Skip to content

Commit 3e02966

Browse files
shiyuhang0shiyuhang0
authored
[close #693] Add test for RangeSplitter (#694)
Signed-off-by: shiyuhang <1136742008@qq.com>
1 parent a459a6e commit 3e02966

2 files changed

Lines changed: 347 additions & 0 deletions

File tree

‎src/test/java/org/tikv/common/MockRegionManager.java‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
/*
2+
* Copyright 2022 TiKV Project Authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*
16+
*/
17+
18+
package org.tikv.common;
19+
20+
import static org.tikv.common.GrpcUtils.encodeKey;
21+
22+
import com.google.common.collect.ImmutableList;
23+
import com.google.protobuf.ByteString;
24+
import java.util.List;
25+
import java.util.Map;
26+
import java.util.stream.Collectors;
27+
import org.tikv.common.apiversion.RequestKeyV1TxnCodec;
28+
import org.tikv.common.key.Key;
29+
import org.tikv.common.region.RegionManager;
30+
import org.tikv.common.region.TiRegion;
31+
import org.tikv.common.region.TiStore;
32+
import org.tikv.common.region.TiStoreType;
33+
import org.tikv.common.util.KeyRangeUtils;
34+
import org.tikv.common.util.Pair;
35+
import org.tikv.kvproto.Coprocessor.KeyRange;
36+
import org.tikv.kvproto.Kvrpcpb.CommandPri;
37+
import org.tikv.kvproto.Kvrpcpb.IsolationLevel;
38+
import org.tikv.kvproto.Metapb;
39+
import org.tikv.kvproto.Metapb.Peer;
40+
import org.tikv.kvproto.Metapb.Region;
41+
42+
public class MockRegionManager extends RegionManager {
43+
44+
private final Map<KeyRange, TiRegion> mockRegionMap;
45+
46+
private static TiRegion region(long id, KeyRange range) {
47+
RequestKeyV1TxnCodec v1 = new RequestKeyV1TxnCodec();
48+
49+
TiConfiguration configuration = new TiConfiguration();
50+
configuration.setIsolationLevel(IsolationLevel.RC);
51+
configuration.setCommandPriority(CommandPri.Low);
52+
Region r =
53+
Metapb.Region.newBuilder()
54+
.setRegionEpoch(Metapb.RegionEpoch.newBuilder().setConfVer(1).setVersion(2))
55+
.setId(id)
56+
.setStartKey(encodeKey(range.getStart().toByteArray()))
57+
.setEndKey(encodeKey(range.getEnd().toByteArray()))
58+
.addPeers(Peer.getDefaultInstance())
59+
.build();
60+
61+
List<Metapb.Store> s = ImmutableList.of(Metapb.Store.newBuilder().setId(id).build());
62+
63+
return new TiRegion(
64+
configuration,
65+
v1.decodeRegion(r),
66+
null,
67+
r.getPeersList(),
68+
s.stream().map(TiStore::new).collect(Collectors.toList()));
69+
}
70+
71+
public MockRegionManager(List<KeyRange> ranges) {
72+
super(null, null);
73+
mockRegionMap =
74+
ranges.stream().collect(Collectors.toMap(kr -> kr, kr -> region(ranges.indexOf(kr), kr)));
75+
}
76+
77+
@Override
78+
public Pair<TiRegion, TiStore> getRegionStorePairByKey(ByteString key, TiStoreType storeType) {
79+
for (Map.Entry<KeyRange, TiRegion> entry : mockRegionMap.entrySet()) {
80+
KeyRange range = entry.getKey();
81+
if (KeyRangeUtils.makeRange(range.getStart(), range.getEnd()).contains(Key.toRawKey(key))) {
82+
TiRegion region = entry.getValue();
83+
return Pair.create(
84+
region, new TiStore(Metapb.Store.newBuilder().setId(region.getId()).build()));
85+
}
86+
}
87+
return null;
88+
}
89+
}

‎src/test/java/org/tikv/util/RangeSplitterTest.java‎

Lines changed: 258 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,258 @@
1+
/*
2+
* Copyright 2022 TiKV Project Authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*
16+
*/
17+
18+
package org.tikv.util;
19+
20+
import static org.junit.Assert.assertEquals;
21+
22+
import com.google.common.collect.ImmutableList;
23+
import com.google.protobuf.ByteString;
24+
import gnu.trove.list.array.TLongArrayList;
25+
import gnu.trove.map.hash.TLongObjectHashMap;
26+
import java.util.ArrayList;
27+
import java.util.List;
28+
import org.junit.Test;
29+
import org.tikv.common.MockRegionManager;
30+
import org.tikv.common.codec.Codec.IntegerCodec;
31+
import org.tikv.common.codec.CodecDataOutput;
32+
import org.tikv.common.key.RowKey;
33+
import org.tikv.common.key.RowKey.DecodeResult.Status;
34+
import org.tikv.common.util.RangeSplitter;
35+
import org.tikv.kvproto.Coprocessor.KeyRange;
36+
37+
public class RangeSplitterTest {
38+
39+
private static KeyRange keyRange(Long s, Long e) {
40+
ByteString sKey = ByteString.EMPTY;
41+
ByteString eKey = ByteString.EMPTY;
42+
if (s != null) {
43+
CodecDataOutput cdo = new CodecDataOutput();
44+
IntegerCodec.writeLongFully(cdo, s, true);
45+
sKey = cdo.toByteString();
46+
}
47+
48+
if (e != null) {
49+
CodecDataOutput cdo = new CodecDataOutput();
50+
IntegerCodec.writeLongFully(cdo, e, true);
51+
eKey = cdo.toByteString();
52+
}
53+
54+
return KeyRange.newBuilder().setStart(sKey).setEnd(eKey).build();
55+
}
56+
57+
private static KeyRange keyRangeByHandle(long tableId, Long s, Long e) {
58+
return keyRangeByHandle(tableId, s, Status.EQUAL, e, Status.EQUAL);
59+
}
60+
61+
private static KeyRange keyRangeByHandle(long tableId, Long s, Status ss, Long e, Status es) {
62+
ByteString sKey = shiftByStatus(handleToByteString(tableId, s), ss);
63+
ByteString eKey = shiftByStatus(handleToByteString(tableId, e), es);
64+
65+
return KeyRange.newBuilder().setStart(sKey).setEnd(eKey).build();
66+
}
67+
68+
private static ByteString shiftByStatus(ByteString v, Status s) {
69+
switch (s) {
70+
case EQUAL:
71+
return v;
72+
case LESS:
73+
return v.substring(0, v.size() - 1);
74+
case GREATER:
75+
return v.concat(ByteString.copyFrom(new byte[] {1, 0}));
76+
default:
77+
throw new IllegalArgumentException("Only EQUAL,LESS,GREATER allowed");
78+
}
79+
}
80+
81+
private static ByteString handleToByteString(long tableId, Long k) {
82+
if (k != null) {
83+
return RowKey.toRowKey(tableId, k).toByteString();
84+
}
85+
return ByteString.EMPTY;
86+
}
87+
88+
@Test
89+
public void splitRangeByRegionTest() {
90+
MockRegionManager mgr =
91+
new MockRegionManager(
92+
ImmutableList.of(keyRange(null, 30L), keyRange(30L, 50L), keyRange(50L, null)));
93+
RangeSplitter s = RangeSplitter.newSplitter(mgr);
94+
List<RangeSplitter.RegionTask> tasks =
95+
s.splitRangeByRegion(
96+
ImmutableList.of(
97+
keyRange(0L, 40L), keyRange(41L, 42L), keyRange(45L, 50L), keyRange(70L, 1000L)));
98+
99+
assertEquals(tasks.get(0).getRegion().getId(), 0);
100+
assertEquals(tasks.get(0).getRanges().size(), 1);
101+
KeyRange range = tasks.get(0).getRanges().get(0);
102+
assertEquals(tasks.get(0).getRanges().get(0), keyRange(0L, 30L));
103+
104+
assertEquals(tasks.get(1).getRegion().getId(), 1);
105+
assertEquals(tasks.get(1).getRanges().get(0), keyRange(30L, 40L));
106+
assertEquals(tasks.get(1).getRanges().get(1), keyRange(41L, 42L));
107+
assertEquals(tasks.get(1).getRanges().get(2), keyRange(45L, 50L));
108+
assertEquals(tasks.get(1).getRanges().size(), 3);
109+
110+
assertEquals(tasks.get(2).getRegion().getId(), 2);
111+
assertEquals(tasks.get(2).getRanges().size(), 1);
112+
assertEquals(tasks.get(2).getRanges().get(0), keyRange(70L, 1000L));
113+
}
114+
115+
@Test
116+
public void splitAndSortHandlesByRegionTest() {
117+
final long tableId = 1;
118+
List<Long> handles = new ArrayList<>();
119+
handles.add(1L);
120+
handles.add(5L);
121+
handles.add(4L);
122+
handles.add(3L);
123+
handles.add(10L);
124+
handles.add(2L);
125+
handles.add(100L);
126+
handles.add(101L);
127+
handles.add(99L);
128+
handles.add(88L);
129+
handles.add(-1L);
130+
handles.add(-255L);
131+
handles.add(-100L);
132+
handles.add(-99L);
133+
handles.add(-98L);
134+
handles.add(Long.MIN_VALUE);
135+
handles.add(8960L);
136+
handles.add(8959L);
137+
handles.add(19999L);
138+
handles.add(15001L);
139+
140+
MockRegionManager mgr =
141+
new MockRegionManager(
142+
ImmutableList.of(
143+
keyRangeByHandle(tableId, null, Status.EQUAL, -100L, Status.EQUAL),
144+
keyRangeByHandle(tableId, -100L, Status.EQUAL, 10L, Status.GREATER),
145+
keyRangeByHandle(tableId, 10L, Status.GREATER, 50L, Status.EQUAL),
146+
keyRangeByHandle(tableId, 50L, Status.EQUAL, 100L, Status.GREATER),
147+
keyRangeByHandle(tableId, 100L, Status.GREATER, 9000L, Status.LESS),
148+
keyRangeByHandle(tableId, 0x2300L /*8960*/, Status.LESS, 16000L, Status.EQUAL),
149+
keyRangeByHandle(tableId, 16000L, Status.EQUAL, null, Status.EQUAL)));
150+
151+
RangeSplitter s = RangeSplitter.newSplitter(mgr);
152+
List<RangeSplitter.RegionTask> tasks =
153+
new ArrayList<>(
154+
s.splitAndSortHandlesByRegion(
155+
ImmutableList.of(tableId),
156+
new TLongArrayList(handles.stream().mapToLong(t -> t).toArray())));
157+
tasks.sort(
158+
(l, r) -> {
159+
Long regionIdLeft = l.getRegion().getId();
160+
Long regionIdRight = r.getRegion().getId();
161+
return regionIdLeft.compareTo(regionIdRight);
162+
});
163+
164+
// [-INF, -100): [Long.MIN_VALUE, Long.MIN_VALUE + 1), [-255, -254)
165+
assertEquals(tasks.get(0).getRegion().getId(), 0);
166+
assertEquals(tasks.get(0).getRanges().size(), 2);
167+
assertEquals(
168+
tasks.get(0).getRanges().get(0),
169+
keyRangeByHandle(tableId, Long.MIN_VALUE, Long.MIN_VALUE + 1));
170+
assertEquals(tasks.get(0).getRanges().get(1), keyRangeByHandle(tableId, -255L, -254L));
171+
172+
// [-100, 10.x): [-100, -97), [-1, 0), [1, 6), [10, 11)
173+
assertEquals(tasks.get(1).getRegion().getId(), 1);
174+
assertEquals(tasks.get(1).getRanges().size(), 4);
175+
assertEquals(tasks.get(1).getRanges().get(0), keyRangeByHandle(tableId, -100L, -97L));
176+
assertEquals(tasks.get(1).getRanges().get(1), keyRangeByHandle(tableId, -1L, 0L));
177+
assertEquals(tasks.get(1).getRanges().get(2), keyRangeByHandle(tableId, 1L, 6L));
178+
assertEquals(tasks.get(1).getRanges().get(3), keyRangeByHandle(tableId, 10L, 11L));
179+
180+
// [10.x, 50): empty
181+
// [50, 100.x): [88, 89) [99, 101)
182+
assertEquals(tasks.get(2).getRegion().getId(), 3);
183+
assertEquals(tasks.get(2).getRanges().size(), 2);
184+
assertEquals(tasks.get(2).getRanges().get(0), keyRangeByHandle(tableId, 88L, 89L));
185+
assertEquals(tasks.get(2).getRanges().get(1), keyRangeByHandle(tableId, 99L, 101L));
186+
187+
// [100.x, less than 8960): [101, 102) [8959, 8960)
188+
assertEquals(tasks.get(3).getRegion().getId(), 4);
189+
assertEquals(tasks.get(3).getRanges().size(), 2);
190+
assertEquals(tasks.get(3).getRanges().get(0), keyRangeByHandle(tableId, 101L, 102L));
191+
assertEquals(tasks.get(3).getRanges().get(1), keyRangeByHandle(tableId, 8959L, 8960L));
192+
193+
// [less than 8960, 16000): [9000, 9001), [15001, 15002)
194+
assertEquals(tasks.get(4).getRegion().getId(), 5);
195+
assertEquals(tasks.get(4).getRanges().size(), 2);
196+
assertEquals(tasks.get(4).getRanges().get(0), keyRangeByHandle(tableId, 8960L, 8961L));
197+
assertEquals(tasks.get(4).getRanges().get(1), keyRangeByHandle(tableId, 15001L, 15002L));
198+
199+
// [16000, INF): [19999, 20000)
200+
assertEquals(tasks.get(5).getRegion().getId(), 6);
201+
assertEquals(tasks.get(5).getRanges().size(), 1);
202+
assertEquals(tasks.get(5).getRanges().get(0), keyRangeByHandle(tableId, 19999L, 20000L));
203+
}
204+
205+
@Test
206+
public void groupByAndSortHandlesByRegionIdTest() {
207+
final long tableId = 1;
208+
List<Long> handles = new ArrayList<>();
209+
handles.add(1L);
210+
handles.add(5L);
211+
handles.add(4L);
212+
handles.add(3L);
213+
handles.add(10L);
214+
handles.add(11L);
215+
handles.add(12L);
216+
handles.add(2L);
217+
handles.add(100L);
218+
handles.add(101L);
219+
handles.add(99L);
220+
handles.add(88L);
221+
handles.add(-1L);
222+
handles.add(-255L);
223+
handles.add(-100L);
224+
handles.add(-99L);
225+
handles.add(-98L);
226+
handles.add(Long.MIN_VALUE);
227+
handles.add(8960L);
228+
handles.add(8959L);
229+
handles.add(19999L);
230+
handles.add(15001L);
231+
handles.add(99999999999L);
232+
handles.add(Long.MAX_VALUE);
233+
234+
MockRegionManager mgr =
235+
new MockRegionManager(
236+
ImmutableList.of(
237+
keyRangeByHandle(tableId, null, Status.EQUAL, -100L, Status.EQUAL),
238+
keyRangeByHandle(tableId, -100L, Status.EQUAL, 10L, Status.GREATER),
239+
keyRangeByHandle(tableId, 10L, Status.GREATER, 50L, Status.EQUAL),
240+
keyRangeByHandle(tableId, 50L, Status.EQUAL, 100L, Status.GREATER),
241+
keyRangeByHandle(tableId, 100L, Status.GREATER, 9000L, Status.LESS),
242+
keyRangeByHandle(tableId, 0x2300L /*8960*/, Status.LESS, 16000L, Status.EQUAL),
243+
keyRangeByHandle(tableId, 16000L, Status.EQUAL, null, Status.EQUAL)));
244+
245+
TLongObjectHashMap<TLongArrayList> result = new TLongObjectHashMap<>();
246+
RangeSplitter.newSplitter(mgr)
247+
.groupByAndSortHandlesByRegionId(
248+
tableId, new TLongArrayList(handles.stream().mapToLong(t -> t).toArray()))
249+
.forEach((k, v) -> result.put(k.first.getId(), v));
250+
assertEquals(2, result.get(0).size());
251+
assertEquals(10, result.get(1).size());
252+
assertEquals(2, result.get(2).size());
253+
assertEquals(3, result.get(3).size());
254+
assertEquals(2, result.get(4).size());
255+
assertEquals(2, result.get(5).size());
256+
assertEquals(3, result.get(6).size());
257+
}
258+
}

0 commit comments

Comments
 (0)