tso: Fix keyspace-group split causing keyspace to fallback to group 0 (#9823)

close #9822

Signed-off-by: ystaticy <y_static_y@sina.com>

Author: ystaticy

pkg/keyspace/tso_keyspace_group.go

@@ -587,9 +587,6 @@ func (m *GroupManager) SplitKeyspaceGroupByID(
 		splitSourceKg.SplitState = &endpoint.SplitState{
 			SplitSource: splitSourceKg.ID,
 		}
-		if err = m.store.SaveKeyspaceGroup(txn, splitSourceKg); err != nil {
-			return err
-		}
 		splitTargetKg = &endpoint.KeyspaceGroup{
 			ID: splitTargetID,
 			// Keep the same user kind and members as the old keyspace group.
@@ -600,8 +597,18 @@ func (m *GroupManager) SplitKeyspaceGroupByID(
 				SplitSource: splitSourceKg.ID,
 			},
 		}
+		// Save the split target keyspace group first, then save the split source keyspace group.
+		// The order matters: if we save splitSourceKg first (which removes some keyspaces from it),
+		// there will be a brief moment where those keyspaces don't belong to any group, causing
+		// them to fallback to group 0. By saving splitTargetKg first (which contains the split
+		// keyspaces), we ensure the keyspaces always belong to a valid group during the transition.
+
 		// Create the new split keyspace group.
-		return m.store.SaveKeyspaceGroup(txn, splitTargetKg)
+		if err = m.store.SaveKeyspaceGroup(txn, splitTargetKg); err != nil {
+			return err
+		}
+		// Update the source keyspace group.
+		return m.store.SaveKeyspaceGroup(txn, splitSourceKg)
 	}); err != nil {
 		return err
 	}
@@ -611,9 +618,9 @@ func (m *GroupManager) SplitKeyspaceGroupByID(
 	return nil
 }
 
+// `old` is the original keyspace list which will be split out,
+// `new` is the keyspace list which will be split from the old keyspace list.
 func buildSplitKeyspaces(
-	// `old` is the original keyspace list which will be split out,
-	// `new` is the keyspace list which will be split from the old keyspace list.
 	old, new []uint32,
 	startKeyspaceID, endKeyspaceID uint32,
 ) (oldSplit, newSplit []uint32, err error) {

pkg/tso/keyspace_group_manager.go

@@ -686,6 +686,16 @@ func (kgm *KeyspaceGroupManager) updateKeyspaceGroup(group *endpoint.KeyspaceGro
 		kgm.metrics.mergeTargetGauge.Dec()
 	}
 
+	failpoint.Inject("delayBeforeCheckingElectionMember", func() {
+		if oldAM == nil && (group.IsSplitTarget() || group.IsSplitSource()) {
+			log.Info("INJECTED DELAY before update the keyspace group",
+				zap.Uint32("target-group-id", group.ID),
+				zap.Uint32("source-group-id", group.SplitState.SplitSource),
+				zap.Int("keyspaces-count", len(group.Keyspaces)))
+			time.Sleep(10 * time.Second) // Adjust this to match client query timing
+		}
+	})
+
 	// If this host is already assigned a replica of this keyspace group, i.e., the member
 	// is already initialized, just update the meta.
 	if oldAM != nil {

tests/integrations/mcs/tso/keyspace_group_manager_test.go

@@ -409,6 +409,13 @@ func (suite *tsoKeyspaceGroupManagerTestSuite) TestTSOKeyspaceGroupSplitClient()
 	re := suite.Require()
 	// Enable the failpoint to slow down the system time to test whether the TSO is monotonic.
 	re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/tso/systemTimeSlow", `return(true)`))
+
+	// Enable the failpoint to delay before checking election member
+	re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/tso/delayBeforeCheckingElectionMember", `return()`))
+	defer func() {
+		re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/tso/delayBeforeCheckingElectionMember"))
+	}()
+
 	// Create the keyspace group `oldID` with keyspaces [777, 888, 999].
 	oldID := suite.allocID()
 	handlersutil.MustCreateKeyspaceGroup(re, suite.pdLeaderServer, &handlers.CreateKeyspaceGroupParams{