From 4693c040568f804d9230be19d2586af3e1204337 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Thu, 9 Jul 2026 18:55:19 +0800 Subject: [PATCH 1/3] This is an automated cherry-pick of #10992 close tikv/pd#10991 Signed-off-by: ti-chi-bot --- pkg/codec/codec.go | 85 +- pkg/codec/codec_test.go | 118 +- pkg/keyspace/util.go | 13 +- pkg/keyspace/util_test.go | 46 + pkg/mock/mockcluster/config.go | 5 + pkg/schedule/checker/merge_checker.go | 11 +- pkg/schedule/checker/merge_checker_test.go | 111 ++ pkg/schedule/checker/split_scatter_group.go | 182 +++ pkg/schedule/checker/split_scatter_test.go | 1061 +++++++++++++++++ .../server/cluster/cross_table_merge_test.go | 103 ++ 10 files changed, 1713 insertions(+), 22 deletions(-) create mode 100644 pkg/schedule/checker/split_scatter_group.go create mode 100644 pkg/schedule/checker/split_scatter_test.go create mode 100644 tests/server/cluster/cross_table_merge_test.go diff --git a/pkg/codec/codec.go b/pkg/codec/codec.go index 481a69acd09..27086b986de 100644 --- a/pkg/codec/codec.go +++ b/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 diff --git a/pkg/codec/codec_test.go b/pkg/codec/codec_test.go index 2121cabf2b8..38d3639d757 100644 --- a/pkg/codec/codec_test.go +++ b/pkg/codec/codec_test.go @@ -38,17 +38,125 @@ 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)...)) + index := EncodeBytes(append(append([]byte{}, prefix...), GenerateIndexKey(tableID, 7)...)) + 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) +} +<<<<<<< HEAD +======= + +func TestGenerateRecordAndIndexKeys(t *testing.T) { + re := require.New(t) + tableID := int64(42) + indexID := int64(7) + + rowKeyPrefix := GenerateRecordKeyPrefix(tableID) + rowKey := GenerateRowKey(tableID, 1) + re.Equal(rowKeyPrefix, rowKey[:len(rowKeyPrefix)]) + re.Equal(tableID, EncodeBytes(rowKey).TableIdentity().TableID) + + indexKey := GenerateIndexKey(tableID, indexID) + _, decodedIndexKey, err := DecodeBytes(EncodeBytes(indexKey)) + re.NoError(err) + re.Equal(indexKey, decodedIndexKey) + re.Equal(tableID, EncodeBytes(indexKey).TableIdentity().TableID) +} +>>>>>>> 2b3abf1483 (codec, checker: fix enable-cross-table-merge for keyspace keys (#10992)) diff --git a/pkg/keyspace/util.go b/pkg/keyspace/util.go index 7df93eefb19..acc833a47bb 100644 --- a/pkg/keyspace/util.go +++ b/pkg/keyspace/util.go @@ -16,7 +16,6 @@ package keyspace import ( "container/heap" - "encoding/binary" "encoding/hex" "regexp" "strconv" @@ -114,10 +113,22 @@ type RegionBound struct { // MakeRegionBound constructs the correct region boundaries of the given keyspace. func MakeRegionBound(id uint32) *RegionBound { +<<<<<<< HEAD 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) + if id == constant.MaxValidKeyspaceID { + // The right bound is an exclusive fencepost, not a real keyspace prefix. + rawRightBound = []byte{'s', 0, 0, 0} + txnRightBound = []byte{'y', 0, 0, 0} + } +>>>>>>> 2b3abf1483 (codec, checker: fix enable-cross-table-merge for keyspace keys (#10992)) return &RegionBound{ RawLeftBound: codec.EncodeBytes(append([]byte{'r'}, keyspaceIDBytes[1:]...)), RawRightBound: codec.EncodeBytes(append([]byte{'r'}, nextKeyspaceIDBytes[1:]...)), diff --git a/pkg/keyspace/util_test.go b/pkg/keyspace/util_test.go index a4b44e2bdce..04eb5ac34ae 100644 --- a/pkg/keyspace/util_test.go +++ b/pkg/keyspace/util_test.go @@ -61,6 +61,52 @@ func TestValidateID(t *testing.T) { } } +<<<<<<< HEAD +======= +func TestMakeRegionBound(t *testing.T) { + re := require.New(t) + encodeKey := func(key []byte) []byte { + return []byte(codec.EncodeBytes(key)) + } + + regionBound := MakeRegionBound(0x010203) + re.Equal(encodeKey([]byte{'r', 0x01, 0x02, 0x03}), regionBound.RawLeftBound) + re.Equal(encodeKey([]byte{'r', 0x01, 0x02, 0x04}), regionBound.RawRightBound) + re.Equal(encodeKey([]byte{'x', 0x01, 0x02, 0x03}), regionBound.TxnLeftBound) + re.Equal(encodeKey([]byte{'x', 0x01, 0x02, 0x04}), regionBound.TxnRightBound) + + carryRegionBound := MakeRegionBound(0x0102ff) + re.Equal(encodeKey([]byte{'r', 0x01, 0x03, 0x00}), carryRegionBound.RawRightBound) + re.Equal(encodeKey([]byte{'x', 0x01, 0x03, 0x00}), carryRegionBound.TxnRightBound) + + maxRegionBound := MakeRegionBound(constant.MaxValidKeyspaceID) + re.Equal(encodeKey([]byte{'r', 0xff, 0xff, 0xff}), maxRegionBound.RawLeftBound) + re.Equal(encodeKey([]byte{'s', 0x00, 0x00, 0x00}), maxRegionBound.RawRightBound) + re.Equal(encodeKey([]byte{'x', 0xff, 0xff, 0xff}), maxRegionBound.TxnLeftBound) + re.Equal(encodeKey([]byte{'y', 0x00, 0x00, 0x00}), maxRegionBound.TxnRightBound) +} + +func TestMaxKeyspaceLabelRuleSplitKeys(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + regionLabeler, err := labeler.NewRegionLabeler(ctx, endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil), time.Hour) + re.NoError(err) + + re.NoError(regionLabeler.SetLabelRule(MakeTxnLabelRule(constant.MaxValidKeyspaceID))) + encodeKey := func(key []byte) []byte { + return []byte(codec.EncodeBytes(key)) + } + re.Equal( + [][]byte{ + encodeKey([]byte{'x', 0xff, 0xff, 0xff}), + encodeKey([]byte{'y', 0x00, 0x00, 0x00}), + }, + regionLabeler.GetSplitKeys(nil, nil), + ) +} + +>>>>>>> 2b3abf1483 (codec, checker: fix enable-cross-table-merge for keyspace keys (#10992)) func TestValidateName(t *testing.T) { re := require.New(t) testCases := []struct { diff --git a/pkg/mock/mockcluster/config.go b/pkg/mock/mockcluster/config.go index c5fdce24f31..33bf17a5afa 100644 --- a/pkg/mock/mockcluster/config.go +++ b/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) }) diff --git a/pkg/schedule/checker/merge_checker.go b/pkg/schedule/checker/merge_checker.go index 571ee134da0..80d607ac6c2 100644 --- a/pkg/schedule/checker/merge_checker.go +++ b/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, diff --git a/pkg/schedule/checker/merge_checker_test.go b/pkg/schedule/checker/merge_checker_test.go index 80a72238ee4..ea7f4ed9817 100644 --- a/pkg/schedule/checker/merge_checker_test.go +++ b/pkg/schedule/checker/merge_checker_test.go @@ -20,11 +20,13 @@ import ( "testing" "time" + "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" "go.uber.org/goleak" "github.com/pingcap/kvproto/pkg/metapb" + "github.com/tikv/pd/pkg/codec" "github.com/tikv/pd/pkg/core" "github.com/tikv/pd/pkg/core/storelimit" "github.com/tikv/pd/pkg/mock/mockcluster" @@ -590,3 +592,112 @@ func newRegionInfo(id uint64, startKey, endKey string, size, keys int64, leader core.SetApproximateKeys(keys), ) } + +func TestAllowMergeCrossTable(t *testing.T) { + ctx := t.Context() + cfg := mockconfig.NewTestOptions() + cluster := mockcluster.NewCluster(ctx, cfg) + // Disable placement rules so AllowMerge is decided only by key type / table ID. + cluster.SetEnablePlacementRules(false) + + const ( + keyspace1 = uint32(42) + keyspace2 = uint32(43) + tableA = int64(100) + tableB = int64(101) + indexID = int64(1) + ) + + // Classic keys. + classicTableA := codec.EncodeBytes(codec.GenerateTableKey(tableA)) + classicTableB := codec.EncodeBytes(codec.GenerateTableKey(tableB)) + classicTableC := codec.EncodeBytes(codec.GenerateTableKey(tableB + 1)) + classicIndexA := codec.EncodeBytes(codec.GenerateIndexKey(tableA, indexID)) + classicRecordA := codec.EncodeBytes(codec.GenerateRowKey(tableA, 1)) + + // Same-keyspace keys. + ks1TableA := encodeKeyspaceRawKey(keyspace1, codec.GenerateTableKey(tableA)) + ks1TableB := encodeKeyspaceRawKey(keyspace1, codec.GenerateTableKey(tableB)) + ks1TableC := encodeKeyspaceRawKey(keyspace1, codec.GenerateTableKey(tableB+1)) + ks1IndexA := encodeKeyspaceRawKey(keyspace1, codec.GenerateIndexKey(tableA, indexID)) + ks1RecordA := encodeKeyspaceRawKey(keyspace1, codec.GenerateRowKey(tableA, 1)) + + // Different-keyspace keys. Adjacent ranges are constructed artificially so + // AllowMerge can reach the table-ID check; production keyspaces are usually + // separated by fence regions and would not be merge candidates. + ks2TableA := encodeKeyspaceRawKey(keyspace2, codec.GenerateTableKey(tableA)) + ks2TableB := encodeKeyspaceRawKey(keyspace2, codec.GenerateTableKey(tableB)) + ks2TableC := encodeKeyspaceRawKey(keyspace2, codec.GenerateTableKey(tableB+1)) + + // Matrix: layout × same/diff table identity × enable-cross-table-merge. + // Logical table identity is (keyspaceID, tableID); classic has no keyspace. + type adjacentPair struct { + startA, endA []byte + startB, endB []byte + } + classicDiffTable := adjacentPair{classicTableA, classicTableB, classicTableB, classicTableC} + classicSameTable := adjacentPair{classicIndexA, classicRecordA, classicRecordA, classicTableB} + sameKSDiffTable := adjacentPair{ks1TableA, ks1TableB, ks1TableB, ks1TableC} + sameKSSameTable := adjacentPair{ks1IndexA, ks1RecordA, ks1RecordA, ks1TableB} + // Different keyspaces with the same numeric table id are still different tables. + diffKSSameTable := adjacentPair{ks1TableA, ks2TableA, ks2TableA, ks2TableB} + diffKSDiffTable := adjacentPair{ks1TableA, ks2TableB, ks2TableB, ks2TableC} + + cases := []struct { + name string + pair adjacentPair + crossTableMerge bool + expectAllow bool + }{ + // Classic: different table IDs. + {"classic/diff-table/cross-enabled", classicDiffTable, true, true}, + {"classic/diff-table/cross-disabled", classicDiffTable, false, false}, + // Classic: same table ID (index + record). + {"classic/same-table/cross-enabled", classicSameTable, true, true}, + {"classic/same-table/cross-disabled", classicSameTable, false, true}, + + // Same keyspace: different table IDs. + {"same-keyspace/diff-table/cross-enabled", sameKSDiffTable, true, true}, + {"same-keyspace/diff-table/cross-disabled", sameKSDiffTable, false, false}, + // Same keyspace: same table ID (index + record). + {"same-keyspace/same-table/cross-enabled", sameKSSameTable, true, true}, + {"same-keyspace/same-table/cross-disabled", sameKSSameTable, false, true}, + + // Different keyspaces: same numeric table ID is not the same logical table. + // With cross-table enabled, merge is still allowed by policy (split keys + // from keyspace labels normally block this path in production). + {"diff-keyspace/same-table/cross-enabled", diffKSSameTable, true, true}, + {"diff-keyspace/same-table/cross-disabled", diffKSSameTable, false, false}, + // Different keyspaces: different table IDs. + {"diff-keyspace/diff-table/cross-enabled", diffKSDiffTable, true, true}, + {"diff-keyspace/diff-table/cross-disabled", diffKSDiffTable, false, false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + re := require.New(t) + cluster.SetEnableCrossTableMerge(tc.crossTableMerge) + regionA := newRegionInfoWithBytes(1, tc.pair.startA, tc.pair.endA) + regionB := newRegionInfoWithBytes(2, tc.pair.startB, tc.pair.endB) + re.Equal(tc.expectAllow, AllowMerge(cluster, regionA, regionB)) + }) + } +} + +func encodeKeyspaceRawKey(keyspaceID uint32, rawKey []byte) []byte { + prefix := codec.MakeKeyspacePrefix(codec.TxnKeyspaceModePrefix, keyspaceID) + return codec.EncodeBytes(append(prefix, rawKey...)) +} + +func newRegionInfoWithBytes(id uint64, startKey, endKey []byte) *core.RegionInfo { + peer := &metapb.Peer{Id: id * 10, StoreId: 1} + return core.NewRegionInfo( + &metapb.Region{ + Id: id, + StartKey: startKey, + EndKey: endKey, + Peers: []*metapb.Peer{peer}, + }, + peer, + ) +} diff --git a/pkg/schedule/checker/split_scatter_group.go b/pkg/schedule/checker/split_scatter_group.go new file mode 100644 index 00000000000..d0ab69f3e7e --- /dev/null +++ b/pkg/schedule/checker/split_scatter_group.go @@ -0,0 +1,182 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package checker + +import ( + "bytes" + "fmt" + + "github.com/tikv/pd/pkg/codec" + "github.com/tikv/pd/pkg/core" +) + +var ( + splitScatterTablePrefix = []byte{'t'} + splitScatterIndexPrefix = []byte("_i") +) + +type splitScatterKeyspaceValidator func(uint32) bool + +type splitScatterDecodedKey struct { + rawKey []byte + keyspacePrefix []byte + keyspaceID uint32 +} + +func (key splitScatterDecodedKey) hasKeyspace() bool { + return len(key.keyspacePrefix) > 0 +} + +func resolveSplitScatterRangeHintWithKeyspaceValidator( + region *core.RegionInfo, + validateKeyspace splitScatterKeyspaceValidator, +) splitScatterRangeHint { + decodedKey, err := decodeSplitScatterRegionKey(region.GetStartKey(), validateKeyspace) + if err != nil { + return splitScatterRangeHint{} + } + rawKey := decodedKey.rawKey + + if !bytes.HasPrefix(rawKey, splitScatterTablePrefix) { + return splitScatterRangeHint{} + } + rest := rawKey[len(splitScatterTablePrefix):] + rest, tableID, err := codec.DecodeInt(rest) + if err != nil { + return splitScatterRangeHint{} + } + + tablePrefix := append([]byte(nil), codec.GenerateTableKey(tableID)...) + tableGroup := makeSplitScatterTableGroup(tableID) + if decodedKey.hasKeyspace() { + tableGroup = makeSplitScatterKeyspaceTableGroup(decodedKey.keyspaceID, tableID) + } + if !bytes.HasPrefix(rest, splitScatterIndexPrefix) { + return splitScatterPrefixRangeWithKeyspaceGroup(decodedKey.keyspacePrefix, tablePrefix, tableGroup) + } + + indexRest := rest[len(splitScatterIndexPrefix):] + _, indexID, err := codec.DecodeInt(indexRest) + if err != nil { + return splitScatterRangeHint{} + } + + indexPrefix := codec.GenerateIndexKey(tableID, indexID) + indexGroup := makeSplitScatterIndexGroup(tableID, indexID) + if decodedKey.hasKeyspace() { + indexGroup = makeSplitScatterKeyspaceIndexGroup(decodedKey.keyspaceID, tableID, indexID) + } + indexRange := splitScatterPrefixRangeWithKeyspaceGroup(decodedKey.keyspacePrefix, indexPrefix, indexGroup) + endKey := region.GetEndKey() + + // We intentionally over-approximate ambiguous table-key ranges. If PD can + // no longer prove the region stays within a single index prefix, it falls + // back to the table-scoped group instead of dropping back to the family + // group, so table-boundary splits and merged ranges still participate in + // the broader scatter continuity/baseline. + // Both endKey and indexRange.endKey are MemComparable-encoded, so + // bytes.Compare correctly reflects the key ordering. + if len(endKey) == 0 || len(indexRange.startKey) == 0 || len(indexRange.endKey) == 0 || bytes.Compare(endKey, indexRange.endKey) > 0 { + return splitScatterPrefixRangeWithKeyspaceGroup(decodedKey.keyspacePrefix, tablePrefix, tableGroup) + } + return indexRange +} + +func decodeSplitScatterRegionKey( + regionKey []byte, + validateKeyspace splitScatterKeyspaceValidator, +) (splitScatterDecodedKey, error) { + _, rawKey, err := codec.DecodeBytes(regionKey) + if err != nil { + return splitScatterDecodedKey{}, err + } + decodedKey := splitScatterDecodedKey{rawKey: rawKey} + mode, keyspaceID, ok := codec.ParseKeyspacePrefix(rawKey) + if !ok || mode != codec.TxnKeyspaceModePrefix { + return decodedKey, nil + } + + // Split-scatter range hints are table/index-scoped, so only TiDB txn + // keyspace keys from a known keyspace range are decoded. With only a + // region key, an API V2 txn prefix can be indistinguishable from a + // classic/raw user key that starts with the same bytes. + if validateKeyspace == nil || !validateKeyspace(keyspaceID) { + return decodedKey, nil + } + decodedKey.rawKey = rawKey[codec.KeyspacePrefixLen:] + decodedKey.keyspacePrefix = codec.MakeKeyspacePrefix(mode, keyspaceID) + decodedKey.keyspaceID = keyspaceID + return decodedKey, nil +} + +func splitScatterPrefixRange(rawPrefix []byte) splitScatterRangeHint { + return splitScatterPrefixRangeWithGroup(rawPrefix, "") +} + +func splitScatterPrefixRangeWithGroup(rawPrefix []byte, scatterGroup string) splitScatterRangeHint { + return splitScatterPrefixRangeWithKeyspaceGroup(nil, rawPrefix, scatterGroup) +} + +func splitScatterPrefixRangeWithKeyspaceGroup(keyspacePrefix, rawPrefix []byte, scatterGroup string) splitScatterRangeHint { + startKey := codec.EncodeBytes(appendKeyspacePrefix(keyspacePrefix, rawPrefix)) + endRawPrefix := splitScatterNextPrefix(rawPrefix) + if len(endRawPrefix) == 0 { + // Current callers use TiDB table/index prefixes, which always start + // with 't' and therefore have a finite next prefix. Keep the fallback + // here so this helper remains safe if it is ever used with an all-0xff + // prefix. + return splitScatterRangeHint{startKey: startKey, scatterGroup: scatterGroup} + } + return splitScatterRangeHint{ + startKey: startKey, + endKey: codec.EncodeBytes(appendKeyspacePrefix(keyspacePrefix, endRawPrefix)), + scatterGroup: scatterGroup, + } +} + +func appendKeyspacePrefix(keyspacePrefix, rawKey []byte) []byte { + key := make([]byte, 0, len(keyspacePrefix)+len(rawKey)) + key = append(key, keyspacePrefix...) + key = append(key, rawKey...) + return key +} + +func makeSplitScatterTableGroup(tableID int64) string { + return fmt.Sprintf("split-scatter-table-%d", tableID) +} + +func makeSplitScatterIndexGroup(tableID, indexID int64) string { + return fmt.Sprintf("split-scatter-index-%d-%d", tableID, indexID) +} + +func makeSplitScatterKeyspaceTableGroup(keyspaceID uint32, tableID int64) string { + return fmt.Sprintf("split-scatter-keyspace-%d-table-%d", keyspaceID, tableID) +} + +func makeSplitScatterKeyspaceIndexGroup(keyspaceID uint32, tableID, indexID int64) string { + return fmt.Sprintf("split-scatter-keyspace-%d-index-%d-%d", keyspaceID, tableID, indexID) +} + +func splitScatterNextPrefix(key []byte) []byte { + next := append([]byte(nil), key...) + for i := len(next) - 1; i >= 0; i-- { + if next[i] == 0xFF { + continue + } + next[i]++ + return next[:i+1] + } + return nil +} diff --git a/pkg/schedule/checker/split_scatter_test.go b/pkg/schedule/checker/split_scatter_test.go new file mode 100644 index 00000000000..0d6076716f1 --- /dev/null +++ b/pkg/schedule/checker/split_scatter_test.go @@ -0,0 +1,1061 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package checker + +import ( + "context" + "testing" + "time" + + "github.com/prometheus/client_golang/prometheus" + promtestutil "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/require" + + "github.com/pingcap/kvproto/pkg/metapb" + + "github.com/tikv/pd/pkg/codec" + "github.com/tikv/pd/pkg/core" + "github.com/tikv/pd/pkg/keyspace" + "github.com/tikv/pd/pkg/keyspace/constant" + "github.com/tikv/pd/pkg/mock/mockcluster" + "github.com/tikv/pd/pkg/mock/mockconfig" + "github.com/tikv/pd/pkg/schedule/hbstream" + "github.com/tikv/pd/pkg/schedule/labeler" + "github.com/tikv/pd/pkg/schedule/operator" + "github.com/tikv/pd/pkg/schedule/scatter" +) + +const ( + splitScatterObservedRegionID uint64 = 101 + splitScatterTestTableID int64 = 42 + splitScatterTestIndexID int64 = 7 + splitScatterTestKeyspaceID uint32 = 4242 + splitScatterTestNextGenKeyspaceID uint32 = constant.SystemKeyspaceID + splitScatterTestSourceWaitVersion = uint64(0) + // CPU usage is only populated to mimic load-split region heartbeat data. + // Current split-scatter dispatch does not rank pending regions by CPU. + splitScatterNoCPUUsage uint64 = 0 + splitScatterReportedCPUUsage uint64 = 1 +) + +func (c *Controller) collectTopPendingSplitScatter(limit int) []splitScatterPendingItem { + return c.splitScatter.collectTopPendingSplitScatter(limit) +} + +func (c *Controller) dispatchSplitScatterRegions() { + c.splitScatter.dispatchSplitScatterRegions() +} + +func TestSplitScatterControllerCleanupResetsPendingGauge(t *testing.T) { + re := require.New(t) + splitScatterPendingGauge.Set(7) + + controller, _, _, cleanup := newTestSplitScatterController(t) + cleanup() + + re.Equal(0, splitScatterPendingCount(controller)) + re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) +} + +func TestRecordSplitScatterBatchCollectsPendingRegions(t *testing.T) { + re := require.New(t) + controller, tc, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102, 103}) + re.Equal(4, splitScatterPendingCount(controller)) + re.Equal(float64(4), promtestutil.ToFloat64(splitScatterPendingGauge)) + + sourceGroup := splitScatterPendingGroup(t, controller, 100) + for _, regionID := range []uint64{101, 102, 103} { + re.Equal(sourceGroup, splitScatterPendingGroup(t, controller, regionID)) + } + + putSplitScatterRegion(tc, 101, "m", "n", splitScatterReportedCPUUsage) + putSplitScatterRegion(tc, 102, "n", "o", splitScatterReportedCPUUsage) + putSplitScatterRegion(tc, 103, "o", "", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + re.ElementsMatch([]uint64{100, 101, 102, 103}, pendingRegionIDs(controller.collectTopPendingSplitScatter(4))) +} + +func TestCheckSplitScatterRegionsCreatesScatterOperator(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) + putSplitScatterRegion(tc, 101, "m", "t", splitScatterReportedCPUUsage) + putSplitScatterRegion(tc, 102, "t", "", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + group := splitScatterPendingGroup(t, controller, 101) + + controller.dispatchSplitScatterRegions() + + var op *operator.Operator + for _, regionID := range []uint64{100, 101, 102} { + op = oc.GetOperator(regionID) + if op != nil { + break + } + } + re.NotNil(op) + re.Equal(scatter.InternalScatterOperatorDesc, op.Desc()) + opGroup, ok := op.GetAdditionalInfo("group") + re.True(ok) + re.Equal(group, opGroup) + batchGroup, ok := op.GetAdditionalInfo("batch-group") + re.True(ok) + re.Equal(group, batchGroup) +} + +func TestDispatchSplitScatterKeepsPendingUntilSplitHeartbeat(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + + controller.dispatchSplitScatterRegions() + + re.Equal(2, splitScatterPendingCount(controller)) + re.Nil(oc.GetOperator(101)) + + putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) + + retrySplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) + re.Empty(controller.collectTopPendingSplitScatter(2)) + advanceSplitScatterSourceVersion(t, tc) + setSplitScatterNextDispatchAt(t, controller, time.Now().Add(-time.Second)) + re.ElementsMatch([]uint64{100, 101}, pendingRegionIDs(controller.collectTopPendingSplitScatter(2))) + + controller.dispatchSplitScatterRegions() + + op := oc.GetOperator(101) + re.NotNil(op) + re.Equal(scatter.InternalScatterOperatorDesc, op.Desc()) +} + +func TestDispatchSplitScatterUsesRequestWaitVersionWhenCacheLags(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + source := tc.GetRegion(100) + re.NotNil(source) + tc.PutRegion(source.Clone(core.SetRegionVersion(4))) + + controller.RecordSplitScatterBatch(100, 6, []uint64{101}) + putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) + advanceSplitScatterRegionVersion(t, tc, 100) + + controller.dispatchSplitScatterRegions() + + re.Empty(oc.GetOperators()) + re.Equal(2, splitScatterPendingCount(controller)) + + advanceSplitScatterRegionVersion(t, tc, 100) + setSplitScatterNextDispatchAt(t, controller, time.Now().Add(-time.Second)) + controller.dispatchSplitScatterRegions() + + re.NotNil(oc.GetOperator(101)) +} + +func TestDispatchSplitScatterRespectsScheduleLimit(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) + putSplitScatterRegion(tc, 101, "m", "t", splitScatterReportedCPUUsage) + putSplitScatterRegion(tc, 102, "t", "", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + tc.SetSplitScatterScheduleLimit(1) + controller.dispatchSplitScatterRegions() + + re.Len(oc.GetOperators(), 1) + re.Equal(uint64(1), oc.OperatorCount(operator.OpSplitScatter)) + + controller.dispatchSplitScatterRegions() + + re.Len(oc.GetOperators(), 1) +} + +func TestRecordSplitScatterBatchSkipsWhenDisabled(t *testing.T) { + re := require.New(t) + controller, tc, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + tc.SetSplitScatterScheduleLimit(0) + + droppedBefore := promtestutil.ToFloat64(splitScatterPendingDroppedCounter) + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) + + re.Equal(0, splitScatterPendingCount(controller)) + re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) + re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingDroppedCounter)-droppedBefore) +} + +func TestDispatchSplitScatterClearsPendingWhenDisabled(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) + re.Equal(3, splitScatterPendingCount(controller)) + + tc.SetSplitScatterScheduleLimit(0) + setSplitScatterNextDispatchAt(t, controller, time.Now().Add(splitScatterRetryBackoff)) + disabledBefore := promtestutil.ToFloat64(splitScatterDispatchDisabledCounter) + controller.dispatchSplitScatterRegions() + + re.Empty(oc.GetOperators()) + re.Equal(0, splitScatterPendingCount(controller)) + re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) + re.Equal(float64(1), promtestutil.ToFloat64(splitScatterDispatchDisabledCounter)-disabledBefore) +} + +func TestDispatchSplitScatterCleansExpiredPendingBeforeEarlyReturn(t *testing.T) { + testCases := []struct { + name string + setupEarlyReturn func(*mockcluster.Cluster, *operator.Controller) + counter prometheus.Counter + }{ + { + name: "disabled", + setupEarlyReturn: func(tc *mockcluster.Cluster, _ *operator.Controller) { + tc.SetSplitScatterScheduleLimit(0) + }, + counter: splitScatterDispatchDisabledCounter, + }, + { + name: "schedule limit", + setupEarlyReturn: func(tc *mockcluster.Cluster, oc *operator.Controller) { + tc.SetSplitScatterScheduleLimit(1) + region := tc.GetRegion(100) + op := operator.NewTestOperator( + region.GetID(), + region.GetRegionEpoch(), + operator.OpSplitScatter|operator.OpRegion, + operator.TransferLeader{FromStore: 1, ToStore: 2}, + ) + require.True(t, oc.AddOperator(op)) + }, + counter: splitScatterDispatchScheduleLimitCounter, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + expireSplitScatterPendingAt(t, controller, 100, time.Now().Add(-time.Second)) + expireSplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) + testCase.setupEarlyReturn(tc, oc) + + expiredBefore := splitScatterPendingExpiredCount("false") + counterBefore := promtestutil.ToFloat64(testCase.counter) + controller.dispatchSplitScatterRegions() + + re.Equal(0, splitScatterPendingCount(controller)) + re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) + re.Equal(float64(2), splitScatterPendingExpiredCount("false")-expiredBefore) + re.Equal(float64(0), promtestutil.ToFloat64(testCase.counter)-counterBefore) + }) + } +} + +func TestCollectTopPendingDelaysMissingRegions(t *testing.T) { + re := require.New(t) + controller, _, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + + missingBefore := promtestutil.ToFloat64(splitScatterDispatchRegionMissingCounter) + re.Empty(controller.collectTopPendingSplitScatter(2)) + + re.Equal(float64(1), promtestutil.ToFloat64(splitScatterDispatchRegionMissingCounter)-missingBefore) + pending := splitScatterPending(t, controller, 101) + re.True(pending.retryAt.After(time.Now())) + + re.Empty(controller.collectTopPendingSplitScatter(2)) + re.Equal(float64(1), promtestutil.ToFloat64(splitScatterDispatchRegionMissingCounter)-missingBefore) +} + +func TestDispatchSplitScatterBacksOffWhenNoCandidates(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + + controller.dispatchSplitScatterRegions() + + re.True(splitScatterNextDispatchAt(t, controller).After(time.Now())) + putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + controller.dispatchSplitScatterRegions() + + re.Empty(oc.GetOperators()) + + retrySplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) + setSplitScatterNextDispatchAt(t, controller, time.Now().Add(-time.Second)) + controller.dispatchSplitScatterRegions() + + re.NotNil(oc.GetOperator(101)) +} + +func TestDispatchSplitScatterRespectsScheduleDeny(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + re.NoError(tc.GetRegionLabeler().SetLabelRule(&labeler.LabelRule{ + ID: "split-scatter-schedule-deny", + Labels: []labeler.RegionLabel{{Key: "schedule", Value: "deny"}}, + RuleType: labeler.KeyRange, + Data: []any{map[string]any{"start_key": "", "end_key": ""}}, + })) + + counterBefore := promtestutil.ToFloat64(splitScatterDispatchScheduleDisabledCounter) + controller.dispatchSplitScatterRegions() + + re.Empty(oc.GetOperators()) + re.Equal(2, splitScatterPendingCount(controller)) + re.Equal(float64(2), promtestutil.ToFloat64(splitScatterDispatchScheduleDisabledCounter)-counterBefore) + for _, regionID := range []uint64{100, 101} { + pending := splitScatterPending(t, controller, regionID) + re.True(pending.retryAt.After(time.Now())) + } +} + +func TestCollectTopPendingRemovesExpiredPending(t *testing.T) { + re := require.New(t) + controller, _, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + expireSplitScatterPendingAt(t, controller, 100, time.Now().Add(-time.Second)) + expireSplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) + + expiredBefore := splitScatterPendingExpiredCount("false") + re.Empty(controller.collectTopPendingSplitScatter(2)) + re.Equal(0, splitScatterPendingCount(controller)) + re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) + re.Equal(float64(2), splitScatterPendingExpiredCount("false")-expiredBefore) +} + +func TestCollectTopPendingMarksAttemptedBeforeExpiration(t *testing.T) { + re := require.New(t) + controller, tc, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + attemptedBefore := splitScatterPendingExpiredCount("true") + unattemptedBefore := splitScatterPendingExpiredCount("false") + re.Len(controller.collectTopPendingSplitScatter(1), 1) + expireSplitScatterPendingAt(t, controller, 100, time.Now().Add(-time.Second)) + expireSplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) + + re.Empty(controller.collectTopPendingSplitScatter(2)) + re.Equal(0, splitScatterPendingCount(controller)) + re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) + re.Equal(float64(1), splitScatterPendingExpiredCount("true")-attemptedBefore) + re.Equal(float64(1), splitScatterPendingExpiredCount("false")-unattemptedBefore) +} + +func TestRecordSplitScatterBatchRespectsPendingLimit(t *testing.T) { + re := require.New(t) + controller, _, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + fillSplitScatterPending(controller, time.Time{}) + + droppedBefore := promtestutil.ToFloat64(splitScatterPendingDroppedCounter) + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + + re.Equal(splitScatterPendingLimit, splitScatterPendingCount(controller)) + re.Equal(float64(2), promtestutil.ToFloat64(splitScatterPendingDroppedCounter)-droppedBefore) + controller.splitScatter.pendingMu.RLock() + _, sourceExists := controller.splitScatter.pending[100] + _, childExists := controller.splitScatter.pending[101] + controller.splitScatter.pendingMu.RUnlock() + re.False(sourceExists) + re.False(childExists) +} + +func TestRecordSplitScatterBatchClearsExpiredPendingBeforeLimitCheck(t *testing.T) { + re := require.New(t) + controller, _, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + fillSplitScatterPending(controller, time.Now().Add(-time.Second)) + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + + re.Equal(2, splitScatterPendingCount(controller)) + re.Equal(makeSplitScatterGroup(100, 101), splitScatterPendingGroup(t, controller, 101)) +} + +func TestCollectTopPendingSortsBeforeLimit(t *testing.T) { + re := require.New(t) + controller, tc, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(200, splitScatterTestSourceWaitVersion, []uint64{201}) + putSplitScatterRegion(tc, 200, "n", "o", splitScatterNoCPUUsage) + putSplitScatterRegion(tc, 201, "o", "p", splitScatterReportedCPUUsage) + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + putSplitScatterRegion(tc, 101, "m", "n", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + pending := controller.collectTopPendingSplitScatter(1) + re.Len(pending, 1) + re.Equal(uint64(100), pending[0].regionID) +} + +func TestCollectTopPendingPrioritizesNearExpiration(t *testing.T) { + re := require.New(t) + controller, tc, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + controller.RecordSplitScatterBatch(200, splitScatterTestSourceWaitVersion, []uint64{201}) + putSplitScatterRegion(tc, 100, "m", "n", splitScatterNoCPUUsage) + putSplitScatterRegion(tc, 101, "n", "o", splitScatterReportedCPUUsage) + putSplitScatterRegion(tc, 200, "o", "p", splitScatterNoCPUUsage) + putSplitScatterRegion(tc, 201, "p", "q", splitScatterReportedCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + advanceSplitScatterRegionVersion(t, tc, 200) + + expireSplitScatterPendingAt(t, controller, 100, time.Now().Add(2*time.Minute)) + expireSplitScatterPendingAt(t, controller, 101, time.Now().Add(2*time.Minute)) + expireSplitScatterPendingAt(t, controller, 200, time.Now().Add(time.Minute)) + expireSplitScatterPendingAt(t, controller, 201, time.Now().Add(time.Minute)) + + pending := controller.collectTopPendingSplitScatter(1) + re.Len(pending, 1) + re.Equal(uint64(200), pending[0].regionID) +} + +func TestCollectTopPendingResolvesRangeHint(t *testing.T) { + testCases := []struct { + name string + startKey []byte + endKey []byte + wantRange splitScatterRangeHint + wantGroup string + keyspaces []uint32 + }{ + { + name: "index region", + startKey: newSplitScatterIndexKey("a"), + endKey: newSplitScatterIndexKey("m"), + wantRange: splitScatterPrefixRange(splitScatterIndexKeyPrefix()), + wantGroup: makeSplitScatterIndexGroup(splitScatterTestTableID, splitScatterTestIndexID), + }, + { + name: "record region", + startKey: newSplitScatterRecordKey(splitScatterTestTableID, "a"), + endKey: newSplitScatterRecordKey(splitScatterTestTableID, "m"), + wantRange: splitScatterPrefixRange(codec.GenerateTableKey(splitScatterTestTableID)), + wantGroup: makeSplitScatterTableGroup(splitScatterTestTableID), + }, + { + name: "bare table boundary", + startKey: newSplitScatterTableBoundaryKey(splitScatterTestTableID), + endKey: newSplitScatterIndexKey("m"), + wantRange: splitScatterPrefixRange(codec.GenerateTableKey(splitScatterTestTableID)), + wantGroup: makeSplitScatterTableGroup(splitScatterTestTableID), + }, + { + name: "cross entity falls back to table", + startKey: newSplitScatterIndexKey("a"), + endKey: newSplitScatterRecordKey(splitScatterTestTableID, "m"), + wantRange: splitScatterPrefixRange(codec.GenerateTableKey(splitScatterTestTableID)), + wantGroup: makeSplitScatterTableGroup(splitScatterTestTableID), + }, + { + name: "cross table uses start table", + startKey: newSplitScatterRecordKey(splitScatterTestTableID, "a"), + endKey: newSplitScatterRecordKey(splitScatterTestTableID+1, "m"), + wantRange: splitScatterPrefixRange(codec.GenerateTableKey(splitScatterTestTableID)), + wantGroup: makeSplitScatterTableGroup(splitScatterTestTableID), + }, + { + name: "nextgen keyspace index region", + startKey: newSplitScatterKeyspaceIndexKey(splitScatterTestNextGenKeyspaceID, "a"), + endKey: newSplitScatterKeyspaceIndexKey(splitScatterTestNextGenKeyspaceID, "m"), + wantRange: splitScatterKeyspacePrefixRange(splitScatterTestNextGenKeyspaceID, splitScatterIndexKeyPrefix()), + wantGroup: makeSplitScatterKeyspaceIndexGroup(splitScatterTestNextGenKeyspaceID, splitScatterTestTableID, splitScatterTestIndexID), + keyspaces: []uint32{splitScatterTestNextGenKeyspaceID}, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + re := require.New(t) + controller, tc, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + putSplitScatterRegionWithKeys(tc, testCase.startKey, testCase.endKey, splitScatterReportedCPUUsage) + putSplitScatterRegion(tc, 100, "z", "", splitScatterNoCPUUsage) + advanceSplitScatterSourceVersion(t, tc) + + re.Equal(makeSplitScatterGroup(100, 101), splitScatterPendingGroup(t, controller, 101)) + re.ElementsMatch([]uint64{100, 101}, pendingRegionIDs(controller.collectTopPendingSplitScatter(2))) + rangeHint := resolveSplitScatterRangeHintWithKeyspaceValidator( + tc.GetRegion(101), + splitScatterKeyspaceValidatorFor(testCase.keyspaces...), + ) + re.Equal(testCase.wantRange.startKey, rangeHint.startKey) + re.Equal(testCase.wantRange.endKey, rangeHint.endKey) + re.Equal(testCase.wantGroup, rangeHint.scatterGroup) + }) + } +} + +func TestResolveSplitScatterRangeHintIgnoresRawLikeKeyspaceKeys(t *testing.T) { + re := require.New(t) + region := core.NewRegionInfo(&metapb.Region{ + Id: 1, + StartKey: newSplitScatterRawKeyspaceRecordKey(splitScatterTestKeyspaceID, splitScatterTestTableID, "a"), + EndKey: newSplitScatterRawKeyspaceRecordKey(splitScatterTestKeyspaceID, splitScatterTestTableID, "m"), + }, nil) + + rangeHint := resolveSplitScatterRangeHintWithKeyspaceValidator( + region, + splitScatterKeyspaceValidatorFor(splitScatterTestKeyspaceID), + ) + re.Equal(splitScatterRangeHint{}, rangeHint) +} + +func TestResolveSplitScatterRangeHintRequiresKnownTxnKeyspaceBounds(t *testing.T) { + testCases := []struct { + name string + keyspaceID uint32 + }{ + {name: "normal keyspace", keyspaceID: splitScatterTestKeyspaceID}, + {name: "max valid keyspace", keyspaceID: constant.MaxValidKeyspaceID}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + re := require.New(t) + controller, tc, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + startKey := newSplitScatterKeyspaceRecordKey(testCase.keyspaceID, "a") + endKey := newSplitScatterKeyspaceRecordKey(testCase.keyspaceID, "m") + region := core.NewRegionInfo(&metapb.Region{ + Id: 1, + StartKey: startKey, + EndKey: endKey, + }, nil) + re.Equal(splitScatterRangeHint{}, resolveSplitScatterRangeHintWithKeyspaceValidator(region, nil)) + + regionBound := keyspace.MakeRegionBound(testCase.keyspaceID) + putSplitScatterRegionWithKeysByID(tc, 90, regionBound.TxnLeftBound, startKey, splitScatterNoCPUUsage) + putSplitScatterRegionWithKeysByID(tc, 91, regionBound.TxnRightBound, nil, splitScatterNoCPUUsage) + + rangeHint := resolveSplitScatterRangeHintWithKeyspaceValidator( + region, + controller.splitScatter.hasSplitScatterTxnKeyspaceBounds, + ) + wantRange := splitScatterKeyspacePrefixRange(testCase.keyspaceID, codec.GenerateTableKey(splitScatterTestTableID)) + wantRange.scatterGroup = makeSplitScatterKeyspaceTableGroup(testCase.keyspaceID, splitScatterTestTableID) + re.Equal(wantRange, rangeHint) + }) + } +} + +func TestDispatchSplitScatterUsesRangeScatterGroup(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + putSplitScatterRegionWithKeysByID(tc, 90, newSplitScatterIndexKey("m"), newSplitScatterIndexKey("z"), splitScatterNoCPUUsage) + putSplitScatterRegionWithKeysByID(tc, 101, newSplitScatterIndexKey("a"), newSplitScatterIndexKey("m"), splitScatterReportedCPUUsage) + advanceSplitScatterRegionVersion(t, tc, 100) + + batchGroup := splitScatterPendingGroup(t, controller, 101) + + controller.dispatchSplitScatterRegions() + + expectedScatterGroup := makeSplitScatterIndexGroup(splitScatterTestTableID, splitScatterTestIndexID) + op := oc.GetOperator(101) + re.NotNil(op) + opGroup, ok := op.GetAdditionalInfo("group") + re.True(ok) + re.Equal(expectedScatterGroup, opGroup) + opBatchGroup, ok := op.GetAdditionalInfo("batch-group") + re.True(ok) + re.Equal(batchGroup, opBatchGroup) +} + +func TestDispatchSplitScatterUsesKeyspaceRangeScatterGroup(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + startKey := newSplitScatterKeyspaceIndexKey(splitScatterTestKeyspaceID, "a") + endKey := newSplitScatterKeyspaceIndexKey(splitScatterTestKeyspaceID, "m") + regionBound := keyspace.MakeRegionBound(splitScatterTestKeyspaceID) + putSplitScatterRegionWithKeysByID(tc, 90, regionBound.TxnLeftBound, startKey, splitScatterNoCPUUsage) + putSplitScatterRegionWithKeysByID(tc, 91, regionBound.TxnRightBound, nil, splitScatterNoCPUUsage) + putSplitScatterRegionWithKeysByID(tc, 101, startKey, endKey, splitScatterReportedCPUUsage) + advanceSplitScatterRegionVersion(t, tc, 100) + + batchGroup := splitScatterPendingGroup(t, controller, 101) + + controller.dispatchSplitScatterRegions() + + expectedScatterGroup := makeSplitScatterKeyspaceIndexGroup( + splitScatterTestKeyspaceID, + splitScatterTestTableID, + splitScatterTestIndexID, + ) + op := oc.GetOperator(101) + re.NotNil(op) + opGroup, ok := op.GetAdditionalInfo("group") + re.True(ok) + re.Equal(expectedScatterGroup, opGroup) + opBatchGroup, ok := op.GetAdditionalInfo("batch-group") + re.True(ok) + re.Equal(batchGroup, opBatchGroup) +} + +func TestDispatchSplitScatterKeepsStableGroupWhenRegionSplitsAgain(t *testing.T) { + re := require.New(t) + controller, tc, oc, cleanup := newTestSplitScatterController(t) + defer cleanup() + + stableGroup := makeSplitScatterIndexGroup(splitScatterTestTableID, splitScatterTestIndexID) + sourceID := uint64(100) + childIDs := []uint64{101, 201, 301} + splitKeys := []string{"m", "t", "x"} + previousBatchGroup := "" + + putSplitScatterRegionWithKeysByID(tc, sourceID, newSplitScatterIndexKey("a"), newSplitScatterIndexKey("z"), splitScatterNoCPUUsage) + for i, childID := range childIDs { + controller.RecordSplitScatterBatch(sourceID, splitScatterTestSourceWaitVersion, []uint64{childID}) + batchGroup := splitScatterPendingGroup(t, controller, childID) + re.NotEqual(previousBatchGroup, batchGroup) + + tc.PutRegion(tc.GetRegion(sourceID).Clone( + core.WithEndKey(newSplitScatterIndexKey(splitKeys[i])), + core.WithIncVersion(), + )) + putSplitScatterRegionWithKeysByID(tc, childID, newSplitScatterIndexKey(splitKeys[i]), newSplitScatterIndexKey("z"), splitScatterReportedCPUUsage) + + controller.dispatchSplitScatterRegions() + + requireInternalScatterOpsUseGroups(t, oc, stableGroup, batchGroup) + removeInternalScatterOps(oc) + previousBatchGroup = batchGroup + } +} + +func TestDispatchSplitScatterBacksOff(t *testing.T) { + testCases := []struct { + name string + putRegion func(*mockcluster.Cluster) + }{ + { + name: "region is not fully replicated", + putRegion: func(tc *mockcluster.Cluster) { + putSplitScatterRegionWithStores(tc, 101, "m", "", splitScatterReportedCPUUsage, 1, 2) + }, + }, + { + name: "scatter internal fails", + putRegion: func(tc *mockcluster.Cluster) { + putSplitScatterRegionWithoutLeader(tc, 101, "m", "", splitScatterReportedCPUUsage) + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + re := require.New(t) + controller, tc, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + testCase.putRegion(tc) + advanceSplitScatterSourceVersion(t, tc) + + controller.dispatchSplitScatterRegions() + + re.Equal(1, splitScatterPendingCount(controller)) + pending := splitScatterObservedPending(t, controller) + re.True(pending.retryAt.After(time.Now())) + re.Empty(pendingRegionIDs(controller.collectTopPendingSplitScatter(2))) + }) + } +} + +func TestDispatchSplitScatterIgnoresStalePendingSnapshot(t *testing.T) { + re := require.New(t) + controller, _, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + stalePending := splitScatterObservedPending(t, controller) + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{102, 101}) + currentPending := splitScatterObservedPending(t, controller) + re.NotEqual(stalePending.group, currentPending.group) + re.Equal(time.Time{}, currentPending.retryAt) + + controller.splitScatter.delayPendingSplitScatter(stalePending) + + currentPending = splitScatterObservedPending(t, controller) + re.Equal(time.Time{}, currentPending.retryAt) + + controller.splitScatter.deletePendingSplitScatter(stalePending) + + currentPending = splitScatterObservedPending(t, controller) + re.Equal(makeSplitScatterGroup(100, 102), currentPending.group) + re.Equal(time.Time{}, currentPending.retryAt) +} + +func TestDispatchSplitScatterIgnoresStalePendingWithSameGroup(t *testing.T) { + re := require.New(t) + controller, _, _, cleanup := newTestSplitScatterController(t) + defer cleanup() + + controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) + stalePending := splitScatterObservedPending(t, controller) + + controller.splitScatter.pendingMu.Lock() + currentPending := controller.splitScatter.pending[splitScatterObservedRegionID] + currentPending.expireAt = currentPending.expireAt.Add(time.Minute) + controller.splitScatter.pending[splitScatterObservedRegionID] = currentPending + controller.splitScatter.pendingMu.Unlock() + + controller.splitScatter.delayPendingSplitScatter(stalePending) + + currentPending = splitScatterObservedPending(t, controller) + re.Equal(time.Time{}, currentPending.retryAt) + + controller.splitScatter.deletePendingSplitScatter(stalePending) + + currentPending = splitScatterObservedPending(t, controller) + re.Equal(stalePending.group, currentPending.group) + re.NotEqual(stalePending.expireAt, currentPending.expireAt) +} + +func newTestSplitScatterController(t *testing.T) (*Controller, *mockcluster.Cluster, *operator.Controller, func()) { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + opt := mockconfig.NewTestOptions() + tc := mockcluster.NewCluster(ctx, opt) + for storeID := uint64(1); storeID <= 4; storeID++ { + tc.AddRegionStore(storeID, 0) + } + putSplitScatterRegion(tc, 100, "", "m", splitScatterNoCPUUsage) + + stream := hbstream.NewTestHeartbeatStreams(ctx, tc, false) + oc := operator.NewController(ctx, tc.GetBasicCluster(), tc.GetSharedConfig(), stream) + controller := NewController(ctx, tc, tc.GetCheckerConfig(), oc) + + cleanup := func() { + controller.splitScatter.clearPendingSplitScatter() + stream.Close() + cancel() + } + return controller, tc, oc, cleanup +} + +func putSplitScatterRegion(tc *mockcluster.Cluster, regionID uint64, startKey, endKey string, cpuUsage uint64) { + tc.AddLeaderRegionWithRange(regionID, startKey, endKey, 1, 2, 3) + region := tc.GetRegion(regionID).Clone(core.SetCPUUsage(cpuUsage)) + tc.PutRegion(region) +} + +func putSplitScatterRegionWithKeys(tc *mockcluster.Cluster, startKey, endKey []byte, cpuUsage uint64) { + putSplitScatterRegionWithKeysByID(tc, splitScatterObservedRegionID, startKey, endKey, cpuUsage) +} + +func putSplitScatterRegionWithKeysByID(tc *mockcluster.Cluster, regionID uint64, startKey, endKey []byte, cpuUsage uint64) { + peers := []*metapb.Peer{ + {Id: regionID*10 + 1, StoreId: 1}, + {Id: regionID*10 + 2, StoreId: 2}, + {Id: regionID*10 + 3, StoreId: 3}, + } + region := core.NewRegionInfo( + &metapb.Region{ + Id: regionID, + StartKey: startKey, + EndKey: endKey, + Peers: peers, + RegionEpoch: &metapb.RegionEpoch{ + ConfVer: 1, + Version: 1, + }, + }, + peers[0], + core.SetCPUUsage(cpuUsage), + ) + tc.PutRegion(region) +} + +func newSplitScatterRegionInfo( + regionID uint64, + startKey, endKey string, + peers []*metapb.Peer, + leader *metapb.Peer, + cpuUsage uint64, +) *core.RegionInfo { + return core.NewRegionInfo( + &metapb.Region{ + Id: regionID, + StartKey: []byte(startKey), + EndKey: []byte(endKey), + Peers: peers, + RegionEpoch: &metapb.RegionEpoch{ + ConfVer: 1, + Version: 1, + }, + }, + leader, + core.SetCPUUsage(cpuUsage), + ) +} + +func putSplitScatterRegionWithStores(tc *mockcluster.Cluster, regionID uint64, startKey, endKey string, cpuUsage uint64, stores ...uint64) { + peers := make([]*metapb.Peer, 0, len(stores)) + for i, storeID := range stores { + peers = append(peers, &metapb.Peer{ + Id: regionID*10 + uint64(i) + 1, + StoreId: storeID, + }) + } + tc.PutRegion(newSplitScatterRegionInfo(regionID, startKey, endKey, peers, peers[0], cpuUsage)) +} + +func putSplitScatterRegionWithoutLeader(tc *mockcluster.Cluster, regionID uint64, startKey, endKey string, cpuUsage uint64) { + peers := []*metapb.Peer{ + {Id: regionID*10 + 1, StoreId: 1}, + {Id: regionID*10 + 2, StoreId: 2}, + {Id: regionID*10 + 3, StoreId: 3}, + } + tc.PutRegion(newSplitScatterRegionInfo(regionID, startKey, endKey, peers, nil, cpuUsage)) +} + +func fillSplitScatterPending(controller *Controller, expireAt time.Time) { + controller.splitScatter.pendingMu.Lock() + defer controller.splitScatter.pendingMu.Unlock() + for regionID := uint64(1000); regionID < 1000+splitScatterPendingLimit; regionID++ { + controller.splitScatter.pending[regionID] = splitScatterPendingItem{ + regionID: regionID, + group: "old", + expireAt: expireAt, + } + } +} + +func advanceSplitScatterSourceVersion(t *testing.T, tc *mockcluster.Cluster) { + advanceSplitScatterRegionVersion(t, tc, 100) +} + +func advanceSplitScatterRegionVersion(t *testing.T, tc *mockcluster.Cluster, regionID uint64) { + t.Helper() + region := tc.GetRegion(regionID) + require.NotNil(t, region) + tc.PutRegion(region.Clone(core.WithIncVersion())) +} + +func splitScatterPendingCount(controller *Controller) int { + controller.splitScatter.pendingMu.RLock() + defer controller.splitScatter.pendingMu.RUnlock() + return len(controller.splitScatter.pending) +} + +func splitScatterPendingExpiredCount(attempted string) float64 { + return promtestutil.ToFloat64(splitScatterPendingExpiredCounter.WithLabelValues(attempted)) +} + +func splitScatterPendingGroup(t *testing.T, controller *Controller, regionID uint64) string { + t.Helper() + return splitScatterPending(t, controller, regionID).group +} + +func splitScatterKeyspaceValidatorFor(keyspaces ...uint32) splitScatterKeyspaceValidator { + return func(keyspaceID uint32) bool { + for _, validKeyspaceID := range keyspaces { + if validKeyspaceID == keyspaceID { + return true + } + } + return false + } +} + +func splitScatterPending(t *testing.T, controller *Controller, regionID uint64) splitScatterPendingItem { + t.Helper() + controller.splitScatter.pendingMu.RLock() + defer controller.splitScatter.pendingMu.RUnlock() + pending, ok := controller.splitScatter.pending[regionID] + require.True(t, ok) + return pending +} + +func splitScatterObservedPending(t *testing.T, controller *Controller) splitScatterPendingItem { + t.Helper() + controller.splitScatter.pendingMu.RLock() + defer controller.splitScatter.pendingMu.RUnlock() + pending, ok := controller.splitScatter.pending[splitScatterObservedRegionID] + require.True(t, ok) + return pending +} + +func expireSplitScatterPendingAt(t *testing.T, controller *Controller, regionID uint64, expireAt time.Time) { + t.Helper() + controller.splitScatter.pendingMu.Lock() + defer controller.splitScatter.pendingMu.Unlock() + pending, ok := controller.splitScatter.pending[regionID] + require.True(t, ok) + pending.expireAt = expireAt + controller.splitScatter.pending[regionID] = pending +} + +func retrySplitScatterPendingAt(t *testing.T, controller *Controller, regionID uint64, retryAt time.Time) { + t.Helper() + controller.splitScatter.pendingMu.Lock() + defer controller.splitScatter.pendingMu.Unlock() + pending, ok := controller.splitScatter.pending[regionID] + require.True(t, ok) + pending.retryAt = retryAt + controller.splitScatter.pending[regionID] = pending +} + +func setSplitScatterNextDispatchAt(t *testing.T, controller *Controller, nextDispatchAt time.Time) { + t.Helper() + controller.splitScatter.pendingMu.Lock() + defer controller.splitScatter.pendingMu.Unlock() + controller.splitScatter.nextDispatchAt = nextDispatchAt +} + +func splitScatterNextDispatchAt(t *testing.T, controller *Controller) time.Time { + t.Helper() + controller.splitScatter.pendingMu.RLock() + defer controller.splitScatter.pendingMu.RUnlock() + return controller.splitScatter.nextDispatchAt +} + +func requireInternalScatterOpsUseGroups(t *testing.T, oc *operator.Controller, scatterGroup, batchGroup string) { + t.Helper() + re := require.New(t) + ops := oc.GetOperators() + re.NotEmpty(ops) + for _, op := range ops { + re.Equal(scatter.InternalScatterOperatorDesc, op.Desc()) + opGroup, ok := op.GetAdditionalInfo("group") + re.True(ok) + re.Equal(scatterGroup, opGroup) + opBatchGroup, ok := op.GetAdditionalInfo("batch-group") + re.True(ok) + re.Equal(batchGroup, opBatchGroup) + } +} + +func removeInternalScatterOps(oc *operator.Controller) { + for _, op := range oc.GetOperators() { + oc.RemoveOperator(op) + } +} + +func pendingRegionIDs(regions []splitScatterPendingItem) []uint64 { + ids := make([]uint64, 0, len(regions)) + for _, region := range regions { + ids = append(ids, region.regionID) + } + return ids +} + +func splitScatterIndexKeyPrefix() []byte { + return codec.GenerateIndexKey(splitScatterTestTableID, splitScatterTestIndexID) +} + +func splitScatterKeyspacePrefixRange(keyspaceID uint32, rawPrefix []byte) splitScatterRangeHint { + startKey := newSplitScatterKeyspaceKey(keyspaceID, codec.TxnKeyspaceModePrefix, rawPrefix) + endRawPrefix := splitScatterNextPrefix(rawPrefix) + if len(endRawPrefix) == 0 { + return splitScatterRangeHint{startKey: startKey} + } + return splitScatterRangeHint{ + startKey: startKey, + endKey: newSplitScatterKeyspaceKey(keyspaceID, codec.TxnKeyspaceModePrefix, endRawPrefix), + } +} + +func newSplitScatterIndexKey(suffix string) []byte { + key := append([]byte(nil), splitScatterIndexKeyPrefix()...) + key = append(key, suffix...) + return codec.EncodeBytes(key) +} + +func newSplitScatterKeyspaceIndexKey(keyspaceID uint32, suffix string) []byte { + key := append([]byte(nil), splitScatterIndexKeyPrefix()...) + key = append(key, suffix...) + return newSplitScatterKeyspaceKey(keyspaceID, codec.TxnKeyspaceModePrefix, key) +} + +func newSplitScatterRecordKey(tableID int64, suffix string) []byte { + key := append([]byte(nil), codec.GenerateRecordKeyPrefix(tableID)...) + key = append(key, suffix...) + return codec.EncodeBytes(key) +} + +func newSplitScatterKeyspaceRecordKey(keyspaceID uint32, suffix string) []byte { + key := append([]byte(nil), codec.GenerateRecordKeyPrefix(splitScatterTestTableID)...) + key = append(key, suffix...) + return newSplitScatterKeyspaceKey(keyspaceID, codec.TxnKeyspaceModePrefix, key) +} + +func newSplitScatterRawKeyspaceRecordKey(keyspaceID uint32, tableID int64, suffix string) []byte { + key := append([]byte(nil), codec.GenerateRecordKeyPrefix(tableID)...) + key = append(key, suffix...) + return newSplitScatterKeyspaceKey(keyspaceID, codec.RawKeyspaceModePrefix, key) +} + +func newSplitScatterTableBoundaryKey(tableID int64) []byte { + return codec.EncodeBytes(codec.GenerateTableKey(tableID)) +} + +func newSplitScatterKeyspaceKey(keyspaceID uint32, mode byte, rawKey []byte) []byte { + key := codec.MakeKeyspacePrefix(mode, keyspaceID) + return codec.EncodeBytes(append(key, rawKey...)) +} diff --git a/tests/server/cluster/cross_table_merge_test.go b/tests/server/cluster/cross_table_merge_test.go new file mode 100644 index 00000000000..21b496d8074 --- /dev/null +++ b/tests/server/cluster/cross_table_merge_test.go @@ -0,0 +1,103 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cluster_test + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/pingcap/kvproto/pkg/metapb" + + "github.com/tikv/pd/pkg/codec" + "github.com/tikv/pd/pkg/core" + "github.com/tikv/pd/pkg/utils/testutil" + "github.com/tikv/pd/pkg/utils/typeutil" + "github.com/tikv/pd/tests" +) + +// TestCrossTableMergeWithKeyspace verifies through a real pd-server that +// enable-cross-table-merge=false blocks merging adjacent regions of different +// tables under a keyspace, while same-table regions still merge. See #10991. +func TestCrossTableMergeWithKeyspace(t *testing.T) { + re := require.New(t) + tc, err := tests.NewTestCluster(t.Context(), 1) + defer tc.Destroy() + re.NoError(err) + re.NoError(tc.RunInitialServers()) + tc.WaitLeader() + leaderServer := tc.GetLeaderServer() + re.NoError(leaderServer.BootstrapCluster()) + tests.MustPutStore(re, tc, &metapb.Store{ + Id: 1, + State: metapb.StoreState_Up, + NodeState: metapb.NodeState_Serving, + LastHeartbeat: time.Now().UnixNano(), + }) + + // Mirror the issue reproduction: pd-ctl config set enable-cross-table-merge false. + svr := leaderServer.GetServer() + schedule := leaderServer.GetConfig().Schedule + schedule.EnableCrossTableMerge = false + schedule.SplitMergeInterval = typeutil.NewDuration(time.Second) + re.NoError(svr.SetScheduleConfig(schedule)) + replication := leaderServer.GetConfig().Replication + replication.MaxReplicas = 1 + re.NoError(svr.SetReplicationConfig(replication)) + + // Keyspace 42, txn mode. + keyspacePrefix := codec.MakeKeyspacePrefix(codec.TxnKeyspaceModePrefix, 42) + tableKey := func(tableID int64) []byte { + return codec.EncodeBytes(append(append([]byte{}, keyspacePrefix...), codec.GenerateTableKey(tableID)...)) + } + rowKey := func(tableID, rowID int64) []byte { + return codec.EncodeBytes(append(append([]byte{}, keyspacePrefix...), codec.GenerateRowKey(tableID, rowID)...)) + } + + // Five contiguous regions inside keyspace 42: three empty single-table + // regions (tables 100..102), then table 103 split at a row key so its two + // halves form a same-table merge control pair. + regions := []struct { + id uint64 + start, end []byte + }{ + {10, tableKey(100), tableKey(101)}, + {11, tableKey(101), tableKey(102)}, + {12, tableKey(102), tableKey(103)}, + {13, tableKey(103), rowKey(103, 500)}, + {14, rowKey(103, 500), tableKey(104)}, + } + for _, r := range regions { + tests.MustPutRegion(re, tc, r.id, 1, r.start, r.end, + core.SetApproximateSize(1), core.SetApproximateKeys(1)) + } + + oc := leaderServer.GetRaftCluster().GetOperatorController() + // The same-table pair must merge: proves the whole merge pipeline + // (patrol -> merge checker -> operator) is live in this setup. + testutil.Eventually(re, func() bool { + op13, op14 := oc.GetOperator(13), oc.GetOperator(14) + return op13 != nil && op14 != nil + }) + // Regions of different tables under the same keyspace must never be + // merged while enable-cross-table-merge is false. + for range 20 { + for _, id := range []uint64{10, 11, 12} { + re.Nil(oc.GetOperator(id), "unexpected operator on cross-table region %d", id) + } + time.Sleep(100 * time.Millisecond) + } +} From 6fc8ba3a27991e550dc046fbf916d13f9912b7fc Mon Sep 17 00:00:00 2001 From: JmPotato Date: Fri, 10 Jul 2026 14:13:22 +0800 Subject: [PATCH 2/3] codec, checker: resolve cherry-pick conflicts in #10995 Signed-off-by: JmPotato --- pkg/codec/codec_test.go | 24 +- pkg/keyspace/util.go | 13 +- pkg/keyspace/util_test.go | 46 - pkg/schedule/checker/merge_checker_test.go | 8 +- pkg/schedule/checker/split_scatter_group.go | 182 ---- pkg/schedule/checker/split_scatter_test.go | 1061 ------------------- 6 files changed, 10 insertions(+), 1324 deletions(-) delete mode 100644 pkg/schedule/checker/split_scatter_group.go delete mode 100644 pkg/schedule/checker/split_scatter_test.go diff --git a/pkg/codec/codec_test.go b/pkg/codec/codec_test.go index 38d3639d757..6107561b68f 100644 --- a/pkg/codec/codec_test.go +++ b/pkg/codec/codec_test.go @@ -73,7 +73,9 @@ func TestTableIDWithKeyspacePrefix(t *testing.T) { // Same table: record and index keys must still resolve to the same identity. record := EncodeBytes(append(append([]byte{}, prefix...), GenerateRowKey(tableID, 1)...)) - index := EncodeBytes(append(append([]byte{}, prefix...), GenerateIndexKey(tableID, 7)...)) + 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()) @@ -140,23 +142,3 @@ func TestParseKeyspacePrefix(t *testing.T) { _, _, ok = ParseKeyspacePrefix([]byte{'t', 0x01, 0x02, 0x03}) re.False(ok) } -<<<<<<< HEAD -======= - -func TestGenerateRecordAndIndexKeys(t *testing.T) { - re := require.New(t) - tableID := int64(42) - indexID := int64(7) - - rowKeyPrefix := GenerateRecordKeyPrefix(tableID) - rowKey := GenerateRowKey(tableID, 1) - re.Equal(rowKeyPrefix, rowKey[:len(rowKeyPrefix)]) - re.Equal(tableID, EncodeBytes(rowKey).TableIdentity().TableID) - - indexKey := GenerateIndexKey(tableID, indexID) - _, decodedIndexKey, err := DecodeBytes(EncodeBytes(indexKey)) - re.NoError(err) - re.Equal(indexKey, decodedIndexKey) - re.Equal(tableID, EncodeBytes(indexKey).TableIdentity().TableID) -} ->>>>>>> 2b3abf1483 (codec, checker: fix enable-cross-table-merge for keyspace keys (#10992)) diff --git a/pkg/keyspace/util.go b/pkg/keyspace/util.go index acc833a47bb..7df93eefb19 100644 --- a/pkg/keyspace/util.go +++ b/pkg/keyspace/util.go @@ -16,6 +16,7 @@ package keyspace import ( "container/heap" + "encoding/binary" "encoding/hex" "regexp" "strconv" @@ -113,22 +114,10 @@ type RegionBound struct { // MakeRegionBound constructs the correct region boundaries of the given keyspace. func MakeRegionBound(id uint32) *RegionBound { -<<<<<<< HEAD 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) - if id == constant.MaxValidKeyspaceID { - // The right bound is an exclusive fencepost, not a real keyspace prefix. - rawRightBound = []byte{'s', 0, 0, 0} - txnRightBound = []byte{'y', 0, 0, 0} - } ->>>>>>> 2b3abf1483 (codec, checker: fix enable-cross-table-merge for keyspace keys (#10992)) return &RegionBound{ RawLeftBound: codec.EncodeBytes(append([]byte{'r'}, keyspaceIDBytes[1:]...)), RawRightBound: codec.EncodeBytes(append([]byte{'r'}, nextKeyspaceIDBytes[1:]...)), diff --git a/pkg/keyspace/util_test.go b/pkg/keyspace/util_test.go index 04eb5ac34ae..a4b44e2bdce 100644 --- a/pkg/keyspace/util_test.go +++ b/pkg/keyspace/util_test.go @@ -61,52 +61,6 @@ func TestValidateID(t *testing.T) { } } -<<<<<<< HEAD -======= -func TestMakeRegionBound(t *testing.T) { - re := require.New(t) - encodeKey := func(key []byte) []byte { - return []byte(codec.EncodeBytes(key)) - } - - regionBound := MakeRegionBound(0x010203) - re.Equal(encodeKey([]byte{'r', 0x01, 0x02, 0x03}), regionBound.RawLeftBound) - re.Equal(encodeKey([]byte{'r', 0x01, 0x02, 0x04}), regionBound.RawRightBound) - re.Equal(encodeKey([]byte{'x', 0x01, 0x02, 0x03}), regionBound.TxnLeftBound) - re.Equal(encodeKey([]byte{'x', 0x01, 0x02, 0x04}), regionBound.TxnRightBound) - - carryRegionBound := MakeRegionBound(0x0102ff) - re.Equal(encodeKey([]byte{'r', 0x01, 0x03, 0x00}), carryRegionBound.RawRightBound) - re.Equal(encodeKey([]byte{'x', 0x01, 0x03, 0x00}), carryRegionBound.TxnRightBound) - - maxRegionBound := MakeRegionBound(constant.MaxValidKeyspaceID) - re.Equal(encodeKey([]byte{'r', 0xff, 0xff, 0xff}), maxRegionBound.RawLeftBound) - re.Equal(encodeKey([]byte{'s', 0x00, 0x00, 0x00}), maxRegionBound.RawRightBound) - re.Equal(encodeKey([]byte{'x', 0xff, 0xff, 0xff}), maxRegionBound.TxnLeftBound) - re.Equal(encodeKey([]byte{'y', 0x00, 0x00, 0x00}), maxRegionBound.TxnRightBound) -} - -func TestMaxKeyspaceLabelRuleSplitKeys(t *testing.T) { - re := require.New(t) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - regionLabeler, err := labeler.NewRegionLabeler(ctx, endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil), time.Hour) - re.NoError(err) - - re.NoError(regionLabeler.SetLabelRule(MakeTxnLabelRule(constant.MaxValidKeyspaceID))) - encodeKey := func(key []byte) []byte { - return []byte(codec.EncodeBytes(key)) - } - re.Equal( - [][]byte{ - encodeKey([]byte{'x', 0xff, 0xff, 0xff}), - encodeKey([]byte{'y', 0x00, 0x00, 0x00}), - }, - regionLabeler.GetSplitKeys(nil, nil), - ) -} - ->>>>>>> 2b3abf1483 (codec, checker: fix enable-cross-table-merge for keyspace keys (#10992)) func TestValidateName(t *testing.T) { re := require.New(t) testCases := []struct { diff --git a/pkg/schedule/checker/merge_checker_test.go b/pkg/schedule/checker/merge_checker_test.go index ea7f4ed9817..75d528bfb1e 100644 --- a/pkg/schedule/checker/merge_checker_test.go +++ b/pkg/schedule/checker/merge_checker_test.go @@ -607,19 +607,23 @@ func TestAllowMergeCrossTable(t *testing.T) { tableB = int64(101) indexID = int64(1) ) + generateIndexKey := func(tableID, indexID int64) []byte { + key := append(codec.GenerateTableKey(tableID), '_', 'i') + return codec.EncodeInt(key, indexID) + } // Classic keys. classicTableA := codec.EncodeBytes(codec.GenerateTableKey(tableA)) classicTableB := codec.EncodeBytes(codec.GenerateTableKey(tableB)) classicTableC := codec.EncodeBytes(codec.GenerateTableKey(tableB + 1)) - classicIndexA := codec.EncodeBytes(codec.GenerateIndexKey(tableA, indexID)) + classicIndexA := codec.EncodeBytes(generateIndexKey(tableA, indexID)) classicRecordA := codec.EncodeBytes(codec.GenerateRowKey(tableA, 1)) // Same-keyspace keys. ks1TableA := encodeKeyspaceRawKey(keyspace1, codec.GenerateTableKey(tableA)) ks1TableB := encodeKeyspaceRawKey(keyspace1, codec.GenerateTableKey(tableB)) ks1TableC := encodeKeyspaceRawKey(keyspace1, codec.GenerateTableKey(tableB+1)) - ks1IndexA := encodeKeyspaceRawKey(keyspace1, codec.GenerateIndexKey(tableA, indexID)) + ks1IndexA := encodeKeyspaceRawKey(keyspace1, generateIndexKey(tableA, indexID)) ks1RecordA := encodeKeyspaceRawKey(keyspace1, codec.GenerateRowKey(tableA, 1)) // Different-keyspace keys. Adjacent ranges are constructed artificially so diff --git a/pkg/schedule/checker/split_scatter_group.go b/pkg/schedule/checker/split_scatter_group.go deleted file mode 100644 index d0ab69f3e7e..00000000000 --- a/pkg/schedule/checker/split_scatter_group.go +++ /dev/null @@ -1,182 +0,0 @@ -// Copyright 2026 TiKV Project Authors. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package checker - -import ( - "bytes" - "fmt" - - "github.com/tikv/pd/pkg/codec" - "github.com/tikv/pd/pkg/core" -) - -var ( - splitScatterTablePrefix = []byte{'t'} - splitScatterIndexPrefix = []byte("_i") -) - -type splitScatterKeyspaceValidator func(uint32) bool - -type splitScatterDecodedKey struct { - rawKey []byte - keyspacePrefix []byte - keyspaceID uint32 -} - -func (key splitScatterDecodedKey) hasKeyspace() bool { - return len(key.keyspacePrefix) > 0 -} - -func resolveSplitScatterRangeHintWithKeyspaceValidator( - region *core.RegionInfo, - validateKeyspace splitScatterKeyspaceValidator, -) splitScatterRangeHint { - decodedKey, err := decodeSplitScatterRegionKey(region.GetStartKey(), validateKeyspace) - if err != nil { - return splitScatterRangeHint{} - } - rawKey := decodedKey.rawKey - - if !bytes.HasPrefix(rawKey, splitScatterTablePrefix) { - return splitScatterRangeHint{} - } - rest := rawKey[len(splitScatterTablePrefix):] - rest, tableID, err := codec.DecodeInt(rest) - if err != nil { - return splitScatterRangeHint{} - } - - tablePrefix := append([]byte(nil), codec.GenerateTableKey(tableID)...) - tableGroup := makeSplitScatterTableGroup(tableID) - if decodedKey.hasKeyspace() { - tableGroup = makeSplitScatterKeyspaceTableGroup(decodedKey.keyspaceID, tableID) - } - if !bytes.HasPrefix(rest, splitScatterIndexPrefix) { - return splitScatterPrefixRangeWithKeyspaceGroup(decodedKey.keyspacePrefix, tablePrefix, tableGroup) - } - - indexRest := rest[len(splitScatterIndexPrefix):] - _, indexID, err := codec.DecodeInt(indexRest) - if err != nil { - return splitScatterRangeHint{} - } - - indexPrefix := codec.GenerateIndexKey(tableID, indexID) - indexGroup := makeSplitScatterIndexGroup(tableID, indexID) - if decodedKey.hasKeyspace() { - indexGroup = makeSplitScatterKeyspaceIndexGroup(decodedKey.keyspaceID, tableID, indexID) - } - indexRange := splitScatterPrefixRangeWithKeyspaceGroup(decodedKey.keyspacePrefix, indexPrefix, indexGroup) - endKey := region.GetEndKey() - - // We intentionally over-approximate ambiguous table-key ranges. If PD can - // no longer prove the region stays within a single index prefix, it falls - // back to the table-scoped group instead of dropping back to the family - // group, so table-boundary splits and merged ranges still participate in - // the broader scatter continuity/baseline. - // Both endKey and indexRange.endKey are MemComparable-encoded, so - // bytes.Compare correctly reflects the key ordering. - if len(endKey) == 0 || len(indexRange.startKey) == 0 || len(indexRange.endKey) == 0 || bytes.Compare(endKey, indexRange.endKey) > 0 { - return splitScatterPrefixRangeWithKeyspaceGroup(decodedKey.keyspacePrefix, tablePrefix, tableGroup) - } - return indexRange -} - -func decodeSplitScatterRegionKey( - regionKey []byte, - validateKeyspace splitScatterKeyspaceValidator, -) (splitScatterDecodedKey, error) { - _, rawKey, err := codec.DecodeBytes(regionKey) - if err != nil { - return splitScatterDecodedKey{}, err - } - decodedKey := splitScatterDecodedKey{rawKey: rawKey} - mode, keyspaceID, ok := codec.ParseKeyspacePrefix(rawKey) - if !ok || mode != codec.TxnKeyspaceModePrefix { - return decodedKey, nil - } - - // Split-scatter range hints are table/index-scoped, so only TiDB txn - // keyspace keys from a known keyspace range are decoded. With only a - // region key, an API V2 txn prefix can be indistinguishable from a - // classic/raw user key that starts with the same bytes. - if validateKeyspace == nil || !validateKeyspace(keyspaceID) { - return decodedKey, nil - } - decodedKey.rawKey = rawKey[codec.KeyspacePrefixLen:] - decodedKey.keyspacePrefix = codec.MakeKeyspacePrefix(mode, keyspaceID) - decodedKey.keyspaceID = keyspaceID - return decodedKey, nil -} - -func splitScatterPrefixRange(rawPrefix []byte) splitScatterRangeHint { - return splitScatterPrefixRangeWithGroup(rawPrefix, "") -} - -func splitScatterPrefixRangeWithGroup(rawPrefix []byte, scatterGroup string) splitScatterRangeHint { - return splitScatterPrefixRangeWithKeyspaceGroup(nil, rawPrefix, scatterGroup) -} - -func splitScatterPrefixRangeWithKeyspaceGroup(keyspacePrefix, rawPrefix []byte, scatterGroup string) splitScatterRangeHint { - startKey := codec.EncodeBytes(appendKeyspacePrefix(keyspacePrefix, rawPrefix)) - endRawPrefix := splitScatterNextPrefix(rawPrefix) - if len(endRawPrefix) == 0 { - // Current callers use TiDB table/index prefixes, which always start - // with 't' and therefore have a finite next prefix. Keep the fallback - // here so this helper remains safe if it is ever used with an all-0xff - // prefix. - return splitScatterRangeHint{startKey: startKey, scatterGroup: scatterGroup} - } - return splitScatterRangeHint{ - startKey: startKey, - endKey: codec.EncodeBytes(appendKeyspacePrefix(keyspacePrefix, endRawPrefix)), - scatterGroup: scatterGroup, - } -} - -func appendKeyspacePrefix(keyspacePrefix, rawKey []byte) []byte { - key := make([]byte, 0, len(keyspacePrefix)+len(rawKey)) - key = append(key, keyspacePrefix...) - key = append(key, rawKey...) - return key -} - -func makeSplitScatterTableGroup(tableID int64) string { - return fmt.Sprintf("split-scatter-table-%d", tableID) -} - -func makeSplitScatterIndexGroup(tableID, indexID int64) string { - return fmt.Sprintf("split-scatter-index-%d-%d", tableID, indexID) -} - -func makeSplitScatterKeyspaceTableGroup(keyspaceID uint32, tableID int64) string { - return fmt.Sprintf("split-scatter-keyspace-%d-table-%d", keyspaceID, tableID) -} - -func makeSplitScatterKeyspaceIndexGroup(keyspaceID uint32, tableID, indexID int64) string { - return fmt.Sprintf("split-scatter-keyspace-%d-index-%d-%d", keyspaceID, tableID, indexID) -} - -func splitScatterNextPrefix(key []byte) []byte { - next := append([]byte(nil), key...) - for i := len(next) - 1; i >= 0; i-- { - if next[i] == 0xFF { - continue - } - next[i]++ - return next[:i+1] - } - return nil -} diff --git a/pkg/schedule/checker/split_scatter_test.go b/pkg/schedule/checker/split_scatter_test.go deleted file mode 100644 index 0d6076716f1..00000000000 --- a/pkg/schedule/checker/split_scatter_test.go +++ /dev/null @@ -1,1061 +0,0 @@ -// Copyright 2026 TiKV Project Authors. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package checker - -import ( - "context" - "testing" - "time" - - "github.com/prometheus/client_golang/prometheus" - promtestutil "github.com/prometheus/client_golang/prometheus/testutil" - "github.com/stretchr/testify/require" - - "github.com/pingcap/kvproto/pkg/metapb" - - "github.com/tikv/pd/pkg/codec" - "github.com/tikv/pd/pkg/core" - "github.com/tikv/pd/pkg/keyspace" - "github.com/tikv/pd/pkg/keyspace/constant" - "github.com/tikv/pd/pkg/mock/mockcluster" - "github.com/tikv/pd/pkg/mock/mockconfig" - "github.com/tikv/pd/pkg/schedule/hbstream" - "github.com/tikv/pd/pkg/schedule/labeler" - "github.com/tikv/pd/pkg/schedule/operator" - "github.com/tikv/pd/pkg/schedule/scatter" -) - -const ( - splitScatterObservedRegionID uint64 = 101 - splitScatterTestTableID int64 = 42 - splitScatterTestIndexID int64 = 7 - splitScatterTestKeyspaceID uint32 = 4242 - splitScatterTestNextGenKeyspaceID uint32 = constant.SystemKeyspaceID - splitScatterTestSourceWaitVersion = uint64(0) - // CPU usage is only populated to mimic load-split region heartbeat data. - // Current split-scatter dispatch does not rank pending regions by CPU. - splitScatterNoCPUUsage uint64 = 0 - splitScatterReportedCPUUsage uint64 = 1 -) - -func (c *Controller) collectTopPendingSplitScatter(limit int) []splitScatterPendingItem { - return c.splitScatter.collectTopPendingSplitScatter(limit) -} - -func (c *Controller) dispatchSplitScatterRegions() { - c.splitScatter.dispatchSplitScatterRegions() -} - -func TestSplitScatterControllerCleanupResetsPendingGauge(t *testing.T) { - re := require.New(t) - splitScatterPendingGauge.Set(7) - - controller, _, _, cleanup := newTestSplitScatterController(t) - cleanup() - - re.Equal(0, splitScatterPendingCount(controller)) - re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) -} - -func TestRecordSplitScatterBatchCollectsPendingRegions(t *testing.T) { - re := require.New(t) - controller, tc, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102, 103}) - re.Equal(4, splitScatterPendingCount(controller)) - re.Equal(float64(4), promtestutil.ToFloat64(splitScatterPendingGauge)) - - sourceGroup := splitScatterPendingGroup(t, controller, 100) - for _, regionID := range []uint64{101, 102, 103} { - re.Equal(sourceGroup, splitScatterPendingGroup(t, controller, regionID)) - } - - putSplitScatterRegion(tc, 101, "m", "n", splitScatterReportedCPUUsage) - putSplitScatterRegion(tc, 102, "n", "o", splitScatterReportedCPUUsage) - putSplitScatterRegion(tc, 103, "o", "", splitScatterReportedCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - - re.ElementsMatch([]uint64{100, 101, 102, 103}, pendingRegionIDs(controller.collectTopPendingSplitScatter(4))) -} - -func TestCheckSplitScatterRegionsCreatesScatterOperator(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) - putSplitScatterRegion(tc, 101, "m", "t", splitScatterReportedCPUUsage) - putSplitScatterRegion(tc, 102, "t", "", splitScatterReportedCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - - group := splitScatterPendingGroup(t, controller, 101) - - controller.dispatchSplitScatterRegions() - - var op *operator.Operator - for _, regionID := range []uint64{100, 101, 102} { - op = oc.GetOperator(regionID) - if op != nil { - break - } - } - re.NotNil(op) - re.Equal(scatter.InternalScatterOperatorDesc, op.Desc()) - opGroup, ok := op.GetAdditionalInfo("group") - re.True(ok) - re.Equal(group, opGroup) - batchGroup, ok := op.GetAdditionalInfo("batch-group") - re.True(ok) - re.Equal(group, batchGroup) -} - -func TestDispatchSplitScatterKeepsPendingUntilSplitHeartbeat(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - - controller.dispatchSplitScatterRegions() - - re.Equal(2, splitScatterPendingCount(controller)) - re.Nil(oc.GetOperator(101)) - - putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) - - retrySplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) - re.Empty(controller.collectTopPendingSplitScatter(2)) - advanceSplitScatterSourceVersion(t, tc) - setSplitScatterNextDispatchAt(t, controller, time.Now().Add(-time.Second)) - re.ElementsMatch([]uint64{100, 101}, pendingRegionIDs(controller.collectTopPendingSplitScatter(2))) - - controller.dispatchSplitScatterRegions() - - op := oc.GetOperator(101) - re.NotNil(op) - re.Equal(scatter.InternalScatterOperatorDesc, op.Desc()) -} - -func TestDispatchSplitScatterUsesRequestWaitVersionWhenCacheLags(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - source := tc.GetRegion(100) - re.NotNil(source) - tc.PutRegion(source.Clone(core.SetRegionVersion(4))) - - controller.RecordSplitScatterBatch(100, 6, []uint64{101}) - putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) - advanceSplitScatterRegionVersion(t, tc, 100) - - controller.dispatchSplitScatterRegions() - - re.Empty(oc.GetOperators()) - re.Equal(2, splitScatterPendingCount(controller)) - - advanceSplitScatterRegionVersion(t, tc, 100) - setSplitScatterNextDispatchAt(t, controller, time.Now().Add(-time.Second)) - controller.dispatchSplitScatterRegions() - - re.NotNil(oc.GetOperator(101)) -} - -func TestDispatchSplitScatterRespectsScheduleLimit(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) - putSplitScatterRegion(tc, 101, "m", "t", splitScatterReportedCPUUsage) - putSplitScatterRegion(tc, 102, "t", "", splitScatterReportedCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - - tc.SetSplitScatterScheduleLimit(1) - controller.dispatchSplitScatterRegions() - - re.Len(oc.GetOperators(), 1) - re.Equal(uint64(1), oc.OperatorCount(operator.OpSplitScatter)) - - controller.dispatchSplitScatterRegions() - - re.Len(oc.GetOperators(), 1) -} - -func TestRecordSplitScatterBatchSkipsWhenDisabled(t *testing.T) { - re := require.New(t) - controller, tc, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - tc.SetSplitScatterScheduleLimit(0) - - droppedBefore := promtestutil.ToFloat64(splitScatterPendingDroppedCounter) - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) - - re.Equal(0, splitScatterPendingCount(controller)) - re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) - re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingDroppedCounter)-droppedBefore) -} - -func TestDispatchSplitScatterClearsPendingWhenDisabled(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101, 102}) - re.Equal(3, splitScatterPendingCount(controller)) - - tc.SetSplitScatterScheduleLimit(0) - setSplitScatterNextDispatchAt(t, controller, time.Now().Add(splitScatterRetryBackoff)) - disabledBefore := promtestutil.ToFloat64(splitScatterDispatchDisabledCounter) - controller.dispatchSplitScatterRegions() - - re.Empty(oc.GetOperators()) - re.Equal(0, splitScatterPendingCount(controller)) - re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) - re.Equal(float64(1), promtestutil.ToFloat64(splitScatterDispatchDisabledCounter)-disabledBefore) -} - -func TestDispatchSplitScatterCleansExpiredPendingBeforeEarlyReturn(t *testing.T) { - testCases := []struct { - name string - setupEarlyReturn func(*mockcluster.Cluster, *operator.Controller) - counter prometheus.Counter - }{ - { - name: "disabled", - setupEarlyReturn: func(tc *mockcluster.Cluster, _ *operator.Controller) { - tc.SetSplitScatterScheduleLimit(0) - }, - counter: splitScatterDispatchDisabledCounter, - }, - { - name: "schedule limit", - setupEarlyReturn: func(tc *mockcluster.Cluster, oc *operator.Controller) { - tc.SetSplitScatterScheduleLimit(1) - region := tc.GetRegion(100) - op := operator.NewTestOperator( - region.GetID(), - region.GetRegionEpoch(), - operator.OpSplitScatter|operator.OpRegion, - operator.TransferLeader{FromStore: 1, ToStore: 2}, - ) - require.True(t, oc.AddOperator(op)) - }, - counter: splitScatterDispatchScheduleLimitCounter, - }, - } - - for _, testCase := range testCases { - t.Run(testCase.name, func(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - expireSplitScatterPendingAt(t, controller, 100, time.Now().Add(-time.Second)) - expireSplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) - testCase.setupEarlyReturn(tc, oc) - - expiredBefore := splitScatterPendingExpiredCount("false") - counterBefore := promtestutil.ToFloat64(testCase.counter) - controller.dispatchSplitScatterRegions() - - re.Equal(0, splitScatterPendingCount(controller)) - re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) - re.Equal(float64(2), splitScatterPendingExpiredCount("false")-expiredBefore) - re.Equal(float64(0), promtestutil.ToFloat64(testCase.counter)-counterBefore) - }) - } -} - -func TestCollectTopPendingDelaysMissingRegions(t *testing.T) { - re := require.New(t) - controller, _, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - - missingBefore := promtestutil.ToFloat64(splitScatterDispatchRegionMissingCounter) - re.Empty(controller.collectTopPendingSplitScatter(2)) - - re.Equal(float64(1), promtestutil.ToFloat64(splitScatterDispatchRegionMissingCounter)-missingBefore) - pending := splitScatterPending(t, controller, 101) - re.True(pending.retryAt.After(time.Now())) - - re.Empty(controller.collectTopPendingSplitScatter(2)) - re.Equal(float64(1), promtestutil.ToFloat64(splitScatterDispatchRegionMissingCounter)-missingBefore) -} - -func TestDispatchSplitScatterBacksOffWhenNoCandidates(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - - controller.dispatchSplitScatterRegions() - - re.True(splitScatterNextDispatchAt(t, controller).After(time.Now())) - putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - - controller.dispatchSplitScatterRegions() - - re.Empty(oc.GetOperators()) - - retrySplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) - setSplitScatterNextDispatchAt(t, controller, time.Now().Add(-time.Second)) - controller.dispatchSplitScatterRegions() - - re.NotNil(oc.GetOperator(101)) -} - -func TestDispatchSplitScatterRespectsScheduleDeny(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - - re.NoError(tc.GetRegionLabeler().SetLabelRule(&labeler.LabelRule{ - ID: "split-scatter-schedule-deny", - Labels: []labeler.RegionLabel{{Key: "schedule", Value: "deny"}}, - RuleType: labeler.KeyRange, - Data: []any{map[string]any{"start_key": "", "end_key": ""}}, - })) - - counterBefore := promtestutil.ToFloat64(splitScatterDispatchScheduleDisabledCounter) - controller.dispatchSplitScatterRegions() - - re.Empty(oc.GetOperators()) - re.Equal(2, splitScatterPendingCount(controller)) - re.Equal(float64(2), promtestutil.ToFloat64(splitScatterDispatchScheduleDisabledCounter)-counterBefore) - for _, regionID := range []uint64{100, 101} { - pending := splitScatterPending(t, controller, regionID) - re.True(pending.retryAt.After(time.Now())) - } -} - -func TestCollectTopPendingRemovesExpiredPending(t *testing.T) { - re := require.New(t) - controller, _, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - expireSplitScatterPendingAt(t, controller, 100, time.Now().Add(-time.Second)) - expireSplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) - - expiredBefore := splitScatterPendingExpiredCount("false") - re.Empty(controller.collectTopPendingSplitScatter(2)) - re.Equal(0, splitScatterPendingCount(controller)) - re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) - re.Equal(float64(2), splitScatterPendingExpiredCount("false")-expiredBefore) -} - -func TestCollectTopPendingMarksAttemptedBeforeExpiration(t *testing.T) { - re := require.New(t) - controller, tc, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - putSplitScatterRegion(tc, 101, "m", "", splitScatterReportedCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - - attemptedBefore := splitScatterPendingExpiredCount("true") - unattemptedBefore := splitScatterPendingExpiredCount("false") - re.Len(controller.collectTopPendingSplitScatter(1), 1) - expireSplitScatterPendingAt(t, controller, 100, time.Now().Add(-time.Second)) - expireSplitScatterPendingAt(t, controller, 101, time.Now().Add(-time.Second)) - - re.Empty(controller.collectTopPendingSplitScatter(2)) - re.Equal(0, splitScatterPendingCount(controller)) - re.Equal(float64(0), promtestutil.ToFloat64(splitScatterPendingGauge)) - re.Equal(float64(1), splitScatterPendingExpiredCount("true")-attemptedBefore) - re.Equal(float64(1), splitScatterPendingExpiredCount("false")-unattemptedBefore) -} - -func TestRecordSplitScatterBatchRespectsPendingLimit(t *testing.T) { - re := require.New(t) - controller, _, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - fillSplitScatterPending(controller, time.Time{}) - - droppedBefore := promtestutil.ToFloat64(splitScatterPendingDroppedCounter) - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - - re.Equal(splitScatterPendingLimit, splitScatterPendingCount(controller)) - re.Equal(float64(2), promtestutil.ToFloat64(splitScatterPendingDroppedCounter)-droppedBefore) - controller.splitScatter.pendingMu.RLock() - _, sourceExists := controller.splitScatter.pending[100] - _, childExists := controller.splitScatter.pending[101] - controller.splitScatter.pendingMu.RUnlock() - re.False(sourceExists) - re.False(childExists) -} - -func TestRecordSplitScatterBatchClearsExpiredPendingBeforeLimitCheck(t *testing.T) { - re := require.New(t) - controller, _, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - fillSplitScatterPending(controller, time.Now().Add(-time.Second)) - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - - re.Equal(2, splitScatterPendingCount(controller)) - re.Equal(makeSplitScatterGroup(100, 101), splitScatterPendingGroup(t, controller, 101)) -} - -func TestCollectTopPendingSortsBeforeLimit(t *testing.T) { - re := require.New(t) - controller, tc, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(200, splitScatterTestSourceWaitVersion, []uint64{201}) - putSplitScatterRegion(tc, 200, "n", "o", splitScatterNoCPUUsage) - putSplitScatterRegion(tc, 201, "o", "p", splitScatterReportedCPUUsage) - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - putSplitScatterRegion(tc, 101, "m", "n", splitScatterReportedCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - - pending := controller.collectTopPendingSplitScatter(1) - re.Len(pending, 1) - re.Equal(uint64(100), pending[0].regionID) -} - -func TestCollectTopPendingPrioritizesNearExpiration(t *testing.T) { - re := require.New(t) - controller, tc, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - controller.RecordSplitScatterBatch(200, splitScatterTestSourceWaitVersion, []uint64{201}) - putSplitScatterRegion(tc, 100, "m", "n", splitScatterNoCPUUsage) - putSplitScatterRegion(tc, 101, "n", "o", splitScatterReportedCPUUsage) - putSplitScatterRegion(tc, 200, "o", "p", splitScatterNoCPUUsage) - putSplitScatterRegion(tc, 201, "p", "q", splitScatterReportedCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - advanceSplitScatterRegionVersion(t, tc, 200) - - expireSplitScatterPendingAt(t, controller, 100, time.Now().Add(2*time.Minute)) - expireSplitScatterPendingAt(t, controller, 101, time.Now().Add(2*time.Minute)) - expireSplitScatterPendingAt(t, controller, 200, time.Now().Add(time.Minute)) - expireSplitScatterPendingAt(t, controller, 201, time.Now().Add(time.Minute)) - - pending := controller.collectTopPendingSplitScatter(1) - re.Len(pending, 1) - re.Equal(uint64(200), pending[0].regionID) -} - -func TestCollectTopPendingResolvesRangeHint(t *testing.T) { - testCases := []struct { - name string - startKey []byte - endKey []byte - wantRange splitScatterRangeHint - wantGroup string - keyspaces []uint32 - }{ - { - name: "index region", - startKey: newSplitScatterIndexKey("a"), - endKey: newSplitScatterIndexKey("m"), - wantRange: splitScatterPrefixRange(splitScatterIndexKeyPrefix()), - wantGroup: makeSplitScatterIndexGroup(splitScatterTestTableID, splitScatterTestIndexID), - }, - { - name: "record region", - startKey: newSplitScatterRecordKey(splitScatterTestTableID, "a"), - endKey: newSplitScatterRecordKey(splitScatterTestTableID, "m"), - wantRange: splitScatterPrefixRange(codec.GenerateTableKey(splitScatterTestTableID)), - wantGroup: makeSplitScatterTableGroup(splitScatterTestTableID), - }, - { - name: "bare table boundary", - startKey: newSplitScatterTableBoundaryKey(splitScatterTestTableID), - endKey: newSplitScatterIndexKey("m"), - wantRange: splitScatterPrefixRange(codec.GenerateTableKey(splitScatterTestTableID)), - wantGroup: makeSplitScatterTableGroup(splitScatterTestTableID), - }, - { - name: "cross entity falls back to table", - startKey: newSplitScatterIndexKey("a"), - endKey: newSplitScatterRecordKey(splitScatterTestTableID, "m"), - wantRange: splitScatterPrefixRange(codec.GenerateTableKey(splitScatterTestTableID)), - wantGroup: makeSplitScatterTableGroup(splitScatterTestTableID), - }, - { - name: "cross table uses start table", - startKey: newSplitScatterRecordKey(splitScatterTestTableID, "a"), - endKey: newSplitScatterRecordKey(splitScatterTestTableID+1, "m"), - wantRange: splitScatterPrefixRange(codec.GenerateTableKey(splitScatterTestTableID)), - wantGroup: makeSplitScatterTableGroup(splitScatterTestTableID), - }, - { - name: "nextgen keyspace index region", - startKey: newSplitScatterKeyspaceIndexKey(splitScatterTestNextGenKeyspaceID, "a"), - endKey: newSplitScatterKeyspaceIndexKey(splitScatterTestNextGenKeyspaceID, "m"), - wantRange: splitScatterKeyspacePrefixRange(splitScatterTestNextGenKeyspaceID, splitScatterIndexKeyPrefix()), - wantGroup: makeSplitScatterKeyspaceIndexGroup(splitScatterTestNextGenKeyspaceID, splitScatterTestTableID, splitScatterTestIndexID), - keyspaces: []uint32{splitScatterTestNextGenKeyspaceID}, - }, - } - - for _, testCase := range testCases { - t.Run(testCase.name, func(t *testing.T) { - re := require.New(t) - controller, tc, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - putSplitScatterRegionWithKeys(tc, testCase.startKey, testCase.endKey, splitScatterReportedCPUUsage) - putSplitScatterRegion(tc, 100, "z", "", splitScatterNoCPUUsage) - advanceSplitScatterSourceVersion(t, tc) - - re.Equal(makeSplitScatterGroup(100, 101), splitScatterPendingGroup(t, controller, 101)) - re.ElementsMatch([]uint64{100, 101}, pendingRegionIDs(controller.collectTopPendingSplitScatter(2))) - rangeHint := resolveSplitScatterRangeHintWithKeyspaceValidator( - tc.GetRegion(101), - splitScatterKeyspaceValidatorFor(testCase.keyspaces...), - ) - re.Equal(testCase.wantRange.startKey, rangeHint.startKey) - re.Equal(testCase.wantRange.endKey, rangeHint.endKey) - re.Equal(testCase.wantGroup, rangeHint.scatterGroup) - }) - } -} - -func TestResolveSplitScatterRangeHintIgnoresRawLikeKeyspaceKeys(t *testing.T) { - re := require.New(t) - region := core.NewRegionInfo(&metapb.Region{ - Id: 1, - StartKey: newSplitScatterRawKeyspaceRecordKey(splitScatterTestKeyspaceID, splitScatterTestTableID, "a"), - EndKey: newSplitScatterRawKeyspaceRecordKey(splitScatterTestKeyspaceID, splitScatterTestTableID, "m"), - }, nil) - - rangeHint := resolveSplitScatterRangeHintWithKeyspaceValidator( - region, - splitScatterKeyspaceValidatorFor(splitScatterTestKeyspaceID), - ) - re.Equal(splitScatterRangeHint{}, rangeHint) -} - -func TestResolveSplitScatterRangeHintRequiresKnownTxnKeyspaceBounds(t *testing.T) { - testCases := []struct { - name string - keyspaceID uint32 - }{ - {name: "normal keyspace", keyspaceID: splitScatterTestKeyspaceID}, - {name: "max valid keyspace", keyspaceID: constant.MaxValidKeyspaceID}, - } - - for _, testCase := range testCases { - t.Run(testCase.name, func(t *testing.T) { - re := require.New(t) - controller, tc, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - startKey := newSplitScatterKeyspaceRecordKey(testCase.keyspaceID, "a") - endKey := newSplitScatterKeyspaceRecordKey(testCase.keyspaceID, "m") - region := core.NewRegionInfo(&metapb.Region{ - Id: 1, - StartKey: startKey, - EndKey: endKey, - }, nil) - re.Equal(splitScatterRangeHint{}, resolveSplitScatterRangeHintWithKeyspaceValidator(region, nil)) - - regionBound := keyspace.MakeRegionBound(testCase.keyspaceID) - putSplitScatterRegionWithKeysByID(tc, 90, regionBound.TxnLeftBound, startKey, splitScatterNoCPUUsage) - putSplitScatterRegionWithKeysByID(tc, 91, regionBound.TxnRightBound, nil, splitScatterNoCPUUsage) - - rangeHint := resolveSplitScatterRangeHintWithKeyspaceValidator( - region, - controller.splitScatter.hasSplitScatterTxnKeyspaceBounds, - ) - wantRange := splitScatterKeyspacePrefixRange(testCase.keyspaceID, codec.GenerateTableKey(splitScatterTestTableID)) - wantRange.scatterGroup = makeSplitScatterKeyspaceTableGroup(testCase.keyspaceID, splitScatterTestTableID) - re.Equal(wantRange, rangeHint) - }) - } -} - -func TestDispatchSplitScatterUsesRangeScatterGroup(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - putSplitScatterRegionWithKeysByID(tc, 90, newSplitScatterIndexKey("m"), newSplitScatterIndexKey("z"), splitScatterNoCPUUsage) - putSplitScatterRegionWithKeysByID(tc, 101, newSplitScatterIndexKey("a"), newSplitScatterIndexKey("m"), splitScatterReportedCPUUsage) - advanceSplitScatterRegionVersion(t, tc, 100) - - batchGroup := splitScatterPendingGroup(t, controller, 101) - - controller.dispatchSplitScatterRegions() - - expectedScatterGroup := makeSplitScatterIndexGroup(splitScatterTestTableID, splitScatterTestIndexID) - op := oc.GetOperator(101) - re.NotNil(op) - opGroup, ok := op.GetAdditionalInfo("group") - re.True(ok) - re.Equal(expectedScatterGroup, opGroup) - opBatchGroup, ok := op.GetAdditionalInfo("batch-group") - re.True(ok) - re.Equal(batchGroup, opBatchGroup) -} - -func TestDispatchSplitScatterUsesKeyspaceRangeScatterGroup(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - startKey := newSplitScatterKeyspaceIndexKey(splitScatterTestKeyspaceID, "a") - endKey := newSplitScatterKeyspaceIndexKey(splitScatterTestKeyspaceID, "m") - regionBound := keyspace.MakeRegionBound(splitScatterTestKeyspaceID) - putSplitScatterRegionWithKeysByID(tc, 90, regionBound.TxnLeftBound, startKey, splitScatterNoCPUUsage) - putSplitScatterRegionWithKeysByID(tc, 91, regionBound.TxnRightBound, nil, splitScatterNoCPUUsage) - putSplitScatterRegionWithKeysByID(tc, 101, startKey, endKey, splitScatterReportedCPUUsage) - advanceSplitScatterRegionVersion(t, tc, 100) - - batchGroup := splitScatterPendingGroup(t, controller, 101) - - controller.dispatchSplitScatterRegions() - - expectedScatterGroup := makeSplitScatterKeyspaceIndexGroup( - splitScatterTestKeyspaceID, - splitScatterTestTableID, - splitScatterTestIndexID, - ) - op := oc.GetOperator(101) - re.NotNil(op) - opGroup, ok := op.GetAdditionalInfo("group") - re.True(ok) - re.Equal(expectedScatterGroup, opGroup) - opBatchGroup, ok := op.GetAdditionalInfo("batch-group") - re.True(ok) - re.Equal(batchGroup, opBatchGroup) -} - -func TestDispatchSplitScatterKeepsStableGroupWhenRegionSplitsAgain(t *testing.T) { - re := require.New(t) - controller, tc, oc, cleanup := newTestSplitScatterController(t) - defer cleanup() - - stableGroup := makeSplitScatterIndexGroup(splitScatterTestTableID, splitScatterTestIndexID) - sourceID := uint64(100) - childIDs := []uint64{101, 201, 301} - splitKeys := []string{"m", "t", "x"} - previousBatchGroup := "" - - putSplitScatterRegionWithKeysByID(tc, sourceID, newSplitScatterIndexKey("a"), newSplitScatterIndexKey("z"), splitScatterNoCPUUsage) - for i, childID := range childIDs { - controller.RecordSplitScatterBatch(sourceID, splitScatterTestSourceWaitVersion, []uint64{childID}) - batchGroup := splitScatterPendingGroup(t, controller, childID) - re.NotEqual(previousBatchGroup, batchGroup) - - tc.PutRegion(tc.GetRegion(sourceID).Clone( - core.WithEndKey(newSplitScatterIndexKey(splitKeys[i])), - core.WithIncVersion(), - )) - putSplitScatterRegionWithKeysByID(tc, childID, newSplitScatterIndexKey(splitKeys[i]), newSplitScatterIndexKey("z"), splitScatterReportedCPUUsage) - - controller.dispatchSplitScatterRegions() - - requireInternalScatterOpsUseGroups(t, oc, stableGroup, batchGroup) - removeInternalScatterOps(oc) - previousBatchGroup = batchGroup - } -} - -func TestDispatchSplitScatterBacksOff(t *testing.T) { - testCases := []struct { - name string - putRegion func(*mockcluster.Cluster) - }{ - { - name: "region is not fully replicated", - putRegion: func(tc *mockcluster.Cluster) { - putSplitScatterRegionWithStores(tc, 101, "m", "", splitScatterReportedCPUUsage, 1, 2) - }, - }, - { - name: "scatter internal fails", - putRegion: func(tc *mockcluster.Cluster) { - putSplitScatterRegionWithoutLeader(tc, 101, "m", "", splitScatterReportedCPUUsage) - }, - }, - } - - for _, testCase := range testCases { - t.Run(testCase.name, func(t *testing.T) { - re := require.New(t) - controller, tc, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - testCase.putRegion(tc) - advanceSplitScatterSourceVersion(t, tc) - - controller.dispatchSplitScatterRegions() - - re.Equal(1, splitScatterPendingCount(controller)) - pending := splitScatterObservedPending(t, controller) - re.True(pending.retryAt.After(time.Now())) - re.Empty(pendingRegionIDs(controller.collectTopPendingSplitScatter(2))) - }) - } -} - -func TestDispatchSplitScatterIgnoresStalePendingSnapshot(t *testing.T) { - re := require.New(t) - controller, _, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - stalePending := splitScatterObservedPending(t, controller) - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{102, 101}) - currentPending := splitScatterObservedPending(t, controller) - re.NotEqual(stalePending.group, currentPending.group) - re.Equal(time.Time{}, currentPending.retryAt) - - controller.splitScatter.delayPendingSplitScatter(stalePending) - - currentPending = splitScatterObservedPending(t, controller) - re.Equal(time.Time{}, currentPending.retryAt) - - controller.splitScatter.deletePendingSplitScatter(stalePending) - - currentPending = splitScatterObservedPending(t, controller) - re.Equal(makeSplitScatterGroup(100, 102), currentPending.group) - re.Equal(time.Time{}, currentPending.retryAt) -} - -func TestDispatchSplitScatterIgnoresStalePendingWithSameGroup(t *testing.T) { - re := require.New(t) - controller, _, _, cleanup := newTestSplitScatterController(t) - defer cleanup() - - controller.RecordSplitScatterBatch(100, splitScatterTestSourceWaitVersion, []uint64{101}) - stalePending := splitScatterObservedPending(t, controller) - - controller.splitScatter.pendingMu.Lock() - currentPending := controller.splitScatter.pending[splitScatterObservedRegionID] - currentPending.expireAt = currentPending.expireAt.Add(time.Minute) - controller.splitScatter.pending[splitScatterObservedRegionID] = currentPending - controller.splitScatter.pendingMu.Unlock() - - controller.splitScatter.delayPendingSplitScatter(stalePending) - - currentPending = splitScatterObservedPending(t, controller) - re.Equal(time.Time{}, currentPending.retryAt) - - controller.splitScatter.deletePendingSplitScatter(stalePending) - - currentPending = splitScatterObservedPending(t, controller) - re.Equal(stalePending.group, currentPending.group) - re.NotEqual(stalePending.expireAt, currentPending.expireAt) -} - -func newTestSplitScatterController(t *testing.T) (*Controller, *mockcluster.Cluster, *operator.Controller, func()) { - t.Helper() - ctx, cancel := context.WithCancel(context.Background()) - opt := mockconfig.NewTestOptions() - tc := mockcluster.NewCluster(ctx, opt) - for storeID := uint64(1); storeID <= 4; storeID++ { - tc.AddRegionStore(storeID, 0) - } - putSplitScatterRegion(tc, 100, "", "m", splitScatterNoCPUUsage) - - stream := hbstream.NewTestHeartbeatStreams(ctx, tc, false) - oc := operator.NewController(ctx, tc.GetBasicCluster(), tc.GetSharedConfig(), stream) - controller := NewController(ctx, tc, tc.GetCheckerConfig(), oc) - - cleanup := func() { - controller.splitScatter.clearPendingSplitScatter() - stream.Close() - cancel() - } - return controller, tc, oc, cleanup -} - -func putSplitScatterRegion(tc *mockcluster.Cluster, regionID uint64, startKey, endKey string, cpuUsage uint64) { - tc.AddLeaderRegionWithRange(regionID, startKey, endKey, 1, 2, 3) - region := tc.GetRegion(regionID).Clone(core.SetCPUUsage(cpuUsage)) - tc.PutRegion(region) -} - -func putSplitScatterRegionWithKeys(tc *mockcluster.Cluster, startKey, endKey []byte, cpuUsage uint64) { - putSplitScatterRegionWithKeysByID(tc, splitScatterObservedRegionID, startKey, endKey, cpuUsage) -} - -func putSplitScatterRegionWithKeysByID(tc *mockcluster.Cluster, regionID uint64, startKey, endKey []byte, cpuUsage uint64) { - peers := []*metapb.Peer{ - {Id: regionID*10 + 1, StoreId: 1}, - {Id: regionID*10 + 2, StoreId: 2}, - {Id: regionID*10 + 3, StoreId: 3}, - } - region := core.NewRegionInfo( - &metapb.Region{ - Id: regionID, - StartKey: startKey, - EndKey: endKey, - Peers: peers, - RegionEpoch: &metapb.RegionEpoch{ - ConfVer: 1, - Version: 1, - }, - }, - peers[0], - core.SetCPUUsage(cpuUsage), - ) - tc.PutRegion(region) -} - -func newSplitScatterRegionInfo( - regionID uint64, - startKey, endKey string, - peers []*metapb.Peer, - leader *metapb.Peer, - cpuUsage uint64, -) *core.RegionInfo { - return core.NewRegionInfo( - &metapb.Region{ - Id: regionID, - StartKey: []byte(startKey), - EndKey: []byte(endKey), - Peers: peers, - RegionEpoch: &metapb.RegionEpoch{ - ConfVer: 1, - Version: 1, - }, - }, - leader, - core.SetCPUUsage(cpuUsage), - ) -} - -func putSplitScatterRegionWithStores(tc *mockcluster.Cluster, regionID uint64, startKey, endKey string, cpuUsage uint64, stores ...uint64) { - peers := make([]*metapb.Peer, 0, len(stores)) - for i, storeID := range stores { - peers = append(peers, &metapb.Peer{ - Id: regionID*10 + uint64(i) + 1, - StoreId: storeID, - }) - } - tc.PutRegion(newSplitScatterRegionInfo(regionID, startKey, endKey, peers, peers[0], cpuUsage)) -} - -func putSplitScatterRegionWithoutLeader(tc *mockcluster.Cluster, regionID uint64, startKey, endKey string, cpuUsage uint64) { - peers := []*metapb.Peer{ - {Id: regionID*10 + 1, StoreId: 1}, - {Id: regionID*10 + 2, StoreId: 2}, - {Id: regionID*10 + 3, StoreId: 3}, - } - tc.PutRegion(newSplitScatterRegionInfo(regionID, startKey, endKey, peers, nil, cpuUsage)) -} - -func fillSplitScatterPending(controller *Controller, expireAt time.Time) { - controller.splitScatter.pendingMu.Lock() - defer controller.splitScatter.pendingMu.Unlock() - for regionID := uint64(1000); regionID < 1000+splitScatterPendingLimit; regionID++ { - controller.splitScatter.pending[regionID] = splitScatterPendingItem{ - regionID: regionID, - group: "old", - expireAt: expireAt, - } - } -} - -func advanceSplitScatterSourceVersion(t *testing.T, tc *mockcluster.Cluster) { - advanceSplitScatterRegionVersion(t, tc, 100) -} - -func advanceSplitScatterRegionVersion(t *testing.T, tc *mockcluster.Cluster, regionID uint64) { - t.Helper() - region := tc.GetRegion(regionID) - require.NotNil(t, region) - tc.PutRegion(region.Clone(core.WithIncVersion())) -} - -func splitScatterPendingCount(controller *Controller) int { - controller.splitScatter.pendingMu.RLock() - defer controller.splitScatter.pendingMu.RUnlock() - return len(controller.splitScatter.pending) -} - -func splitScatterPendingExpiredCount(attempted string) float64 { - return promtestutil.ToFloat64(splitScatterPendingExpiredCounter.WithLabelValues(attempted)) -} - -func splitScatterPendingGroup(t *testing.T, controller *Controller, regionID uint64) string { - t.Helper() - return splitScatterPending(t, controller, regionID).group -} - -func splitScatterKeyspaceValidatorFor(keyspaces ...uint32) splitScatterKeyspaceValidator { - return func(keyspaceID uint32) bool { - for _, validKeyspaceID := range keyspaces { - if validKeyspaceID == keyspaceID { - return true - } - } - return false - } -} - -func splitScatterPending(t *testing.T, controller *Controller, regionID uint64) splitScatterPendingItem { - t.Helper() - controller.splitScatter.pendingMu.RLock() - defer controller.splitScatter.pendingMu.RUnlock() - pending, ok := controller.splitScatter.pending[regionID] - require.True(t, ok) - return pending -} - -func splitScatterObservedPending(t *testing.T, controller *Controller) splitScatterPendingItem { - t.Helper() - controller.splitScatter.pendingMu.RLock() - defer controller.splitScatter.pendingMu.RUnlock() - pending, ok := controller.splitScatter.pending[splitScatterObservedRegionID] - require.True(t, ok) - return pending -} - -func expireSplitScatterPendingAt(t *testing.T, controller *Controller, regionID uint64, expireAt time.Time) { - t.Helper() - controller.splitScatter.pendingMu.Lock() - defer controller.splitScatter.pendingMu.Unlock() - pending, ok := controller.splitScatter.pending[regionID] - require.True(t, ok) - pending.expireAt = expireAt - controller.splitScatter.pending[regionID] = pending -} - -func retrySplitScatterPendingAt(t *testing.T, controller *Controller, regionID uint64, retryAt time.Time) { - t.Helper() - controller.splitScatter.pendingMu.Lock() - defer controller.splitScatter.pendingMu.Unlock() - pending, ok := controller.splitScatter.pending[regionID] - require.True(t, ok) - pending.retryAt = retryAt - controller.splitScatter.pending[regionID] = pending -} - -func setSplitScatterNextDispatchAt(t *testing.T, controller *Controller, nextDispatchAt time.Time) { - t.Helper() - controller.splitScatter.pendingMu.Lock() - defer controller.splitScatter.pendingMu.Unlock() - controller.splitScatter.nextDispatchAt = nextDispatchAt -} - -func splitScatterNextDispatchAt(t *testing.T, controller *Controller) time.Time { - t.Helper() - controller.splitScatter.pendingMu.RLock() - defer controller.splitScatter.pendingMu.RUnlock() - return controller.splitScatter.nextDispatchAt -} - -func requireInternalScatterOpsUseGroups(t *testing.T, oc *operator.Controller, scatterGroup, batchGroup string) { - t.Helper() - re := require.New(t) - ops := oc.GetOperators() - re.NotEmpty(ops) - for _, op := range ops { - re.Equal(scatter.InternalScatterOperatorDesc, op.Desc()) - opGroup, ok := op.GetAdditionalInfo("group") - re.True(ok) - re.Equal(scatterGroup, opGroup) - opBatchGroup, ok := op.GetAdditionalInfo("batch-group") - re.True(ok) - re.Equal(batchGroup, opBatchGroup) - } -} - -func removeInternalScatterOps(oc *operator.Controller) { - for _, op := range oc.GetOperators() { - oc.RemoveOperator(op) - } -} - -func pendingRegionIDs(regions []splitScatterPendingItem) []uint64 { - ids := make([]uint64, 0, len(regions)) - for _, region := range regions { - ids = append(ids, region.regionID) - } - return ids -} - -func splitScatterIndexKeyPrefix() []byte { - return codec.GenerateIndexKey(splitScatterTestTableID, splitScatterTestIndexID) -} - -func splitScatterKeyspacePrefixRange(keyspaceID uint32, rawPrefix []byte) splitScatterRangeHint { - startKey := newSplitScatterKeyspaceKey(keyspaceID, codec.TxnKeyspaceModePrefix, rawPrefix) - endRawPrefix := splitScatterNextPrefix(rawPrefix) - if len(endRawPrefix) == 0 { - return splitScatterRangeHint{startKey: startKey} - } - return splitScatterRangeHint{ - startKey: startKey, - endKey: newSplitScatterKeyspaceKey(keyspaceID, codec.TxnKeyspaceModePrefix, endRawPrefix), - } -} - -func newSplitScatterIndexKey(suffix string) []byte { - key := append([]byte(nil), splitScatterIndexKeyPrefix()...) - key = append(key, suffix...) - return codec.EncodeBytes(key) -} - -func newSplitScatterKeyspaceIndexKey(keyspaceID uint32, suffix string) []byte { - key := append([]byte(nil), splitScatterIndexKeyPrefix()...) - key = append(key, suffix...) - return newSplitScatterKeyspaceKey(keyspaceID, codec.TxnKeyspaceModePrefix, key) -} - -func newSplitScatterRecordKey(tableID int64, suffix string) []byte { - key := append([]byte(nil), codec.GenerateRecordKeyPrefix(tableID)...) - key = append(key, suffix...) - return codec.EncodeBytes(key) -} - -func newSplitScatterKeyspaceRecordKey(keyspaceID uint32, suffix string) []byte { - key := append([]byte(nil), codec.GenerateRecordKeyPrefix(splitScatterTestTableID)...) - key = append(key, suffix...) - return newSplitScatterKeyspaceKey(keyspaceID, codec.TxnKeyspaceModePrefix, key) -} - -func newSplitScatterRawKeyspaceRecordKey(keyspaceID uint32, tableID int64, suffix string) []byte { - key := append([]byte(nil), codec.GenerateRecordKeyPrefix(tableID)...) - key = append(key, suffix...) - return newSplitScatterKeyspaceKey(keyspaceID, codec.RawKeyspaceModePrefix, key) -} - -func newSplitScatterTableBoundaryKey(tableID int64) []byte { - return codec.EncodeBytes(codec.GenerateTableKey(tableID)) -} - -func newSplitScatterKeyspaceKey(keyspaceID uint32, mode byte, rawKey []byte) []byte { - key := codec.MakeKeyspacePrefix(mode, keyspaceID) - return codec.EncodeBytes(append(key, rawKey...)) -} From b88b6fd4d82cf56d9edf0ae3f40d1b9a280dc167 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Mon, 13 Jul 2026 16:44:32 +0800 Subject: [PATCH 3/3] keyspace: complete cherry-pick resolution in #10995 Signed-off-by: JmPotato --- pkg/keyspace/util.go | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/pkg/keyspace/util.go b/pkg/keyspace/util.go index 7df93eefb19..c0b40a0facb 100644 --- a/pkg/keyspace/util.go +++ b/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), } }