codec, checker: fix enable-cross-table-merge for keyspace keys (#10992)

pkg/codec/codec.go

@@ -33,42 +33,103 @@ const (
 	encGroupSize = 8
 	encMarker    = byte(0xFF)
 	encPad       = byte(0x0)
+
+	// RawKeyspaceModePrefix is the raw keyspace prefix mode byte.
+	RawKeyspaceModePrefix = byte('r')
+	// TxnKeyspaceModePrefix is the txn keyspace prefix mode byte.
+	TxnKeyspaceModePrefix = byte('x')
+	// KeyspacePrefixLen is the raw keyspace prefix length before memcomparable encoding.
+	KeyspacePrefixLen = 4
 )
 
 // Key represents high-level Key type.
 type Key []byte
 
-// TableID returns the table ID of the key, if the key is not table key, returns 0.
-func (k Key) TableID() int64 {
+// MakeKeyspacePrefix constructs the raw keyspace prefix for the given mode and keyspace ID.
+// Keyspace keys encode the lower 24 bits of the keyspace ID after the mode byte.
+func MakeKeyspacePrefix(mode byte, id uint32) []byte {
+	prefix := make([]byte, KeyspacePrefixLen)
+	binary.BigEndian.PutUint32(prefix, id)
+	prefix[0] = mode
+	return prefix
+}
+
+// ParseKeyspacePrefix parses a raw keyspace prefix from key.
+// It returns false for keys that do not start with a known keyspace mode byte.
+func ParseKeyspacePrefix(key []byte) (mode byte, id uint32, ok bool) {
+	if len(key) < KeyspacePrefixLen {
+		return 0, 0, false
+	}
+	mode = key[0]
+	if mode != RawKeyspaceModePrefix && mode != TxnKeyspaceModePrefix {
+		return 0, 0, false
+	}
+	idBytes := [KeyspacePrefixLen]byte{0, key[1], key[2], key[3]}
+	id = binary.BigEndian.Uint32(idBytes[:])
+	return mode, id, true
+}
+
+// unwrapKeyspace strips the API v2 txn keyspace prefix (mode byte + 24-bit id)
+// when the remainder is a TiDB meta/table key. TiDB data only lives under the
+// txn ('x') mode; raw-mode payloads are arbitrary user bytes, so raw keys and
+// keys that only happen to start with 'x' are left unchanged with hasKeyspace
+// false.
+func unwrapKeyspace(key []byte) (payload []byte, keyspaceID uint32, hasKeyspace bool) {
+	mode, keyspaceID, ok := ParseKeyspacePrefix(key)
+	if !ok || mode != TxnKeyspaceModePrefix {
+		return key, 0, false
+	}
+	rest := key[KeyspacePrefixLen:]
+	if !bytes.HasPrefix(rest, tablePrefix) && !bytes.HasPrefix(rest, metaPrefix) {
+		return key, 0, false
+	}
+	return rest, keyspaceID, true
+}
+
+// TableIdentity identifies the logical table a key belongs to. HasKeyspace is
+// false for classic TiDB keys, distinguishing them from keyspace 0. TableID is
+// 0 when the key is not a table key (including meta keys), so all non-table
+// keys of one keyspace share a single identity. Two table keys belong to the
+// same logical table iff their TableIdentity values are equal.
+type TableIdentity struct {
+	KeyspaceID  uint32
+	TableID     int64
+	HasKeyspace bool
+}
+
+// TableIdentity returns the keyspace-qualified table identity of an encoded key.
+func (k Key) TableIdentity() TableIdentity {
 	_, key, err := DecodeBytes(k)
 	if err != nil {
-		// should never happen
-		return 0
+		// should never happen for region boundary keys produced by TiKV
+		return TableIdentity{}
 	}
-	if !bytes.HasPrefix(key, tablePrefix) {
-		return 0
+	key, keyspaceID, hasKeyspace := unwrapKeyspace(key)
+	identity := TableIdentity{KeyspaceID: keyspaceID, HasKeyspace: hasKeyspace}
+	if bytes.HasPrefix(key, tablePrefix) {
+		// A truncated table key fails to decode and keeps TableID 0, i.e. it
+		// is treated as a non-table key, matching the historical semantics.
+		_, identity.TableID, _ = DecodeInt(key[len(tablePrefix):])
 	}
-	key = key[len(tablePrefix):]
-
-	_, tableID, _ := DecodeInt(key)
-	return tableID
+	return identity
 }
 
 // MetaOrTable checks if the key is a meta key or table key.
 // If the key is a meta key, it returns true and 0.
 // If the key is a table key, it returns false and table ID.
 // Otherwise, it returns false and 0.
+// It supports both classic TiDB keys and API v2 keyspace-prefixed keys.
 func (k Key) MetaOrTable() (bool, int64) {
 	_, key, err := DecodeBytes(k)
 	if err != nil {
 		return false, 0
 	}
+	key, _, _ = unwrapKeyspace(key)
 	if bytes.HasPrefix(key, metaPrefix) {
 		return true, 0
 	}
 	if bytes.HasPrefix(key, tablePrefix) {
-		key = key[len(tablePrefix):]
-		_, tableID, _ := DecodeInt(key)
+		_, tableID, _ := DecodeInt(key[len(tablePrefix):])
 		return false, tableID
 	}
 	return false, 0

pkg/codec/codec_test.go

@@ -38,17 +38,107 @@ func TestDecodeBytes(t *testing.T) {
 func TestTableID(t *testing.T) {
 	re := require.New(t)
 	key := EncodeBytes([]byte("t\x80\x00\x00\x00\x00\x00\x00\xff"))
-	re.Equal(int64(0xff), key.TableID())
+	re.Equal(int64(0xff), key.TableIdentity().TableID)
 
 	key = EncodeBytes([]byte("t\x80\x00\x00\x00\x00\x00\x00\xff_i\x01\x02"))
-	re.Equal(int64(0xff), key.TableID())
+	re.Equal(int64(0xff), key.TableIdentity().TableID)
 
 	key = []byte("t\x80\x00\x00\x00\x00\x00\x00\xff")
-	re.Equal(int64(0), key.TableID())
+	re.Equal(int64(0), key.TableIdentity().TableID)
 
 	key = EncodeBytes([]byte("T\x00\x00\x00\x00\x00\x00\x00\xff"))
-	re.Equal(int64(0), key.TableID())
+	re.Equal(int64(0), key.TableIdentity().TableID)
 
 	key = EncodeBytes([]byte("t\x80\x00\x00\x00\x00\x00\xff"))
-	re.Equal(int64(0), key.TableID())
+	re.Equal(int64(0), key.TableIdentity().TableID)
+}
+
+func TestTableIDWithKeyspacePrefix(t *testing.T) {
+	re := require.New(t)
+	tableID := int64(100)
+	otherTableID := int64(200)
+	keyspaceID := uint32(42)
+
+	classic := EncodeBytes(GenerateTableKey(tableID))
+	re.Equal(TableIdentity{TableID: tableID}, classic.TableIdentity())
+
+	prefix := MakeKeyspacePrefix(TxnKeyspaceModePrefix, keyspaceID)
+	identity := TableIdentity{KeyspaceID: keyspaceID, TableID: tableID, HasKeyspace: true}
+	encoded := EncodeBytes(append(append([]byte{}, prefix...), GenerateTableKey(tableID)...))
+	re.Equal(identity, encoded.TableIdentity())
+
+	other := EncodeBytes(append(append([]byte{}, prefix...), GenerateTableKey(otherTableID)...))
+	re.Equal(otherTableID, other.TableIdentity().TableID)
+	re.NotEqual(encoded.TableIdentity(), other.TableIdentity())
+
+	// Same table: record and index keys must still resolve to the same identity.
+	record := EncodeBytes(append(append([]byte{}, prefix...), GenerateRowKey(tableID, 1)...))
+	indexKey := append(GenerateTableKey(tableID), '_', 'i')
+	indexKey = EncodeInt(indexKey, 7)
+	index := EncodeBytes(append(append([]byte{}, prefix...), indexKey...))
+	re.Equal(identity, record.TableIdentity())
+	re.Equal(identity, index.TableIdentity())
+
+	// Same numeric table id under different keyspaces is a different identity.
+	ks1 := EncodeBytes(append(MakeKeyspacePrefix(TxnKeyspaceModePrefix, 1), GenerateTableKey(tableID)...))
+	ks2 := EncodeBytes(append(MakeKeyspacePrefix(TxnKeyspaceModePrefix, 2), GenerateTableKey(tableID)...))
+	re.Equal(ks1.TableIdentity().TableID, ks2.TableIdentity().TableID)
+	re.NotEqual(ks1.TableIdentity(), ks2.TableIdentity())
+
+	// A raw key that only happens to start with the txn mode byte but is not
+	// followed by a TiDB table/meta payload must not be treated as a table key.
+	ambiguous := EncodeBytes([]byte{'x', 0x00, 0x00, 0x2a, 'u', 's', 'e', 'r'})
+	re.Equal(TableIdentity{}, ambiguous.TableIdentity())
+
+	// TiDB data only lives under the txn mode: a raw-mode keyspace key whose
+	// payload happens to look like a table key gets no table identity.
+	rawMode := EncodeBytes(append(MakeKeyspacePrefix(RawKeyspaceModePrefix, keyspaceID), GenerateTableKey(tableID)...))
+	re.Equal(TableIdentity{}, rawMode.TableIdentity())
+}
+
+func TestMetaOrTableWithKeyspacePrefix(t *testing.T) {
+	re := require.New(t)
+	tableID := int64(55)
+	keyspaceID := uint32(7)
+	prefix := MakeKeyspacePrefix(TxnKeyspaceModePrefix, keyspaceID)
+
+	isMeta, id := EncodeBytes(append(append([]byte{}, prefix...), metaPrefix...)).MetaOrTable()
+	re.True(isMeta)
+	re.Equal(int64(0), id)
+
+	isMeta, id = EncodeBytes(append(append([]byte{}, prefix...), GenerateTableKey(tableID)...)).MetaOrTable()
+	re.False(isMeta)
+	re.Equal(tableID, id)
+
+	isMeta, id = EncodeBytes([]byte("hello")).MetaOrTable()
+	re.False(isMeta)
+	re.Equal(int64(0), id)
+}
+
+func TestMakeKeyspacePrefix(t *testing.T) {
+	re := require.New(t)
+	re.Equal([]byte{'r', 0x01, 0x02, 0x03}, MakeKeyspacePrefix(RawKeyspaceModePrefix, 0x010203))
+	// Only the lower 24 bits of the keyspace ID are encoded.
+	re.Equal([]byte{'x', 0xff, 0xff, 0xff}, MakeKeyspacePrefix(TxnKeyspaceModePrefix, 0xffffff))
+}
+
+func TestParseKeyspacePrefix(t *testing.T) {
+	re := require.New(t)
+
+	mode, id, ok := ParseKeyspacePrefix([]byte{'r', 0x01, 0x02, 0x03})
+	re.True(ok)
+	re.Equal(RawKeyspaceModePrefix, mode)
+	re.Equal(uint32(0x010203), id)
+
+	mode, id, ok = ParseKeyspacePrefix([]byte{'x', 0xff, 0xff, 0xff, 't'})
+	re.True(ok)
+	re.Equal(TxnKeyspaceModePrefix, mode)
+	re.Equal(uint32(0xffffff), id)
+
+	// Too short.
+	_, _, ok = ParseKeyspacePrefix([]byte{'x', 0x01, 0x02})
+	re.False(ok)
+	// Unknown mode byte.
+	_, _, ok = ParseKeyspacePrefix([]byte{'t', 0x01, 0x02, 0x03})
+	re.False(ok)
 }

pkg/keyspace/util.go

@@ -16,7 +16,6 @@ package keyspace
 
 import (
 	"container/heap"
-	"encoding/binary"
 	"encoding/hex"
 	"regexp"
 	"strconv"
@@ -114,15 +113,15 @@ type RegionBound struct {
 
 // MakeRegionBound constructs the correct region boundaries of the given keyspace.
 func MakeRegionBound(id uint32) *RegionBound {
-	keyspaceIDBytes := make([]byte, 4)
-	nextKeyspaceIDBytes := make([]byte, 4)
-	binary.BigEndian.PutUint32(keyspaceIDBytes, id)
-	binary.BigEndian.PutUint32(nextKeyspaceIDBytes, id+1)
+	rawLeftBound := codec.MakeKeyspacePrefix(codec.RawKeyspaceModePrefix, id)
+	rawRightBound := codec.MakeKeyspacePrefix(codec.RawKeyspaceModePrefix, id+1)
+	txnLeftBound := codec.MakeKeyspacePrefix(codec.TxnKeyspaceModePrefix, id)
+	txnRightBound := codec.MakeKeyspacePrefix(codec.TxnKeyspaceModePrefix, id+1)
 	return &RegionBound{
-		RawLeftBound:  codec.EncodeBytes(append([]byte{'r'}, keyspaceIDBytes[1:]...)),
-		RawRightBound: codec.EncodeBytes(append([]byte{'r'}, nextKeyspaceIDBytes[1:]...)),
-		TxnLeftBound:  codec.EncodeBytes(append([]byte{'x'}, keyspaceIDBytes[1:]...)),
-		TxnRightBound: codec.EncodeBytes(append([]byte{'x'}, nextKeyspaceIDBytes[1:]...)),
+		RawLeftBound:  codec.EncodeBytes(rawLeftBound),
+		RawRightBound: codec.EncodeBytes(rawRightBound),
+		TxnLeftBound:  codec.EncodeBytes(txnLeftBound),
+		TxnRightBound: codec.EncodeBytes(txnRightBound),
 	}
 }
 

pkg/mock/mockcluster/config.go

@@ -47,6 +47,11 @@ func (mc *Cluster) SetEnableOneWayMerge(v bool) {
 	mc.updateScheduleConfig(func(s *sc.ScheduleConfig) { s.EnableOneWayMerge = v })
 }
 
+// SetEnableCrossTableMerge updates the EnableCrossTableMerge configuration.
+func (mc *Cluster) SetEnableCrossTableMerge(v bool) {
+	mc.updateScheduleConfig(func(s *sc.ScheduleConfig) { s.EnableCrossTableMerge = v })
+}
+
 // SetMaxSnapshotCount updates the MaxSnapshotCount configuration.
 func (mc *Cluster) SetMaxSnapshotCount(v int) {
 	mc.updateScheduleConfig(func(s *sc.ScheduleConfig) { s.MaxSnapshotCount = uint64(v) })

pkg/schedule/checker/merge_checker.go

@@ -269,18 +269,21 @@ func AllowMerge(cluster sche.SharedCluster, region, adjacent *core.RegionInfo) b
 		if cluster.GetSharedConfig().IsCrossTableMergeEnabled() {
 			return true
 		}
-		return isTableIDSame(region, adjacent)
+		return isSameTableIdentity(region, adjacent)
 	case constant.Raw:
 		return true
 	case constant.Txn:
 		return true
 	default:
-		return isTableIDSame(region, adjacent)
+		return isSameTableIdentity(region, adjacent)
 	}
 }
 
-func isTableIDSame(region, adjacent *core.RegionInfo) bool {
-	return codec.Key(region.GetStartKey()).TableID() == codec.Key(adjacent.GetStartKey()).TableID()
+// isSameTableIdentity reports whether two regions belong to the same logical
+// table, i.e. the same table ID within the same keyspace (if any).
+func isSameTableIdentity(region, adjacent *core.RegionInfo) bool {
+	return codec.Key(region.GetStartKey()).TableIdentity() ==
+		codec.Key(adjacent.GetStartKey()).TableIdentity()
 }
 
 // Check whether there is a peer of the adjacent region on an offline store,