Skip to content

Commit 65701f0

Browse files
lhy1024HunDunDM
authored andcommitted
affinity: add base models and manager skeleton (tikv#9997)
ref tikv#9764 Signed-off-by: lhy1024 <admin@liudos.us>
1 parent ed8bbfa commit 65701f0

File tree

6 files changed

+1316
-0
lines changed

6 files changed

+1316
-0
lines changed
Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
// Copyright 2025 TiKV Project Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package affinity
16+
17+
import (
18+
"fmt"
19+
"maps"
20+
"testing"
21+
22+
"github.com/stretchr/testify/require"
23+
"go.uber.org/goleak"
24+
25+
"github.com/pingcap/kvproto/pkg/metapb"
26+
27+
"github.com/tikv/pd/pkg/core"
28+
"github.com/tikv/pd/pkg/utils/keyutil"
29+
"github.com/tikv/pd/pkg/utils/testutil"
30+
)
31+
32+
func TestMain(m *testing.M) {
33+
goleak.VerifyTestMain(m, testutil.LeakOptions...)
34+
}
35+
36+
//nolint:unused
37+
func getGroupForTest(re *require.Assertions, m *Manager, id string) *runtimeGroupInfo {
38+
m.RLock()
39+
defer m.RUnlock()
40+
group, ok := m.groups[id]
41+
re.True(ok)
42+
// Return a cloned runtimeGroupInfo to ensure lock safety.
43+
newGroup := *group
44+
newGroup.Regions = maps.Clone(group.Regions)
45+
return &newGroup
46+
}
47+
48+
//nolint:unused
49+
func createGroupForTest(re *require.Assertions, m *Manager, id string, rangeCount int) []keyutil.KeyRange {
50+
gkr := GroupKeyRanges{
51+
KeyRanges: make([]keyutil.KeyRange, rangeCount),
52+
GroupID: id,
53+
}
54+
for i := range rangeCount {
55+
gkr.KeyRanges[i] = keyutil.KeyRange{
56+
StartKey: []byte(fmt.Sprintf("test-%s-%04d", id, i)),
57+
EndKey: []byte(fmt.Sprintf("test-%s-%04d", id, i+1)),
58+
}
59+
}
60+
re.NoError(m.CreateAffinityGroups([]GroupKeyRanges{gkr}))
61+
return gkr.KeyRanges
62+
}
63+
64+
//nolint:unused
65+
func testCacheStale(re *require.Assertions, m *Manager, region *core.RegionInfo) {
66+
cache, group := m.getCache(region)
67+
if cache != nil && group != nil {
68+
re.NotEqual(cache.affinityVer, group.affinityVer)
69+
}
70+
}
71+
72+
// generateRegionForTest generates a test Region from the given information,
73+
// where voterStoreIDs[0] is used as the leaderStoreID.
74+
//
75+
//nolint:unparam
76+
func generateRegionForTest(id uint64, voterStoreIDs []uint64, keyRange keyutil.KeyRange) *core.RegionInfo {
77+
peers := make([]*metapb.Peer, len(voterStoreIDs))
78+
for i, storeID := range voterStoreIDs {
79+
peers[i] = &metapb.Peer{
80+
Id: id*10 + uint64(i),
81+
StoreId: storeID,
82+
}
83+
}
84+
meta := &metapb.Region{
85+
Id: id,
86+
StartKey: keyRange.StartKey,
87+
EndKey: keyRange.EndKey,
88+
Peers: peers,
89+
}
90+
return core.NewRegionInfo(meta, peers[0])
91+
}
92+
93+
var nonOverlappingRange = keyutil.KeyRange{
94+
StartKey: []byte("non-overlapping"),
95+
EndKey: []byte("noop"),
96+
}

0 commit comments

Comments
 (0)