Skip to content

Commit

Permalink
[close #693] Add test for RangeSplitter (#694) #698
Browse files Browse the repository at this point in the history 
Signed-off-by: ti-srebot <[email protected]>

Signed-off-by: ti-srebot <[email protected]>
Co-authored-by: shi yuhang <[email protected]>
  • Loading branch information
@ti-srebot @shiyuhang0
ti-srebot and shiyuhang0 authored Dec 29, 2022
1 parent a345184 commit 30930e2
Show file tree
Hide file tree
Showing 2 changed files with 347 additions and 0 deletions.
89 changes: 89 additions & 0 deletions src/test/java/org/tikv/common/MockRegionManager.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
/*
* Copyright 2022 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 org.tikv.common;

import static org.tikv.common.GrpcUtils.encodeKey;

import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteString;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.tikv.common.apiversion.RequestKeyV1TxnCodec;
import org.tikv.common.key.Key;
import org.tikv.common.region.RegionManager;
import org.tikv.common.region.TiRegion;
import org.tikv.common.region.TiStore;
import org.tikv.common.region.TiStoreType;
import org.tikv.common.util.KeyRangeUtils;
import org.tikv.common.util.Pair;
import org.tikv.kvproto.Coprocessor.KeyRange;
import org.tikv.kvproto.Kvrpcpb.CommandPri;
import org.tikv.kvproto.Kvrpcpb.IsolationLevel;
import org.tikv.kvproto.Metapb;
import org.tikv.kvproto.Metapb.Peer;
import org.tikv.kvproto.Metapb.Region;

public class MockRegionManager extends RegionManager {

private final Map<KeyRange, TiRegion> mockRegionMap;

private static TiRegion region(long id, KeyRange range) {
RequestKeyV1TxnCodec v1 = new RequestKeyV1TxnCodec();

TiConfiguration configuration = new TiConfiguration();
configuration.setIsolationLevel(IsolationLevel.RC);
configuration.setCommandPriority(CommandPri.Low);
Region r =
Metapb.Region.newBuilder()
.setRegionEpoch(Metapb.RegionEpoch.newBuilder().setConfVer(1).setVersion(2))
.setId(id)
.setStartKey(encodeKey(range.getStart().toByteArray()))
.setEndKey(encodeKey(range.getEnd().toByteArray()))
.addPeers(Peer.getDefaultInstance())
.build();

List<Metapb.Store> s = ImmutableList.of(Metapb.Store.newBuilder().setId(id).build());

return new TiRegion(
configuration,
v1.decodeRegion(r),
null,
r.getPeersList(),
s.stream().map(TiStore::new).collect(Collectors.toList()));
}

public MockRegionManager(List<KeyRange> ranges) {
super(null, null);
mockRegionMap =
ranges.stream().collect(Collectors.toMap(kr -> kr, kr -> region(ranges.indexOf(kr), kr)));
}

@Override
public Pair<TiRegion, TiStore> getRegionStorePairByKey(ByteString key, TiStoreType storeType) {
for (Map.Entry<KeyRange, TiRegion> entry : mockRegionMap.entrySet()) {
KeyRange range = entry.getKey();
if (KeyRangeUtils.makeRange(range.getStart(), range.getEnd()).contains(Key.toRawKey(key))) {
TiRegion region = entry.getValue();
return Pair.create(
region, new TiStore(Metapb.Store.newBuilder().setId(region.getId()).build()));
}
}
return null;
}
}
258 changes: 258 additions & 0 deletions src/test/java/org/tikv/util/RangeSplitterTest.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,258 @@
/*
* Copyright 2022 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 org.tikv.util;

import static org.junit.Assert.assertEquals;

import com.google.common.collect.ImmutableList;
import com.google.protobuf.ByteString;
import gnu.trove.list.array.TLongArrayList;
import gnu.trove.map.hash.TLongObjectHashMap;
import java.util.ArrayList;
import java.util.List;
import org.junit.Test;
import org.tikv.common.MockRegionManager;
import org.tikv.common.codec.Codec.IntegerCodec;
import org.tikv.common.codec.CodecDataOutput;
import org.tikv.common.key.RowKey;
import org.tikv.common.key.RowKey.DecodeResult.Status;
import org.tikv.common.util.RangeSplitter;
import org.tikv.kvproto.Coprocessor.KeyRange;

public class RangeSplitterTest {

private static KeyRange keyRange(Long s, Long e) {
ByteString sKey = ByteString.EMPTY;
ByteString eKey = ByteString.EMPTY;
if (s != null) {
CodecDataOutput cdo = new CodecDataOutput();
IntegerCodec.writeLongFully(cdo, s, true);
sKey = cdo.toByteString();
}

if (e != null) {
CodecDataOutput cdo = new CodecDataOutput();
IntegerCodec.writeLongFully(cdo, e, true);
eKey = cdo.toByteString();
}

return KeyRange.newBuilder().setStart(sKey).setEnd(eKey).build();
}

private static KeyRange keyRangeByHandle(long tableId, Long s, Long e) {
return keyRangeByHandle(tableId, s, Status.EQUAL, e, Status.EQUAL);
}

private static KeyRange keyRangeByHandle(long tableId, Long s, Status ss, Long e, Status es) {
ByteString sKey = shiftByStatus(handleToByteString(tableId, s), ss);
ByteString eKey = shiftByStatus(handleToByteString(tableId, e), es);

return KeyRange.newBuilder().setStart(sKey).setEnd(eKey).build();
}

private static ByteString shiftByStatus(ByteString v, Status s) {
switch (s) {
case EQUAL:
return v;
case LESS:
return v.substring(0, v.size() - 1);
case GREATER:
return v.concat(ByteString.copyFrom(new byte[] {1, 0}));
default:
throw new IllegalArgumentException("Only EQUAL,LESS,GREATER allowed");
}
}

private static ByteString handleToByteString(long tableId, Long k) {
if (k != null) {
return RowKey.toRowKey(tableId, k).toByteString();
}
return ByteString.EMPTY;
}

@Test
public void splitRangeByRegionTest() {
MockRegionManager mgr =
new MockRegionManager(
ImmutableList.of(keyRange(null, 30L), keyRange(30L, 50L), keyRange(50L, null)));
RangeSplitter s = RangeSplitter.newSplitter(mgr);
List<RangeSplitter.RegionTask> tasks =
s.splitRangeByRegion(
ImmutableList.of(
keyRange(0L, 40L), keyRange(41L, 42L), keyRange(45L, 50L), keyRange(70L, 1000L)));

assertEquals(tasks.get(0).getRegion().getId(), 0);
assertEquals(tasks.get(0).getRanges().size(), 1);
KeyRange range = tasks.get(0).getRanges().get(0);
assertEquals(tasks.get(0).getRanges().get(0), keyRange(0L, 30L));

assertEquals(tasks.get(1).getRegion().getId(), 1);
assertEquals(tasks.get(1).getRanges().get(0), keyRange(30L, 40L));
assertEquals(tasks.get(1).getRanges().get(1), keyRange(41L, 42L));
assertEquals(tasks.get(1).getRanges().get(2), keyRange(45L, 50L));
assertEquals(tasks.get(1).getRanges().size(), 3);

assertEquals(tasks.get(2).getRegion().getId(), 2);
assertEquals(tasks.get(2).getRanges().size(), 1);
assertEquals(tasks.get(2).getRanges().get(0), keyRange(70L, 1000L));
}

@Test
public void splitAndSortHandlesByRegionTest() {
final long tableId = 1;
List<Long> handles = new ArrayList<>();
handles.add(1L);
handles.add(5L);
handles.add(4L);
handles.add(3L);
handles.add(10L);
handles.add(2L);
handles.add(100L);
handles.add(101L);
handles.add(99L);
handles.add(88L);
handles.add(-1L);
handles.add(-255L);
handles.add(-100L);
handles.add(-99L);
handles.add(-98L);
handles.add(Long.MIN_VALUE);
handles.add(8960L);
handles.add(8959L);
handles.add(19999L);
handles.add(15001L);

MockRegionManager mgr =
new MockRegionManager(
ImmutableList.of(
keyRangeByHandle(tableId, null, Status.EQUAL, -100L, Status.EQUAL),
keyRangeByHandle(tableId, -100L, Status.EQUAL, 10L, Status.GREATER),
keyRangeByHandle(tableId, 10L, Status.GREATER, 50L, Status.EQUAL),
keyRangeByHandle(tableId, 50L, Status.EQUAL, 100L, Status.GREATER),
keyRangeByHandle(tableId, 100L, Status.GREATER, 9000L, Status.LESS),
keyRangeByHandle(tableId, 0x2300L /*8960*/, Status.LESS, 16000L, Status.EQUAL),
keyRangeByHandle(tableId, 16000L, Status.EQUAL, null, Status.EQUAL)));

RangeSplitter s = RangeSplitter.newSplitter(mgr);
List<RangeSplitter.RegionTask> tasks =
new ArrayList<>(
s.splitAndSortHandlesByRegion(
ImmutableList.of(tableId),
new TLongArrayList(handles.stream().mapToLong(t -> t).toArray())));
tasks.sort(
(l, r) -> {
Long regionIdLeft = l.getRegion().getId();
Long regionIdRight = r.getRegion().getId();
return regionIdLeft.compareTo(regionIdRight);
});

// [-INF, -100): [Long.MIN_VALUE, Long.MIN_VALUE + 1), [-255, -254)
assertEquals(tasks.get(0).getRegion().getId(), 0);
assertEquals(tasks.get(0).getRanges().size(), 2);
assertEquals(
tasks.get(0).getRanges().get(0),
keyRangeByHandle(tableId, Long.MIN_VALUE, Long.MIN_VALUE + 1));
assertEquals(tasks.get(0).getRanges().get(1), keyRangeByHandle(tableId, -255L, -254L));

// [-100, 10.x): [-100, -97), [-1, 0), [1, 6), [10, 11)
assertEquals(tasks.get(1).getRegion().getId(), 1);
assertEquals(tasks.get(1).getRanges().size(), 4);
assertEquals(tasks.get(1).getRanges().get(0), keyRangeByHandle(tableId, -100L, -97L));
assertEquals(tasks.get(1).getRanges().get(1), keyRangeByHandle(tableId, -1L, 0L));
assertEquals(tasks.get(1).getRanges().get(2), keyRangeByHandle(tableId, 1L, 6L));
assertEquals(tasks.get(1).getRanges().get(3), keyRangeByHandle(tableId, 10L, 11L));

// [10.x, 50): empty
// [50, 100.x): [88, 89) [99, 101)
assertEquals(tasks.get(2).getRegion().getId(), 3);
assertEquals(tasks.get(2).getRanges().size(), 2);
assertEquals(tasks.get(2).getRanges().get(0), keyRangeByHandle(tableId, 88L, 89L));
assertEquals(tasks.get(2).getRanges().get(1), keyRangeByHandle(tableId, 99L, 101L));

// [100.x, less than 8960): [101, 102) [8959, 8960)
assertEquals(tasks.get(3).getRegion().getId(), 4);
assertEquals(tasks.get(3).getRanges().size(), 2);
assertEquals(tasks.get(3).getRanges().get(0), keyRangeByHandle(tableId, 101L, 102L));
assertEquals(tasks.get(3).getRanges().get(1), keyRangeByHandle(tableId, 8959L, 8960L));

// [less than 8960, 16000): [9000, 9001), [15001, 15002)
assertEquals(tasks.get(4).getRegion().getId(), 5);
assertEquals(tasks.get(4).getRanges().size(), 2);
assertEquals(tasks.get(4).getRanges().get(0), keyRangeByHandle(tableId, 8960L, 8961L));
assertEquals(tasks.get(4).getRanges().get(1), keyRangeByHandle(tableId, 15001L, 15002L));

// [16000, INF): [19999, 20000)
assertEquals(tasks.get(5).getRegion().getId(), 6);
assertEquals(tasks.get(5).getRanges().size(), 1);
assertEquals(tasks.get(5).getRanges().get(0), keyRangeByHandle(tableId, 19999L, 20000L));
}

@Test
public void groupByAndSortHandlesByRegionIdTest() {
final long tableId = 1;
List<Long> handles = new ArrayList<>();
handles.add(1L);
handles.add(5L);
handles.add(4L);
handles.add(3L);
handles.add(10L);
handles.add(11L);
handles.add(12L);
handles.add(2L);
handles.add(100L);
handles.add(101L);
handles.add(99L);
handles.add(88L);
handles.add(-1L);
handles.add(-255L);
handles.add(-100L);
handles.add(-99L);
handles.add(-98L);
handles.add(Long.MIN_VALUE);
handles.add(8960L);
handles.add(8959L);
handles.add(19999L);
handles.add(15001L);
handles.add(99999999999L);
handles.add(Long.MAX_VALUE);

MockRegionManager mgr =
new MockRegionManager(
ImmutableList.of(
keyRangeByHandle(tableId, null, Status.EQUAL, -100L, Status.EQUAL),
keyRangeByHandle(tableId, -100L, Status.EQUAL, 10L, Status.GREATER),
keyRangeByHandle(tableId, 10L, Status.GREATER, 50L, Status.EQUAL),
keyRangeByHandle(tableId, 50L, Status.EQUAL, 100L, Status.GREATER),
keyRangeByHandle(tableId, 100L, Status.GREATER, 9000L, Status.LESS),
keyRangeByHandle(tableId, 0x2300L /*8960*/, Status.LESS, 16000L, Status.EQUAL),
keyRangeByHandle(tableId, 16000L, Status.EQUAL, null, Status.EQUAL)));

TLongObjectHashMap<TLongArrayList> result = new TLongObjectHashMap<>();
RangeSplitter.newSplitter(mgr)
.groupByAndSortHandlesByRegionId(
tableId, new TLongArrayList(handles.stream().mapToLong(t -> t).toArray()))
.forEach((k, v) -> result.put(k.first.getId(), v));
assertEquals(2, result.get(0).size());
assertEquals(10, result.get(1).size());
assertEquals(2, result.get(2).size());
assertEquals(3, result.get(3).size());
assertEquals(2, result.get(4).size());
assertEquals(2, result.get(5).size());
assertEquals(3, result.get(6).size());
}
}

0 comments on commit 30930e2

Please sign in to comment.