fix startWatchLoop leak (#6352)

Signed-off-by: Ryan Leung <[email protected]>
Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

pkg/keyspace/keyspace_test.go

@@ -61,7 +61,7 @@ func (suite *keyspaceTestSuite) SetupTest() {
 	allocator := mockid.NewIDAllocator()
 	kgm := NewKeyspaceGroupManager(suite.ctx, store, nil, 0)
 	suite.manager = NewKeyspaceManager(store, nil, allocator, &mockConfig{}, kgm)
-	suite.NoError(kgm.Bootstrap())
+	suite.NoError(kgm.Bootstrap(suite.ctx))
 	suite.NoError(suite.manager.Bootstrap())
 }
 

pkg/keyspace/tso_keyspace_group.go

@@ -101,7 +101,7 @@ func NewKeyspaceGroupManager(ctx context.Context, store endpoint.KeyspaceGroupSt
 }
 
 // Bootstrap saves default keyspace group info and init group mapping in the memory.
-func (m *GroupManager) Bootstrap() error {
+func (m *GroupManager) Bootstrap(ctx context.Context) error {
 	// Force the membership restriction that the default keyspace must belong to default keyspace group.
 	// Have no information to specify the distribution of the default keyspace group replicas, so just
 	// leave the replica/member list empty. The TSO service will assign the default keyspace group replica
@@ -135,7 +135,7 @@ func (m *GroupManager) Bootstrap() error {
 		m.nodesBalancer = balancer.GenByPolicy[string](m.policy)
 		m.serviceRegistryMap = make(map[string]string)
 		m.wg.Add(1)
-		go m.startWatchLoop()
+		go m.startWatchLoop(ctx)
 	}
 	return nil
 }
@@ -146,10 +146,10 @@ func (m *GroupManager) Close() {
 	m.wg.Wait()
 }
 
-func (m *GroupManager) startWatchLoop() {
+func (m *GroupManager) startWatchLoop(parentCtx context.Context) {
 	defer logutil.LogPanic()
 	defer m.wg.Done()
-	ctx, cancel := context.WithCancel(m.ctx)
+	ctx, cancel := context.WithCancel(parentCtx)
 	defer cancel()
 	var (
 		resp     *clientv3.GetResponse

pkg/keyspace/tso_keyspace_group_test.go

@@ -47,7 +47,7 @@ func (suite *keyspaceGroupTestSuite) SetupTest() {
 	idAllocator := mockid.NewIDAllocator()
 	cluster := mockcluster.NewCluster(suite.ctx, mockconfig.NewTestOptions())
 	suite.kg = NewKeyspaceManager(store, cluster, idAllocator, &mockConfig{}, suite.kgm)
-	suite.NoError(suite.kgm.Bootstrap())
+	suite.NoError(suite.kgm.Bootstrap(suite.ctx))
 }
 
 func (suite *keyspaceGroupTestSuite) TearDownTest() {

server/cluster/cluster.go

@@ -275,7 +275,7 @@ func (c *RaftCluster) Start(s Server) error {
 		return nil
 	}
 	if s.IsAPIServiceMode() {
-		err = c.keyspaceGroupManager.Bootstrap()
+		err = c.keyspaceGroupManager.Bootstrap(c.ctx)
 		if err != nil {
 			return err
 		}