server: Add feature flags that helps TiKV to implement feature gate that doesn't depend on incremental semver

client/go.mod

@@ -43,3 +43,5 @@ require (
 	gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+replace github.com/pingcap/kvproto => github.com/MyonKeminta/kvproto v0.0.0-20260203075250-99a2286adbee

client/go.sum

@@ -1,5 +1,7 @@
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/MyonKeminta/kvproto v0.0.0-20260203075250-99a2286adbee h1:MmEK5DwZJPhRMfIu4X1SxM+Qamqtkz7/b/AiZ53l8e8=
+github.com/MyonKeminta/kvproto v0.0.0-20260203075250-99a2286adbee/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
 github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
@@ -53,8 +55,6 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c h1:xpW9bvK+HuuTm
 github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ueLEUZgtx2cIogM0v4Zj5uvvzhuuiu7Pn8HzMPg=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
-github.com/pingcap/kvproto v0.0.0-20260106110113-438649d89ee7 h1:ENZsbCKjZe5GyNVuKHTGYiYd4LnM8PRphBsYU6M19g0=
-github.com/pingcap/kvproto v0.0.0-20260106110113-438649d89ee7/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

go.mod

@@ -250,3 +250,5 @@ require (
 // which will cause several different tests to fail. So this is a temporary workaround to use the old version of `testify`.
 // TODO: fix those flasky tests introduced by the behavior change of `Eventually` and `EventuallyWithT` assertions.
 replace github.com/stretchr/testify => github.com/stretchr/testify v1.10.0
+
+replace github.com/pingcap/kvproto => github.com/MyonKeminta/kvproto v0.0.0-20260203075250-99a2286adbee

pkg/core/store_option.go

@@ -67,6 +67,15 @@ func SetStoreVersion(githash, version string) StoreCreateOption {
 	}
 }
 
+// SetStoreFeatureFlags sets the feature flags for the store.
+func SetStoreFeatureFlags(featureFlags uint64) StoreCreateOption {
+	return func(store *StoreInfo) {
+		meta := typeutil.DeepClone(store.meta, StoreFactory)
+		meta.FeatureFlags = featureFlags
+		store.meta = meta
+	}
+}
+
 // SetStoreDeployPath sets the deploy path for the store.
 func SetStoreDeployPath(deployPath string) StoreCreateOption {
 	return func(store *StoreInfo) {

server/cluster/cluster.go

@@ -125,6 +125,23 @@ const (
 	syncRegionTaskRunner = "sync-region-async"
 )
 
+// StoreFeatureFlags represents a set of flags indicating whether a feature of stores is available in a cluster.
+// This type is used for internally deciding whether a feature can be safely enabled. This is used to manages
+// features that:
+//
+//   - a single store node cannot enable it without confirming that the whole cluster has supported it otherwise there
+//     would be compatibility issue;
+//   - the feature gate based on version number is for some reason not applicable.
+//
+// This type directly equals to the field `metapb.Store.FeatureFlags`.
+type StoreFeatureFlags uint64
+
+const (
+	// FeatureFlagInMemoryPessimisticLocksNextGen is the flag representing the feature in-memory pessimistic
+	// locks for next gen or TiDB-X clusters.
+	FeatureFlagInMemoryPessimisticLocksNextGen StoreFeatureFlags = 1 << iota
+)
+
 // Server is the interface for cluster.
 type Server interface {
 	GetAllocator() id.Allocator
@@ -200,6 +217,12 @@ type RaftCluster struct {
 	logRunner ratelimit.Runner
 	// syncRegionRunner is used to sync region asynchronously.
 	syncRegionRunner ratelimit.Runner
+
+	clusterVersionUpdateMu syncutil.Mutex
+	// clusterFeatureFlags is the current feature flags indicating what features (in a specific set) can be safely
+	// enabled in the cluster.
+	// This field is only updated with clusterVersionUpdateMu acquired.
+	clusterFeatureFlags atomic.Uint64
 }
 
 // Status saves some state information.
@@ -808,6 +831,8 @@ func (c *RaftCluster) LoadClusterInfo() (*RaftCluster, error) {
 		zap.Duration("cost", time.Since(start)),
 	)
 
+	c.updateClusterFeatureFlags(c.GetStores())
+
 	start = time.Now()
 
 	// used to load region from kv storage to cache storage.
@@ -1464,6 +1489,7 @@ func (c *RaftCluster) putStoreImpl(store *metapb.Store, force bool) error {
 		opts = append(opts,
 			core.SetStoreAddress(store.Address, store.StatusAddress, store.PeerAddress),
 			core.SetStoreVersion(store.GitHash, store.Version),
+			core.SetStoreFeatureFlags(store.FeatureFlags),
 			core.SetStoreLabels(labels),
 			core.SetStoreStartTime(store.StartTimestamp),
 			core.SetStoreDeployPath(store.DeployPath))
@@ -2051,6 +2077,9 @@ func (*RaftCluster) resetHealthStatus() {
 
 // OnStoreVersionChange changes the version of the cluster when needed.
 func (c *RaftCluster) OnStoreVersionChange() {
+	c.clusterVersionUpdateMu.Lock()
+	defer c.clusterVersionUpdateMu.Unlock()
+
 	var minVersion *semver.Version
 	stores := c.GetStores()
 	for _, s := range stores {
@@ -2069,20 +2098,46 @@ func (c *RaftCluster) OnStoreVersionChange() {
 	failpoint.Inject("versionChangeConcurrency", func() {
 		time.Sleep(500 * time.Millisecond)
 	})
-	if minVersion == nil || clusterVersion.Equal(*minVersion) {
-		return
+	if minVersion != nil && !clusterVersion.Equal(*minVersion) {
+		if !c.opt.CASClusterVersion(clusterVersion, minVersion) {
+			log.Error("cluster version changed by API at the same time")
+		}
+		err := c.opt.Persist(c.storage)
+		if err != nil {
+			log.Error("persist cluster version meet error", errs.ZapError(err))
+		}
+		log.Info("cluster version changed",
+			zap.Stringer("old-cluster-version", clusterVersion),
+			zap.Stringer("new-cluster-version", minVersion))
 	}
 
-	if !c.opt.CASClusterVersion(clusterVersion, minVersion) {
-		log.Error("cluster version changed by API at the same time")
-	}
-	err := c.opt.Persist(c.storage)
-	if err != nil {
-		log.Error("persist cluster version meet error", errs.ZapError(err))
+	c.updateClusterFeatureFlags(stores)
+}
+
+func (c *RaftCluster) updateClusterFeatureFlags(stores []*core.StoreInfo) {
+	newFeatureFlags := calcClusterFeatureFlags(stores)
+	c.clusterFeatureFlags.Store(uint64(newFeatureFlags))
+}
+
+// calcClusterFeatureFlags merges feature flags of all stores and outputs feature flags of the whole cluster.
+func calcClusterFeatureFlags(stores []*core.StoreInfo) StoreFeatureFlags {
+	var res StoreFeatureFlags
+	// Add conjunction flags, and exclude it later if some of the nodes does not satisfy.
+	res |= FeatureFlagInMemoryPessimisticLocksNextGen
+
+	for _, store := range stores {
+		if store.IsRemoved() {
+			continue
+		}
+
+		nodeFeatureFlags := StoreFeatureFlags(store.GetMeta().GetFeatureFlags())
+		// In-memory pessimistic lock can be enabled if all TiKV nodes supports it,
+		if store.IsTiKV() && nodeFeatureFlags&FeatureFlagInMemoryPessimisticLocksNextGen == 0 {
+			res &= ^FeatureFlagInMemoryPessimisticLocksNextGen
+		}
 	}
-	log.Info("cluster version changed",
-		zap.Stringer("old-cluster-version", clusterVersion),
-		zap.Stringer("new-cluster-version", minVersion))
+
+	return res
 }
 
 func (c *RaftCluster) changedRegionNotifier() <-chan *core.RegionInfo {
@@ -2472,6 +2527,11 @@ func (c *RaftCluster) GetClusterVersion() string {
 	return c.opt.GetClusterVersion().String()
 }
 
+// GetClusterFeatureFlags returns the current cluster feature flags.
+func (c *RaftCluster) GetClusterFeatureFlags() StoreFeatureFlags {
+	return StoreFeatureFlags(c.clusterFeatureFlags.Load())
+}
+
 // GetEtcdClient returns the current etcd client
 func (c *RaftCluster) GetEtcdClient() *clientv3.Client {
 	return c.etcdClient

server/cluster/cluster_test.go

@@ -4204,3 +4204,173 @@ func BenchmarkHandleRegionHeartbeat(b *testing.B) {
 		re.NoError(err)
 	}
 }
+
+func TestCalcClusterFeatureFlags_InMemoryPessimisticLocksForNextGen(t *testing.T) {
+	re := require.New(t)
+
+	newStore := func(id uint64, addr string, flags uint64) *core.StoreInfo {
+		return core.NewStoreInfo(&metapb.Store{
+			Id:           id,
+			Address:      addr,
+			Version:      "8.5.0",
+			FeatureFlags: flags,
+		})
+	}
+	newTiFlashStore := func(id uint64, addr string, flags uint64) *core.StoreInfo {
+		return core.NewStoreInfo(&metapb.Store{
+			Id:           id,
+			Address:      addr,
+			Version:      "8.5.0",
+			FeatureFlags: flags,
+			Labels:       []*metapb.StoreLabel{{Key: core.EngineKey, Value: core.EngineTiFlash}},
+		})
+	}
+	newTombstoneStore := func(id uint64, addr string, flags uint64) *core.StoreInfo {
+		return core.NewStoreInfo(&metapb.Store{
+			Id:           id,
+			Address:      addr,
+			Version:      "8.5.0",
+			FeatureFlags: flags,
+			State:        metapb.StoreState_Tombstone,
+			NodeState:    metapb.NodeState_Removed,
+		})
+	}
+
+	flag := uint64(FeatureFlagInMemoryPessimisticLocksNextGen)
+
+	// All TiKV stores support the flag.
+	stores := []*core.StoreInfo{
+		newStore(1, "tikv1", flag),
+		newStore(2, "tikv2", flag),
+		newStore(3, "tikv3", flag),
+	}
+	re.Equal(FeatureFlagInMemoryPessimisticLocksNextGen, calcClusterFeatureFlags(stores))
+
+	// One TiKV store doesn't support the flag.
+	stores = []*core.StoreInfo{
+		newStore(1, "tikv1", flag),
+		newStore(2, "tikv2", 0),
+		newStore(3, "tikv3", flag),
+	}
+	re.Equal(StoreFeatureFlags(0), calcClusterFeatureFlags(stores))
+
+	// Tombstone stores should be ignored.
+	stores = []*core.StoreInfo{
+		newStore(1, "tikv1", flag),
+		newTombstoneStore(2, "tikv2", 0),
+		newStore(3, "tikv3", flag),
+	}
+	re.Equal(FeatureFlagInMemoryPessimisticLocksNextGen, calcClusterFeatureFlags(stores))
+
+	// In-memory pessimistic locks don't have the requirement on TiFlash nodes.
+	stores = []*core.StoreInfo{
+		newStore(1, "tikv1", flag),
+		newStore(2, "tikv2", flag),
+		newTiFlashStore(3, "tiflash1", 0),
+	}
+	re.Equal(FeatureFlagInMemoryPessimisticLocksNextGen, calcClusterFeatureFlags(stores))
+}
+
+func TestFeatureFlagsPersistence(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	_, opt, err := newTestScheduleConfig()
+	re.NoError(err)
+	s := storage.NewStorageWithMemoryBackend()
+	c := newTestRaftCluster(ctx, mockid.NewIDAllocator(), opt, s)
+
+	// A store with no feature flag set.
+	store := &metapb.Store{
+		Id:           1,
+		Address:      "tikv1",
+		Version:      "8.5.0",
+		FeatureFlags: 0,
+	}
+	re.NoError(c.PutMetaStore(store))
+
+	var loaded metapb.Store
+	ok, err := s.LoadStoreMeta(1, &loaded)
+	re.NoError(err)
+	re.True(ok)
+	re.Equal(uint64(0), loaded.FeatureFlags)
+
+	// A store with a flag set.
+	store = &metapb.Store{
+		Id:           2,
+		Address:      "tikv2",
+		Version:      "8.5.0",
+		FeatureFlags: uint64(FeatureFlagInMemoryPessimisticLocksNextGen),
+	}
+	re.NoError(c.PutMetaStore(store))
+
+	ok, err = s.LoadStoreMeta(2, &loaded)
+	re.NoError(err)
+	re.True(ok)
+	re.Equal(uint64(FeatureFlagInMemoryPessimisticLocksNextGen), loaded.FeatureFlags)
+
+	// Updating an existing store and set a flag, simulating upgrading.
+	store = &metapb.Store{
+		Id:           1,
+		Address:      "tikv1",
+		Version:      "8.5.0",
+		FeatureFlags: uint64(FeatureFlagInMemoryPessimisticLocksNextGen),
+	}
+	re.NoError(c.PutMetaStore(store))
+
+	ok, err = s.LoadStoreMeta(1, &loaded)
+	re.NoError(err)
+	re.True(ok)
+	re.Equal(uint64(FeatureFlagInMemoryPessimisticLocksNextGen), loaded.FeatureFlags)
+}
+
+func TestGetClusterFeatureFlags(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	_, opt, err := newTestScheduleConfig()
+	re.NoError(err)
+	c := newTestRaftCluster(ctx, mockid.NewIDAllocator(), opt, storage.NewStorageWithMemoryBackend())
+
+	flag := uint64(FeatureFlagInMemoryPessimisticLocksNextGen)
+
+	// Initialize with mixed states: store 1 without flag, store 2 with flag.
+	re.NoError(c.PutMetaStore(&metapb.Store{
+		Id: 1, Address: "tikv1", Version: "8.5.0", FeatureFlags: 0,
+	}))
+	re.NoError(c.PutMetaStore(&metapb.Store{
+		Id: 2, Address: "tikv2", Version: "8.5.0", FeatureFlags: flag,
+	}))
+	re.Equal(StoreFeatureFlags(0), c.GetClusterFeatureFlags())
+
+	// Add TiFlash without the flag - should not affect cluster flags.
+	re.NoError(c.PutMetaStore(&metapb.Store{
+		Id: 3, Address: "tiflash1", Version: "8.5.0", FeatureFlags: 0,
+		Labels: []*metapb.StoreLabel{{Key: core.EngineKey, Value: core.EngineTiFlash}},
+	}))
+	re.Equal(StoreFeatureFlags(0), c.GetClusterFeatureFlags())
+
+	// Upgrade store 1 to support the flag.
+	re.NoError(c.PutMetaStore(&metapb.Store{
+		Id: 1, Address: "tikv1", Version: "8.5.0", FeatureFlags: flag,
+	}))
+	re.Equal(FeatureFlagInMemoryPessimisticLocksNextGen, c.GetClusterFeatureFlags())
+
+	// Add another TiKV with the flag - cluster flags should remain set.
+	re.NoError(c.PutMetaStore(&metapb.Store{
+		Id: 4, Address: "tikv3", Version: "8.5.0", FeatureFlags: flag,
+	}))
+	re.Equal(FeatureFlagInMemoryPessimisticLocksNextGen, c.GetClusterFeatureFlags())
+
+	// Add a TiKV without the flag - cluster flags should be cleared.
+	re.NoError(c.PutMetaStore(&metapb.Store{
+		Id: 5, Address: "tikv4", Version: "8.5.0", FeatureFlags: 0,
+	}))
+	re.Equal(StoreFeatureFlags(0), c.GetClusterFeatureFlags())
+
+	// Bury the store without flag - cluster flags should be set again.
+	re.NoError(c.BuryStore(5, true))
+	re.Equal(FeatureFlagInMemoryPessimisticLocksNextGen, c.GetClusterFeatureFlags())
+}

server/grpc_service.go

@@ -975,6 +975,7 @@ func (s *GrpcServer) StoreHeartbeat(ctx context.Context, request *pdpb.StoreHear
 
 	resp.ReplicationStatus = rc.GetReplicationMode().GetReplicationStatus()
 	resp.ClusterVersion = rc.GetClusterVersion()
+	resp.FeatureFlags = uint64(rc.GetClusterFeatureFlags())
 	rc.GetUnsafeRecoveryController().HandleStoreHeartbeat(request, resp)
 
 	return resp, nil

tests/integrations/go.mod

@@ -239,3 +239,5 @@ require (
 // which will cause several different tests to fail. So this is a temporary workaround to use the old version of `testify`.
 // TODO: fix those flasky tests introduced by the behavior change of `Eventually` and `EventuallyWithT` assertions.
 replace github.com/stretchr/testify => github.com/stretchr/testify v1.10.0
+
+replace github.com/pingcap/kvproto => github.com/MyonKeminta/kvproto v0.0.0-20260203075250-99a2286adbee