tso: update keyspace group revision after initial load

pkg/keyspace/tso_keyspace_group.go

@@ -58,6 +58,7 @@ const (
 	defaultKeyspaceCountSplitThreshold = 40000
 	// autoSplitKeyspaceGroupPatrolInterval is the interval for patrolling keyspace group size for auto-split.
 	autoSplitKeyspaceGroupPatrolInterval = 15 * time.Minute
+	keyspaceGroupRevisionValue           = "1"
 )
 
 const (
@@ -90,6 +91,58 @@ type GroupManager struct {
 	tsoNodesWatcher *etcdutil.LoopWatcher
 }
 
+type keyspaceGroupRevisionStorage struct {
+	endpoint.KeyspaceGroupStorage
+}
+
+type keyspaceGroupRevisionTxn struct {
+	kv.Txn
+	changed bool
+}
+
+// Save records keyspace group writes so the transaction can advance the revision marker.
+func (txn *keyspaceGroupRevisionTxn) Save(key, value string) error {
+	if err := txn.Txn.Save(key, value); err != nil {
+		return err
+	}
+	if strings.HasPrefix(key, keypath.KeyspaceGroupIDPrefix()) {
+		txn.changed = true
+	}
+	return nil
+}
+
+// Remove records keyspace group deletions so the transaction can advance the revision marker.
+func (txn *keyspaceGroupRevisionTxn) Remove(key string) error {
+	if err := txn.Txn.Remove(key); err != nil {
+		return err
+	}
+	if strings.HasPrefix(key, keypath.KeyspaceGroupIDPrefix()) {
+		txn.changed = true
+	}
+	return nil
+}
+
+// RunInTxn atomically advances the revision marker when the transaction changes a keyspace group.
+func (s *keyspaceGroupRevisionStorage) RunInTxn(ctx context.Context, f func(txn kv.Txn) error) error {
+	return s.KeyspaceGroupStorage.RunInTxn(ctx, func(txn kv.Txn) error {
+		revisionTxn := &keyspaceGroupRevisionTxn{Txn: txn}
+		if err := f(revisionTxn); err != nil {
+			return err
+		}
+		if !revisionTxn.changed {
+			return nil
+		}
+		return txn.Save(keypath.KeyspaceGroupRevisionPath(), keyspaceGroupRevisionValue)
+	})
+}
+
+func withKeyspaceGroupRevision(store endpoint.KeyspaceGroupStorage) endpoint.KeyspaceGroupStorage {
+	if _, ok := store.(*keyspaceGroupRevisionStorage); ok {
+		return store
+	}
+	return &keyspaceGroupRevisionStorage{KeyspaceGroupStorage: store}
+}
+
 // NewKeyspaceGroupManager creates a Manager of keyspace group related data.
 func NewKeyspaceGroupManager(
 	ctx context.Context,
@@ -104,7 +157,7 @@ func NewKeyspaceGroupManager(
 	m := &GroupManager{
 		ctx:                ctx,
 		cancel:             cancel,
-		store:              store,
+		store:              withKeyspaceGroupRevision(store),
 		groups:             groups,
 		client:             client,
 		nodesBalancer:      balancer.GenByPolicy[string](defaultBalancerPolicy),
@@ -141,6 +194,13 @@ func (m *GroupManager) Bootstrap(ctx context.Context) error {
 	if err != nil && err != errs.ErrKeyspaceGroupExists {
 		return err
 	}
+	// Persist a marker on every bootstrap so upgraded clusters retain a
+	// revision that is at least as new as all existing keyspace group state.
+	if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error {
+		return txn.Save(keypath.KeyspaceGroupRevisionPath(), keyspaceGroupRevisionValue)
+	}); err != nil {
+		return err
+	}
 
 	// Load all the keyspace groups from the storage and add to the respective userKind groups.
 	groups, err := m.store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0)

pkg/keyspace/tso_keyspace_group_test.go

@@ -33,6 +33,7 @@ import (
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/storage/kv"
 	"github.com/tikv/pd/pkg/utils/etcdutil"
+	"github.com/tikv/pd/pkg/utils/keypath"
 	"github.com/tikv/pd/pkg/versioninfo/kerneltype"
 )
 
@@ -131,6 +132,37 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupOperations() {
 	re.Error(err)
 }
 
+func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupChangesUpdateRevisionMarker() {
+	re := suite.Require()
+
+	removeMarker := func() {
+		re.NoError(suite.kgm.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
+			return txn.Remove(keypath.KeyspaceGroupRevisionPath())
+		}))
+	}
+	checkMarker := func() {
+		var marker string
+		re.NoError(suite.kgm.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
+			var err error
+			marker, err = txn.Load(keypath.KeyspaceGroupRevisionPath())
+			return err
+		}))
+		re.NotEmpty(marker)
+	}
+
+	removeMarker()
+	re.NoError(suite.kgm.CreateKeyspaceGroups([]*endpoint.KeyspaceGroup{{
+		ID:       1,
+		UserKind: endpoint.Standard.String(),
+	}}))
+	checkMarker()
+
+	removeMarker()
+	_, err := suite.kgm.DeleteKeyspaceGroupByID(1)
+	re.NoError(err)
+	checkMarker()
+}
+
 func (suite *keyspaceGroupTestSuite) TestKeyspaceAssignment() {
 	re := suite.Require()
 

pkg/tso/keyspace_group_manager.go

@@ -534,9 +534,25 @@ func (kgm *KeyspaceGroupManager) InitializeTSOServerWatchLoop() error {
 // membership/distribution metadata.
 // Key: /pd/{cluster_id}/tso/keyspace_groups/membership/{group}
 // Value: endpoint.KeyspaceGroup
+// Revision marker: /pd/{cluster_id}/tso/keyspace_groups/revision
 func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error {
 	defaultKGConfigured := false
+	maxLoadedModRevision := uint64(0)
+	preEventsFn := func([]*clientv3.Event) error {
+		maxLoadedModRevision = 0
+		return nil
+	}
 	putFn := func(kv *mvccpb.KeyValue) error {
+		if string(kv.Key) == keypath.KeyspaceGroupRevisionPath() {
+			failpoint.Inject("SkipKeyspaceWatch", func(val failpoint.Value) {
+				addr, ok := val.(string)
+				if ok && addr == kgm.electionNamePrefix {
+					failpoint.Return(nil)
+				}
+			})
+			maxLoadedModRevision = max(maxLoadedModRevision, uint64(kv.ModRevision))
+			return nil
+		}
 		group := &endpoint.KeyspaceGroup{}
 		if err := json.Unmarshal(kv.Value, group); err != nil {
 			return errs.ErrJSONUnmarshal.Wrap(err)
@@ -547,13 +563,17 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error {
 				failpoint.Return(nil)
 			}
 		})
+		maxLoadedModRevision = max(maxLoadedModRevision, uint64(kv.ModRevision))
 		kgm.updateKeyspaceGroup(group)
 		if group.ID == constant.DefaultKeyspaceGroupID {
 			defaultKGConfigured = true
 		}
 		return nil
 	}
 	deleteFn := func(kv *mvccpb.KeyValue) error {
+		if string(kv.Key) == keypath.KeyspaceGroupRevisionPath() {
+			return nil
+		}
 		groupID, err := ExtractKeyspaceGroupIDFromPath(kgm.compiledKGMembershipIDRegexp, string(kv.Key))
 		if err != nil {
 			return err
@@ -581,6 +601,8 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error {
 					zap.Uint64("new-mod-revision", uint64(last.Kv.ModRevision)),
 				)
 			}
+		} else if maxLoadedModRevision > 0 {
+			kgm.SetModRevision(maxLoadedModRevision)
 		}
 		return nil
 	}
@@ -589,9 +611,8 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error {
 		&kgm.wg,
 		kgm.etcdClient,
 		"keyspace-watcher",
-		// To keep the consistency with the previous code, we should trim the suffix `/`.
-		strings.TrimSuffix(keypath.KeyspaceGroupIDPrefix(), "/"),
-		func([]*clientv3.Event) error { return nil },
+		keypath.KeyspaceGroupPrefix(),
+		preEventsFn,
 		putFn,
 		deleteFn,
 		postEventsFn,

pkg/tso/keyspace_group_manager_test.go

@@ -210,6 +210,140 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsAssignment() {
 	suite.runTestLoadKeyspaceGroupsAssignment(re, maxCountInUse+1, 0, 10)
 }
 
+func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSetsModRevision() {
+	re := suite.Require()
+
+	mgr := suite.newUniqueKeyspaceGroupManager(1)
+	re.NotNil(mgr)
+	defer mgr.Close()
+
+	const (
+		groupID    = uint32(1)
+		keyspaceID = uint32(101)
+	)
+	err := addKeyspaceGroupAssignment(
+		suite.ctx,
+		suite.etcdClient,
+		groupID,
+		[]string{mgr.tsoServiceID.ServiceAddr},
+		[]int{mcs.DefaultKeyspaceGroupReplicaPriority},
+		[]uint32{keyspaceID},
+	)
+	re.NoError(err)
+
+	resp, err := suite.etcdClient.Get(suite.ctx, keypath.KeyspaceGroupIDPath(groupID))
+	re.NoError(err)
+	re.Len(resp.Kvs, 1)
+	groupRevision := uint64(resp.Kvs[0].ModRevision)
+
+	const deletedGroupID = uint32(2)
+	err = addKeyspaceGroupAssignment(
+		suite.ctx,
+		suite.etcdClient,
+		deletedGroupID,
+		[]string{mgr.tsoServiceID.ServiceAddr},
+		[]int{mcs.DefaultKeyspaceGroupReplicaPriority},
+		[]uint32{keyspaceID + 1},
+	)
+	re.NoError(err)
+	deleteResp, err := suite.etcdClient.Txn(suite.ctx).Then(
+		clientv3.OpDelete(keypath.KeyspaceGroupIDPath(deletedGroupID)),
+		clientv3.OpPut(keypath.KeyspaceGroupRevisionPath(), "1"),
+	).Commit()
+	re.NoError(err)
+	deletedRevision := uint64(deleteResp.Header.Revision)
+	re.Greater(deletedRevision, groupRevision)
+
+	err = mgr.Initialize()
+	re.NoError(err)
+
+	_, kg, loadedGroupID, loadedRevision, err := mgr.FindGroupByKeyspaceID(keyspaceID)
+	re.NoError(err)
+	re.NotNil(kg)
+	re.Equal(groupID, loadedGroupID)
+	re.Equal(deletedRevision, loadedRevision)
+}
+
+func (suite *keyspaceGroupManagerTestSuite) TestSnapshotRevisionRemainsComparableAcrossManagers() {
+	re := suite.Require()
+	keypath.SetClusterID(rand.Uint64())
+
+	cfg1 := suite.createConfig()
+	cfg2 := suite.createConfig()
+	mgr1 := suite.newKeyspaceGroupManager(1, cfg1)
+	mgr2 := suite.newKeyspaceGroupManager(1, cfg2)
+	defer mgr1.Close()
+	defer mgr2.Close()
+
+	const (
+		groupID    = uint32(1)
+		keyspaceID = uint32(101)
+	)
+	re.NoError(addKeyspaceGroupAssignment(
+		suite.ctx,
+		suite.etcdClient,
+		groupID,
+		[]string{
+			mgr1.tsoServiceID.ServiceAddr,
+			mgr2.tsoServiceID.ServiceAddr,
+		},
+		[]int{
+			mcs.DefaultKeyspaceGroupReplicaPriority,
+			mcs.DefaultKeyspaceGroupReplicaPriority,
+		},
+		[]uint32{keyspaceID},
+	))
+	re.NoError(mgr1.Initialize())
+
+	_, _, _, oldRevision, err := mgr1.FindGroupByKeyspaceID(keyspaceID)
+	re.NoError(err)
+	re.NotZero(oldRevision)
+
+	_, err = suite.etcdClient.Put(suite.ctx, "/unrelated/revision-gap", "1")
+	re.NoError(err)
+
+	re.NoError(mgr2.Initialize())
+	_, _, _, newRevision, err := mgr2.FindGroupByKeyspaceID(keyspaceID)
+	re.NoError(err)
+	re.Equal(oldRevision, newRevision)
+}
+
+func (suite *keyspaceGroupManagerTestSuite) TestInitialSkipKeyspaceWatchDoesNotAdvanceRevision() {
+	re := suite.Require()
+
+	mgr := suite.newUniqueKeyspaceGroupManager(1)
+	re.NotNil(mgr)
+	defer mgr.Close()
+
+	point := fmt.Sprintf("return(\"%s\")", mgr.electionNamePrefix)
+	re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/tso/SkipKeyspaceWatch", point))
+	defer func() {
+		re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/tso/SkipKeyspaceWatch"))
+	}()
+
+	const (
+		groupID    = uint32(1)
+		keyspaceID = uint32(101)
+	)
+	re.NoError(addKeyspaceGroupAssignment(
+		suite.ctx,
+		suite.etcdClient,
+		groupID,
+		[]string{mgr.tsoServiceID.ServiceAddr},
+		[]int{mcs.DefaultKeyspaceGroupReplicaPriority},
+		[]uint32{keyspaceID},
+	))
+	re.NoError(mgr.Initialize())
+
+	mgr.RLock()
+	loadedGroup := mgr.kgs[groupID]
+	loadedRevision := mgr.modRevision
+	mgr.RUnlock()
+
+	re.Nil(loadedGroup)
+	re.Zero(loadedRevision, "a skipped initial watch must not advertise an unapplied revision")
+}
+
 // TestLoadWithDifferentBatchSize tests the loading of the keyspace group assignment with the different batch size.
 func (suite *keyspaceGroupManagerTestSuite) TestLoadWithDifferentBatchSize() {
 	re := suite.Require()

pkg/utils/etcdutil/etcdutil.go

@@ -374,7 +374,6 @@ type LoopWatcher struct {
 	postEventsFn func([]*clientv3.Event) error
 	// preEventsFn is used to call before handling all events.
 	preEventsFn func([]*clientv3.Event) error
-
 	// forceLoadMu is used to ensure two force loads have minimal interval.
 	forceLoadMu syncutil.RWMutex
 	// lastTimeForceLoad is used to record the last time force loading data from etcd.
@@ -639,6 +638,7 @@ func (lw *LoopWatcher) watch(ctx context.Context, revision int64) (nextRevision
 func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error) {
 	startKey := lw.key
 	limit := lw.loadBatchSize
+	snapshotRevision := int64(0)
 	opts := lw.buildLoadingOpts(limit)
 
 	if err := lw.preEventsFn([]*clientv3.Event{}); err != nil {
@@ -677,10 +677,19 @@ func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error)
 					return 0, err
 				}
 				opts = lw.buildLoadingOpts(limit)
+				if snapshotRevision > 0 {
+					opts = append(opts, clientv3.WithRev(snapshotRevision))
+				}
 				continue
 			}
 			return 0, err
 		}
+		if snapshotRevision == 0 {
+			snapshotRevision = resp.Header.Revision
+			// Keep all remaining pages on the same snapshot. Otherwise a write
+			// between pages could be skipped when the watch starts.
+			opts = append(opts, clientv3.WithRev(snapshotRevision))
+		}
 		for i, item := range resp.Kvs {
 			if i == len(resp.Kvs)-1 && resp.More {
 				// If there are more keys, we need to load the next batch.
@@ -701,7 +710,7 @@ func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error)
 		}
 		// Note: if there are no keys in etcd, the resp.More is false. It also means the load is finished.
 		if !resp.More {
-			return resp.Header.Revision + 1, err
+			return snapshotRevision + 1, err
 		}
 	}
 }