tikv: support CSE keyspace-level transaction safepoint

tikv/compatible_txn_safe_point_loader.go

@@ -10,6 +10,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/pingcap/kvproto/pkg/keyspacepb"
 	"github.com/pkg/errors"
 	"github.com/tikv/client-go/v2/internal/apicodec"
 	"github.com/tikv/client-go/v2/util"
@@ -23,6 +24,19 @@ const (
 	keyspaceLevelTxnSafePointPath = "/keyspaces/tidb/%d/tidb/store/gcworker/saved_safe_point"
 )
 
+// IsCSEKeyspaceLevelGC reports whether a keyspace uses the CSE keyspace-level GC metadata format.
+// TODO: Replace this implementation with pd.IsCSEKeyspaceLevelGC after API v3 support is merged into client-go.
+func IsCSEKeyspaceLevelGC(keyspaceMeta *keyspacepb.KeyspaceMeta) bool {
+	return keyspaceMeta != nil && keyspaceMeta.Config != nil && keyspaceMeta.Config["safe_point_version"] == "v2"
+}
+
+func compatibleTxnSafePointPath(keyspaceMeta *keyspacepb.KeyspaceMeta) string {
+	if pd.IsKeyspaceUsingKeyspaceLevelGC(keyspaceMeta) || IsCSEKeyspaceLevelGC(keyspaceMeta) {
+		return fmt.Sprintf(keyspaceLevelTxnSafePointPath, keyspaceMeta.Id)
+	}
+	return unifiedTxnSafePointPath
+}
+
 // compatibleTxnSafePointLoader is used to load txn safe point from etcd for old versions where the GetGCState API
 // is not yet supported.
 //
@@ -101,11 +115,7 @@ func (l *compatibleTxnSafePointLoader) loadTxnSafePoint(ctx context.Context) (ui
 		}
 	}
 
-	key := unifiedTxnSafePointPath
-	keyspaceMeta := l.codec.GetKeyspaceMeta()
-	if pd.IsKeyspaceUsingKeyspaceLevelGC(keyspaceMeta) {
-		key = fmt.Sprintf(keyspaceLevelTxnSafePointPath, keyspaceMeta.Id)
-	}
+	key := compatibleTxnSafePointPath(l.codec.GetKeyspaceMeta())
 
 	// Follow the same implementation as the EtcdSafePointKV by setting the timeout 5 seconds.
 	ctx, cancel := context.WithTimeout(ctx, time.Second*5)

tikv/compatible_txn_safe_point_loader_test.go

@@ -0,0 +1,61 @@
+package tikv
+
+import (
+	"testing"
+
+	"github.com/pingcap/kvproto/pkg/keyspacepb"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsCSEKeyspaceLevelGC(t *testing.T) {
+	testCases := []struct {
+		name string
+		meta *keyspacepb.KeyspaceMeta
+		want bool
+	}{
+		{name: "nil keyspace"},
+		{name: "missing config", meta: &keyspacepb.KeyspaceMeta{}},
+		{name: "CSE keyspace-level GC", meta: &keyspacepb.KeyspaceMeta{Config: map[string]string{"safe_point_version": "v2"}}, want: true},
+		{name: "case-sensitive version", meta: &keyspacepb.KeyspaceMeta{Config: map[string]string{"safe_point_version": "V2"}}},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.want, IsCSEKeyspaceLevelGC(tc.meta))
+		})
+	}
+}
+
+func TestCompatibleTxnSafePointPath(t *testing.T) {
+	testCases := []struct {
+		name string
+		meta *keyspacepb.KeyspaceMeta
+		want string
+	}{
+		{
+			name: "null keyspace",
+			want: unifiedTxnSafePointPath,
+		},
+		{
+			name: "native keyspace-level GC",
+			meta: &keyspacepb.KeyspaceMeta{Id: 1, Config: map[string]string{"gc_management_type": "keyspace_level"}},
+			want: "/keyspaces/tidb/1/tidb/store/gcworker/saved_safe_point",
+		},
+		{
+			name: "CSE keyspace-level GC",
+			meta: &keyspacepb.KeyspaceMeta{Id: 2, Config: map[string]string{"safe_point_version": "v2"}},
+			want: "/keyspaces/tidb/2/tidb/store/gcworker/saved_safe_point",
+		},
+		{
+			name: "unified GC",
+			meta: &keyspacepb.KeyspaceMeta{Id: 3, Config: map[string]string{"gc_management_type": "unified"}},
+			want: unifiedTxnSafePointPath,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.want, compatibleTxnSafePointPath(tc.meta))
+		})
+	}
+}