tso: fix bugs to make split test case to pass (tikv#6389)

ref tikv#6232

fix bugs to make split test case to pass

Signed-off-by: Bin Shi <binshi.bing@gmail.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

pkg/tso/keyspace_group_manager.go

@@ -131,8 +131,10 @@ func (s *state) getKeyspaceGroupMetaWithCheck(
 	// The keyspace doesn't belong to this keyspace group, we should check if it belongs to any other
 	// keyspace groups, and return the correct keyspace group meta to the client.
 	if kgid, ok := s.keyspaceLookupTable[keyspaceID]; ok {
-		return s.ams[kgid], s.kgs[kgid], kgid,
-			genNotServedErr(errs.ErrGetAllocatorManager, keyspaceGroupID)
+		if s.ams[kgid] != nil {
+			return s.ams[kgid], s.kgs[kgid], kgid, nil
+		}
+		return nil, s.kgs[kgid], kgid, genNotServedErr(errs.ErrGetAllocatorManager, keyspaceGroupID)
 	}
 
 	// The keyspace doesn't belong to any keyspace group but the keyspace has been assigned to a
@@ -463,6 +465,7 @@ func (kgm *KeyspaceGroupManager) watchKeyspaceGroupsMetaChange(revision int64) (
 	defer watcher.Close()
 
 	ksgPrefix := strings.Join([]string{kgm.legacySvcRootPath, endpoint.KeyspaceGroupIDPrefix()}, "/")
+	log.Info("start to watch keyspace group meta change", zap.Int64("revision", revision), zap.String("prefix", ksgPrefix))
 
 	for {
 		watchChan := watcher.Watch(kgm.ctx, ksgPrefix, clientv3.WithPrefix(), clientv3.WithRev(revision))
@@ -840,7 +843,7 @@ func (kgm *KeyspaceGroupManager) HandleTSORequest(
 		return pdpb.Timestamp{}, curKeyspaceGroupID, err
 	}
 	ts, err = am.HandleRequest(dcLocation, count)
-	return ts, keyspaceGroupID, err
+	return ts, curKeyspaceGroupID, err
 }
 
 func (kgm *KeyspaceGroupManager) checkKeySpaceGroupID(id uint32) error {

pkg/tso/keyspace_group_manager_test.go

@@ -420,21 +420,20 @@ func (suite *keyspaceGroupManagerTestSuite) TestGetKeyspaceGroupMetaWithCheck()
 	re.Equal(uint32(0), kgid)
 	re.NotNil(am)
 	re.NotNil(kg)
+	// Should succeed and get the meta of keyspace group 0, because keyspace 0
+	// belongs to group 0, though the specified group 1 doesn't exist.
+	am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(0, 1)
+	re.NoError(err)
+	re.Equal(uint32(0), kgid)
+	re.NotNil(am)
+	re.NotNil(kg)
 	// Should fail because keyspace 3 isn't explicitly assigned to any keyspace
 	// group, and the specified group isn't the default keyspace group.
 	am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(3, 100)
 	re.Error(err)
 	re.Equal(uint32(100), kgid)
 	re.Nil(am)
 	re.Nil(kg)
-	// Should fail but still be able to get the meta of keyspace group 0,
-	// because keyspace 0 belongs to group 0, though the specified group 1
-	// doesn't exist.
-	am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(0, 1)
-	re.Error(err)
-	re.Equal(uint32(0), kgid)
-	re.NotNil(am)
-	re.NotNil(kg)
 }
 
 // TestDefaultMembershipRestriction tests the restriction of default keyspace always
@@ -498,9 +497,9 @@ func (suite *keyspaceGroupManagerTestSuite) TestDefaultMembershipRestriction() {
 	re.Equal(mcsutils.DefaultKeyspaceGroupID, kgid)
 	re.NotNil(am)
 	re.NotNil(kg)
-	// Should fail but still be able to get the keyspace group meta from the default keyspace group
+	// Should succeed and return the keyspace group meta from the default keyspace group
 	am, kg, kgid, err = mgr.getKeyspaceGroupMetaWithCheck(mcsutils.DefaultKeyspaceID, 3)
-	re.Error(err)
+	re.NoError(err)
 	re.Equal(mcsutils.DefaultKeyspaceGroupID, kgid)
 	re.NotNil(am)
 	re.NotNil(kg)