From 30930e22e484969a61e1b547b118e60a9dd93187 Mon Sep 17 00:00:00 2001 From: ti-srebot <66930949+ti-srebot@users.noreply.github.com> Date: Thu, 29 Dec 2022 15:47:56 +0800 Subject: [PATCH] [close #693] Add test for RangeSplitter (#694) #698 Signed-off-by: ti-srebot Signed-off-by: ti-srebot Co-authored-by: shi yuhang <52435083+shiyuhang0@users.noreply.github.com> --- .../org/tikv/common/MockRegionManager.java | 89 ++++++ .../java/org/tikv/util/RangeSplitterTest.java | 258 ++++++++++++++++++ 2 files changed, 347 insertions(+) create mode 100644 src/test/java/org/tikv/common/MockRegionManager.java create mode 100644 src/test/java/org/tikv/util/RangeSplitterTest.java diff --git a/src/test/java/org/tikv/common/MockRegionManager.java b/src/test/java/org/tikv/common/MockRegionManager.java new file mode 100644 index 00000000000..cac65312d0b --- /dev/null +++ b/src/test/java/org/tikv/common/MockRegionManager.java @@ -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 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 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 ranges) { + super(null, null); + mockRegionMap = + ranges.stream().collect(Collectors.toMap(kr -> kr, kr -> region(ranges.indexOf(kr), kr))); + } + + @Override + public Pair getRegionStorePairByKey(ByteString key, TiStoreType storeType) { + for (Map.Entry 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; + } +} diff --git a/src/test/java/org/tikv/util/RangeSplitterTest.java b/src/test/java/org/tikv/util/RangeSplitterTest.java new file mode 100644 index 00000000000..7207f959d43 --- /dev/null +++ b/src/test/java/org/tikv/util/RangeSplitterTest.java @@ -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 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 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 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 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 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()); + } +}