From 1c97d4919057dfa4ab18f1167217799fe7fadbc1 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Thu, 24 Mar 2022 19:26:52 +0800 Subject: [PATCH 01/32] api-v2: enable api-v2 Signed-off-by: iosmanthus --- .../java/org/tikv/common/ConfigUtils.java | 5 + src/main/java/org/tikv/common/PDClient.java | 32 ++- .../java/org/tikv/common/TiConfiguration.java | 183 +++++++++++++++++- src/main/java/org/tikv/common/TiSession.java | 5 +- .../region/AbstractRegionStoreClient.java | 37 ++++ .../tikv/common/region/RegionStoreClient.java | 110 +++++++++-- .../org/tikv/common/TiConfigurationTest.java | 9 + 7 files changed, 337 insertions(+), 44 deletions(-) diff --git a/src/main/java/org/tikv/common/ConfigUtils.java b/src/main/java/org/tikv/common/ConfigUtils.java index 600a81d109b..6e7ea07f9e4 100644 --- a/src/main/java/org/tikv/common/ConfigUtils.java +++ b/src/main/java/org/tikv/common/ConfigUtils.java @@ -119,6 +119,9 @@ public class ConfigUtils { public static final String TIFLASH_ENABLE = "tiflash.enable"; public static final String TIKV_WARM_UP_ENABLE = "tikv.warm_up.enable"; + + public static final String TIKV_API_VERSION = "tikv.api_version";; + public static final String DEF_PD_ADDRESSES = "127.0.0.1:2379"; public static final String DEF_TIMEOUT = "200ms"; public static final String DEF_TIKV_GRPC_INGEST_TIMEOUT = "200s"; @@ -200,4 +203,6 @@ public class ConfigUtils { public static final int DEF_TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT = 10; public static final int DEF_TIKV_SCAN_REGIONS_LIMIT = 1000; + + public static final int DEF_TIKV_API_VERSION = 1; } diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index 4b8a283c40c..064491a9aa3 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -311,7 +311,7 @@ private boolean isScatterRegionFinish(GetOperatorResponse resp) { public Pair getRegionByKey(BackOffer backOffer, ByteString key) { Histogram.Timer requestTimer = PD_GET_REGION_BY_KEY_REQUEST_LATENCY.startTimer(); try { - if (conf.isTxnKVMode()) { + if (conf.isTxnKVMode() || conf.getApiVersion().isV2()) { CodecDataOutput cdo = new CodecDataOutput(); BytesCodec.writeBytes(cdo, key.toByteArray()); key = cdo.toByteString(); @@ -326,7 +326,7 @@ public Pair getRegionByKey(BackOffer backOffer, Byte GetRegionResponse resp = callWithRetry(backOffer, PDGrpc.getGetRegionMethod(), request, handler); - return new Pair(decodeRegion(resp.getRegion()), resp.getLeader()); + return new Pair<>(decodeRegion(resp.getRegion()), resp.getLeader()); } finally { requestTimer.observeDuration(); } @@ -806,35 +806,31 @@ private Metapb.Region decodeRegion(Metapb.Region region) { .setRegionEpoch(region.getRegionEpoch()) .addAllPeers(region.getPeersList()); - if (region.getStartKey().isEmpty() || isRawRegion) { + if (conf.getApiVersion().isV1() && (region.getStartKey().isEmpty() || isRawRegion)) { builder.setStartKey(region.getStartKey()); } else { - if (!conf.isTest()) { + try { byte[] decodedStartKey = BytesCodec.readBytes(new CodecDataInput(region.getStartKey())); builder.setStartKey(ByteString.copyFrom(decodedStartKey)); - } else { - try { - byte[] decodedStartKey = BytesCodec.readBytes(new CodecDataInput(region.getStartKey())); - builder.setStartKey(ByteString.copyFrom(decodedStartKey)); - } catch (Exception e) { - builder.setStartKey(region.getStartKey()); + } catch (Exception e) { + if (!conf.isTest()) { + throw e; } + builder.setStartKey(region.getStartKey()); } } - if (region.getEndKey().isEmpty() || isRawRegion) { + if (conf.getApiVersion().isV1() && (region.getEndKey().isEmpty() || isRawRegion)) { builder.setEndKey(region.getEndKey()); } else { - if (!conf.isTest()) { + try { byte[] decodedEndKey = BytesCodec.readBytes(new CodecDataInput(region.getEndKey())); builder.setEndKey(ByteString.copyFrom(decodedEndKey)); - } else { - try { - byte[] decodedEndKey = BytesCodec.readBytes(new CodecDataInput(region.getEndKey())); - builder.setEndKey(ByteString.copyFrom(decodedEndKey)); - } catch (Exception e) { - builder.setEndKey(region.getEndKey()); + } catch (Exception e) { + if (!conf.isTest()) { + throw e; } + builder.setEndKey(region.getEndKey()); } } diff --git a/src/main/java/org/tikv/common/TiConfiguration.java b/src/main/java/org/tikv/common/TiConfiguration.java index bc64a7ce4b5..a35b2e6208d 100644 --- a/src/main/java/org/tikv/common/TiConfiguration.java +++ b/src/main/java/org/tikv/common/TiConfiguration.java @@ -17,14 +17,155 @@ package org.tikv.common; -import static org.tikv.common.ConfigUtils.*; +import static org.tikv.common.ConfigUtils.DEF_BATCH_DELETE_CONCURRENCY; +import static org.tikv.common.ConfigUtils.DEF_BATCH_GET_CONCURRENCY; +import static org.tikv.common.ConfigUtils.DEF_BATCH_PUT_CONCURRENCY; +import static org.tikv.common.ConfigUtils.DEF_BATCH_SCAN_CONCURRENCY; +import static org.tikv.common.ConfigUtils.DEF_CHECK_HEALTH_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_DB_PREFIX; +import static org.tikv.common.ConfigUtils.DEF_DELETE_RANGE_CONCURRENCY; +import static org.tikv.common.ConfigUtils.DEF_FORWARD_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_GRPC_FORWARD_ENABLE; +import static org.tikv.common.ConfigUtils.DEF_HEALTH_CHECK_PERIOD_DURATION; +import static org.tikv.common.ConfigUtils.DEF_INDEX_SCAN_BATCH_SIZE; +import static org.tikv.common.ConfigUtils.DEF_INDEX_SCAN_CONCURRENCY; +import static org.tikv.common.ConfigUtils.DEF_KV_CLIENT_CONCURRENCY; +import static org.tikv.common.ConfigUtils.DEF_MAX_FRAME_SIZE; +import static org.tikv.common.ConfigUtils.DEF_METRICS_ENABLE; +import static org.tikv.common.ConfigUtils.DEF_METRICS_PORT; +import static org.tikv.common.ConfigUtils.DEF_PD_ADDRESSES; +import static org.tikv.common.ConfigUtils.DEF_REPLICA_READ; +import static org.tikv.common.ConfigUtils.DEF_SCAN_BATCH_SIZE; +import static org.tikv.common.ConfigUtils.DEF_SCAN_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_SHOW_ROWID; +import static org.tikv.common.ConfigUtils.DEF_TABLE_SCAN_CONCURRENCY; +import static org.tikv.common.ConfigUtils.DEF_TIFLASH_ENABLE; +import static org.tikv.common.ConfigUtils.DEF_TIKV_API_VERSION; +import static org.tikv.common.ConfigUtils.DEF_TIKV_BO_REGION_MISS_BASE_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_ENABLE_ATOMIC_FOR_CAS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_GRPC_IDLE_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_TIKV_GRPC_INGEST_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_TIKV_GRPC_KEEPALIVE_TIME; +import static org.tikv.common.ConfigUtils.DEF_TIKV_GRPC_KEEPALIVE_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_TIKV_GRPC_WARM_UP_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_TIKV_IMPORTER_MAX_KV_BATCH_BYTES; +import static org.tikv.common.ConfigUtils.DEF_TIKV_IMPORTER_MAX_KV_BATCH_SIZE; +import static org.tikv.common.ConfigUtils.DEF_TIKV_NETWORK_MAPPING_NAME; +import static org.tikv.common.ConfigUtils.DEF_TIKV_PD_FIRST_GET_MEMBER_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_TIKV_RAWKV_BATCH_READ_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_RAWKV_BATCH_WRITE_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_RAWKV_CLEAN_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_RAWKV_DEFAULT_BACKOFF_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_RAWKV_READ_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_RAWKV_SCAN_SLOWLOG_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_RAWKV_SCAN_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_RAWKV_WRITE_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_SCAN_REGIONS_LIMIT; +import static org.tikv.common.ConfigUtils.DEF_TIKV_SCATTER_WAIT_SECONDS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_TLS_ENABLE; +import static org.tikv.common.ConfigUtils.DEF_TIKV_USE_JKS; +import static org.tikv.common.ConfigUtils.DEF_TIKV_WARM_UP_ENABLE; +import static org.tikv.common.ConfigUtils.DEF_TIMEOUT; +import static org.tikv.common.ConfigUtils.DEF_TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT; +import static org.tikv.common.ConfigUtils.DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_ERROR_THRESHOLD_PERCENTAGE; +import static org.tikv.common.ConfigUtils.DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_REQUST_VOLUMN_THRESHOLD; +import static org.tikv.common.ConfigUtils.DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_WINDOW_IN_SECONDS; +import static org.tikv.common.ConfigUtils.DEF_TiKV_CIRCUIT_BREAK_ENABLE; +import static org.tikv.common.ConfigUtils.DEF_TiKV_CIRCUIT_BREAK_SLEEP_WINDOW_IN_SECONDS; +import static org.tikv.common.ConfigUtils.FOLLOWER; +import static org.tikv.common.ConfigUtils.HIGH_COMMAND_PRIORITY; +import static org.tikv.common.ConfigUtils.LEADER_AND_FOLLOWER; +import static org.tikv.common.ConfigUtils.LOW_COMMAND_PRIORITY; +import static org.tikv.common.ConfigUtils.NORMAL_COMMAND_PRIORITY; +import static org.tikv.common.ConfigUtils.RAW_KV_MODE; +import static org.tikv.common.ConfigUtils.READ_COMMITTED_ISOLATION_LEVEL; +import static org.tikv.common.ConfigUtils.SNAPSHOT_ISOLATION_LEVEL; +import static org.tikv.common.ConfigUtils.TIFLASH_ENABLE; +import static org.tikv.common.ConfigUtils.TIKV_API_VERSION; +import static org.tikv.common.ConfigUtils.TIKV_BATCH_DELETE_CONCURRENCY; +import static org.tikv.common.ConfigUtils.TIKV_BATCH_GET_CONCURRENCY; +import static org.tikv.common.ConfigUtils.TIKV_BATCH_PUT_CONCURRENCY; +import static org.tikv.common.ConfigUtils.TIKV_BATCH_SCAN_CONCURRENCY; +import static org.tikv.common.ConfigUtils.TIKV_BO_REGION_MISS_BASE_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_DB_PREFIX; +import static org.tikv.common.ConfigUtils.TIKV_DELETE_RANGE_CONCURRENCY; +import static org.tikv.common.ConfigUtils.TIKV_ENABLE_ATOMIC_FOR_CAS; +import static org.tikv.common.ConfigUtils.TIKV_ENABLE_GRPC_FORWARD; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_FORWARD_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_HEALTH_CHECK_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_IDLE_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_INGEST_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_KEEPALIVE_TIME; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_KEEPALIVE_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_MAX_FRAME_SIZE; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_SCAN_BATCH_SIZE; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_SCAN_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_GRPC_WARM_UP_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_HEALTH_CHECK_PERIOD_DURATION; +import static org.tikv.common.ConfigUtils.TIKV_IMPORTER_MAX_KV_BATCH_BYTES; +import static org.tikv.common.ConfigUtils.TIKV_IMPORTER_MAX_KV_BATCH_SIZE; +import static org.tikv.common.ConfigUtils.TIKV_INDEX_SCAN_BATCH_SIZE; +import static org.tikv.common.ConfigUtils.TIKV_INDEX_SCAN_CONCURRENCY; +import static org.tikv.common.ConfigUtils.TIKV_JKS_KEY_PASSWORD; +import static org.tikv.common.ConfigUtils.TIKV_JKS_KEY_PATH; +import static org.tikv.common.ConfigUtils.TIKV_JKS_TRUST_PASSWORD; +import static org.tikv.common.ConfigUtils.TIKV_JKS_TRUST_PATH; +import static org.tikv.common.ConfigUtils.TIKV_KEY_CERT_CHAIN; +import static org.tikv.common.ConfigUtils.TIKV_KEY_FILE; +import static org.tikv.common.ConfigUtils.TIKV_KV_CLIENT_CONCURRENCY; +import static org.tikv.common.ConfigUtils.TIKV_KV_MODE; +import static org.tikv.common.ConfigUtils.TIKV_METRICS_ENABLE; +import static org.tikv.common.ConfigUtils.TIKV_METRICS_PORT; +import static org.tikv.common.ConfigUtils.TIKV_NETWORK_MAPPING_NAME; +import static org.tikv.common.ConfigUtils.TIKV_PD_ADDRESSES; +import static org.tikv.common.ConfigUtils.TIKV_PD_FIRST_GET_MEMBER_TIMEOUT; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_BATCH_READ_SLOWLOG_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_BATCH_READ_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_BATCH_WRITE_SLOWLOG_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_BATCH_WRITE_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_CLEAN_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_DEFAULT_BACKOFF_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_READ_SLOWLOG_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_READ_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_SCAN_SLOWLOG_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_SCAN_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_SERVER_SLOWLOG_FACTOR; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_WRITE_SLOWLOG_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_RAWKV_WRITE_TIMEOUT_IN_MS; +import static org.tikv.common.ConfigUtils.TIKV_REPLICA_READ; +import static org.tikv.common.ConfigUtils.TIKV_REQUEST_COMMAND_PRIORITY; +import static org.tikv.common.ConfigUtils.TIKV_REQUEST_ISOLATION_LEVEL; +import static org.tikv.common.ConfigUtils.TIKV_SCAN_REGIONS_LIMIT; +import static org.tikv.common.ConfigUtils.TIKV_SCATTER_WAIT_SECONDS; +import static org.tikv.common.ConfigUtils.TIKV_SHOW_ROWID; +import static org.tikv.common.ConfigUtils.TIKV_TABLE_SCAN_CONCURRENCY; +import static org.tikv.common.ConfigUtils.TIKV_TLS_ENABLE; +import static org.tikv.common.ConfigUtils.TIKV_TRUST_CERT_COLLECTION; +import static org.tikv.common.ConfigUtils.TIKV_USE_JKS; +import static org.tikv.common.ConfigUtils.TIKV_WARM_UP_ENABLE; +import static org.tikv.common.ConfigUtils.TXN_KV_MODE; +import static org.tikv.common.ConfigUtils.TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT; +import static org.tikv.common.ConfigUtils.TiKV_CIRCUIT_BREAK_AVAILABILITY_ERROR_THRESHOLD_PERCENTAGE; +import static org.tikv.common.ConfigUtils.TiKV_CIRCUIT_BREAK_AVAILABILITY_REQUEST_VOLUMN_THRESHOLD; +import static org.tikv.common.ConfigUtils.TiKV_CIRCUIT_BREAK_AVAILABILITY_WINDOW_IN_SECONDS; +import static org.tikv.common.ConfigUtils.TiKV_CIRCUIT_BREAK_ENABLE; +import static org.tikv.common.ConfigUtils.TiKV_CIRCUIT_BREAK_SLEEP_WINDOW_IN_SECONDS; import io.grpc.Metadata; import java.io.IOException; import java.io.InputStream; import java.io.Serializable; import java.net.URI; -import java.util.*; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Optional; +import java.util.Properties; import java.util.concurrent.ConcurrentHashMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -34,6 +175,29 @@ import org.tikv.kvproto.Kvrpcpb.IsolationLevel; public class TiConfiguration implements Serializable { + public static enum ApiVersion { + V1, + V2; + + public static ApiVersion fromInt(int version) { + switch (version) { + case 1: + return V1; + case 2: + return V2; + default: + throw new IllegalArgumentException("unknown api version " + version); + } + } + + public boolean isV1() { + return this == V1; + } + + public boolean isV2() { + return this == V2; + } + } private static final Logger logger = LoggerFactory.getLogger(TiConfiguration.class); private static final ConcurrentHashMap settings = new ConcurrentHashMap<>(); @@ -41,6 +205,8 @@ public class TiConfiguration implements Serializable { Metadata.Key.of("tikv-forwarded-host", Metadata.ASCII_STRING_MARSHALLER); public static final Metadata.Key PD_FORWARD_META_DATA_KEY = Metadata.Key.of("pd-forwarded-host", Metadata.ASCII_STRING_MARSHALLER); + public static final String API_V2_RAW_PREFIX = "r"; + public static final String API_V2_TXN_PREFIX = "x"; static { // priority: system environment > config file > default @@ -151,6 +317,8 @@ private static void loadFromDefaultProperties() { setIfMissing( TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT, DEF_TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT); setIfMissing(TIKV_SCAN_REGIONS_LIMIT, DEF_TIKV_SCAN_REGIONS_LIMIT); + + setIfMissing(TIKV_API_VERSION, DEF_TIKV_API_VERSION); } public static void listAll() { @@ -403,6 +571,8 @@ private static ReplicaRead getReplicaRead(String key) { private int scanRegionsLimit = getInt(TIKV_SCAN_REGIONS_LIMIT); + private ApiVersion apiVersion = ApiVersion.fromInt(getInt(TIKV_API_VERSION)); + public enum KVMode { TXN, RAW @@ -1075,4 +1245,13 @@ public int getScanRegionsLimit() { public void setScanRegionsLimit(int scanRegionsLimit) { this.scanRegionsLimit = scanRegionsLimit; } + + public ApiVersion getApiVersion() { + return apiVersion; + } + + public TiConfiguration setApiVersion(ApiVersion version) { + this.apiVersion = version; + return this; + } } diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java index 8dd1793cbca..48960117511 100644 --- a/src/main/java/org/tikv/common/TiSession.java +++ b/src/main/java/org/tikv/common/TiSession.java @@ -67,7 +67,6 @@ * contention */ public class TiSession implements AutoCloseable { - private static final Logger logger = LoggerFactory.getLogger(TiSession.class); private static final Map sessionCachedMap = new HashMap<>(); private final TiConfiguration conf; @@ -156,7 +155,7 @@ public TiSession(TiConfiguration conf) { logger.info("enable grpc forward for high available"); } if (conf.isWarmUpEnable() && conf.isRawKVMode()) { - warmUp(); + warmup(); } this.circuitBreaker = new CircuitBreakerImpl(conf); logger.info("TiSession initialized in " + conf.getKvMode() + " mode"); @@ -177,7 +176,7 @@ private static VersionInfo getVersionInfo() { return info; } - private synchronized void warmUp() { + private synchronized void warmup() { long warmUpStartTime = System.nanoTime(); BackOffer backOffer = ConcreteBackOffer.newRawKVBackOff(); try { diff --git a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java index 1dc6c321fc4..1008d7db6c4 100644 --- a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java @@ -34,6 +34,7 @@ import org.slf4j.LoggerFactory; import org.tikv.common.AbstractGRPCClient; import org.tikv.common.TiConfiguration; +import org.tikv.common.TiConfiguration.KVMode; import org.tikv.common.exception.GrpcException; import org.tikv.common.log.SlowLog; import org.tikv.common.log.SlowLogSpan; @@ -155,6 +156,42 @@ public boolean onStoreUnreachable(BackOffer backOffer) { return false; } + public ByteString buildRequestKey(ByteString key) { + switch (conf.getApiVersion()) { + case V1: + return key; + case V2: + if (conf.getKvMode() == KVMode.RAW) { + return ByteString.copyFromUtf8(TiConfiguration.API_V2_RAW_PREFIX).concat(key); + } else if (conf.getKvMode() == KVMode.TXN) { + return ByteString.copyFromUtf8(TiConfiguration.API_V2_TXN_PREFIX).concat(key); + } + default: + throw new IllegalArgumentException("unknown api version or kv mode"); + } + } + + public ByteString unwrapResponseKey(ByteString key) { + switch (conf.getApiVersion()) { + case V1: + return key; + case V2: + if (conf.getKvMode() == KVMode.RAW) { + if (!key.startsWith(ByteString.copyFromUtf8(TiConfiguration.API_V2_RAW_PREFIX))) { + throw new IllegalArgumentException("key corrupted, wrong prefix"); + } + return key.substring(1); + } else if (conf.getKvMode() == KVMode.TXN) { + if (!key.startsWith(ByteString.copyFromUtf8(TiConfiguration.API_V2_TXN_PREFIX))) { + throw new IllegalArgumentException("key corrupted, wrong prefix"); + } + return key.substring(1); + } + default: + throw new IllegalArgumentException("unknown api version or kv mode"); + } + } + private Kvrpcpb.Context addTraceId(Kvrpcpb.Context context, SlowLog slowLog) { if (slowLog.getThresholdMS() < 0) { // disable tikv tracing diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index 1c090a92a0f..fd35f0b4848 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -18,7 +18,9 @@ package org.tikv.common.region; import static org.tikv.common.region.RegionStoreClient.RequestTypes.REQ_TYPE_DAG; -import static org.tikv.common.util.BackOffFunction.BackOffFuncType.*; +import static org.tikv.common.util.BackOffFunction.BackOffFuncType.BoRegionMiss; +import static org.tikv.common.util.BackOffFunction.BackOffFuncType.BoTxnLock; +import static org.tikv.common.util.BackOffFunction.BackOffFuncType.BoTxnLockFast; import com.google.common.annotations.VisibleForTesting; import com.google.protobuf.ByteString; @@ -29,23 +31,83 @@ import io.grpc.Metadata; import io.grpc.stub.MetadataUtils; import io.prometheus.client.Histogram; -import java.util.*; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Queue; +import java.util.Set; import java.util.function.Supplier; +import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.tikv.common.PDClient; import org.tikv.common.StoreVersion; import org.tikv.common.TiConfiguration; import org.tikv.common.Version; -import org.tikv.common.exception.*; +import org.tikv.common.exception.GrpcException; +import org.tikv.common.exception.KeyException; +import org.tikv.common.exception.RawCASConflictException; +import org.tikv.common.exception.RegionException; +import org.tikv.common.exception.SelectException; +import org.tikv.common.exception.TiClientInternalException; +import org.tikv.common.exception.TiKVException; import org.tikv.common.log.SlowLogEmptyImpl; import org.tikv.common.operation.KVErrorHandler; import org.tikv.common.operation.RegionErrorHandler; import org.tikv.common.streaming.StreamingResponse; -import org.tikv.common.util.*; +import org.tikv.common.util.BackOffFunction; +import org.tikv.common.util.BackOffer; +import org.tikv.common.util.Batch; +import org.tikv.common.util.ChannelFactory; +import org.tikv.common.util.ConcreteBackOffer; +import org.tikv.common.util.HistogramUtils; +import org.tikv.common.util.Pair; +import org.tikv.common.util.RangeSplitter; import org.tikv.kvproto.Coprocessor; import org.tikv.kvproto.Errorpb; -import org.tikv.kvproto.Kvrpcpb.*; +import org.tikv.kvproto.Kvrpcpb.BatchGetRequest; +import org.tikv.kvproto.Kvrpcpb.BatchGetResponse; +import org.tikv.kvproto.Kvrpcpb.CommitRequest; +import org.tikv.kvproto.Kvrpcpb.CommitResponse; +import org.tikv.kvproto.Kvrpcpb.GetRequest; +import org.tikv.kvproto.Kvrpcpb.GetResponse; +import org.tikv.kvproto.Kvrpcpb.KeyError; +import org.tikv.kvproto.Kvrpcpb.KvPair; +import org.tikv.kvproto.Kvrpcpb.Mutation; +import org.tikv.kvproto.Kvrpcpb.PrewriteRequest; +import org.tikv.kvproto.Kvrpcpb.PrewriteResponse; +import org.tikv.kvproto.Kvrpcpb.RawBatchDeleteRequest; +import org.tikv.kvproto.Kvrpcpb.RawBatchDeleteResponse; +import org.tikv.kvproto.Kvrpcpb.RawBatchGetRequest; +import org.tikv.kvproto.Kvrpcpb.RawBatchGetResponse; +import org.tikv.kvproto.Kvrpcpb.RawBatchPutRequest; +import org.tikv.kvproto.Kvrpcpb.RawBatchPutResponse; +import org.tikv.kvproto.Kvrpcpb.RawCASRequest; +import org.tikv.kvproto.Kvrpcpb.RawCASResponse; +import org.tikv.kvproto.Kvrpcpb.RawDeleteRangeRequest; +import org.tikv.kvproto.Kvrpcpb.RawDeleteRangeResponse; +import org.tikv.kvproto.Kvrpcpb.RawDeleteRequest; +import org.tikv.kvproto.Kvrpcpb.RawDeleteResponse; +import org.tikv.kvproto.Kvrpcpb.RawGetKeyTTLRequest; +import org.tikv.kvproto.Kvrpcpb.RawGetKeyTTLResponse; +import org.tikv.kvproto.Kvrpcpb.RawGetRequest; +import org.tikv.kvproto.Kvrpcpb.RawGetResponse; +import org.tikv.kvproto.Kvrpcpb.RawPutRequest; +import org.tikv.kvproto.Kvrpcpb.RawPutResponse; +import org.tikv.kvproto.Kvrpcpb.RawScanRequest; +import org.tikv.kvproto.Kvrpcpb.RawScanResponse; +import org.tikv.kvproto.Kvrpcpb.ScanRequest; +import org.tikv.kvproto.Kvrpcpb.ScanResponse; +import org.tikv.kvproto.Kvrpcpb.SplitRegionRequest; +import org.tikv.kvproto.Kvrpcpb.SplitRegionResponse; +import org.tikv.kvproto.Kvrpcpb.TxnHeartBeatRequest; +import org.tikv.kvproto.Kvrpcpb.TxnHeartBeatResponse; import org.tikv.kvproto.Metapb; import org.tikv.kvproto.TikvGrpc; import org.tikv.kvproto.TikvGrpc.TikvBlockingStub; @@ -64,7 +126,6 @@ /** Note that RegionStoreClient itself is not thread-safe */ public class RegionStoreClient extends AbstractRegionStoreClient { - private static final Logger logger = LoggerFactory.getLogger(RegionStoreClient.class); @VisibleForTesting public final AbstractLockResolverClient lockResolverClient; private final TiStoreType storeType; @@ -175,7 +236,7 @@ public ByteString get(BackOffer backOffer, ByteString key, long version) GetRequest.newBuilder() .setContext( makeContext(getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .setKey(key) + .setKey(buildRequestKey(key)) .setVersion(version) .build(); @@ -286,7 +347,7 @@ public List scan( .setContext( makeContext( getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .setStartKey(startKey) + .setStartKey(buildRequestKey(startKey)) .setVersion(version) .setKeyOnly(keyOnly) .setLimit(getConf().getScanBatchSize()) @@ -333,7 +394,7 @@ private List doScan(ScanResponse resp) { KvPair.newBuilder() .setError(kvPair.getError()) .setValue(kvPair.getValue()) - .setKey(lock.getKey()) + .setKey(unwrapResponseKey(lock.getKey())) .build()); } else { newKvPairs.add(kvPair); @@ -788,8 +849,7 @@ public List splitRegion(Iterable splitKeys) { if (resp.hasRegionError()) { throw new TiClientInternalException( String.format( - "failed to split region %d because %s", - region.getId(), resp.getRegionError().toString())); + "failed to split region %d because %s", region.getId(), resp.getRegionError())); } return resp.getRegionsList(); @@ -805,7 +865,7 @@ public Optional rawGet(BackOffer backOffer, ByteString key) { () -> RawGetRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(key) + .setKey(buildRequestKey(key)) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -844,7 +904,7 @@ public Optional rawGetKeyTTL(BackOffer backOffer, ByteString key) { () -> RawGetKeyTTLRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(key) + .setKey(buildRequestKey(key)) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -883,7 +943,7 @@ public void rawDelete(BackOffer backOffer, ByteString key, boolean atomicForCAS) () -> RawDeleteRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(key) + .setKey(buildRequestKey(key)) .setForCas(atomicForCAS) .build(); @@ -921,7 +981,7 @@ public void rawPut( () -> RawPutRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(key) + .setKey(buildRequestKey(key)) .setValue(value) .setTtl(ttl) .setForCas(atomicForCAS) @@ -965,7 +1025,7 @@ public void rawCompareAndSet( () -> RawCASRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(key) + .setKey(buildRequestKey(key)) .setValue(value) .setPreviousValue(prevValue.orElse(ByteString.EMPTY)) .setPreviousNotExist(!prevValue.isPresent()) @@ -1075,7 +1135,7 @@ public void rawBatchPut(BackOffer backOffer, Batch batch, long ttl, boolean atom for (int i = 0; i < batch.getKeys().size(); i++) { pairs.add( KvPair.newBuilder() - .setKey(batch.getKeys().get(i)) + .setKey(buildRequestKey(batch.getKeys().get(i))) .setValue(batch.getValues().get(i)) .build()); } @@ -1152,7 +1212,7 @@ public List rawScan(BackOffer backOffer, ByteString key, int limit, bool () -> RawScanRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setStartKey(key) + .setStartKey(buildRequestKey(key)) .setKeyOnly(keyOnly) .setLimit(limit) .build(); @@ -1180,7 +1240,15 @@ private List rawScanHelper(RawScanResponse resp) { if (resp.hasRegionError()) { throw new RegionException(resp.getRegionError()); } - return resp.getKvsList(); + return resp.getKvsList() + .stream() + .map( + kvPair -> + KvPair.newBuilder() + .setKey(unwrapResponseKey(kvPair.getKey())) + .setValue(kvPair.getValue()) + .build()) + .collect(Collectors.toList()); } /** @@ -1198,8 +1266,8 @@ public void rawDeleteRange(BackOffer backOffer, ByteString startKey, ByteString () -> RawDeleteRangeRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setStartKey(startKey) - .setEndKey(endKey) + .setStartKey(buildRequestKey(startKey)) + .setEndKey(buildRequestKey(endKey)) .build(); RegionErrorHandler handler = diff --git a/src/test/java/org/tikv/common/TiConfigurationTest.java b/src/test/java/org/tikv/common/TiConfigurationTest.java index f4968c9b1bc..9d2ac595ff4 100644 --- a/src/test/java/org/tikv/common/TiConfigurationTest.java +++ b/src/test/java/org/tikv/common/TiConfigurationTest.java @@ -19,6 +19,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; import static org.tikv.common.ConfigUtils.TIKV_GRPC_HEALTH_CHECK_TIMEOUT; import static org.tikv.common.ConfigUtils.TIKV_HEALTH_CHECK_PERIOD_DURATION; @@ -27,6 +28,7 @@ import java.io.ObjectOutputStream; import org.junit.Assert; import org.junit.Test; +import org.tikv.common.TiConfiguration.ApiVersion; public class TiConfigurationTest { @@ -102,4 +104,11 @@ public void serializeTest() throws IOException { oos.flush(); } } + + @Test + public void testApiVersion() { + TiConfiguration conf = TiConfiguration.createDefault(); + assertTrue(conf.getApiVersion().isV1()); + assertTrue(conf.setApiVersion(ApiVersion.V2).getApiVersion().isV1()); + } } From f45d5838c9578e33d7ca6ba62896264f49f1a0d1 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Fri, 25 Mar 2022 18:18:38 +0800 Subject: [PATCH 02/32] attach api version to context and fix lock resolver api version Signed-off-by: iosmanthus --- .../java/org/tikv/common/TiConfiguration.java | 12 +++++++++ .../region/AbstractRegionStoreClient.java | 27 ++++++++++++++----- .../org/tikv/txn/LockResolverClientV4.java | 12 ++++----- 3 files changed, 38 insertions(+), 13 deletions(-) diff --git a/src/main/java/org/tikv/common/TiConfiguration.java b/src/main/java/org/tikv/common/TiConfiguration.java index a35b2e6208d..fa180d4a3e2 100644 --- a/src/main/java/org/tikv/common/TiConfiguration.java +++ b/src/main/java/org/tikv/common/TiConfiguration.java @@ -171,6 +171,7 @@ import org.slf4j.LoggerFactory; import org.tikv.common.pd.PDUtils; import org.tikv.common.replica.ReplicaSelector; +import org.tikv.kvproto.Kvrpcpb; import org.tikv.kvproto.Kvrpcpb.CommandPri; import org.tikv.kvproto.Kvrpcpb.IsolationLevel; @@ -197,6 +198,17 @@ public boolean isV1() { public boolean isV2() { return this == V2; } + + public Kvrpcpb.APIVersion toPb() { + switch (this) { + case V1: + return Kvrpcpb.APIVersion.V1; + case V2: + return Kvrpcpb.APIVersion.V2; + default: + throw new IllegalArgumentException("unknown api version " + this); + } + } } private static final Logger logger = LoggerFactory.getLogger(TiConfiguration.class); diff --git a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java index 1008d7db6c4..f8ef94d14b1 100644 --- a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java @@ -104,7 +104,8 @@ protected TikvGrpc.TikvFutureStub getAsyncStub() { } @Override - public void close() throws GrpcException {} + public void close() throws GrpcException { + } /** * onNotLeader deals with NotLeaderError and returns whether re-splitting key range is needed @@ -207,15 +208,27 @@ private Kvrpcpb.Context addTraceId(Kvrpcpb.Context context, SlowLog slowLog) { .build(); } + protected Kvrpcpb.Context withApiVersion(Kvrpcpb.Context context) { + return Kvrpcpb.Context.newBuilder(context).setApiVersion(conf.getApiVersion().toPb()).build(); + } + protected Kvrpcpb.Context makeContext(TiStoreType storeType, SlowLog slowLog) { Kvrpcpb.Context context = region.getReplicaContext(java.util.Collections.emptySet(), storeType); - return addTraceId(context, slowLog); + return withApiVersion(addTraceId(context, slowLog)); } protected Kvrpcpb.Context makeContext( Set resolvedLocks, TiStoreType storeType, SlowLog slowLog) { Kvrpcpb.Context context = region.getReplicaContext(resolvedLocks, storeType); - return addTraceId(context, slowLog); + return withApiVersion(addTraceId(context, slowLog)); + } + + protected Kvrpcpb.Context makeContext() { + return withApiVersion(region.getLeaderContext()); + } + + protected Kvrpcpb.Context makeContext(Metapb.Peer peer) { + return withApiVersion(region.getReplicaContext(peer)); } private void updateClientStub() { @@ -323,8 +336,8 @@ private Metapb.Peer switchLeaderStore(BackOffer backOffer) { TikvGrpc.newFutureStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS); Kvrpcpb.RawGetRequest rawGetRequest = Kvrpcpb.RawGetRequest.newBuilder() - .setContext(region.getReplicaContext(peer)) - .setKey(key) + .setContext(makeContext(peer)) + .setKey(buildRequestKey(key)) .build(); ListenableFuture task = stub.rawGet(rawGetRequest); responses.add(new SwitchLeaderTask(task, peer)); @@ -385,8 +398,8 @@ private TiStore switchProxyStore(BackOffer backOffer) { header.put(TiConfiguration.FORWARD_META_DATA_KEY, store.getStore().getAddress()); Kvrpcpb.RawGetRequest rawGetRequest = Kvrpcpb.RawGetRequest.newBuilder() - .setContext(region.getReplicaContext(region.getLeader())) - .setKey(key) + .setContext(makeContext()) + .setKey(buildRequestKey(key)) .build(); ListenableFuture task = MetadataUtils.attachHeaders(stub, header).rawGet(rawGetRequest); diff --git a/src/main/java/org/tikv/txn/LockResolverClientV4.java b/src/main/java/org/tikv/txn/LockResolverClientV4.java index 4a723537c51..1aaa0221cb4 100644 --- a/src/main/java/org/tikv/txn/LockResolverClientV4.java +++ b/src/main/java/org/tikv/txn/LockResolverClientV4.java @@ -167,10 +167,10 @@ private void resolvePessimisticLock(BackOffer bo, Lock lock, Set cl Supplier factory = () -> Kvrpcpb.PessimisticRollbackRequest.newBuilder() - .setContext(region.getLeaderContext()) + .setContext(makeContext()) + .addKeys(buildRequestKey(lock.getKey())) .setStartVersion(lock.getTxnID()) .setForUpdateTs(forUpdateTS) - .addKeys(lock.getKey()) .build(); KVErrorHandler handler = @@ -285,8 +285,8 @@ private TxnStatus getTxnStatus( () -> { TiRegion primaryKeyRegion = regionManager.getRegionByKey(primary); return Kvrpcpb.CheckTxnStatusRequest.newBuilder() - .setContext(primaryKeyRegion.getLeaderContext()) - .setPrimaryKey(primary) + .setContext(withApiVersion(primaryKeyRegion.getLeaderContext())) + .setPrimaryKey(buildRequestKey(primary)) .setLockTs(txnID) .setCallerStartTs(callerStartTS) .setCurrentTs(currentTS) @@ -362,7 +362,7 @@ private void resolveLock( Kvrpcpb.ResolveLockRequest.Builder builder = Kvrpcpb.ResolveLockRequest.newBuilder() - .setContext(region.getLeaderContext()) + .setContext(makeContext()) .setStartVersion(lock.getTxnID()); if (txnStatus.isCommitted()) { @@ -373,7 +373,7 @@ private void resolveLock( if (lock.getTxnSize() < BIG_TXN_THRESHOLD) { // Only resolve specified keys when it is a small transaction, // prevent from scanning the whole region in this case. - builder.addKeys(lock.getKey()); + builder.addKeys(buildRequestKey(lock.getKey())); } Supplier factory = builder::build; From 8ff27f29a6e4e4fe8453d721c60a2efe1ea703d1 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Mon, 28 Mar 2022 14:48:08 +0800 Subject: [PATCH 03/32] ./dev/javafmt Signed-off-by: iosmanthus --- .../java/org/tikv/common/region/AbstractRegionStoreClient.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java index f8ef94d14b1..0a84b8b0103 100644 --- a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java @@ -104,8 +104,7 @@ protected TikvGrpc.TikvFutureStub getAsyncStub() { } @Override - public void close() throws GrpcException { - } + public void close() throws GrpcException {} /** * onNotLeader deals with NotLeaderError and returns whether re-splitting key range is needed From 7bb666b0a94c6bc52f60fc7974ab144036a30b6d Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Mon, 28 Mar 2022 19:03:54 +0800 Subject: [PATCH 04/32] fix api v2 config test Signed-off-by: iosmanthus --- src/test/java/org/tikv/common/TiConfigurationTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/test/java/org/tikv/common/TiConfigurationTest.java b/src/test/java/org/tikv/common/TiConfigurationTest.java index 9d2ac595ff4..cb7bb1060ee 100644 --- a/src/test/java/org/tikv/common/TiConfigurationTest.java +++ b/src/test/java/org/tikv/common/TiConfigurationTest.java @@ -109,6 +109,6 @@ public void serializeTest() throws IOException { public void testApiVersion() { TiConfiguration conf = TiConfiguration.createDefault(); assertTrue(conf.getApiVersion().isV1()); - assertTrue(conf.setApiVersion(ApiVersion.V2).getApiVersion().isV1()); + assertTrue(conf.setApiVersion(ApiVersion.V2).getApiVersion().isV2()); } } From 790ffd2c29b47eb2307b4a1130534cd99c4ee63f Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Mon, 28 Mar 2022 19:06:03 +0800 Subject: [PATCH 05/32] add topb test Signed-off-by: iosmanthus --- src/test/java/org/tikv/common/TiConfigurationTest.java | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/test/java/org/tikv/common/TiConfigurationTest.java b/src/test/java/org/tikv/common/TiConfigurationTest.java index cb7bb1060ee..184b50708fe 100644 --- a/src/test/java/org/tikv/common/TiConfigurationTest.java +++ b/src/test/java/org/tikv/common/TiConfigurationTest.java @@ -29,6 +29,7 @@ import org.junit.Assert; import org.junit.Test; import org.tikv.common.TiConfiguration.ApiVersion; +import org.tikv.kvproto.Kvrpcpb.APIVersion; public class TiConfigurationTest { @@ -111,4 +112,10 @@ public void testApiVersion() { assertTrue(conf.getApiVersion().isV1()); assertTrue(conf.setApiVersion(ApiVersion.V2).getApiVersion().isV2()); } + + @Test + public void testApiVersionToPb() { + assertEquals(APIVersion.V1, ApiVersion.V1.toPb()); + assertEquals(APIVersion.V2, ApiVersion.V2.toPb()); + } } From c5e502e45c99ad2cc319359d2d3e2a04480adf8e Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Wed, 6 Apr 2022 15:28:55 +0800 Subject: [PATCH 06/32] fix ingest interface for rawkv Signed-off-by: iosmanthus --- dev/proto.sh | 4 +- src/main/java/org/tikv/common/PDClient.java | 33 ++++- .../java/org/tikv/common/StoreVersion.java | 1 - .../java/org/tikv/common/TiConfiguration.java | 134 +++++++++++++----- src/main/java/org/tikv/common/TiSession.java | 8 +- .../tikv/common/importer/ImporterClient.java | 39 +++-- .../operation/iterator/RawScanIterator.java | 5 +- .../region/AbstractRegionStoreClient.java | 53 +------ .../tikv/common/region/RegionStoreClient.java | 113 ++++++++------- .../java/org/tikv/common/region/TiRegion.java | 1 + .../org/tikv/txn/LockResolverClientV4.java | 8 +- src/test/java/org/tikv/BaseRawKVTest.java | 12 +- src/test/java/org/tikv/BaseTxnKVTest.java | 9 ++ .../java/org/tikv/common/StoreApiVersion.java | 52 +++++++ .../tikv/common/importer/RawKVIngestTest.java | 2 +- .../java/org/tikv/raw/RawKVClientTest.java | 3 +- 16 files changed, 302 insertions(+), 175 deletions(-) create mode 100644 src/test/java/org/tikv/common/StoreApiVersion.java diff --git a/dev/proto.sh b/dev/proto.sh index 5d8dbd4b171..f1a85b5bf67 100755 --- a/dev/proto.sh +++ b/dev/proto.sh @@ -14,7 +14,7 @@ # limitations under the License. # -kvproto_hash=58f2ac94aa38f49676dd628fbcc1d669a77a62ac +kvproto_hash=afd3c854a8bec2cfe0582c914fda8d9754f1c40a raft_rs_hash=b9891b673573fad77ebcf9bbe0969cf945841926 tipb_hash=c4d518eb1d60c21f05b028b36729e64610346dac @@ -32,7 +32,7 @@ cd "$TIKV_CLIENT_HOME" || exit if [ -d "$kvproto_dir" ]; then git -C ${kvproto_dir} fetch -p else - git clone https://github.com/pingcap/kvproto ${kvproto_dir} + git clone https://github.com/iosmanthus/kvproto ${kvproto_dir} fi git -C ${kvproto_dir} checkout ${kvproto_hash} diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index 064491a9aa3..2a87d8c38a5 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -311,6 +311,7 @@ private boolean isScatterRegionFinish(GetOperatorResponse resp) { public Pair getRegionByKey(BackOffer backOffer, ByteString key) { Histogram.Timer requestTimer = PD_GET_REGION_BY_KEY_REQUEST_LATENCY.startTimer(); try { + key = conf.buildRequestKey(key); if (conf.isTxnKVMode() || conf.getApiVersion().isV2()) { CodecDataOutput cdo = new CodecDataOutput(); BytesCodec.writeBytes(cdo, key.toByteArray()); @@ -806,12 +807,22 @@ private Metapb.Region decodeRegion(Metapb.Region region) { .setRegionEpoch(region.getRegionEpoch()) .addAllPeers(region.getPeersList()); - if (conf.getApiVersion().isV1() && (region.getStartKey().isEmpty() || isRawRegion)) { + if (region.getStartKey().isEmpty() || (conf.getApiVersion().isV1() && isRawRegion)) { builder.setStartKey(region.getStartKey()); } else { try { - byte[] decodedStartKey = BytesCodec.readBytes(new CodecDataInput(region.getStartKey())); - builder.setStartKey(ByteString.copyFrom(decodedStartKey)); + ByteString decodedStartKey = + ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getStartKey()))); + if (conf.getApiVersion().isV2()) { + if (ByteString.unsignedLexicographicalComparator() + .compare(decodedStartKey, conf.getKeyPrefix()) + < 0) { + decodedStartKey = ByteString.EMPTY; + } else { + decodedStartKey = decodedStartKey.substring(conf.getKeyPrefix().size()); + } + } + builder.setStartKey(decodedStartKey); } catch (Exception e) { if (!conf.isTest()) { throw e; @@ -820,12 +831,22 @@ private Metapb.Region decodeRegion(Metapb.Region region) { } } - if (conf.getApiVersion().isV1() && (region.getEndKey().isEmpty() || isRawRegion)) { + if (region.getEndKey().isEmpty() || (conf.getApiVersion().isV1() && isRawRegion)) { builder.setEndKey(region.getEndKey()); } else { try { - byte[] decodedEndKey = BytesCodec.readBytes(new CodecDataInput(region.getEndKey())); - builder.setEndKey(ByteString.copyFrom(decodedEndKey)); + ByteString decodedEndKey = + ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getEndKey()))); + if (conf.getApiVersion().isV2()) { + if (ByteString.unsignedLexicographicalComparator() + .compare(decodedEndKey, conf.getEndKey()) + >= 0) { + decodedEndKey = ByteString.EMPTY; + } else { + decodedEndKey = decodedEndKey.substring(conf.getKeyPrefix().size()); + } + } + builder.setEndKey(decodedEndKey); } catch (Exception e) { if (!conf.isTest()) { throw e; diff --git a/src/main/java/org/tikv/common/StoreVersion.java b/src/main/java/org/tikv/common/StoreVersion.java index c23ce496840..046eca61bd2 100644 --- a/src/main/java/org/tikv/common/StoreVersion.java +++ b/src/main/java/org/tikv/common/StoreVersion.java @@ -27,7 +27,6 @@ import org.tikv.kvproto.Metapb; public class StoreVersion { - private static final int SCALE = 10000; private final Logger logger = LoggerFactory.getLogger(this.getClass()); private int v0 = 9999; diff --git a/src/main/java/org/tikv/common/TiConfiguration.java b/src/main/java/org/tikv/common/TiConfiguration.java index fa180d4a3e2..3d32762a25f 100644 --- a/src/main/java/org/tikv/common/TiConfiguration.java +++ b/src/main/java/org/tikv/common/TiConfiguration.java @@ -152,6 +152,7 @@ import static org.tikv.common.ConfigUtils.TiKV_CIRCUIT_BREAK_ENABLE; import static org.tikv.common.ConfigUtils.TiKV_CIRCUIT_BREAK_SLEEP_WINDOW_IN_SECONDS; +import com.google.protobuf.ByteString; import io.grpc.Metadata; import java.io.IOException; import java.io.InputStream; @@ -176,49 +177,14 @@ import org.tikv.kvproto.Kvrpcpb.IsolationLevel; public class TiConfiguration implements Serializable { - public static enum ApiVersion { - V1, - V2; - - public static ApiVersion fromInt(int version) { - switch (version) { - case 1: - return V1; - case 2: - return V2; - default: - throw new IllegalArgumentException("unknown api version " + version); - } - } - - public boolean isV1() { - return this == V1; - } - - public boolean isV2() { - return this == V2; - } - - public Kvrpcpb.APIVersion toPb() { - switch (this) { - case V1: - return Kvrpcpb.APIVersion.V1; - case V2: - return Kvrpcpb.APIVersion.V2; - default: - throw new IllegalArgumentException("unknown api version " + this); - } - } - } - private static final Logger logger = LoggerFactory.getLogger(TiConfiguration.class); private static final ConcurrentHashMap settings = new ConcurrentHashMap<>(); public static final Metadata.Key FORWARD_META_DATA_KEY = Metadata.Key.of("tikv-forwarded-host", Metadata.ASCII_STRING_MARSHALLER); public static final Metadata.Key PD_FORWARD_META_DATA_KEY = Metadata.Key.of("pd-forwarded-host", Metadata.ASCII_STRING_MARSHALLER); - public static final String API_V2_RAW_PREFIX = "r"; - public static final String API_V2_TXN_PREFIX = "x"; + public static final ByteString API_V2_RAW_PREFIX = ByteString.copyFromUtf8("r"); + public static final ByteString API_V2_TXN_PREFIX = ByteString.copyFromUtf8("x"); static { // priority: system environment > config file > default @@ -334,7 +300,7 @@ private static void loadFromDefaultProperties() { } public static void listAll() { - logger.info("static configurations are:" + new ArrayList<>(settings.entrySet()).toString()); + logger.info("static configurations are:" + new ArrayList<>(settings.entrySet())); } private static void set(String key, String value) { @@ -562,7 +528,7 @@ private static ReplicaRead getReplicaRead(String key) { private String jksTrustPath = getOption(TIKV_JKS_TRUST_PATH).orElse(null); private String jksTrustPassword = getOption(TIKV_JKS_TRUST_PASSWORD).orElse(null); - private boolean tiFlashEnable = getBoolean(TIFLASH_ENABLE); + private final boolean tiFlashEnable = getBoolean(TIFLASH_ENABLE); private boolean warmUpEnable = getBoolean(TIKV_WARM_UP_ENABLE); private boolean isTest = false; @@ -1266,4 +1232,94 @@ public TiConfiguration setApiVersion(ApiVersion version) { this.apiVersion = version; return this; } + + public ByteString getKeyPrefix() { + if (apiVersion.isV2()) { + return isRawKVMode() ? API_V2_RAW_PREFIX : API_V2_TXN_PREFIX; + } + return ByteString.EMPTY; + } + + public ByteString getEndKey() { + if (apiVersion.isV2()) { + byte end = (byte) (getKeyPrefix().toByteArray()[0] + 1); + return ByteString.copyFrom(new byte[] {end}); + } + return ByteString.EMPTY; + } + + public ByteString buildRequestKey(ByteString key, boolean isEndKey) { + switch (apiVersion) { + case V1: + return key; + case V2: + if (isEndKey && key.isEmpty()) { + return getEndKey(); + } + return getKeyPrefix().concat(key); + default: + throw new IllegalArgumentException("unknown api version or kv mode"); + } + } + + public ByteString buildRequestKey(ByteString key) { + return buildRequestKey(key, false); + } + + public ByteString unwrapResponseKey(ByteString key) { + return unwrapResponseKey(key, false); + } + + public ByteString unwrapResponseKey(ByteString key, boolean isEndKey) { + switch (apiVersion) { + case V1: + return key; + case V2: + if (isEndKey && key.equals(getEndKey())) { + return ByteString.EMPTY; + } + // TODO: empty key is not always scan back!!!! + if (!key.isEmpty() && !key.startsWith(getKeyPrefix())) { + throw new IllegalArgumentException("key corrupted, wrong prefix"); + } + return key.substring(key.isEmpty() ? 0 : 1); + default: + throw new IllegalArgumentException("unknown api version or kv mode"); + } + } + + public enum ApiVersion { + V1, + V2; + + public static ApiVersion fromInt(int version) { + switch (version) { + case 1: + return V1; + case 2: + return V2; + default: + throw new IllegalArgumentException("unknown api version " + version); + } + } + + public boolean isV1() { + return this == V1; + } + + public boolean isV2() { + return this == V2; + } + + public Kvrpcpb.APIVersion toPb() { + switch (this) { + case V1: + return Kvrpcpb.APIVersion.V1; + case V2: + return Kvrpcpb.APIVersion.V2; + default: + throw new IllegalArgumentException("unknown api version " + this); + } + } + } } diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java index 48960117511..dd6a710564e 100644 --- a/src/main/java/org/tikv/common/TiSession.java +++ b/src/main/java/org/tikv/common/TiSession.java @@ -158,7 +158,8 @@ public TiSession(TiConfiguration conf) { warmup(); } this.circuitBreaker = new CircuitBreakerImpl(conf); - logger.info("TiSession initialized in " + conf.getKvMode() + " mode"); + logger.info("TiSession initialized in " + conf.getKvMode() + " mode in API version: " + + conf.getApiVersion()); } private static VersionInfo getVersionInfo() { @@ -182,10 +183,7 @@ private synchronized void warmup() { try { // let JVM ClassLoader load gRPC error related classes // this operation may cost 100ms - Errorpb.Error.newBuilder() - .setNotLeader(Errorpb.NotLeader.newBuilder().build()) - .build() - .toString(); + Errorpb.Error.newBuilder().setNotLeader(Errorpb.NotLeader.newBuilder().build()).build(); this.client = getPDClient(); this.regionManager = getRegionManager(); diff --git a/src/main/java/org/tikv/common/importer/ImporterClient.java b/src/main/java/org/tikv/common/importer/ImporterClient.java index 054e85caa12..71482a9847c 100644 --- a/src/main/java/org/tikv/common/importer/ImporterClient.java +++ b/src/main/java/org/tikv/common/importer/ImporterClient.java @@ -43,18 +43,19 @@ import org.tikv.common.util.Pair; import org.tikv.kvproto.Errorpb.Error; import org.tikv.kvproto.ImportSstpb; +import org.tikv.kvproto.ImportSstpb.RawWriteBatch; import org.tikv.kvproto.Metapb; public class ImporterClient { private static final Logger logger = LoggerFactory.getLogger(ImporterClient.class); - private TiConfiguration tiConf; - private TiSession tiSession; - private ByteString uuid; - private Key minKey; - private Key maxKey; + private final TiConfiguration tiConf; + private final TiSession tiSession; + private final ByteString uuid; + private final Key minKey; + private final Key maxKey; private TiRegion region; - private Long ttl; + private final Long ttl; private boolean deduplicate = false; @@ -108,9 +109,13 @@ public void write(Iterator> iterator) throws TiKVEx String.format("duplicate key found, key = %s", preKey.toStringUtf8())); } } else { + ByteString key = tiConf.buildRequestKey(pair.first); pairs.add( - ImportSstpb.Pair.newBuilder().setKey(pair.first).setValue(pair.second).build()); - totalBytes += (pair.first.size() + pair.second.size()); + ImportSstpb.Pair.newBuilder() + .setKey(key) + .setValue(pair.second) + .build()); + totalBytes += (key.size() + pair.second.size()); preKey = pair.first; } } @@ -138,10 +143,10 @@ private void init() { long regionId = region.getId(); Metapb.RegionEpoch regionEpoch = region.getRegionEpoch(); ImportSstpb.Range range = - tiConf.isTxnKVMode() + tiConf.isTxnKVMode() || tiConf.getApiVersion().isV2() ? ImportSstpb.Range.newBuilder() - .setStart(encode(minKey.toByteString())) - .setEnd(encode(maxKey.toByteString())) + .setStart(encode(tiConf.buildRequestKey(minKey.toByteString()))) + .setEnd(encode(tiConf.buildRequestKey(maxKey.toByteString(), true))) .build() : ImportSstpb.Range.newBuilder() .setStart(minKey.toByteString()) @@ -150,6 +155,7 @@ private void init() { sstMeta = ImportSstpb.SSTMeta.newBuilder() + .setApiVersion(tiConf.getApiVersion().toPb()) .setUuid(uuid) .setRegionId(regionId) .setRegionEpoch(regionEpoch) @@ -216,11 +222,14 @@ private void writeBatch(List pairs) { } else { ImportSstpb.RawWriteBatch batch; - if (ttl == null || ttl <= 0) { - batch = ImportSstpb.RawWriteBatch.newBuilder().addAllPairs(pairs).build(); - } else { - batch = ImportSstpb.RawWriteBatch.newBuilder().addAllPairs(pairs).setTtl(ttl).build(); + RawWriteBatch.Builder batchBuilder = RawWriteBatch.newBuilder().addAllPairs(pairs); + if (ttl != null && ttl > 0) { + batchBuilder.setTtl(ttl); + } + if (tiConf.getApiVersion().isV2()) { + batchBuilder.setTs(tiSession.getTimestamp().getVersion()); } + batch = batchBuilder.build(); ImportSstpb.RawWriteRequest request = ImportSstpb.RawWriteRequest.newBuilder().setBatch(batch).build(); diff --git a/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java b/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java index 09f17c99fb5..9c2225188db 100644 --- a/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java +++ b/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java @@ -30,7 +30,6 @@ import org.tikv.kvproto.Kvrpcpb; public class RawScanIterator extends ScanIterator { - private final BackOffer scanBackOffer; public RawScanIterator( @@ -73,7 +72,9 @@ private boolean endOfScan() { return false; } ByteString lastKey = currentCache.get(index).getKey(); - return !lastKey.isEmpty() && Key.toRawKey(lastKey).compareTo(endKey) >= 0; + boolean a = !lastKey.isEmpty(); + boolean b = Key.toRawKey(lastKey).compareTo(endKey) >= 0; + return a && b; } boolean isCacheDrained() { diff --git a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java index 0a84b8b0103..117615ca10a 100644 --- a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java @@ -34,7 +34,6 @@ import org.slf4j.LoggerFactory; import org.tikv.common.AbstractGRPCClient; import org.tikv.common.TiConfiguration; -import org.tikv.common.TiConfiguration.KVMode; import org.tikv.common.exception.GrpcException; import org.tikv.common.log.SlowLog; import org.tikv.common.log.SlowLogSpan; @@ -156,42 +155,6 @@ public boolean onStoreUnreachable(BackOffer backOffer) { return false; } - public ByteString buildRequestKey(ByteString key) { - switch (conf.getApiVersion()) { - case V1: - return key; - case V2: - if (conf.getKvMode() == KVMode.RAW) { - return ByteString.copyFromUtf8(TiConfiguration.API_V2_RAW_PREFIX).concat(key); - } else if (conf.getKvMode() == KVMode.TXN) { - return ByteString.copyFromUtf8(TiConfiguration.API_V2_TXN_PREFIX).concat(key); - } - default: - throw new IllegalArgumentException("unknown api version or kv mode"); - } - } - - public ByteString unwrapResponseKey(ByteString key) { - switch (conf.getApiVersion()) { - case V1: - return key; - case V2: - if (conf.getKvMode() == KVMode.RAW) { - if (!key.startsWith(ByteString.copyFromUtf8(TiConfiguration.API_V2_RAW_PREFIX))) { - throw new IllegalArgumentException("key corrupted, wrong prefix"); - } - return key.substring(1); - } else if (conf.getKvMode() == KVMode.TXN) { - if (!key.startsWith(ByteString.copyFromUtf8(TiConfiguration.API_V2_TXN_PREFIX))) { - throw new IllegalArgumentException("key corrupted, wrong prefix"); - } - return key.substring(1); - } - default: - throw new IllegalArgumentException("unknown api version or kv mode"); - } - } - private Kvrpcpb.Context addTraceId(Kvrpcpb.Context context, SlowLog slowLog) { if (slowLog.getThresholdMS() < 0) { // disable tikv tracing @@ -207,27 +170,23 @@ private Kvrpcpb.Context addTraceId(Kvrpcpb.Context context, SlowLog slowLog) { .build(); } - protected Kvrpcpb.Context withApiVersion(Kvrpcpb.Context context) { - return Kvrpcpb.Context.newBuilder(context).setApiVersion(conf.getApiVersion().toPb()).build(); - } - protected Kvrpcpb.Context makeContext(TiStoreType storeType, SlowLog slowLog) { Kvrpcpb.Context context = region.getReplicaContext(java.util.Collections.emptySet(), storeType); - return withApiVersion(addTraceId(context, slowLog)); + return addTraceId(context, slowLog); } protected Kvrpcpb.Context makeContext( Set resolvedLocks, TiStoreType storeType, SlowLog slowLog) { Kvrpcpb.Context context = region.getReplicaContext(resolvedLocks, storeType); - return withApiVersion(addTraceId(context, slowLog)); + return addTraceId(context, slowLog); } protected Kvrpcpb.Context makeContext() { - return withApiVersion(region.getLeaderContext()); + return region.getLeaderContext(); } protected Kvrpcpb.Context makeContext(Metapb.Peer peer) { - return withApiVersion(region.getReplicaContext(peer)); + return region.getReplicaContext(peer); } private void updateClientStub() { @@ -336,7 +295,7 @@ private Metapb.Peer switchLeaderStore(BackOffer backOffer) { Kvrpcpb.RawGetRequest rawGetRequest = Kvrpcpb.RawGetRequest.newBuilder() .setContext(makeContext(peer)) - .setKey(buildRequestKey(key)) + .setKey(conf.buildRequestKey(key)) .build(); ListenableFuture task = stub.rawGet(rawGetRequest); responses.add(new SwitchLeaderTask(task, peer)); @@ -398,7 +357,7 @@ private TiStore switchProxyStore(BackOffer backOffer) { Kvrpcpb.RawGetRequest rawGetRequest = Kvrpcpb.RawGetRequest.newBuilder() .setContext(makeContext()) - .setKey(buildRequestKey(key)) + .setKey(conf.buildRequestKey(key)) .build(); ListenableFuture task = MetadataUtils.attachHeaders(stub, header).rawGet(rawGetRequest); diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index fd35f0b4848..1c4547d4ceb 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -124,12 +124,17 @@ // need to be re-split across regions/stores, region info outdated, e.t.c., you // should retry it in an upper client logic (KVClient, TxnClient, e.t.c.) -/** Note that RegionStoreClient itself is not thread-safe */ +/** + * Note that RegionStoreClient itself is not thread-safe + */ public class RegionStoreClient extends AbstractRegionStoreClient { private static final Logger logger = LoggerFactory.getLogger(RegionStoreClient.class); - @VisibleForTesting public final AbstractLockResolverClient lockResolverClient; + @VisibleForTesting + public final AbstractLockResolverClient lockResolverClient; private final TiStoreType storeType; - /** startTS -> List(locks) */ + /** + * startTS -> List(locks) + */ private final Map> resolvedLocks = new HashMap<>(); private final PDClient pdClient; @@ -226,7 +231,7 @@ public synchronized Set getResolvedLocks(Long version) { * @param version key version * @return value * @throws TiClientInternalException TiSpark Client exception, unexpected - * @throws KeyException Key may be locked + * @throws KeyException Key may be locked */ public ByteString get(BackOffer backOffer, ByteString key, long version) throws TiClientInternalException, KeyException { @@ -236,7 +241,7 @@ public ByteString get(BackOffer backOffer, ByteString key, long version) GetRequest.newBuilder() .setContext( makeContext(getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .setKey(buildRequestKey(key)) + .setKey(conf.buildRequestKey(key)) .setVersion(version) .build(); @@ -260,7 +265,7 @@ public ByteString get(BackOffer backOffer, ByteString key, long version) /** * @param resp GetResponse * @throws TiClientInternalException TiSpark Client exception, unexpected - * @throws KeyException Key may be locked + * @throws KeyException Key may be locked */ private void handleGetResponse(GetResponse resp) throws TiClientInternalException, KeyException { if (resp == null) { @@ -275,14 +280,14 @@ private void handleGetResponse(GetResponse resp) throws TiClientInternalExceptio } } - public List batchGet(BackOffer backOffer, Iterable keys, long version) { + public List batchGet(BackOffer backOffer, List keys, long version) { boolean forWrite = false; Supplier request = () -> BatchGetRequest.newBuilder() .setContext( makeContext(getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .addAllKeys(keys) + .addAllKeys(keys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) .setVersion(version) .build(); KVErrorHandler handler = @@ -347,7 +352,7 @@ public List scan( .setContext( makeContext( getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .setStartKey(buildRequestKey(startKey)) + .setStartKey(conf.buildRequestKey(startKey)) .setVersion(version) .setKeyOnly(keyOnly) .setLimit(getConf().getScanBatchSize()) @@ -394,7 +399,7 @@ private List doScan(ScanResponse resp) { KvPair.newBuilder() .setError(kvPair.getError()) .setValue(kvPair.getValue()) - .setKey(unwrapResponseKey(lock.getKey())) + .setKey(conf.unwrapResponseKey(lock.getKey())) .build()); } else { newKvPairs.add(kvPair); @@ -416,15 +421,11 @@ public List scan(BackOffer backOffer, ByteString startKey, long version) * @param startTs startTs of prewrite * @param lockTTL lock ttl * @throws TiClientInternalException TiSpark Client exception, unexpected - * @throws KeyException Key may be locked - * @throws RegionException region error occurs + * @throws KeyException Key may be locked + * @throws RegionException region error occurs */ public void prewrite( - BackOffer backOffer, - ByteString primary, - Iterable mutations, - long startTs, - long lockTTL) + BackOffer backOffer, ByteString primary, List mutations, long startTs, long lockTTL) throws TiClientInternalException, KeyException, RegionException { this.prewrite(backOffer, primary, mutations, startTs, lockTTL, false); } @@ -437,7 +438,7 @@ public void prewrite( public void prewrite( BackOffer bo, ByteString primaryLock, - Iterable mutations, + List mutations, long startTs, long ttl, boolean skipConstraintCheck) @@ -448,15 +449,23 @@ public void prewrite( () -> getIsV4() ? PrewriteRequest.newBuilder() - .setContext(makeContext(storeType, bo.getSlowLog())) - .setStartVersion(startTs) - .setPrimaryLock(primaryLock) - .addAllMutations(mutations) - .setLockTtl(ttl) - .setSkipConstraintCheck(skipConstraintCheck) - .setMinCommitTs(startTs) - .setTxnSize(16) - .build() + .setContext(makeContext(storeType, bo.getSlowLog())) + .setStartVersion(startTs) + .setPrimaryLock(conf.buildRequestKey(primaryLock)) + .addAllMutations( + mutations + .stream() + .map( + mutation -> + Mutation.newBuilder(mutation) + .setKey(conf.buildRequestKey(mutation.getKey())) + .build()) + .collect(Collectors.toList())) + .setLockTtl(ttl) + .setSkipConstraintCheck(skipConstraintCheck) + .setMinCommitTs(startTs) + .setTxnSize(16) + .build() : PrewriteRequest.newBuilder() .setContext(makeContext(storeType, bo.getSlowLog())) .setStartVersion(startTs) @@ -488,8 +497,8 @@ public void prewrite( * @param backOffer backOffer * @param resp response * @return Return true means the rpc call success. Return false means the rpc call fail, - * RegionStoreClient should retry. Throw an Exception means the rpc call fail, - * RegionStoreClient cannot handle this kind of error + * RegionStoreClient should retry. Throw an Exception means the rpc call fail, RegionStoreClient + * cannot handle this kind of error * @throws TiClientInternalException * @throws RegionException * @throws KeyException @@ -531,7 +540,9 @@ private boolean isPrewriteSuccess(BackOffer backOffer, PrewriteResponse resp, lo return false; } - /** TXN Heart Beat: update primary key ttl */ + /** + * TXN Heart Beat: update primary key ttl + */ public void txnHeartBeat(BackOffer bo, ByteString primaryLock, long startTs, long ttl) { boolean forWrite = false; while (true) { @@ -587,7 +598,7 @@ private boolean isTxnHeartBeatSuccess(TxnHeartBeatResponse resp) * @param startTs start version * @param commitTs commit version */ - public void commit(BackOffer backOffer, Iterable keys, long startTs, long commitTs) + public void commit(BackOffer backOffer, List keys, long startTs, long commitTs) throws KeyException { boolean forWrite = true; Supplier factory = @@ -595,7 +606,7 @@ public void commit(BackOffer backOffer, Iterable keys, long startTs, CommitRequest.newBuilder() .setStartVersion(startTs) .setCommitVersion(commitTs) - .addAllKeys(keys) + .addAllKeys(keys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) .setContext(makeContext(storeType, backOffer.getSlowLog())) .build(); KVErrorHandler handler = @@ -817,12 +828,13 @@ public Iterator coprocessStreaming( * @param splitKeys is the split points for a specific region. * @return a split region info. */ - public List splitRegion(Iterable splitKeys) { + public List splitRegion(List splitKeys) { Supplier request = () -> SplitRegionRequest.newBuilder() .setContext(makeContext(storeType, SlowLogEmptyImpl.INSTANCE)) - .addAllSplitKeys(splitKeys) + .addAllSplitKeys( + splitKeys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) .setIsRawKv(conf.isRawKVMode()) .build(); @@ -865,7 +877,7 @@ public Optional rawGet(BackOffer backOffer, ByteString key) { () -> RawGetRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(buildRequestKey(key)) + .setKey(conf.buildRequestKey(key)) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -904,7 +916,7 @@ public Optional rawGetKeyTTL(BackOffer backOffer, ByteString key) { () -> RawGetKeyTTLRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(buildRequestKey(key)) + .setKey(conf.buildRequestKey(key)) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -943,7 +955,7 @@ public void rawDelete(BackOffer backOffer, ByteString key, boolean atomicForCAS) () -> RawDeleteRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(buildRequestKey(key)) + .setKey(conf.buildRequestKey(key)) .setForCas(atomicForCAS) .build(); @@ -981,7 +993,7 @@ public void rawPut( () -> RawPutRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(buildRequestKey(key)) + .setKey(conf.buildRequestKey(key)) .setValue(value) .setTtl(ttl) .setForCas(atomicForCAS) @@ -1025,7 +1037,7 @@ public void rawCompareAndSet( () -> RawCASRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(buildRequestKey(key)) + .setKey(conf.buildRequestKey(key)) .setValue(value) .setPreviousValue(prevValue.orElse(ByteString.EMPTY)) .setPreviousNotExist(!prevValue.isPresent()) @@ -1078,7 +1090,7 @@ public List rawBatchGet(BackOffer backoffer, List keys) { () -> RawBatchGetRequest.newBuilder() .setContext(makeContext(storeType, backoffer.getSlowLog())) - .addAllKeys(keys) + .addAllKeys(keys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -1135,7 +1147,7 @@ public void rawBatchPut(BackOffer backOffer, Batch batch, long ttl, boolean atom for (int i = 0; i < batch.getKeys().size(); i++) { pairs.add( KvPair.newBuilder() - .setKey(buildRequestKey(batch.getKeys().get(i))) + .setKey(conf.buildRequestKey(batch.getKeys().get(i))) .setValue(batch.getValues().get(i)) .build()); } @@ -1167,7 +1179,7 @@ public void rawBatchDelete(BackOffer backoffer, List keys, boolean a () -> RawBatchDeleteRequest.newBuilder() .setContext(makeContext(storeType, backoffer.getSlowLog())) - .addAllKeys(keys) + .addAllKeys(keys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) .setForCas(atomicForCAS) .build(); RegionErrorHandler handler = @@ -1212,7 +1224,8 @@ public List rawScan(BackOffer backOffer, ByteString key, int limit, bool () -> RawScanRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setStartKey(buildRequestKey(key)) + .setStartKey(conf.buildRequestKey(key)) + .setEndKey(conf.getEndKey()) .setKeyOnly(keyOnly) .setLimit(limit) .build(); @@ -1243,11 +1256,11 @@ private List rawScanHelper(RawScanResponse resp) { return resp.getKvsList() .stream() .map( - kvPair -> - KvPair.newBuilder() - .setKey(unwrapResponseKey(kvPair.getKey())) - .setValue(kvPair.getValue()) - .build()) + kvPair -> KvPair.newBuilder() + .setKey(conf.unwrapResponseKey(kvPair.getKey())) + .setValue(kvPair.getValue()) + .build() + ) .collect(Collectors.toList()); } @@ -1266,8 +1279,8 @@ public void rawDeleteRange(BackOffer backOffer, ByteString startKey, ByteString () -> RawDeleteRangeRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setStartKey(buildRequestKey(startKey)) - .setEndKey(buildRequestKey(endKey)) + .setStartKey(conf.buildRequestKey(startKey)) + .setEndKey(conf.buildRequestKey(endKey, true)) .build(); RegionErrorHandler handler = diff --git a/src/main/java/org/tikv/common/region/TiRegion.java b/src/main/java/org/tikv/common/region/TiRegion.java index b24c5f66eb4..3c0ce8e48d4 100644 --- a/src/main/java/org/tikv/common/region/TiRegion.java +++ b/src/main/java/org/tikv/common/region/TiRegion.java @@ -169,6 +169,7 @@ private Kvrpcpb.Context getContext( Kvrpcpb.Context.Builder builder = Kvrpcpb.Context.newBuilder(); builder + .setApiVersion(conf.getApiVersion().toPb()) .setIsolationLevel(this.isolationLevel) .setPriority(this.commandPri) .setRegionId(meta.getId()) diff --git a/src/main/java/org/tikv/txn/LockResolverClientV4.java b/src/main/java/org/tikv/txn/LockResolverClientV4.java index 1aaa0221cb4..7f428cfdb42 100644 --- a/src/main/java/org/tikv/txn/LockResolverClientV4.java +++ b/src/main/java/org/tikv/txn/LockResolverClientV4.java @@ -168,7 +168,7 @@ private void resolvePessimisticLock(BackOffer bo, Lock lock, Set cl () -> Kvrpcpb.PessimisticRollbackRequest.newBuilder() .setContext(makeContext()) - .addKeys(buildRequestKey(lock.getKey())) + .addKeys(conf.buildRequestKey(lock.getKey())) .setStartVersion(lock.getTxnID()) .setForUpdateTs(forUpdateTS) .build(); @@ -285,8 +285,8 @@ private TxnStatus getTxnStatus( () -> { TiRegion primaryKeyRegion = regionManager.getRegionByKey(primary); return Kvrpcpb.CheckTxnStatusRequest.newBuilder() - .setContext(withApiVersion(primaryKeyRegion.getLeaderContext())) - .setPrimaryKey(buildRequestKey(primary)) + .setContext(primaryKeyRegion.getLeaderContext()) + .setPrimaryKey(conf.buildRequestKey(primary)) .setLockTs(txnID) .setCallerStartTs(callerStartTS) .setCurrentTs(currentTS) @@ -373,7 +373,7 @@ private void resolveLock( if (lock.getTxnSize() < BIG_TXN_THRESHOLD) { // Only resolve specified keys when it is a small transaction, // prevent from scanning the whole region in this case. - builder.addKeys(buildRequestKey(lock.getKey())); + builder.addKeys(conf.buildRequestKey(lock.getKey())); } Supplier factory = builder::build; diff --git a/src/test/java/org/tikv/BaseRawKVTest.java b/src/test/java/org/tikv/BaseRawKVTest.java index 9b8a9042f3c..25b0b777901 100644 --- a/src/test/java/org/tikv/BaseRawKVTest.java +++ b/src/test/java/org/tikv/BaseRawKVTest.java @@ -18,13 +18,14 @@ package org.tikv; import org.tikv.common.PDClient; +import org.tikv.common.StoreApiVersion; import org.tikv.common.StoreVersion; import org.tikv.common.TiConfiguration; +import org.tikv.common.TiConfiguration.ApiVersion; import org.tikv.common.TiSession; import org.tikv.util.TestUtils; public class BaseRawKVTest { - protected boolean tikvVersionNewerThan(String expectedVersion) { TiConfiguration conf = createTiConfiguration(); TiSession session = TiSession.create(conf); @@ -43,6 +44,15 @@ protected TiConfiguration createTiConfiguration() { conf.setEnableAtomicForCAS(true); conf.setEnableGrpcForward(false); conf.setEnableAtomicForCAS(true); + conf.setRawKVScanTimeoutInMS(1000000000); + + conf.setWarmUpEnable(false); + try (TiSession session = TiSession.create(conf)) { + PDClient pdClient = session.getPDClient(); + conf.setApiVersion(ApiVersion.fromInt(StoreApiVersion.acquire(pdClient).get())); + } catch (Exception ignore) { + } + return conf; } } diff --git a/src/test/java/org/tikv/BaseTxnKVTest.java b/src/test/java/org/tikv/BaseTxnKVTest.java index 82fa45135f9..302d44b21cc 100644 --- a/src/test/java/org/tikv/BaseTxnKVTest.java +++ b/src/test/java/org/tikv/BaseTxnKVTest.java @@ -17,7 +17,10 @@ package org.tikv; +import org.tikv.common.StoreApiVersion; import org.tikv.common.TiConfiguration; +import org.tikv.common.TiConfiguration.ApiVersion; +import org.tikv.common.TiSession; import org.tikv.util.TestUtils; public class BaseTxnKVTest { @@ -30,6 +33,12 @@ protected TiConfiguration createTiConfiguration() { : TiConfiguration.createDefault(pdAddrsStr); conf.setTest(true); conf.setEnableGrpcForward(false); + + try (TiSession session = TiSession.create(conf)) { + int version = StoreApiVersion.acquire(session.getPDClient()).get(); + conf.setApiVersion(ApiVersion.fromInt(version)); + } catch (Exception ignore) { + } return conf; } } diff --git a/src/test/java/org/tikv/common/StoreApiVersion.java b/src/test/java/org/tikv/common/StoreApiVersion.java new file mode 100644 index 00000000000..44f1f7791af --- /dev/null +++ b/src/test/java/org/tikv/common/StoreApiVersion.java @@ -0,0 +1,52 @@ +package org.tikv.common; + +import com.google.gson.Gson; +import com.google.gson.JsonObject; +import org.apache.http.HttpEntity; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpGet; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; +import org.apache.http.util.EntityUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.tikv.common.util.BackOffer; +import org.tikv.common.util.ConcreteBackOffer; +import org.tikv.kvproto.Metapb.Store; + +public class StoreApiVersion { + private static final Logger logger = LoggerFactory.getLogger(StoreApiVersion.class); + private final int version; + + private StoreApiVersion(int version) { + this.version = version; + } + + public int get() { + return version; + } + + public static StoreApiVersion acquire(PDClient client) { + int version = 1; + BackOffer backOffer = ConcreteBackOffer.newCustomBackOff(BackOffer.PD_INFO_BACKOFF); + for (Store store : client.getAllStores(backOffer)) { + String statusAddr = store.getStatusAddress(); + String api = "http://" + statusAddr + "/config"; + try (CloseableHttpClient httpClient = HttpClients.createDefault()) { + HttpGet request = new HttpGet(api); + try (CloseableHttpResponse response = httpClient.execute(request)) { + HttpEntity entity = response.getEntity(); + String content = EntityUtils.toString(entity); + JsonObject object = new Gson().fromJson(content, JsonObject.class); + try { + version = object.get("storage").getAsJsonObject().get("api-version").getAsInt(); + } catch (Exception ignored) { + } + } + } catch (Exception e) { + logger.error("Failed to get store api version: ", e); + } + } + return new StoreApiVersion(version); + } +} diff --git a/src/test/java/org/tikv/common/importer/RawKVIngestTest.java b/src/test/java/org/tikv/common/importer/RawKVIngestTest.java index a8e5f94cf08..a76afffeb6e 100644 --- a/src/test/java/org/tikv/common/importer/RawKVIngestTest.java +++ b/src/test/java/org/tikv/common/importer/RawKVIngestTest.java @@ -34,11 +34,11 @@ import org.tikv.common.TiSession; import org.tikv.common.key.Key; import org.tikv.common.util.Pair; +import org.tikv.kvproto.Kvrpcpb.KvPair; import org.tikv.raw.RawKVClient; import org.tikv.util.TestUtils; public class RawKVIngestTest extends BaseRawKVTest { - private TiSession session; private static final int KEY_NUMBER = 16; diff --git a/src/test/java/org/tikv/raw/RawKVClientTest.java b/src/test/java/org/tikv/raw/RawKVClientTest.java index d4ad7bdbe06..faa17a5f73e 100644 --- a/src/test/java/org/tikv/raw/RawKVClientTest.java +++ b/src/test/java/org/tikv/raw/RawKVClientTest.java @@ -325,7 +325,7 @@ public void batchPutTest() { @Test public void deleteRangeTest() { - client.deleteRange(ByteString.EMPTY, ByteString.EMPTY); + checkDeleteRange(ByteString.EMPTY, ByteString.EMPTY); } @Test @@ -534,7 +534,6 @@ private void baseTest( rawDeleteTest(deleteCases, benchmark); } - // TODO: check whether cluster supports ttl long ttl = 10; rawTTLTest(10, ttl, benchmark); From 8a11a888362d24b1be7396e9cd3918f1d7bba4b5 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 12 Apr 2022 11:27:12 +0800 Subject: [PATCH 07/32] fix ingest interface for rawkv Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/TiSession.java | 7 +- .../tikv/common/importer/ImporterClient.java | 6 +- .../tikv/common/region/RegionStoreClient.java | 71 +++++++++---------- .../tikv/common/importer/RawKVIngestTest.java | 1 - 4 files changed, 38 insertions(+), 47 deletions(-) diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java index dd6a710564e..35d26508888 100644 --- a/src/main/java/org/tikv/common/TiSession.java +++ b/src/main/java/org/tikv/common/TiSession.java @@ -158,8 +158,11 @@ public TiSession(TiConfiguration conf) { warmup(); } this.circuitBreaker = new CircuitBreakerImpl(conf); - logger.info("TiSession initialized in " + conf.getKvMode() + " mode in API version: " - + conf.getApiVersion()); + logger.info( + "TiSession initialized in " + + conf.getKvMode() + + " mode in API version: " + + conf.getApiVersion()); } private static VersionInfo getVersionInfo() { diff --git a/src/main/java/org/tikv/common/importer/ImporterClient.java b/src/main/java/org/tikv/common/importer/ImporterClient.java index 71482a9847c..ac048ff53bd 100644 --- a/src/main/java/org/tikv/common/importer/ImporterClient.java +++ b/src/main/java/org/tikv/common/importer/ImporterClient.java @@ -110,11 +110,7 @@ public void write(Iterator> iterator) throws TiKVEx } } else { ByteString key = tiConf.buildRequestKey(pair.first); - pairs.add( - ImportSstpb.Pair.newBuilder() - .setKey(key) - .setValue(pair.second) - .build()); + pairs.add(ImportSstpb.Pair.newBuilder().setKey(key).setValue(pair.second).build()); totalBytes += (key.size() + pair.second.size()); preKey = pair.first; } diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index 1c4547d4ceb..2e94722dc0e 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -124,17 +124,12 @@ // need to be re-split across regions/stores, region info outdated, e.t.c., you // should retry it in an upper client logic (KVClient, TxnClient, e.t.c.) -/** - * Note that RegionStoreClient itself is not thread-safe - */ +/** Note that RegionStoreClient itself is not thread-safe */ public class RegionStoreClient extends AbstractRegionStoreClient { private static final Logger logger = LoggerFactory.getLogger(RegionStoreClient.class); - @VisibleForTesting - public final AbstractLockResolverClient lockResolverClient; + @VisibleForTesting public final AbstractLockResolverClient lockResolverClient; private final TiStoreType storeType; - /** - * startTS -> List(locks) - */ + /** startTS -> List(locks) */ private final Map> resolvedLocks = new HashMap<>(); private final PDClient pdClient; @@ -231,7 +226,7 @@ public synchronized Set getResolvedLocks(Long version) { * @param version key version * @return value * @throws TiClientInternalException TiSpark Client exception, unexpected - * @throws KeyException Key may be locked + * @throws KeyException Key may be locked */ public ByteString get(BackOffer backOffer, ByteString key, long version) throws TiClientInternalException, KeyException { @@ -265,7 +260,7 @@ public ByteString get(BackOffer backOffer, ByteString key, long version) /** * @param resp GetResponse * @throws TiClientInternalException TiSpark Client exception, unexpected - * @throws KeyException Key may be locked + * @throws KeyException Key may be locked */ private void handleGetResponse(GetResponse resp) throws TiClientInternalException, KeyException { if (resp == null) { @@ -421,8 +416,8 @@ public List scan(BackOffer backOffer, ByteString startKey, long version) * @param startTs startTs of prewrite * @param lockTTL lock ttl * @throws TiClientInternalException TiSpark Client exception, unexpected - * @throws KeyException Key may be locked - * @throws RegionException region error occurs + * @throws KeyException Key may be locked + * @throws RegionException region error occurs */ public void prewrite( BackOffer backOffer, ByteString primary, List mutations, long startTs, long lockTTL) @@ -449,23 +444,23 @@ public void prewrite( () -> getIsV4() ? PrewriteRequest.newBuilder() - .setContext(makeContext(storeType, bo.getSlowLog())) - .setStartVersion(startTs) - .setPrimaryLock(conf.buildRequestKey(primaryLock)) - .addAllMutations( - mutations - .stream() - .map( - mutation -> - Mutation.newBuilder(mutation) - .setKey(conf.buildRequestKey(mutation.getKey())) - .build()) - .collect(Collectors.toList())) - .setLockTtl(ttl) - .setSkipConstraintCheck(skipConstraintCheck) - .setMinCommitTs(startTs) - .setTxnSize(16) - .build() + .setContext(makeContext(storeType, bo.getSlowLog())) + .setStartVersion(startTs) + .setPrimaryLock(conf.buildRequestKey(primaryLock)) + .addAllMutations( + mutations + .stream() + .map( + mutation -> + Mutation.newBuilder(mutation) + .setKey(conf.buildRequestKey(mutation.getKey())) + .build()) + .collect(Collectors.toList())) + .setLockTtl(ttl) + .setSkipConstraintCheck(skipConstraintCheck) + .setMinCommitTs(startTs) + .setTxnSize(16) + .build() : PrewriteRequest.newBuilder() .setContext(makeContext(storeType, bo.getSlowLog())) .setStartVersion(startTs) @@ -497,8 +492,8 @@ public void prewrite( * @param backOffer backOffer * @param resp response * @return Return true means the rpc call success. Return false means the rpc call fail, - * RegionStoreClient should retry. Throw an Exception means the rpc call fail, RegionStoreClient - * cannot handle this kind of error + * RegionStoreClient should retry. Throw an Exception means the rpc call fail, + * RegionStoreClient cannot handle this kind of error * @throws TiClientInternalException * @throws RegionException * @throws KeyException @@ -540,9 +535,7 @@ private boolean isPrewriteSuccess(BackOffer backOffer, PrewriteResponse resp, lo return false; } - /** - * TXN Heart Beat: update primary key ttl - */ + /** TXN Heart Beat: update primary key ttl */ public void txnHeartBeat(BackOffer bo, ByteString primaryLock, long startTs, long ttl) { boolean forWrite = false; while (true) { @@ -1256,11 +1249,11 @@ private List rawScanHelper(RawScanResponse resp) { return resp.getKvsList() .stream() .map( - kvPair -> KvPair.newBuilder() - .setKey(conf.unwrapResponseKey(kvPair.getKey())) - .setValue(kvPair.getValue()) - .build() - ) + kvPair -> + KvPair.newBuilder() + .setKey(conf.unwrapResponseKey(kvPair.getKey())) + .setValue(kvPair.getValue()) + .build()) .collect(Collectors.toList()); } diff --git a/src/test/java/org/tikv/common/importer/RawKVIngestTest.java b/src/test/java/org/tikv/common/importer/RawKVIngestTest.java index a76afffeb6e..e25567dd32f 100644 --- a/src/test/java/org/tikv/common/importer/RawKVIngestTest.java +++ b/src/test/java/org/tikv/common/importer/RawKVIngestTest.java @@ -34,7 +34,6 @@ import org.tikv.common.TiSession; import org.tikv.common.key.Key; import org.tikv.common.util.Pair; -import org.tikv.kvproto.Kvrpcpb.KvPair; import org.tikv.raw.RawKVClient; import org.tikv.util.TestUtils; From f9f4cfe004cfdb212c89fc6421b9a976ef39ec3b Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 19 Apr 2022 21:48:39 +0800 Subject: [PATCH 08/32] add api version test Signed-off-by: iosmanthus --- src/test/java/org/tikv/BaseRawKVTest.java | 5 +- src/test/java/org/tikv/BaseTxnKVTest.java | 6 +- .../java/org/tikv/common/ApiVersionTest.java | 136 ++++++++++++++++++ .../java/org/tikv/common/StoreApiVersion.java | 52 ------- .../java/org/tikv/common/StoreConfig.java | 54 +++++++ .../java/org/tikv/raw/RawKVClientTest.java | 3 +- 6 files changed, 196 insertions(+), 60 deletions(-) create mode 100644 src/test/java/org/tikv/common/ApiVersionTest.java delete mode 100644 src/test/java/org/tikv/common/StoreApiVersion.java create mode 100644 src/test/java/org/tikv/common/StoreConfig.java diff --git a/src/test/java/org/tikv/BaseRawKVTest.java b/src/test/java/org/tikv/BaseRawKVTest.java index 25b0b777901..f2d12dc5cbd 100644 --- a/src/test/java/org/tikv/BaseRawKVTest.java +++ b/src/test/java/org/tikv/BaseRawKVTest.java @@ -18,10 +18,9 @@ package org.tikv; import org.tikv.common.PDClient; -import org.tikv.common.StoreApiVersion; +import org.tikv.common.StoreConfig; import org.tikv.common.StoreVersion; import org.tikv.common.TiConfiguration; -import org.tikv.common.TiConfiguration.ApiVersion; import org.tikv.common.TiSession; import org.tikv.util.TestUtils; @@ -49,7 +48,7 @@ protected TiConfiguration createTiConfiguration() { conf.setWarmUpEnable(false); try (TiSession session = TiSession.create(conf)) { PDClient pdClient = session.getPDClient(); - conf.setApiVersion(ApiVersion.fromInt(StoreApiVersion.acquire(pdClient).get())); + conf.setApiVersion(StoreConfig.acquireApiVersion(pdClient)); } catch (Exception ignore) { } diff --git a/src/test/java/org/tikv/BaseTxnKVTest.java b/src/test/java/org/tikv/BaseTxnKVTest.java index 302d44b21cc..965cd61a1d5 100644 --- a/src/test/java/org/tikv/BaseTxnKVTest.java +++ b/src/test/java/org/tikv/BaseTxnKVTest.java @@ -17,9 +17,8 @@ package org.tikv; -import org.tikv.common.StoreApiVersion; +import org.tikv.common.StoreConfig; import org.tikv.common.TiConfiguration; -import org.tikv.common.TiConfiguration.ApiVersion; import org.tikv.common.TiSession; import org.tikv.util.TestUtils; @@ -35,8 +34,7 @@ protected TiConfiguration createTiConfiguration() { conf.setEnableGrpcForward(false); try (TiSession session = TiSession.create(conf)) { - int version = StoreApiVersion.acquire(session.getPDClient()).get(); - conf.setApiVersion(ApiVersion.fromInt(version)); + conf.setApiVersion(StoreConfig.acquireApiVersion(session.getPDClient())); } catch (Exception ignore) { } return conf; diff --git a/src/test/java/org/tikv/common/ApiVersionTest.java b/src/test/java/org/tikv/common/ApiVersionTest.java new file mode 100644 index 00000000000..0aecaf3becf --- /dev/null +++ b/src/test/java/org/tikv/common/ApiVersionTest.java @@ -0,0 +1,136 @@ +package org.tikv.common; + +import com.google.protobuf.ByteString; +import java.util.Optional; +import org.junit.Assert; +import org.junit.Assume; +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.tikv.common.TiConfiguration.ApiVersion; +import org.tikv.raw.RawKVClient; + +public class ApiVersionTest { + private static final Logger logger = LoggerFactory.getLogger(ApiVersionTest.class); + + private TiConfiguration createConfiguration() { + TiConfiguration conf = TiConfiguration.createRawDefault(); + + conf.setTest(true); + conf.setEnableAtomicForCAS(true); + conf.setWarmUpEnable(false); + + return conf; + } + + private RawKVClient createRawClient(ApiVersion apiVersion) { + TiConfiguration conf = createConfiguration(); + conf.setApiVersion(apiVersion); + return TiSession.create(conf).createRawClient(); + } + + private ApiVersion getClusterApiVersion() { + return StoreConfig.acquireApiVersion(TiSession.create(createConfiguration()).getPDClient()); + } + + private boolean getClusterEnabledTtl() { + return StoreConfig.ifTllEnable(TiSession.create(createConfiguration()).getPDClient()); + } + + private boolean minTiKVVersion(String version) { + return StoreVersion.minTiKVVersion( + version, TiSession.create(createConfiguration()).getPDClient()); + } + + @Test + public void testAccessV2Cluster() { + Assume.assumeTrue(getClusterApiVersion().isV2()); + + Assert.assertTrue(getClusterEnabledTtl()); + + // V1 client can't access V2 cluster + RawKVClient client = createRawClient(ApiVersion.V1); + try { + client.get(ByteString.EMPTY); + Assert.fail("Should not be able to access V2 cluster with V1 client"); + } catch (Exception e) { + Assert.assertNotNull(e); + } + + try { + client.put(ByteString.EMPTY, ByteString.EMPTY, 10); + Assert.fail("Should not be able to access V2 cluster with V1 client using TTL"); + } catch (Exception e) { + Assert.assertNotNull(e); + } + + // V2 client can access V2 cluster + client = createRawClient(ApiVersion.V2); + client.putIfAbsent(ByteString.EMPTY, ByteString.EMPTY); + client.put(ByteString.EMPTY, ByteString.EMPTY, 10); + Optional result = client.get(ByteString.EMPTY); + Assert.assertTrue(result.isPresent()); + result.ifPresent(value -> Assert.assertEquals(ByteString.EMPTY, value)); + + client.delete(ByteString.EMPTY); + } + + @Test + public void testAccessV1Cluster() { + Assume.assumeTrue(minTiKVVersion("6.0.0")); + Assume.assumeTrue(getClusterApiVersion().isV1()); + Assume.assumeFalse(getClusterEnabledTtl()); + + // V1 client can access V1 cluster's raw data, no ttl allowed + RawKVClient client = createRawClient(ApiVersion.V1); + client.put(ByteString.EMPTY, ByteString.EMPTY); + Optional result = client.get(ByteString.EMPTY); + Assert.assertTrue(result.isPresent()); + result.ifPresent(value -> Assert.assertEquals(ByteString.EMPTY, value)); + client.delete(ByteString.EMPTY); + + try { + client.put(ByteString.EMPTY, ByteString.EMPTY, 10); + Assert.fail("Should not be able to access V1 cluster without TTL"); + } catch (Exception e) { + Assert.assertNotNull(e); + } + + // V2 client can't access V1 cluster + client = createRawClient(ApiVersion.V2); + try { + client.put(ByteString.EMPTY, ByteString.EMPTY); + Assert.fail("Should not be able to access V1 cluster with V2 Client"); + } catch (Exception e) { + Assert.assertNotNull(e); + } + } + + @Test + public void testAccessV1ClusterWithTtl() throws InterruptedException { + Assume.assumeTrue(minTiKVVersion("6.0.0")); + Assume.assumeTrue(getClusterApiVersion().isV1()); + Assume.assumeTrue(getClusterEnabledTtl()); + + // V1 client can access V1 cluster's raw data, ttl allowed + RawKVClient client = createRawClient(ApiVersion.V1); + client.put(ByteString.EMPTY, ByteString.EMPTY, 5); + Optional result = client.get(ByteString.EMPTY); + Assert.assertTrue(result.isPresent()); + result.ifPresent(value -> Assert.assertEquals(ByteString.EMPTY, value)); + + logger.info("Waiting for ttl to expire"); + Thread.sleep(5000); + + Assert.assertFalse(client.get(ByteString.EMPTY).isPresent()); + + // V2 client can't access V1 cluster with TTL + client = createRawClient(ApiVersion.V2); + try { + client.put(ByteString.EMPTY, ByteString.EMPTY, 5); + Assert.fail("Should not be able to access V1 cluster with TTL with V2 Client"); + } catch (Exception e) { + Assert.assertNotNull(e); + } + } +} diff --git a/src/test/java/org/tikv/common/StoreApiVersion.java b/src/test/java/org/tikv/common/StoreApiVersion.java deleted file mode 100644 index 44f1f7791af..00000000000 --- a/src/test/java/org/tikv/common/StoreApiVersion.java +++ /dev/null @@ -1,52 +0,0 @@ -package org.tikv.common; - -import com.google.gson.Gson; -import com.google.gson.JsonObject; -import org.apache.http.HttpEntity; -import org.apache.http.client.methods.CloseableHttpResponse; -import org.apache.http.client.methods.HttpGet; -import org.apache.http.impl.client.CloseableHttpClient; -import org.apache.http.impl.client.HttpClients; -import org.apache.http.util.EntityUtils; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.tikv.common.util.BackOffer; -import org.tikv.common.util.ConcreteBackOffer; -import org.tikv.kvproto.Metapb.Store; - -public class StoreApiVersion { - private static final Logger logger = LoggerFactory.getLogger(StoreApiVersion.class); - private final int version; - - private StoreApiVersion(int version) { - this.version = version; - } - - public int get() { - return version; - } - - public static StoreApiVersion acquire(PDClient client) { - int version = 1; - BackOffer backOffer = ConcreteBackOffer.newCustomBackOff(BackOffer.PD_INFO_BACKOFF); - for (Store store : client.getAllStores(backOffer)) { - String statusAddr = store.getStatusAddress(); - String api = "http://" + statusAddr + "/config"; - try (CloseableHttpClient httpClient = HttpClients.createDefault()) { - HttpGet request = new HttpGet(api); - try (CloseableHttpResponse response = httpClient.execute(request)) { - HttpEntity entity = response.getEntity(); - String content = EntityUtils.toString(entity); - JsonObject object = new Gson().fromJson(content, JsonObject.class); - try { - version = object.get("storage").getAsJsonObject().get("api-version").getAsInt(); - } catch (Exception ignored) { - } - } - } catch (Exception e) { - logger.error("Failed to get store api version: ", e); - } - } - return new StoreApiVersion(version); - } -} diff --git a/src/test/java/org/tikv/common/StoreConfig.java b/src/test/java/org/tikv/common/StoreConfig.java new file mode 100644 index 00000000000..97d345c83a8 --- /dev/null +++ b/src/test/java/org/tikv/common/StoreConfig.java @@ -0,0 +1,54 @@ +package org.tikv.common; + +import com.google.gson.Gson; +import com.google.gson.JsonObject; +import java.util.List; +import org.apache.http.HttpEntity; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpGet; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; +import org.apache.http.util.EntityUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.tikv.common.TiConfiguration.ApiVersion; +import org.tikv.common.util.BackOffer; +import org.tikv.common.util.ConcreteBackOffer; +import org.tikv.kvproto.Metapb.Store; + +public class StoreConfig { + private static final Logger logger = LoggerFactory.getLogger(StoreConfig.class); + + private static JsonObject getConfig(PDClient client) { + BackOffer backOffer = ConcreteBackOffer.newCustomBackOff(BackOffer.PD_INFO_BACKOFF); + List stores = client.getAllStores(backOffer); + if (stores.isEmpty()) { + throw new IllegalStateException("No store found"); + } + + Store store = stores.get(0); + String statusAddr = store.getStatusAddress(); + String api = "http://" + statusAddr + "/config"; + try (CloseableHttpClient httpClient = HttpClients.createDefault()) { + HttpGet request = new HttpGet(api); + try (CloseableHttpResponse response = httpClient.execute(request)) { + HttpEntity entity = response.getEntity(); + String content = EntityUtils.toString(entity); + return new Gson().fromJson(content, JsonObject.class); + } + } catch (Exception e) { + logger.error("Failed to get store api version: ", e); + throw new IllegalStateException(e); + } + } + + public static ApiVersion acquireApiVersion(PDClient client) { + return getConfig(client).get("storage").getAsJsonObject().get("api-version").getAsInt() == 1 + ? ApiVersion.V1 + : ApiVersion.V2; + } + + public static boolean ifTllEnable(PDClient client) { + return getConfig(client).get("storage").getAsJsonObject().get("enable-ttl").getAsBoolean(); + } +} diff --git a/src/test/java/org/tikv/raw/RawKVClientTest.java b/src/test/java/org/tikv/raw/RawKVClientTest.java index e157dd8382c..ee20df3ee55 100644 --- a/src/test/java/org/tikv/raw/RawKVClientTest.java +++ b/src/test/java/org/tikv/raw/RawKVClientTest.java @@ -50,6 +50,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.tikv.BaseRawKVTest; +import org.tikv.common.StoreConfig; import org.tikv.common.TiConfiguration; import org.tikv.common.TiSession; import org.tikv.common.codec.KeyUtils; @@ -899,7 +900,7 @@ private void rawDeleteRangeTest(boolean benchmark) { } public void rawTTLTest(int cases, long ttl, boolean benchmark) { - Assume.assumeTrue(tikvVersionNewerThan("v5.0.0")); + Assume.assumeTrue(StoreConfig.ifTllEnable(session.getPDClient())); logger.info("ttl testing"); if (benchmark) { for (int i = 0; i < cases; i++) { From e5ff7e03c9d96e501e8cd7b86ba84d1c29a20fc8 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 19 Apr 2022 21:50:44 +0800 Subject: [PATCH 09/32] add license header for new files Signed-off-by: iosmanthus --- .../java/org/tikv/common/ApiVersionTest.java | 17 +++++++++++++++++ src/test/java/org/tikv/common/StoreConfig.java | 17 +++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/src/test/java/org/tikv/common/ApiVersionTest.java b/src/test/java/org/tikv/common/ApiVersionTest.java index 0aecaf3becf..b6ab167f77f 100644 --- a/src/test/java/org/tikv/common/ApiVersionTest.java +++ b/src/test/java/org/tikv/common/ApiVersionTest.java @@ -1,3 +1,20 @@ +/* + * 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 com.google.protobuf.ByteString; diff --git a/src/test/java/org/tikv/common/StoreConfig.java b/src/test/java/org/tikv/common/StoreConfig.java index 97d345c83a8..eb6f0d9278e 100644 --- a/src/test/java/org/tikv/common/StoreConfig.java +++ b/src/test/java/org/tikv/common/StoreConfig.java @@ -1,3 +1,20 @@ +/* + * 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 com.google.gson.Gson; From 3e63dc4a4695552c2e7cf848d6a7de15e7b60028 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Wed, 20 Apr 2022 13:08:47 +0800 Subject: [PATCH 10/32] fix ApiVersionTest for older version of TiKV Signed-off-by: iosmanthus --- src/test/java/org/tikv/common/StoreConfig.java | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/test/java/org/tikv/common/StoreConfig.java b/src/test/java/org/tikv/common/StoreConfig.java index eb6f0d9278e..48acb7880d9 100644 --- a/src/test/java/org/tikv/common/StoreConfig.java +++ b/src/test/java/org/tikv/common/StoreConfig.java @@ -18,6 +18,7 @@ package org.tikv.common; import com.google.gson.Gson; +import com.google.gson.JsonElement; import com.google.gson.JsonObject; import java.util.List; import org.apache.http.HttpEntity; @@ -60,12 +61,12 @@ private static JsonObject getConfig(PDClient client) { } public static ApiVersion acquireApiVersion(PDClient client) { - return getConfig(client).get("storage").getAsJsonObject().get("api-version").getAsInt() == 1 - ? ApiVersion.V1 - : ApiVersion.V2; + JsonElement version = getConfig(client).get("storage").getAsJsonObject().get("api-version"); + return version == null ? ApiVersion.V1 : ApiVersion.fromInt(version.getAsInt()); } public static boolean ifTllEnable(PDClient client) { - return getConfig(client).get("storage").getAsJsonObject().get("enable-ttl").getAsBoolean(); + JsonElement ttlEnabled = getConfig(client).get("storage").getAsJsonObject().get("enable-ttl"); + return ttlEnabled != null && ttlEnabled.getAsBoolean(); } } From 04fa138f89fbfc63d06de2a0d093d04f4989ceb9 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 17 May 2022 18:47:40 +0800 Subject: [PATCH 11/32] fix kvproto wrong commit Signed-off-by: iosmanthus --- dev/proto.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/proto.sh b/dev/proto.sh index 71e8c125530..a988ce8cd51 100755 --- a/dev/proto.sh +++ b/dev/proto.sh @@ -21,7 +21,7 @@ if [ -d $proto_dir ]; then fi repos=("https://github.com/pingcap/kvproto" "https://github.com/pingcap/raft-rs" "https://github.com/pingcap/tipb") -commits=(58f2ac94aa38f49676dd628fbcc1d669a77a62ac b9891b673573fad77ebcf9bbe0969cf945841926 c4d518eb1d60c21f05b028b36729e64610346dac) +commits=(3056ca36e6f2a71a9fc7ba7135e6b119fd977553 b9891b673573fad77ebcf9bbe0969cf945841926 c4d518eb1d60c21f05b028b36729e64610346dac) for i in "${!repos[@]}"; do repo_name=$(basename ${repos[$i]}) From 66d07c4b8d53b211f8d9228b7f388cc5494d7465 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Mon, 23 May 2022 19:11:09 +0800 Subject: [PATCH 12/32] wip: add mock test for error handler for EpochNotMatch region error in api v2 Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/PDClient.java | 3 ++- src/main/java/org/tikv/common/ReadOnlyPDClient.java | 2 ++ .../java/org/tikv/common/operation/RegionErrorHandler.java | 2 ++ src/test/java/org/tikv/common/RegionErrorTest.java | 5 +++++ 4 files changed, 11 insertions(+), 1 deletion(-) create mode 100644 src/test/java/org/tikv/common/RegionErrorTest.java diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index 990e5739bcd..86a18c13bd1 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -801,7 +801,8 @@ public String toString() { } } - private Metapb.Region decodeRegion(Metapb.Region region) { + @Override + public Metapb.Region decodeRegion(Metapb.Region region) { final boolean isRawRegion = conf.isRawKVMode(); Metapb.Region.Builder builder = Metapb.Region.newBuilder() diff --git a/src/main/java/org/tikv/common/ReadOnlyPDClient.java b/src/main/java/org/tikv/common/ReadOnlyPDClient.java index d5a7c3aea87..7ae029b80ef 100644 --- a/src/main/java/org/tikv/common/ReadOnlyPDClient.java +++ b/src/main/java/org/tikv/common/ReadOnlyPDClient.java @@ -64,6 +64,8 @@ List scanRegions( */ Store getStore(BackOffer backOffer, long storeId); + Metapb.Region decodeRegion(Metapb.Region region); + List getAllStores(BackOffer backOffer); TiConfiguration.ReplicaRead getReplicaRead(); diff --git a/src/main/java/org/tikv/common/operation/RegionErrorHandler.java b/src/main/java/org/tikv/common/operation/RegionErrorHandler.java index a809d304412..23ce76101c7 100644 --- a/src/main/java/org/tikv/common/operation/RegionErrorHandler.java +++ b/src/main/java/org/tikv/common/operation/RegionErrorHandler.java @@ -212,6 +212,8 @@ private boolean onRegionEpochNotMatch(BackOffer backOffer, List c List newRegions = new ArrayList<>(currentRegions.size()); // If the region epoch is not ahead of TiKV's, replace region meta in region cache. for (Metapb.Region meta : currentRegions) { + // The region needs to be decoded to plain format. + meta = regionManager.getPDClient().decodeRegion(meta); TiRegion region = regionManager.createRegion(meta, backOffer); newRegions.add(region); if (recv.getRegion().getVerID() == region.getVerID()) { diff --git a/src/test/java/org/tikv/common/RegionErrorTest.java b/src/test/java/org/tikv/common/RegionErrorTest.java new file mode 100644 index 00000000000..d3b3042da3c --- /dev/null +++ b/src/test/java/org/tikv/common/RegionErrorTest.java @@ -0,0 +1,5 @@ +package org.tikv.common; + +public class RegionErrorTest extends MockThreeStoresTest { + +} From 6771d13257c787edd38cc08773ed54fcf8b74984 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 24 May 2022 18:31:00 +0800 Subject: [PATCH 13/32] add mock test for EpochNotMatch in API v2 Signed-off-by: iosmanthus --- .../java/org/tikv/common/KVMockServer.java | 218 +++++++++--------- .../java/org/tikv/common/RegionErrorTest.java | 53 ++++- .../tikv/common/RegionStoreClientTest.java | 23 +- 3 files changed, 179 insertions(+), 115 deletions(-) diff --git a/src/test/java/org/tikv/common/KVMockServer.java b/src/test/java/org/tikv/common/KVMockServer.java index 18de79819ef..c87b8a62450 100644 --- a/src/test/java/org/tikv/common/KVMockServer.java +++ b/src/test/java/org/tikv/common/KVMockServer.java @@ -34,12 +34,14 @@ import io.grpc.stub.StreamObserver; import java.io.IOException; import java.net.ServerSocket; +import java.nio.charset.StandardCharsets; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.SortedMap; import java.util.TreeMap; +import java.util.function.Supplier; import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -49,8 +51,6 @@ import org.tikv.kvproto.Errorpb; import org.tikv.kvproto.Errorpb.EpochNotMatch; import org.tikv.kvproto.Errorpb.Error; -import org.tikv.kvproto.Errorpb.NotLeader; -import org.tikv.kvproto.Errorpb.ServerIsBusy; import org.tikv.kvproto.Kvrpcpb; import org.tikv.kvproto.Kvrpcpb.Context; import org.tikv.kvproto.TikvGrpc; @@ -63,7 +63,9 @@ public class KVMockServer extends TikvGrpc.TikvImplBase { private TiRegion region; private State state = State.Normal; private final TreeMap dataMap = new TreeMap<>(); - private final Map errorMap = new HashMap<>(); + private final Map> regionErrMap = new HashMap<>(); + + private final Map> keyErrMap = new HashMap<>(); // for KV error public static final int ABORT = 1; @@ -111,21 +113,28 @@ public void put(String key, ByteString data) { put(ByteString.copyFromUtf8(key), data); } - public void putError(String key, int code) { - errorMap.put(ByteString.copyFromUtf8(key), code); + public void putError(String key, Supplier builder) { + regionErrMap.put(toRawKey(key.getBytes(StandardCharsets.UTF_8)), builder); } public void clearAllMap() { dataMap.clear(); - errorMap.clear(); + regionErrMap.clear(); } - private void verifyContext(Context context) throws Exception { - if (context.getRegionId() != region.getId() - || !context.getRegionEpoch().equals(region.getRegionEpoch()) - || !context.getPeer().equals(region.getLeader())) { + private Errorpb.Error verifyContext(Context context) throws Exception { + if (context.getRegionId() != region.getId() || !context.getPeer().equals(region.getLeader())) { throw new Exception("context doesn't match"); } + + Errorpb.Error.Builder errBuilder = Errorpb.Error.newBuilder(); + + if (!context.getRegionEpoch().equals(region.getRegionEpoch())) { + return errBuilder + .setEpochNotMatch(EpochNotMatch.newBuilder().addCurrentRegions(region.getMeta()).build()) + .build(); + } + return null; } @Override @@ -138,18 +147,21 @@ public void rawGet( throw new Exception(State.Fail.toString()); default: } - verifyContext(request.getContext()); - ByteString key = request.getKey(); - + Key key = toRawKey(request.getKey()); Kvrpcpb.RawGetResponse.Builder builder = Kvrpcpb.RawGetResponse.newBuilder(); - Integer errorCode = errorMap.remove(key); - Errorpb.Error.Builder errBuilder = Errorpb.Error.newBuilder(); - if (errorCode != null) { - setErrorInfo(errorCode, errBuilder); - builder.setRegionError(errBuilder.build()); + + Error e = verifyContext(request.getContext()); + if (e != null) { + responseObserver.onNext(builder.setRegionError(e).build()); + responseObserver.onCompleted(); + return; + } + + Supplier errProvider = regionErrMap.remove(key); + if (errProvider != null) { + builder.setRegionError(errProvider.get().build()); } else { - Key rawKey = toRawKey(key); - ByteString value = dataMap.get(rawKey); + ByteString value = dataMap.get(key); if (value == null) { value = ByteString.EMPTY; } @@ -158,24 +170,29 @@ public void rawGet( responseObserver.onNext(builder.build()); responseObserver.onCompleted(); } catch (Exception e) { + logger.error("internal error", e); responseObserver.onError(Status.INTERNAL.asRuntimeException()); } } - /** */ + @Override public void rawPut( org.tikv.kvproto.Kvrpcpb.RawPutRequest request, io.grpc.stub.StreamObserver responseObserver) { try { - verifyContext(request.getContext()); - ByteString key = request.getKey(); - + Key key = toRawKey(request.getKey()); Kvrpcpb.RawPutResponse.Builder builder = Kvrpcpb.RawPutResponse.newBuilder(); - Integer errorCode = errorMap.remove(key); - Errorpb.Error.Builder errBuilder = Errorpb.Error.newBuilder(); - if (errorCode != null) { - setErrorInfo(errorCode, errBuilder); - builder.setRegionError(errBuilder.build()); + + Error e = verifyContext(request.getContext()); + if (e != null) { + responseObserver.onNext(builder.setRegionError(e).build()); + responseObserver.onCompleted(); + return; + } + + Supplier errProvider = regionErrMap.remove(key); + if (errProvider != null) { + builder.setRegionError(errProvider.get().build()); } responseObserver.onNext(builder.build()); responseObserver.onCompleted(); @@ -184,40 +201,24 @@ public void rawPut( } } - private void setErrorInfo(int errorCode, Errorpb.Error.Builder errBuilder) { - if (errorCode == NOT_LEADER) { - errBuilder.setNotLeader(Errorpb.NotLeader.getDefaultInstance()); - } else if (errorCode == REGION_NOT_FOUND) { - errBuilder.setRegionNotFound(Errorpb.RegionNotFound.getDefaultInstance()); - } else if (errorCode == KEY_NOT_IN_REGION) { - errBuilder.setKeyNotInRegion(Errorpb.KeyNotInRegion.getDefaultInstance()); - } else if (errorCode == STALE_EPOCH) { - errBuilder.setEpochNotMatch(Errorpb.EpochNotMatch.getDefaultInstance()); - } else if (errorCode == STALE_COMMAND) { - errBuilder.setStaleCommand(Errorpb.StaleCommand.getDefaultInstance()); - } else if (errorCode == SERVER_IS_BUSY) { - errBuilder.setServerIsBusy(Errorpb.ServerIsBusy.getDefaultInstance()); - } else if (errorCode == STORE_NOT_MATCH) { - errBuilder.setStoreNotMatch(Errorpb.StoreNotMatch.getDefaultInstance()); - } else if (errorCode == RAFT_ENTRY_TOO_LARGE) { - errBuilder.setRaftEntryTooLarge(Errorpb.RaftEntryTooLarge.getDefaultInstance()); - } - } - - /** */ + @Override public void rawDelete( org.tikv.kvproto.Kvrpcpb.RawDeleteRequest request, io.grpc.stub.StreamObserver responseObserver) { try { - verifyContext(request.getContext()); - ByteString key = request.getKey(); - + Key key = toRawKey(request.getKey()); Kvrpcpb.RawDeleteResponse.Builder builder = Kvrpcpb.RawDeleteResponse.newBuilder(); - Integer errorCode = errorMap.remove(key); - Errorpb.Error.Builder errBuilder = Errorpb.Error.newBuilder(); - if (errorCode != null) { - setErrorInfo(errorCode, errBuilder); - builder.setRegionError(errBuilder.build()); + + Error e = verifyContext(request.getContext()); + if (e != null) { + responseObserver.onNext(builder.setRegionError(e).build()); + responseObserver.onCompleted(); + return; + } + + Supplier errProvider = regionErrMap.remove(key); + if (errProvider != null) { + builder.setRegionError(errProvider.get().build()); } responseObserver.onNext(builder.build()); responseObserver.onCompleted(); @@ -231,24 +232,24 @@ public void kvGet( org.tikv.kvproto.Kvrpcpb.GetRequest request, io.grpc.stub.StreamObserver responseObserver) { try { - verifyContext(request.getContext()); if (request.getVersion() == 0) { throw new Exception(); } - ByteString key = request.getKey(); - + Key key = toRawKey(request.getKey()); Kvrpcpb.GetResponse.Builder builder = Kvrpcpb.GetResponse.newBuilder(); - Integer errorCode = errorMap.remove(key); - Kvrpcpb.KeyError.Builder errBuilder = Kvrpcpb.KeyError.newBuilder(); - if (errorCode != null) { - if (errorCode == ABORT) { - errBuilder.setAbort("ABORT"); - } else if (errorCode == RETRY) { - errBuilder.setRetryable("Retry"); - } - builder.setError(errBuilder); + + Error e = verifyContext(request.getContext()); + if (e != null) { + responseObserver.onNext(builder.setRegionError(e).build()); + responseObserver.onCompleted(); + return; + } + + Supplier errProvider = keyErrMap.remove(key); + if (errProvider != null) { + builder.setError(errProvider.get().build()); } else { - ByteString value = dataMap.get(toRawKey(key)); + ByteString value = dataMap.get(key); builder.setValue(value); } responseObserver.onNext(builder.build()); @@ -263,23 +264,24 @@ public void kvScan( org.tikv.kvproto.Kvrpcpb.ScanRequest request, io.grpc.stub.StreamObserver responseObserver) { try { - verifyContext(request.getContext()); if (request.getVersion() == 0) { throw new Exception(); } - ByteString key = request.getStartKey(); - + Key key = toRawKey(request.getStartKey()); Kvrpcpb.ScanResponse.Builder builder = Kvrpcpb.ScanResponse.newBuilder(); - Error.Builder errBuilder = Error.newBuilder(); - Integer errorCode = errorMap.remove(key); - if (errorCode != null) { - if (errorCode == ABORT) { - errBuilder.setServerIsBusy(Errorpb.ServerIsBusy.getDefaultInstance()); - } - builder.setRegionError(errBuilder.build()); + + Error e = verifyContext(request.getContext()); + if (e != null) { + responseObserver.onNext(builder.setRegionError(e).build()); + responseObserver.onCompleted(); + return; + } + + Supplier errProvider = regionErrMap.remove(key); + if (errProvider != null) { + builder.setRegionError(errProvider.get().build()); } else { - ByteString startKey = request.getStartKey(); - SortedMap kvs = dataMap.tailMap(toRawKey(startKey)); + SortedMap kvs = dataMap.tailMap(key); builder.addAllPairs( kvs.entrySet() .stream() @@ -303,25 +305,28 @@ public void kvBatchGet( org.tikv.kvproto.Kvrpcpb.BatchGetRequest request, io.grpc.stub.StreamObserver responseObserver) { try { - verifyContext(request.getContext()); if (request.getVersion() == 0) { throw new Exception(); } List keys = request.getKeysList(); Kvrpcpb.BatchGetResponse.Builder builder = Kvrpcpb.BatchGetResponse.newBuilder(); - Error.Builder errBuilder = Error.newBuilder(); + Error e = verifyContext(request.getContext()); + if (e != null) { + responseObserver.onNext(builder.setRegionError(e).build()); + responseObserver.onCompleted(); + return; + } + ImmutableList.Builder resultList = ImmutableList.builder(); for (ByteString key : keys) { - Integer errorCode = errorMap.remove(key); - if (errorCode != null) { - if (errorCode == ABORT) { - errBuilder.setServerIsBusy(Errorpb.ServerIsBusy.getDefaultInstance()); - } - builder.setRegionError(errBuilder.build()); + Key rawKey = toRawKey(key); + Supplier errProvider = regionErrMap.remove(rawKey); + if (errProvider != null) { + builder.setRegionError(errProvider.get().build()); break; } else { - ByteString value = dataMap.get(toRawKey(key)); + ByteString value = dataMap.get(rawKey); resultList.add(Kvrpcpb.KvPair.newBuilder().setKey(key).setValue(value).build()); } } @@ -338,8 +343,6 @@ public void coprocessor( org.tikv.kvproto.Coprocessor.Request requestWrap, io.grpc.stub.StreamObserver responseObserver) { try { - verifyContext(requestWrap.getContext()); - DAGRequest request = DAGRequest.parseFrom(requestWrap.getData()); if (request.getStartTsFallback() == 0) { throw new Exception(); @@ -348,25 +351,22 @@ public void coprocessor( List keyRanges = requestWrap.getRangesList(); Coprocessor.Response.Builder builderWrap = Coprocessor.Response.newBuilder(); - SelectResponse.Builder builder = SelectResponse.newBuilder(); - org.tikv.kvproto.Errorpb.Error.Builder errBuilder = - org.tikv.kvproto.Errorpb.Error.newBuilder(); + Error e = verifyContext(requestWrap.getContext()); + if (e != null) { + responseObserver.onNext(builderWrap.setRegionError(e).build()); + responseObserver.onCompleted(); + return; + } + SelectResponse.Builder builder = SelectResponse.newBuilder(); for (Coprocessor.KeyRange keyRange : keyRanges) { - Integer errorCode = errorMap.remove(keyRange.getStart()); - if (errorCode != null) { - if (STALE_EPOCH == errorCode) { - errBuilder.setEpochNotMatch(EpochNotMatch.getDefaultInstance()); - } else if (NOT_LEADER == errorCode) { - errBuilder.setNotLeader(NotLeader.getDefaultInstance()); - } else { - errBuilder.setServerIsBusy(ServerIsBusy.getDefaultInstance()); - } - builderWrap.setRegionError(errBuilder.build()); + Key startKey = toRawKey(keyRange.getStart()); + Supplier errProvider = regionErrMap.remove(startKey); + if (errProvider != null) { + builderWrap.setRegionError(errProvider.get().build()); break; } else { - ByteString startKey = keyRange.getStart(); - SortedMap kvs = dataMap.tailMap(toRawKey(startKey)); + SortedMap kvs = dataMap.tailMap(startKey); builder.addAllChunks( kvs.entrySet() .stream() diff --git a/src/test/java/org/tikv/common/RegionErrorTest.java b/src/test/java/org/tikv/common/RegionErrorTest.java index d3b3042da3c..07ce3af262e 100644 --- a/src/test/java/org/tikv/common/RegionErrorTest.java +++ b/src/test/java/org/tikv/common/RegionErrorTest.java @@ -1,5 +1,56 @@ package org.tikv.common; +import com.google.protobuf.ByteString; +import java.util.Optional; +import java.util.stream.Collectors; +import org.junit.Assert; +import org.junit.Test; +import org.tikv.common.TiConfiguration.ApiVersion; +import org.tikv.common.codec.Codec.BytesCodec; +import org.tikv.common.codec.CodecDataOutput; +import org.tikv.common.region.TiRegion; +import org.tikv.common.region.TiStore; +import org.tikv.kvproto.Metapb; +import org.tikv.raw.RawKVClient; + public class RegionErrorTest extends MockThreeStoresTest { - + private RawKVClient createClient() { + return TiSession.create(session.getConf().setApiVersion(ApiVersion.V2)).createRawClient(); + } + + private ByteString encode(ByteString key) { + CodecDataOutput cdo = new CodecDataOutput(); + BytesCodec.writeBytes(cdo, key.toByteArray()); + return cdo.toByteString(); + } + + @Test + public void testOnEpochNotMatch() { + try (RawKVClient client = createClient()) { + ByteString key = ByteString.copyFromUtf8("test-epoch-not-match"); + ByteString value = ByteString.copyFromUtf8("value"); + + ByteString requestKey = client.getSession().getConf().buildRequestKey(key); + put(requestKey, value); + + TiRegion newRegion = + new TiRegion( + this.region.getConf(), + Metapb.Region.newBuilder() + .mergeFrom(this.region.getMeta()) + .setRegionEpoch(Metapb.RegionEpoch.newBuilder().setConfVer(2).setVersion(3)) + .build(), + this.region.getLeader(), + this.region.getPeersList(), + stores.stream().map(TiStore::new).collect(Collectors.toList())); + + Assert.assertEquals(Optional.of(value), client.get(key)); + + for (KVMockServer server : servers) { + server.setRegion(newRegion); + } + + Assert.assertEquals(Optional.of(value), client.get(key)); + } + } } diff --git a/src/test/java/org/tikv/common/RegionStoreClientTest.java b/src/test/java/org/tikv/common/RegionStoreClientTest.java index e3aff1243e1..1a03ad80e26 100644 --- a/src/test/java/org/tikv/common/RegionStoreClientTest.java +++ b/src/test/java/org/tikv/common/RegionStoreClientTest.java @@ -32,6 +32,10 @@ import org.tikv.common.region.TiStore; import org.tikv.common.util.BackOffer; import org.tikv.common.util.ConcreteBackOffer; +import org.tikv.kvproto.Errorpb; +import org.tikv.kvproto.Errorpb.EpochNotMatch; +import org.tikv.kvproto.Errorpb.NotLeader; +import org.tikv.kvproto.Errorpb.ServerIsBusy; import org.tikv.kvproto.Kvrpcpb; import org.tikv.kvproto.Metapb; @@ -75,12 +79,15 @@ public void doRawGetTest(RegionStoreClient client) { Optional value = client.rawGet(defaultBackOff(), ByteString.copyFromUtf8("key1")); assertEquals(ByteString.copyFromUtf8("value1"), value.get()); - server.putError("error1", KVMockServer.NOT_LEADER); + server.putError( + "error1", () -> Errorpb.Error.newBuilder().setNotLeader(NotLeader.getDefaultInstance())); // since not_leader is retryable, so the result should be correct. value = client.rawGet(defaultBackOff(), ByteString.copyFromUtf8("key1")); assertEquals(ByteString.copyFromUtf8("value1"), value.get()); - server.putError("failure", KVMockServer.STALE_EPOCH); + server.putError( + "failure", + () -> Errorpb.Error.newBuilder().setEpochNotMatch(EpochNotMatch.getDefaultInstance())); try { // since stale epoch is not retryable, so the test should fail. client.rawGet(defaultBackOff(), ByteString.copyFromUtf8("failure")); @@ -102,7 +109,9 @@ public void doGetTest(RegionStoreClient client) { ByteString value = client.get(defaultBackOff(), ByteString.copyFromUtf8("key1"), 1); assertEquals(ByteString.copyFromUtf8("value1"), value); - server.putError("error1", KVMockServer.ABORT); + server.putError( + "error1", + () -> Errorpb.Error.newBuilder().setServerIsBusy(ServerIsBusy.getDefaultInstance())); try { client.get(defaultBackOff(), ByteString.copyFromUtf8("error1"), 1); fail(); @@ -134,7 +143,9 @@ public void doBatchGetTest(RegionStoreClient client) { assertEquals( kv.getKey().toStringUtf8().replace("key", "value"), kv.getValue().toStringUtf8())); - server.putError("error1", KVMockServer.ABORT); + server.putError( + "error1", + () -> Errorpb.Error.newBuilder().setServerIsBusy(ServerIsBusy.getDefaultInstance())); try { client.batchGet( defaultBackOff(), @@ -165,7 +176,9 @@ public void doScanTest(RegionStoreClient client) { assertEquals( kv.getKey().toStringUtf8().replace("key", "value"), kv.getValue().toStringUtf8())); - server.putError("error1", KVMockServer.ABORT); + server.putError( + "error1", + () -> Errorpb.Error.newBuilder().setServerIsBusy(ServerIsBusy.getDefaultInstance())); try { client.scan(defaultBackOff(), ByteString.copyFromUtf8("error1"), 1); fail(); From b7c54f8036161bdd7f2902e399d64e6a8947fce7 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 24 May 2022 18:34:52 +0800 Subject: [PATCH 14/32] fix license header Signed-off-by: iosmanthus --- .../java/org/tikv/common/RegionErrorTest.java | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/test/java/org/tikv/common/RegionErrorTest.java b/src/test/java/org/tikv/common/RegionErrorTest.java index 07ce3af262e..6990dab842b 100644 --- a/src/test/java/org/tikv/common/RegionErrorTest.java +++ b/src/test/java/org/tikv/common/RegionErrorTest.java @@ -1,3 +1,20 @@ +/* + * 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 com.google.protobuf.ByteString; From 72eb3711ac9014bf1b920b5bea3be6999fd6677e Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 24 May 2022 19:11:20 +0800 Subject: [PATCH 15/32] add some comments for epoch not match mock test Signed-off-by: iosmanthus --- src/test/java/org/tikv/common/RegionErrorTest.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/test/java/org/tikv/common/RegionErrorTest.java b/src/test/java/org/tikv/common/RegionErrorTest.java index 6990dab842b..5724d772fac 100644 --- a/src/test/java/org/tikv/common/RegionErrorTest.java +++ b/src/test/java/org/tikv/common/RegionErrorTest.java @@ -50,6 +50,9 @@ public void testOnEpochNotMatch() { ByteString requestKey = client.getSession().getConf().buildRequestKey(key); put(requestKey, value); + Assert.assertEquals(Optional.of(value), client.get(key)); + + // Increase the region epoch for the cluster TiRegion newRegion = new TiRegion( this.region.getConf(), @@ -61,8 +64,6 @@ public void testOnEpochNotMatch() { this.region.getPeersList(), stores.stream().map(TiStore::new).collect(Collectors.toList())); - Assert.assertEquals(Optional.of(value), client.get(key)); - for (KVMockServer server : servers) { server.setRegion(newRegion); } From 844358827653ecfa07ac8ee288fec72f171054fa Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Wed, 25 May 2022 18:16:47 +0800 Subject: [PATCH 16/32] git fire!: add tests for scanRegions Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/PDClient.java | 28 +++--- src/main/java/org/tikv/common/TiSession.java | 3 +- .../java/org/tikv/common/KVMockServer.java | 71 +++++++++------ .../org/tikv/common/PDClientV2MockTest.java | 88 +++++++++++++++++++ .../java/org/tikv/common/PDMockServer.java | 20 ++++- .../org/tikv/common/PDMockServerTest.java | 17 ++++ .../java/org/tikv/common/RegionErrorTest.java | 40 +++++---- 7 files changed, 211 insertions(+), 56 deletions(-) create mode 100644 src/test/java/org/tikv/common/PDClientV2MockTest.java diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index 33daf4988fe..56a3c3edada 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -219,7 +219,8 @@ public Boolean isCheckerPaused(PDChecker checker) { try { ObjectMapper mapper = new ObjectMapper(); HashMap status = - mapper.readValue(new URL(api), new TypeReference>() {}); + mapper.readValue(new URL(api), new TypeReference>() { + }); return status.get("paused"); } catch (Exception e) { logger.error(String.format("failed to get %s checker status.", checker.apiName()), e); @@ -357,15 +358,18 @@ public List scanRegions( Pdpb.ScanRegionsRequest request = Pdpb.ScanRegionsRequest.newBuilder() .setHeader(header) - .setStartKey(startKey) - .setEndKey(endKey) + .setStartKey(conf.buildRequestKey(startKey)) + .setEndKey(conf.buildRequestKey(startKey, true)) .setLimit(limit) .build(); Pdpb.ScanRegionsResponse resp = stub.scanRegions(request); if (resp == null) { return null; } - return resp.getRegionsList(); + return resp.getRegionsList().stream().map(r -> { + Metapb.Region decoded = decodeRegion(r.getRegion()); + return Pdpb.Region.newBuilder().mergeFrom(r).setRegion(decoded).build(); + }).collect(Collectors.toList()); } private Supplier buildGetStoreReq(long storeId) { @@ -398,12 +402,12 @@ public Store getStore(BackOffer backOffer, long storeId) { @Override public List getAllStores(BackOffer backOffer) { return callWithRetry( - backOffer, - PDGrpc.getGetAllStoresMethod(), - buildGetAllStoresReq(), - new PDErrorHandler<>( - r -> r.getHeader().hasError() ? buildFromPdpbError(r.getHeader().getError()) : null, - this)) + backOffer, + PDGrpc.getGetAllStoresMethod(), + buildGetAllStoresReq(), + new PDErrorHandler<>( + r -> r.getHeader().hasError() ? buildFromPdpbError(r.getHeader().getError()) : null, + this)) .getStoresList(); } @@ -829,7 +833,7 @@ public Metapb.Region decodeRegion(Metapb.Region region) { ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getStartKey()))); if (conf.getApiVersion().isV2()) { if (ByteString.unsignedLexicographicalComparator() - .compare(decodedStartKey, conf.getKeyPrefix()) + .compare(decodedStartKey, conf.getKeyPrefix()) < 0) { decodedStartKey = ByteString.EMPTY; } else { @@ -853,7 +857,7 @@ public Metapb.Region decodeRegion(Metapb.Region region) { ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getEndKey()))); if (conf.getApiVersion().isV2()) { if (ByteString.unsignedLexicographicalComparator() - .compare(decodedEndKey, conf.getEndKey()) + .compare(decodedEndKey, conf.getEndKey()) >= 0) { decodedEndKey = ByteString.EMPTY; } else { diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java index 7c5467cd583..0c38117211a 100644 --- a/src/main/java/org/tikv/common/TiSession.java +++ b/src/main/java/org/tikv/common/TiSession.java @@ -184,7 +184,8 @@ private static VersionInfo getVersionInfo() { return info; } - private synchronized void warmup() { + @VisibleForTesting + public synchronized void warmup() { long warmUpStartTime = System.nanoTime(); BackOffer backOffer = ConcreteBackOffer.newRawKVBackOff(getPDClient().getClusterId()); try { diff --git a/src/test/java/org/tikv/common/KVMockServer.java b/src/test/java/org/tikv/common/KVMockServer.java index c87b8a62450..69d8a55ee0e 100644 --- a/src/test/java/org/tikv/common/KVMockServer.java +++ b/src/test/java/org/tikv/common/KVMockServer.java @@ -157,9 +157,12 @@ public void rawGet( return; } - Supplier errProvider = regionErrMap.remove(key); + Supplier errProvider = regionErrMap.get(key); if (errProvider != null) { - builder.setRegionError(errProvider.get().build()); + Error.Builder eb = errProvider.get(); + if (eb != null) { + builder.setRegionError(eb.build()); + } } else { ByteString value = dataMap.get(key); if (value == null) { @@ -190,10 +193,14 @@ public void rawPut( return; } - Supplier errProvider = regionErrMap.remove(key); + Supplier errProvider = regionErrMap.get(key); if (errProvider != null) { - builder.setRegionError(errProvider.get().build()); + Error.Builder eb = errProvider.get(); + if (eb != null) { + builder.setRegionError(eb.build()); + } } + responseObserver.onNext(builder.build()); responseObserver.onCompleted(); } catch (Exception e) { @@ -216,9 +223,12 @@ public void rawDelete( return; } - Supplier errProvider = regionErrMap.remove(key); + Supplier errProvider = regionErrMap.get(key); if (errProvider != null) { - builder.setRegionError(errProvider.get().build()); + Error.Builder eb = errProvider.get(); + if (eb != null) { + builder.setRegionError(eb.build()); + } } responseObserver.onNext(builder.build()); responseObserver.onCompleted(); @@ -277,9 +287,12 @@ public void kvScan( return; } - Supplier errProvider = regionErrMap.remove(key); + Supplier errProvider = regionErrMap.get(key); if (errProvider != null) { - builder.setRegionError(errProvider.get().build()); + Error.Builder eb = errProvider.get(); + if (eb != null) { + builder.setRegionError(eb.build()); + } } else { SortedMap kvs = dataMap.tailMap(key); builder.addAllPairs( @@ -321,14 +334,17 @@ public void kvBatchGet( ImmutableList.Builder resultList = ImmutableList.builder(); for (ByteString key : keys) { Key rawKey = toRawKey(key); - Supplier errProvider = regionErrMap.remove(rawKey); + Supplier errProvider = regionErrMap.get(rawKey); if (errProvider != null) { - builder.setRegionError(errProvider.get().build()); - break; - } else { - ByteString value = dataMap.get(rawKey); - resultList.add(Kvrpcpb.KvPair.newBuilder().setKey(key).setValue(value).build()); + Error.Builder eb = errProvider.get(); + if (eb != null) { + builder.setRegionError(eb.build()); + break; + } } + + ByteString value = dataMap.get(rawKey); + resultList.add(Kvrpcpb.KvPair.newBuilder().setKey(key).setValue(value).build()); } builder.addAllPairs(resultList.build()); responseObserver.onNext(builder.build()); @@ -361,20 +377,23 @@ public void coprocessor( SelectResponse.Builder builder = SelectResponse.newBuilder(); for (Coprocessor.KeyRange keyRange : keyRanges) { Key startKey = toRawKey(keyRange.getStart()); - Supplier errProvider = regionErrMap.remove(startKey); + Supplier errProvider = regionErrMap.get(startKey); if (errProvider != null) { - builderWrap.setRegionError(errProvider.get().build()); - break; - } else { - SortedMap kvs = dataMap.tailMap(startKey); - builder.addAllChunks( - kvs.entrySet() - .stream() - .filter(Objects::nonNull) - .filter(kv -> kv.getKey().compareTo(toRawKey(keyRange.getEnd())) <= 0) - .map(kv -> Chunk.newBuilder().setRowsData(kv.getValue()).build()) - .collect(Collectors.toList())); + Error.Builder eb = errProvider.get(); + if (eb != null) { + builderWrap.setRegionError(eb.build()); + break; + } } + + SortedMap kvs = dataMap.tailMap(startKey); + builder.addAllChunks( + kvs.entrySet() + .stream() + .filter(Objects::nonNull) + .filter(kv -> kv.getKey().compareTo(toRawKey(keyRange.getEnd())) <= 0) + .map(kv -> Chunk.newBuilder().setRowsData(kv.getValue()).build()) + .collect(Collectors.toList())); } responseObserver.onNext(builderWrap.setData(builder.build().toByteString()).build()); diff --git a/src/test/java/org/tikv/common/PDClientV2MockTest.java b/src/test/java/org/tikv/common/PDClientV2MockTest.java new file mode 100644 index 00000000000..cfd13f36edd --- /dev/null +++ b/src/test/java/org/tikv/common/PDClientV2MockTest.java @@ -0,0 +1,88 @@ +/* + * Copyright 2017 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 com.google.protobuf.ByteString; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.tikv.common.codec.Codec.BytesCodec; +import org.tikv.common.codec.CodecDataOutput; +import org.tikv.common.util.ConcreteBackOffer; +import org.tikv.kvproto.Metapb; + +public class PDClientV2MockTest extends PDClientMockTest { + @Before + public void init() throws Exception { + upgradeToV2Cluster(); + } + + private PDClient createClient() { + return session.getPDClient(); + } + + + public static ByteString encode(ByteString key) { + CodecDataOutput cdo = new CodecDataOutput(); + BytesCodec.writeBytes(cdo, key.toByteArray()); + return cdo.toByteString(); + } + + @Test + public void testGetRegionById() throws Exception { + ByteString start = ByteString.copyFromUtf8("regionById"); + ByteString end = ByteString.copyFromUtf8("regionById0"); + leader.addGetRegionByIDListener( + request -> + GrpcUtils.makeGetRegionResponse( + leader.getClusterId(), + GrpcUtils.makeRegion( + 1, + encode(session.getConf().buildRequestKey(start)), + encode(session.getConf().buildRequestKey(end)), + GrpcUtils.makeRegionEpoch(2, 3), + GrpcUtils.makePeer(1, 10), + GrpcUtils.makePeer(2, 20)))); + try (PDClient client = createClient()) { + Metapb.Region r = client.getRegionByID(ConcreteBackOffer.newRawKVBackOff(), 1).first; + Assert.assertEquals(start, r.getStartKey()); + Assert.assertEquals(end, r.getEndKey()); + } + + leader.addGetRegionByIDListener( + request -> + GrpcUtils.makeGetRegionResponse( + leader.getClusterId(), + GrpcUtils.makeRegion( + 1, + encode(session.getConf().buildRequestKey(start)), + encode(session.getConf().getEndKey()), + GrpcUtils.makeRegionEpoch(2, 3), + GrpcUtils.makePeer(1, 10), + GrpcUtils.makePeer(2, 20)))); + try (PDClient client = createClient()) { + Metapb.Region r = client.getRegionByID(ConcreteBackOffer.newRawKVBackOff(), 1).first; + Assert.assertEquals(start, r.getStartKey()); + Assert.assertEquals(ByteString.EMPTY, r.getEndKey()); + } + } + + @Test + public void testScanRegions() { + } +} diff --git a/src/test/java/org/tikv/common/PDMockServer.java b/src/test/java/org/tikv/common/PDMockServer.java index 7980e0703f3..5bda86612f4 100644 --- a/src/test/java/org/tikv/common/PDMockServer.java +++ b/src/test/java/org/tikv/common/PDMockServer.java @@ -37,6 +37,8 @@ import org.tikv.kvproto.Pdpb.GetRegionResponse; import org.tikv.kvproto.Pdpb.GetStoreRequest; import org.tikv.kvproto.Pdpb.GetStoreResponse; +import org.tikv.kvproto.Pdpb.ScanRegionsRequest; +import org.tikv.kvproto.Pdpb.ScanRegionsResponse; import org.tikv.kvproto.Pdpb.TsoRequest; import org.tikv.kvproto.Pdpb.TsoResponse; @@ -50,6 +52,8 @@ public class PDMockServer extends PDGrpc.PDImplBase { private Function getRegionListener; private Function getRegionByIDListener; + private Function scanRegionsListener; + public void addGetMembersListener(Function func) { getMembersListener = func; } @@ -71,10 +75,12 @@ public StreamObserver tso(StreamObserver resp) { private int logical = 0; @Override - public void onNext(TsoRequest value) {} + public void onNext(TsoRequest value) { + } @Override - public void onError(Throwable t) {} + public void onError(Throwable t) { + } @Override public void onCompleted() { @@ -126,6 +132,16 @@ public void getStore(GetStoreRequest request, StreamObserver r } } + @Override + public void scanRegions(ScanRegionsRequest request, StreamObserver resp) { + try { + resp.onNext(Optional.ofNullable(scanRegionsListener.apply(request)).get()); + resp.onCompleted(); + } catch (Exception e) { + resp.onError(Status.INTERNAL.asRuntimeException()); + } + } + public void start(long clusterId) throws IOException { int port; try (ServerSocket s = new ServerSocket(0)) { diff --git a/src/test/java/org/tikv/common/PDMockServerTest.java b/src/test/java/org/tikv/common/PDMockServerTest.java index 81821885b06..08018d0d495 100644 --- a/src/test/java/org/tikv/common/PDMockServerTest.java +++ b/src/test/java/org/tikv/common/PDMockServerTest.java @@ -22,6 +22,7 @@ import java.util.List; import org.junit.After; import org.junit.Before; +import org.tikv.common.TiConfiguration.ApiVersion; public abstract class PDMockServerTest { protected static final String LOCAL_ADDR = "127.0.0.1"; @@ -35,6 +36,21 @@ public void setup() throws IOException { setup(LOCAL_ADDR); } + void upgradeToV2Cluster() throws Exception { + if (session == null) { + throw new IllegalStateException("Cluster is not initialized"); + } + + if (session.getConf().getApiVersion().isV2()) { + return; + } + + TiConfiguration conf = session.getConf().setApiVersion(ApiVersion.V2); + session.close(); + + session = TiSession.create(conf); + } + void setup(String addr) throws IOException { int[] ports = new int[3]; for (int i = 0; i < ports.length; i++) { @@ -62,6 +78,7 @@ void setup(String addr) throws IOException { conf.setWarmUpEnable(false); conf.setTimeout(2000); conf.setEnableGrpcForward(true); + session = TiSession.create(conf); } diff --git a/src/test/java/org/tikv/common/RegionErrorTest.java b/src/test/java/org/tikv/common/RegionErrorTest.java index 5724d772fac..94f70115abc 100644 --- a/src/test/java/org/tikv/common/RegionErrorTest.java +++ b/src/test/java/org/tikv/common/RegionErrorTest.java @@ -21,30 +21,28 @@ import java.util.Optional; import java.util.stream.Collectors; import org.junit.Assert; +import org.junit.Before; import org.junit.Test; -import org.tikv.common.TiConfiguration.ApiVersion; -import org.tikv.common.codec.Codec.BytesCodec; -import org.tikv.common.codec.CodecDataOutput; import org.tikv.common.region.TiRegion; import org.tikv.common.region.TiStore; import org.tikv.kvproto.Metapb; import org.tikv.raw.RawKVClient; public class RegionErrorTest extends MockThreeStoresTest { - private RawKVClient createClient() { - return TiSession.create(session.getConf().setApiVersion(ApiVersion.V2)).createRawClient(); + @Before + public void init() throws Exception { + upgradeToV2Cluster(); } - private ByteString encode(ByteString key) { - CodecDataOutput cdo = new CodecDataOutput(); - BytesCodec.writeBytes(cdo, key.toByteArray()); - return cdo.toByteString(); + private RawKVClient createClient() { + return session.createRawClient(); } @Test public void testOnEpochNotMatch() { try (RawKVClient client = createClient()) { - ByteString key = ByteString.copyFromUtf8("test-epoch-not-match"); + // Construct a key that is less than the prefix of RAW API v2; + ByteString key = ByteString.copyFromUtf8("key-test-epoch-not-match"); ByteString value = ByteString.copyFromUtf8("value"); ByteString requestKey = client.getSession().getConf().buildRequestKey(key); @@ -52,22 +50,34 @@ public void testOnEpochNotMatch() { Assert.assertEquals(Optional.of(value), client.get(key)); - // Increase the region epoch for the cluster + Metapb.Region newMeta = + Metapb.Region.newBuilder() + .mergeFrom(this.region.getMeta()) + .setRegionEpoch(Metapb.RegionEpoch.newBuilder().setConfVer(2).setVersion(3)) + .setStartKey(PDClientV2MockTest.encode(requestKey)) + .setEndKey(PDClientV2MockTest.encode(requestKey.concat(ByteString.copyFromUtf8("0")))) + .build(); + + // Increase the region epoch for the cluster, + // this will cause the cluster return an EpochNotMatch region error. TiRegion newRegion = new TiRegion( this.region.getConf(), - Metapb.Region.newBuilder() - .mergeFrom(this.region.getMeta()) - .setRegionEpoch(Metapb.RegionEpoch.newBuilder().setConfVer(2).setVersion(3)) - .build(), + newMeta, this.region.getLeader(), this.region.getPeersList(), stores.stream().map(TiStore::new).collect(Collectors.toList())); + // Update the region of each server for (KVMockServer server : servers) { server.setRegion(newRegion); } + // Forbid the client get region from PD leader. + leader.addGetRegionListener(request -> null); + + // The get should success since the region cache + // will be updated the currentRegions of `EpochNotMatch` error. Assert.assertEquals(Optional.of(value), client.get(key)); } } From 62879c6a28c67c85f0fedf77f85e5fd9f6c283e8 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Thu, 26 May 2022 20:19:28 +0800 Subject: [PATCH 17/32] add tests for scanRegions Signed-off-by: iosmanthus --- .../org/tikv/common/PDClientV2MockTest.java | 74 ++++++++++++------- .../java/org/tikv/common/PDMockServer.java | 4 + 2 files changed, 50 insertions(+), 28 deletions(-) diff --git a/src/test/java/org/tikv/common/PDClientV2MockTest.java b/src/test/java/org/tikv/common/PDClientV2MockTest.java index cfd13f36edd..57eb5840ff1 100644 --- a/src/test/java/org/tikv/common/PDClientV2MockTest.java +++ b/src/test/java/org/tikv/common/PDClientV2MockTest.java @@ -18,6 +18,7 @@ package org.tikv.common; import com.google.protobuf.ByteString; +import java.util.List; import org.junit.Assert; import org.junit.Before; import org.junit.Test; @@ -25,6 +26,10 @@ import org.tikv.common.codec.CodecDataOutput; import org.tikv.common.util.ConcreteBackOffer; import org.tikv.kvproto.Metapb; +import org.tikv.kvproto.Pdpb; +import org.tikv.kvproto.Pdpb.GetRegionResponse; +import org.tikv.kvproto.Pdpb.Region; +import org.tikv.kvproto.Pdpb.ScanRegionsResponse; public class PDClientV2MockTest extends PDClientMockTest { @Before @@ -43,46 +48,59 @@ public static ByteString encode(ByteString key) { return cdo.toByteString(); } + private GetRegionResponse makeGetRegionResponse(String start, String end) { + return GrpcUtils.makeGetRegionResponse(leader.getClusterId(), makeRegion(start, end)); + } + + private Metapb.Region makeRegion(String start, String end) { + return GrpcUtils.makeRegion( + 1, + encode(session.getConf().buildRequestKey(ByteString.copyFromUtf8(start))), + encode(session.getConf().buildRequestKey(ByteString.copyFromUtf8(end), true)), + GrpcUtils.makeRegionEpoch(2, 3), + GrpcUtils.makePeer(1, 10), + GrpcUtils.makePeer(2, 20)); + } + @Test public void testGetRegionById() throws Exception { - ByteString start = ByteString.copyFromUtf8("regionById"); - ByteString end = ByteString.copyFromUtf8("regionById0"); + String start = "getRegionById"; + String end = "getRegionByIdEnd"; leader.addGetRegionByIDListener( - request -> - GrpcUtils.makeGetRegionResponse( - leader.getClusterId(), - GrpcUtils.makeRegion( - 1, - encode(session.getConf().buildRequestKey(start)), - encode(session.getConf().buildRequestKey(end)), - GrpcUtils.makeRegionEpoch(2, 3), - GrpcUtils.makePeer(1, 10), - GrpcUtils.makePeer(2, 20)))); + request -> makeGetRegionResponse(start, end)); try (PDClient client = createClient()) { Metapb.Region r = client.getRegionByID(ConcreteBackOffer.newRawKVBackOff(), 1).first; - Assert.assertEquals(start, r.getStartKey()); - Assert.assertEquals(end, r.getEndKey()); + Assert.assertEquals(start, r.getStartKey().toStringUtf8()); + Assert.assertEquals(end, r.getEndKey().toStringUtf8()); } - leader.addGetRegionByIDListener( - request -> - GrpcUtils.makeGetRegionResponse( - leader.getClusterId(), - GrpcUtils.makeRegion( - 1, - encode(session.getConf().buildRequestKey(start)), - encode(session.getConf().getEndKey()), - GrpcUtils.makeRegionEpoch(2, 3), - GrpcUtils.makePeer(1, 10), - GrpcUtils.makePeer(2, 20)))); + leader.addGetRegionByIDListener(request -> makeGetRegionResponse(start, "")); try (PDClient client = createClient()) { Metapb.Region r = client.getRegionByID(ConcreteBackOffer.newRawKVBackOff(), 1).first; - Assert.assertEquals(start, r.getStartKey()); - Assert.assertEquals(ByteString.EMPTY, r.getEndKey()); + Assert.assertEquals(start, r.getStartKey().toStringUtf8()); + Assert.assertEquals("", r.getEndKey().toStringUtf8()); } } @Test - public void testScanRegions() { + public void testScanRegions() throws Exception { + String start = "scanRegions"; + String end = "scanRegionsEnd"; + + leader.addScanRegionsListener(request -> + ScanRegionsResponse.newBuilder() + .addRegions(Pdpb.Region.newBuilder().setRegion(makeRegion(start, end)).build()) + .build()); + + try (PDClient client = createClient()) { + List regions = client. + scanRegions(ConcreteBackOffer.newRawKVBackOff(), ByteString.EMPTY, ByteString.EMPTY, + 1); + + for (Region r : regions) { + Assert.assertEquals(start, r.getRegion().getStartKey().toStringUtf8()); + Assert.assertEquals(end, r.getRegion().getEndKey().toStringUtf8()); + } + } } } diff --git a/src/test/java/org/tikv/common/PDMockServer.java b/src/test/java/org/tikv/common/PDMockServer.java index 5bda86612f4..8396a681ecc 100644 --- a/src/test/java/org/tikv/common/PDMockServer.java +++ b/src/test/java/org/tikv/common/PDMockServer.java @@ -132,6 +132,10 @@ public void getStore(GetStoreRequest request, StreamObserver r } } + public void addScanRegionsListener(Function func) { + scanRegionsListener = func; + } + @Override public void scanRegions(ScanRegionsRequest request, StreamObserver resp) { try { From edbb9f57b0d9fd752f1c38ae99793e19427e2aa6 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Thu, 26 May 2022 20:23:45 +0800 Subject: [PATCH 18/32] add mock test for ApiV2 PD client Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/PDClient.java | 31 ++++++++++--------- .../org/tikv/common/PDClientV2MockTest.java | 19 ++++++------ .../java/org/tikv/common/PDMockServer.java | 6 ++-- 3 files changed, 28 insertions(+), 28 deletions(-) diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index 56a3c3edada..9df9e7c860c 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -219,8 +219,7 @@ public Boolean isCheckerPaused(PDChecker checker) { try { ObjectMapper mapper = new ObjectMapper(); HashMap status = - mapper.readValue(new URL(api), new TypeReference>() { - }); + mapper.readValue(new URL(api), new TypeReference>() {}); return status.get("paused"); } catch (Exception e) { logger.error(String.format("failed to get %s checker status.", checker.apiName()), e); @@ -366,10 +365,14 @@ public List scanRegions( if (resp == null) { return null; } - return resp.getRegionsList().stream().map(r -> { - Metapb.Region decoded = decodeRegion(r.getRegion()); - return Pdpb.Region.newBuilder().mergeFrom(r).setRegion(decoded).build(); - }).collect(Collectors.toList()); + return resp.getRegionsList() + .stream() + .map( + r -> { + Metapb.Region decoded = decodeRegion(r.getRegion()); + return Pdpb.Region.newBuilder().mergeFrom(r).setRegion(decoded).build(); + }) + .collect(Collectors.toList()); } private Supplier buildGetStoreReq(long storeId) { @@ -402,12 +405,12 @@ public Store getStore(BackOffer backOffer, long storeId) { @Override public List getAllStores(BackOffer backOffer) { return callWithRetry( - backOffer, - PDGrpc.getGetAllStoresMethod(), - buildGetAllStoresReq(), - new PDErrorHandler<>( - r -> r.getHeader().hasError() ? buildFromPdpbError(r.getHeader().getError()) : null, - this)) + backOffer, + PDGrpc.getGetAllStoresMethod(), + buildGetAllStoresReq(), + new PDErrorHandler<>( + r -> r.getHeader().hasError() ? buildFromPdpbError(r.getHeader().getError()) : null, + this)) .getStoresList(); } @@ -833,7 +836,7 @@ public Metapb.Region decodeRegion(Metapb.Region region) { ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getStartKey()))); if (conf.getApiVersion().isV2()) { if (ByteString.unsignedLexicographicalComparator() - .compare(decodedStartKey, conf.getKeyPrefix()) + .compare(decodedStartKey, conf.getKeyPrefix()) < 0) { decodedStartKey = ByteString.EMPTY; } else { @@ -857,7 +860,7 @@ public Metapb.Region decodeRegion(Metapb.Region region) { ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getEndKey()))); if (conf.getApiVersion().isV2()) { if (ByteString.unsignedLexicographicalComparator() - .compare(decodedEndKey, conf.getEndKey()) + .compare(decodedEndKey, conf.getEndKey()) >= 0) { decodedEndKey = ByteString.EMPTY; } else { diff --git a/src/test/java/org/tikv/common/PDClientV2MockTest.java b/src/test/java/org/tikv/common/PDClientV2MockTest.java index 57eb5840ff1..d5ffcdcb005 100644 --- a/src/test/java/org/tikv/common/PDClientV2MockTest.java +++ b/src/test/java/org/tikv/common/PDClientV2MockTest.java @@ -41,7 +41,6 @@ private PDClient createClient() { return session.getPDClient(); } - public static ByteString encode(ByteString key) { CodecDataOutput cdo = new CodecDataOutput(); BytesCodec.writeBytes(cdo, key.toByteArray()); @@ -66,8 +65,7 @@ private Metapb.Region makeRegion(String start, String end) { public void testGetRegionById() throws Exception { String start = "getRegionById"; String end = "getRegionByIdEnd"; - leader.addGetRegionByIDListener( - request -> makeGetRegionResponse(start, end)); + leader.addGetRegionByIDListener(request -> makeGetRegionResponse(start, end)); try (PDClient client = createClient()) { Metapb.Region r = client.getRegionByID(ConcreteBackOffer.newRawKVBackOff(), 1).first; Assert.assertEquals(start, r.getStartKey().toStringUtf8()); @@ -87,15 +85,16 @@ public void testScanRegions() throws Exception { String start = "scanRegions"; String end = "scanRegionsEnd"; - leader.addScanRegionsListener(request -> - ScanRegionsResponse.newBuilder() - .addRegions(Pdpb.Region.newBuilder().setRegion(makeRegion(start, end)).build()) - .build()); + leader.addScanRegionsListener( + request -> + ScanRegionsResponse.newBuilder() + .addRegions(Pdpb.Region.newBuilder().setRegion(makeRegion(start, end)).build()) + .build()); try (PDClient client = createClient()) { - List regions = client. - scanRegions(ConcreteBackOffer.newRawKVBackOff(), ByteString.EMPTY, ByteString.EMPTY, - 1); + List regions = + client.scanRegions( + ConcreteBackOffer.newRawKVBackOff(), ByteString.EMPTY, ByteString.EMPTY, 1); for (Region r : regions) { Assert.assertEquals(start, r.getRegion().getStartKey().toStringUtf8()); diff --git a/src/test/java/org/tikv/common/PDMockServer.java b/src/test/java/org/tikv/common/PDMockServer.java index 8396a681ecc..c081eb382c0 100644 --- a/src/test/java/org/tikv/common/PDMockServer.java +++ b/src/test/java/org/tikv/common/PDMockServer.java @@ -75,12 +75,10 @@ public StreamObserver tso(StreamObserver resp) { private int logical = 0; @Override - public void onNext(TsoRequest value) { - } + public void onNext(TsoRequest value) {} @Override - public void onError(Throwable t) { - } + public void onError(Throwable t) {} @Override public void onCompleted() { From 8b22d7706bd10266fdbcd6bef12f5c999c0d89b1 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Thu, 26 May 2022 20:29:17 +0800 Subject: [PATCH 19/32] revert src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java Signed-off-by: iosmanthus --- .../org/tikv/common/operation/iterator/RawScanIterator.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java b/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java index 9c2225188db..5b27132699d 100644 --- a/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java +++ b/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java @@ -72,9 +72,7 @@ private boolean endOfScan() { return false; } ByteString lastKey = currentCache.get(index).getKey(); - boolean a = !lastKey.isEmpty(); - boolean b = Key.toRawKey(lastKey).compareTo(endKey) >= 0; - return a && b; + return !lastKey.isEmpty() && Key.toRawKey(lastKey).compareTo(endKey) >= 0; } boolean isCacheDrained() { From 34a81a072ce13ac7198ee11c546555b9dedc7979 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Mon, 30 May 2022 20:18:21 +0800 Subject: [PATCH 20/32] refactor: introduce RequestKeyCodec to reduce if Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/PDClient.java | 174 +++++++++--------- .../org/tikv/common/ReadOnlyPDClient.java | 5 +- src/main/java/org/tikv/common/TiSession.java | 25 ++- .../tikv/common/apiversion/CodecUtils.java | 18 ++ .../common/apiversion/RequestKeyCodec.java | 23 +++ .../apiversion/RequestKeyV1RawCodec.java | 34 ++++ .../apiversion/RequestKeyV1TxnCodec.java | 46 +++++ .../common/apiversion/RequestKeyV2Codec.java | 79 ++++++++ .../apiversion/RequestKeyV2RawCodec.java | 10 + .../apiversion/RequestKeyV2TxnCodec.java | 10 + .../tikv/common/importer/ImporterClient.java | 28 ++- .../tikv/common/operation/KVErrorHandler.java | 6 +- .../common/operation/RegionErrorHandler.java | 2 +- .../region/AbstractRegionStoreClient.java | 7 +- .../tikv/common/region/RegionStoreClient.java | 159 ++++++++++------ .../tikv/txn/AbstractLockResolverClient.java | 5 +- src/main/java/org/tikv/txn/Lock.java | 7 +- .../org/tikv/txn/LockResolverClientV4.java | 12 +- .../org/tikv/common/PDClientV2MockTest.java | 10 +- .../java/org/tikv/common/RegionErrorTest.java | 2 +- 20 files changed, 485 insertions(+), 177 deletions(-) create mode 100644 src/main/java/org/tikv/common/apiversion/CodecUtils.java create mode 100644 src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java create mode 100644 src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java create mode 100644 src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java create mode 100644 src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java create mode 100644 src/main/java/org/tikv/common/apiversion/RequestKeyV2RawCodec.java create mode 100644 src/main/java/org/tikv/common/apiversion/RequestKeyV2TxnCodec.java diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index 9df9e7c860c..c957054ca71 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -62,9 +62,7 @@ import org.apache.http.impl.client.HttpClients; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import org.tikv.common.codec.Codec.BytesCodec; -import org.tikv.common.codec.CodecDataInput; -import org.tikv.common.codec.CodecDataOutput; +import org.tikv.common.apiversion.RequestKeyCodec; import org.tikv.common.codec.KeyUtils; import org.tikv.common.exception.GrpcException; import org.tikv.common.exception.TiClientInternalException; @@ -110,7 +108,9 @@ public class PDClient extends AbstractGRPCClient private static final long MIN_TRY_UPDATE_DURATION = 50; private static final int PAUSE_CHECKER_TIMEOUT = 300; // in seconds private static final int KEEP_CHECKER_PAUSE_PERIOD = PAUSE_CHECKER_TIMEOUT / 5; // in seconds - private final Logger logger = LoggerFactory.getLogger(PDClient.class); + private static final Logger logger = LoggerFactory.getLogger(PDClient.class); + + private final RequestKeyCodec codec; private RequestHeader header; private TsoRequest tsoReq; private volatile PDClientWrapper pdClientWrapper; @@ -130,19 +130,22 @@ public class PDClient extends AbstractGRPCClient .labelNames("cluster") .register(); - private PDClient(TiConfiguration conf, ChannelFactory channelFactory) { + private PDClient(TiConfiguration conf, RequestKeyCodec codec, ChannelFactory channelFactory) { super(conf, channelFactory); initCluster(); + this.codec = codec; this.blockingStub = getBlockingStub(); this.asyncStub = getAsyncStub(); } - public static ReadOnlyPDClient create(TiConfiguration conf, ChannelFactory channelFactory) { - return createRaw(conf, channelFactory); + public static ReadOnlyPDClient create( + TiConfiguration conf, RequestKeyCodec codec, ChannelFactory channelFactory) { + return createRaw(conf, codec, channelFactory); } - static PDClient createRaw(TiConfiguration conf, ChannelFactory channelFactory) { - return new PDClient(conf, channelFactory); + static PDClient createRaw( + TiConfiguration conf, RequestKeyCodec codec, ChannelFactory channelFactory) { + return new PDClient(conf, codec, channelFactory); } public HostMapping getHostMapping() { @@ -313,23 +316,19 @@ public Pair getRegionByKey(BackOffer backOffer, Byte Histogram.Timer requestTimer = PD_GET_REGION_BY_KEY_REQUEST_LATENCY.labels(getClusterId().toString()).startTimer(); try { - key = conf.buildRequestKey(key); - if (conf.isTxnKVMode() || conf.getApiVersion().isV2()) { - CodecDataOutput cdo = new CodecDataOutput(); - BytesCodec.writeBytes(cdo, key.toByteArray()); - key = cdo.toByteString(); - } - ByteString queryKey = key; - Supplier request = - () -> GetRegionRequest.newBuilder().setHeader(header).setRegionKey(queryKey).build(); + () -> + GetRegionRequest.newBuilder() + .setHeader(header) + .setRegionKey(codec.encodePdQuery(key)) + .build(); PDErrorHandler handler = new PDErrorHandler<>(getRegionResponseErrorExtractor, this); GetRegionResponse resp = callWithRetry(backOffer, PDGrpc.getGetRegionMethod(), request, handler); - return new Pair<>(decodeRegion(resp.getRegion()), resp.getLeader()); + return new Pair<>(codec.decodeRegion(resp.getRegion()), resp.getLeader()); } finally { requestTimer.observeDuration(); } @@ -344,7 +343,8 @@ public Pair getRegionByID(BackOffer backOffer, long GetRegionResponse resp = callWithRetry(backOffer, PDGrpc.getGetRegionByIDMethod(), request, handler); - return new Pair(decodeRegion(resp.getRegion()), resp.getLeader()); + return new Pair( + codec.decodeRegion(resp.getRegion()), resp.getLeader()); } @Override @@ -354,22 +354,29 @@ public List scanRegions( // introduce a warm-up timeout for ScanRegions requests PDGrpc.PDBlockingStub stub = getBlockingStub().withDeadlineAfter(conf.getWarmUpTimeout(), TimeUnit.MILLISECONDS); + org.apache.commons.lang3.tuple.Pair range = + codec.encodePdQueryRange(startKey, endKey); Pdpb.ScanRegionsRequest request = Pdpb.ScanRegionsRequest.newBuilder() .setHeader(header) - .setStartKey(conf.buildRequestKey(startKey)) - .setEndKey(conf.buildRequestKey(startKey, true)) + .setStartKey(range.getLeft()) + .setEndKey(range.getRight()) .setLimit(limit) .build(); Pdpb.ScanRegionsResponse resp = stub.scanRegions(request); if (resp == null) { return null; } + + if (conf.getApiVersion().isV1() && conf.isRawKVMode()) { + return resp.getRegionsList(); + } + return resp.getRegionsList() .stream() .map( r -> { - Metapb.Region decoded = decodeRegion(r.getRegion()); + Metapb.Region decoded = codec.decodeRegion(r.getRegion()); return Pdpb.Region.newBuilder().mergeFrom(r).setRegion(decoded).build(); }) .collect(Collectors.toList()); @@ -819,65 +826,64 @@ public String toString() { } } - @Override - public Metapb.Region decodeRegion(Metapb.Region region) { - final boolean isRawRegion = conf.isRawKVMode(); - Metapb.Region.Builder builder = - Metapb.Region.newBuilder() - .setId(region.getId()) - .setRegionEpoch(region.getRegionEpoch()) - .addAllPeers(region.getPeersList()); - - if (region.getStartKey().isEmpty() || (conf.getApiVersion().isV1() && isRawRegion)) { - builder.setStartKey(region.getStartKey()); - } else { - try { - ByteString decodedStartKey = - ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getStartKey()))); - if (conf.getApiVersion().isV2()) { - if (ByteString.unsignedLexicographicalComparator() - .compare(decodedStartKey, conf.getKeyPrefix()) - < 0) { - decodedStartKey = ByteString.EMPTY; - } else { - decodedStartKey = decodedStartKey.substring(conf.getKeyPrefix().size()); - } - } - builder.setStartKey(decodedStartKey); - } catch (Exception e) { - if (!conf.isTest()) { - throw e; - } - builder.setStartKey(region.getStartKey()); - } - } - - if (region.getEndKey().isEmpty() || (conf.getApiVersion().isV1() && isRawRegion)) { - builder.setEndKey(region.getEndKey()); - } else { - try { - ByteString decodedEndKey = - ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getEndKey()))); - if (conf.getApiVersion().isV2()) { - if (ByteString.unsignedLexicographicalComparator() - .compare(decodedEndKey, conf.getEndKey()) - >= 0) { - decodedEndKey = ByteString.EMPTY; - } else { - decodedEndKey = decodedEndKey.substring(conf.getKeyPrefix().size()); - } - } - builder.setEndKey(decodedEndKey); - } catch (Exception e) { - if (!conf.isTest()) { - throw e; - } - builder.setEndKey(region.getEndKey()); - } - } - - return builder.build(); - } + // public Metapb.Region decodeRegion(Metapb.Region region) { + // final boolean isRawRegion = conf.isRawKVMode(); + // Metapb.Region.Builder builder = + // Metapb.Region.newBuilder() + // .setId(region.getId()) + // .setRegionEpoch(region.getRegionEpoch()) + // .addAllPeers(region.getPeersList()); + + // if (region.getStartKey().isEmpty() || (conf.getApiVersion().isV1() && isRawRegion)) { + // builder.setStartKey(region.getStartKey()); + // } else { + // try { + // ByteString decodedStartKey = + // ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getStartKey()))); + // if (conf.getApiVersion().isV2()) { + // if (ByteString.unsignedLexicographicalComparator() + // .compare(decodedStartKey, conf.getKeyPrefix()) + // < 0) { + // decodedStartKey = ByteString.EMPTY; + // } else { + // decodedStartKey = decodedStartKey.substring(conf.getKeyPrefix().size()); + // } + // } + // builder.setStartKey(decodedStartKey); + // } catch (Exception e) { + // if (!conf.isTest()) { + // throw e; + // } + // builder.setStartKey(region.getStartKey()); + // } + // } + + // if (region.getEndKey().isEmpty() || (conf.getApiVersion().isV1() && isRawRegion)) { + // builder.setEndKey(region.getEndKey()); + // } else { + // try { + // ByteString decodedEndKey = + // ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getEndKey()))); + // if (conf.getApiVersion().isV2()) { + // if (ByteString.unsignedLexicographicalComparator() + // .compare(decodedEndKey, conf.getEndKey()) + // >= 0) { + // decodedEndKey = ByteString.EMPTY; + // } else { + // decodedEndKey = decodedEndKey.substring(conf.getKeyPrefix().size()); + // } + // } + // builder.setEndKey(decodedEndKey); + // } catch (Exception e) { + // if (!conf.isTest()) { + // throw e; + // } + // builder.setEndKey(region.getEndKey()); + // } + // } + + // return builder.build(); + // } public Long getClusterId() { return header.getClusterId(); @@ -886,4 +892,8 @@ public Long getClusterId() { public List getPdAddrs() { return pdAddrs; } + + public RequestKeyCodec getCodec() { + return codec; + } } diff --git a/src/main/java/org/tikv/common/ReadOnlyPDClient.java b/src/main/java/org/tikv/common/ReadOnlyPDClient.java index 7ae029b80ef..ddf1855e614 100644 --- a/src/main/java/org/tikv/common/ReadOnlyPDClient.java +++ b/src/main/java/org/tikv/common/ReadOnlyPDClient.java @@ -19,6 +19,7 @@ import com.google.protobuf.ByteString; import java.util.List; +import org.tikv.common.apiversion.RequestKeyCodec; import org.tikv.common.meta.TiTimestamp; import org.tikv.common.util.BackOffer; import org.tikv.common.util.Pair; @@ -64,11 +65,11 @@ List scanRegions( */ Store getStore(BackOffer backOffer, long storeId); - Metapb.Region decodeRegion(Metapb.Region region); - List getAllStores(BackOffer backOffer); TiConfiguration.ReplicaRead getReplicaRead(); Long getClusterId(); + + RequestKeyCodec getCodec(); } diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java index 0c38117211a..bf2365390f6 100644 --- a/src/main/java/org/tikv/common/TiSession.java +++ b/src/main/java/org/tikv/common/TiSession.java @@ -33,6 +33,11 @@ import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.tikv.common.apiversion.RequestKeyCodec; +import org.tikv.common.apiversion.RequestKeyV1RawCodec; +import org.tikv.common.apiversion.RequestKeyV1TxnCodec; +import org.tikv.common.apiversion.RequestKeyV2RawCodec; +import org.tikv.common.apiversion.RequestKeyV2TxnCodec; import org.tikv.common.catalog.Catalog; import org.tikv.common.exception.TiKVException; import org.tikv.common.importer.ImporterStoreClient; @@ -70,6 +75,7 @@ public class TiSession implements AutoCloseable { private static final Logger logger = LoggerFactory.getLogger(TiSession.class); private static final Map sessionCachedMap = new HashMap<>(); private final TiConfiguration conf; + private final RequestKeyCodec keyCodec; private final ChannelFactory channelFactory; // below object creation is either heavy or making connection (pd), pending for lazy loading private volatile PDClient client; @@ -117,6 +123,21 @@ public TiSession(TiConfiguration conf) { this.metricsServer = MetricsServer.getInstance(conf); this.conf = conf; + + if (conf.getApiVersion().isV1()) { + if (conf.isRawKVMode()) { + keyCodec = new RequestKeyV1RawCodec(); + } else { + keyCodec = new RequestKeyV1TxnCodec(); + } + } else { + if (conf.isRawKVMode()) { + keyCodec = new RequestKeyV2RawCodec(); + } else { + keyCodec = new RequestKeyV2TxnCodec(); + } + } + if (conf.isTlsEnable()) { if (conf.isJksEnable()) { this.channelFactory = @@ -153,7 +174,7 @@ public TiSession(TiConfiguration conf) { conf.getIdleTimeout()); } - this.client = PDClient.createRaw(conf, channelFactory); + this.client = PDClient.createRaw(conf, keyCodec, channelFactory); this.enableGrpcForward = conf.getEnableGrpcForward(); if (this.enableGrpcForward) { logger.info("enable grpc forward for high available"); @@ -357,7 +378,7 @@ public PDClient getPDClient() { if (res == null) { synchronized (this) { if (client == null) { - client = PDClient.createRaw(this.getConf(), channelFactory); + client = PDClient.createRaw(this.getConf(), keyCodec, channelFactory); } res = client; } diff --git a/src/main/java/org/tikv/common/apiversion/CodecUtils.java b/src/main/java/org/tikv/common/apiversion/CodecUtils.java new file mode 100644 index 00000000000..803007bbc78 --- /dev/null +++ b/src/main/java/org/tikv/common/apiversion/CodecUtils.java @@ -0,0 +1,18 @@ +package org.tikv.common.apiversion; + +import com.google.protobuf.ByteString; +import org.tikv.common.codec.Codec.BytesCodec; +import org.tikv.common.codec.CodecDataInput; +import org.tikv.common.codec.CodecDataOutput; + +public class CodecUtils { + public static ByteString encode(ByteString key) { + CodecDataOutput cdo = new CodecDataOutput(); + BytesCodec.writeBytes(cdo, key.toByteArray()); + return cdo.toByteString(); + } + + public static ByteString decode(ByteString key) { + return ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(key))); + } +} diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java new file mode 100644 index 00000000000..3fbbea82b33 --- /dev/null +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java @@ -0,0 +1,23 @@ +package org.tikv.common.apiversion; + +import com.google.protobuf.ByteString; +import org.apache.commons.lang3.tuple.Pair; +import org.tikv.kvproto.Metapb; + +public interface RequestKeyCodec { + ByteString encodeKey(ByteString key); + + ByteString decodeKey(ByteString key); + + Pair encodeRange(ByteString start, ByteString end); + + + ByteString encodePdQuery(ByteString key); + + default Pair encodePdQueryRange(ByteString start, ByteString end) { + Pair range = encodeRange(start, end); + return Pair.of(CodecUtils.encode(range.getLeft()), CodecUtils.encode(range.getRight())); + } + + Metapb.Region decodeRegion(Metapb.Region region); +} diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java new file mode 100644 index 00000000000..2611f30b924 --- /dev/null +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java @@ -0,0 +1,34 @@ +package org.tikv.common.apiversion; + +import com.google.protobuf.ByteString; +import org.apache.commons.lang3.tuple.Pair; +import org.tikv.kvproto.Metapb; + +public class RequestKeyV1RawCodec implements RequestKeyCodec { + public RequestKeyV1RawCodec() {} + + @Override + public ByteString encodeKey(ByteString key) { + return key; + } + + @Override + public ByteString decodeKey(ByteString key) { + return key; + } + + @Override + public Pair encodeRange(ByteString start, ByteString end) { + return Pair.of(start, end); + } + + @Override + public ByteString encodePdQuery(ByteString key) { + return key; + } + + @Override + public Metapb.Region decodeRegion(Metapb.Region region) { + return region; + } +} diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java new file mode 100644 index 00000000000..6cf3cdcb0a1 --- /dev/null +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java @@ -0,0 +1,46 @@ +package org.tikv.common.apiversion; + +import com.google.protobuf.ByteString; +import org.apache.commons.lang3.tuple.Pair; +import org.tikv.kvproto.Metapb; + +public class RequestKeyV1TxnCodec implements RequestKeyCodec { + public RequestKeyV1TxnCodec() {} + + @Override + public ByteString encodeKey(ByteString key) { + return key; + } + + @Override + public ByteString decodeKey(ByteString key) { + return key; + } + + @Override + public Pair encodeRange(ByteString start, ByteString end) { + return Pair.of(start, end); + } + + @Override + public ByteString encodePdQuery(ByteString key) { + return CodecUtils.encode(key); + } + + @Override + public Metapb.Region decodeRegion(Metapb.Region region) { + Metapb.Region.Builder builder = Metapb.Region.newBuilder().mergeFrom(region); + ByteString start = region.getStartKey(); + ByteString end = region.getEndKey(); + + if (!start.isEmpty()) { + start = CodecUtils.decode(start); + } + + if (!end.isEmpty()) { + end = CodecUtils.decode(end); + } + + return builder.setStartKey(start).setEndKey(end).build(); + } +} diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java new file mode 100644 index 00000000000..22000afc87e --- /dev/null +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java @@ -0,0 +1,79 @@ +package org.tikv.common.apiversion; + +import com.google.protobuf.ByteString; +import org.apache.commons.lang3.tuple.Pair; +import org.tikv.kvproto.Metapb; +import org.tikv.kvproto.Metapb.Region; + +public class RequestKeyV2Codec implements RequestKeyCodec { + protected static ByteString RAW_KEY_PREFIX = ByteString.copyFromUtf8("r"); + protected static ByteString RAW_END_KEY = + ByteString.copyFrom(new byte[]{(byte) (RAW_KEY_PREFIX.toByteArray()[0] + 1)}); + + protected static ByteString TXN_KEY_PREFIX = ByteString.copyFromUtf8("x"); + protected static ByteString TXN_END_KEY = + ByteString.copyFrom(new byte[]{(byte) (TXN_KEY_PREFIX.toByteArray()[0] + 1)}); + protected ByteString keyPrefix; + + protected ByteString infiniteEndKey; + + @Override + public ByteString encodeKey(ByteString key) { + return keyPrefix.concat(key); + } + + @Override + public ByteString decodeKey(ByteString key) { + if (key.isEmpty()) { + return key; + } + + if (!key.startsWith(keyPrefix)) { + throw new IllegalArgumentException("key corrupted, wrong prefix"); + } + + return key.substring(1); + } + + @Override + public Pair encodeRange(ByteString start, ByteString end) { + start = encodeKey(start); + + end = end.isEmpty() ? infiniteEndKey : encodeKey(end); + + return Pair.of(start, end); + } + + @Override + public ByteString encodePdQuery(ByteString key) { + return CodecUtils.encode(encodeKey(key)); + } + + @Override + public Region decodeRegion(Region region) { + Metapb.Region.Builder builder = Metapb.Region.newBuilder().mergeFrom(region); + + ByteString start = region.getStartKey(); + ByteString end = region.getEndKey(); + + if (!start.isEmpty()) { + start = CodecUtils.decode(start); + if (ByteString.unsignedLexicographicalComparator().compare(start, keyPrefix) < 0) { + start = ByteString.EMPTY; + } else { + start = decodeKey(start); + } + } + + if (!end.isEmpty()) { + end = CodecUtils.decode(end); + if (ByteString.unsignedLexicographicalComparator().compare(end, infiniteEndKey) >= 0) { + end = ByteString.EMPTY; + } else { + end = decodeKey(end); + } + } + + return builder.setStartKey(start).setEndKey(end).build(); + } +} diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2RawCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2RawCodec.java new file mode 100644 index 00000000000..3f6868a9877 --- /dev/null +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2RawCodec.java @@ -0,0 +1,10 @@ +package org.tikv.common.apiversion; + +public class RequestKeyV2RawCodec extends RequestKeyV2Codec { + public RequestKeyV2RawCodec() { + super(); + + this.keyPrefix = RAW_KEY_PREFIX; + this.infiniteEndKey = RAW_END_KEY; + } +} diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2TxnCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2TxnCodec.java new file mode 100644 index 00000000000..47b7ce2427a --- /dev/null +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2TxnCodec.java @@ -0,0 +1,10 @@ +package org.tikv.common.apiversion; + +public class RequestKeyV2TxnCodec extends RequestKeyV2Codec { + public RequestKeyV2TxnCodec() { + super(); + + this.keyPrefix = TXN_KEY_PREFIX; + this.infiniteEndKey = TXN_END_KEY; + } +} diff --git a/src/main/java/org/tikv/common/importer/ImporterClient.java b/src/main/java/org/tikv/common/importer/ImporterClient.java index da8ab53d6d4..f8ddce88242 100644 --- a/src/main/java/org/tikv/common/importer/ImporterClient.java +++ b/src/main/java/org/tikv/common/importer/ImporterClient.java @@ -29,6 +29,7 @@ import org.slf4j.LoggerFactory; import org.tikv.common.TiConfiguration; import org.tikv.common.TiSession; +import org.tikv.common.apiversion.RequestKeyCodec; import org.tikv.common.codec.Codec; import org.tikv.common.codec.CodecDataOutput; import org.tikv.common.exception.GrpcException; @@ -64,6 +65,8 @@ public class ImporterClient { private List clientList; private ImporterStoreClient clientLeader; + private final RequestKeyCodec codec; + public ImporterClient( TiSession tiSession, ByteString uuid, Key minKey, Key maxKey, TiRegion region, Long ttl) { this.uuid = uuid; @@ -73,6 +76,7 @@ public ImporterClient( this.maxKey = maxKey; this.region = region; this.ttl = ttl; + this.codec = tiSession.getPDClient().getCodec(); } public boolean isDeduplicate() { @@ -109,7 +113,7 @@ public void write(Iterator> iterator) throws TiKVEx String.format("duplicate key found, key = %s", preKey.toStringUtf8())); } } else { - ByteString key = tiConf.buildRequestKey(pair.first); + ByteString key = codec.encodeKey(pair.first); pairs.add(ImportSstpb.Pair.newBuilder().setKey(key).setValue(pair.second).build()); totalBytes += (key.size() + pair.second.size()); preKey = pair.first; @@ -138,16 +142,14 @@ public void write(Iterator> iterator) throws TiKVEx private void init() { long regionId = region.getId(); Metapb.RegionEpoch regionEpoch = region.getRegionEpoch(); + org.apache.commons.lang3.tuple.Pair keyRange = + codec.encodePdQueryRange(minKey.toByteString(), maxKey.toByteString()); + ImportSstpb.Range range = - tiConf.isTxnKVMode() || tiConf.getApiVersion().isV2() - ? ImportSstpb.Range.newBuilder() - .setStart(encode(tiConf.buildRequestKey(minKey.toByteString()))) - .setEnd(encode(tiConf.buildRequestKey(maxKey.toByteString(), true))) - .build() - : ImportSstpb.Range.newBuilder() - .setStart(minKey.toByteString()) - .setEnd(maxKey.toByteString()) - .build(); + ImportSstpb.Range.newBuilder() + .setStart(keyRange.getLeft()) + .setEnd(keyRange.getRight()) + .build(); sstMeta = ImportSstpb.SSTMeta.newBuilder() @@ -172,12 +174,6 @@ private void init() { } } - private ByteString encode(ByteString key) { - CodecDataOutput cdo = new CodecDataOutput(); - Codec.BytesCodec.writeBytes(cdo, key.toByteArray()); - return cdo.toByteString(); - } - private void startWrite() { for (ImporterStoreClient client : clientList) { client.startWrite(); diff --git a/src/main/java/org/tikv/common/operation/KVErrorHandler.java b/src/main/java/org/tikv/common/operation/KVErrorHandler.java index 135664530bf..22d05658425 100644 --- a/src/main/java/org/tikv/common/operation/KVErrorHandler.java +++ b/src/main/java/org/tikv/common/operation/KVErrorHandler.java @@ -23,6 +23,7 @@ import java.util.function.Function; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.tikv.common.apiversion.RequestKeyCodec; import org.tikv.common.exception.GrpcException; import org.tikv.common.exception.KeyException; import org.tikv.common.region.RegionErrorReceiver; @@ -44,6 +45,8 @@ public class KVErrorHandler implements ErrorHandler { private final boolean forWrite; private final RegionErrorHandler regionHandler; + private final RequestKeyCodec codec; + public KVErrorHandler( RegionManager regionManager, RegionErrorReceiver recv, @@ -59,6 +62,7 @@ public KVErrorHandler( this.resolveLockResultCallback = resolveLockResultCallback; this.callerStartTS = callerStartTS; this.forWrite = forWrite; + this.codec = regionManager.getPDClient().getCodec(); } private void resolveLock(BackOffer backOffer, Lock lock) { @@ -100,7 +104,7 @@ public boolean handleResponseError(BackOffer backOffer, RespT resp) { Kvrpcpb.KeyError keyError = getKeyError.apply(resp); if (keyError != null) { try { - Lock lock = AbstractLockResolverClient.extractLockFromKeyErr(keyError); + Lock lock = AbstractLockResolverClient.extractLockFromKeyErr(keyError, codec); resolveLock(backOffer, lock); return true; } catch (KeyException e) { diff --git a/src/main/java/org/tikv/common/operation/RegionErrorHandler.java b/src/main/java/org/tikv/common/operation/RegionErrorHandler.java index 23ce76101c7..7e7eff2b9dc 100644 --- a/src/main/java/org/tikv/common/operation/RegionErrorHandler.java +++ b/src/main/java/org/tikv/common/operation/RegionErrorHandler.java @@ -213,7 +213,7 @@ private boolean onRegionEpochNotMatch(BackOffer backOffer, List c // If the region epoch is not ahead of TiKV's, replace region meta in region cache. for (Metapb.Region meta : currentRegions) { // The region needs to be decoded to plain format. - meta = regionManager.getPDClient().decodeRegion(meta); + meta = regionManager.getPDClient().getCodec().decodeRegion(meta); TiRegion region = regionManager.createRegion(meta, backOffer); newRegions.add(region); if (recv.getRegion().getVerID() == region.getVerID()) { diff --git a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java index b64258f3472..f848de8798d 100644 --- a/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java @@ -34,6 +34,7 @@ import org.slf4j.LoggerFactory; import org.tikv.common.AbstractGRPCClient; import org.tikv.common.TiConfiguration; +import org.tikv.common.apiversion.RequestKeyCodec; import org.tikv.common.exception.GrpcException; import org.tikv.common.log.SlowLog; import org.tikv.common.log.SlowLogSpan; @@ -66,6 +67,7 @@ public abstract class AbstractRegionStoreClient .register(); protected final RegionManager regionManager; + protected final RequestKeyCodec codec; protected TiRegion region; protected TiStore store; @@ -84,6 +86,7 @@ protected AbstractRegionStoreClient( this.region = region; this.regionManager = regionManager; this.store = store; + this.codec = regionManager.getPDClient().getCodec(); if (this.store.getProxyStore() != null) { this.timeout = conf.getForwardTimeout(); } @@ -303,7 +306,7 @@ private Metapb.Peer switchLeaderStore(BackOffer backOffer) { Kvrpcpb.RawGetRequest rawGetRequest = Kvrpcpb.RawGetRequest.newBuilder() .setContext(makeContext(peer)) - .setKey(conf.buildRequestKey(key)) + .setKey(codec.encodeKey(key)) .build(); ListenableFuture task = stub.rawGet(rawGetRequest); responses.add(new SwitchLeaderTask(task, peer)); @@ -365,7 +368,7 @@ private TiStore switchProxyStore(BackOffer backOffer) { Kvrpcpb.RawGetRequest rawGetRequest = Kvrpcpb.RawGetRequest.newBuilder() .setContext(makeContext()) - .setKey(conf.buildRequestKey(key)) + .setKey(codec.encodeKey(key)) .build(); ListenableFuture task = MetadataUtils.attachHeaders(stub, header).rawGet(rawGetRequest); diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index c3d49c9b318..9d0039e28ca 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -124,12 +124,17 @@ // need to be re-split across regions/stores, region info outdated, e.t.c., you // should retry it in an upper client logic (KVClient, TxnClient, e.t.c.) -/** Note that RegionStoreClient itself is not thread-safe */ +/** + * Note that RegionStoreClient itself is not thread-safe + */ public class RegionStoreClient extends AbstractRegionStoreClient { private static final Logger logger = LoggerFactory.getLogger(RegionStoreClient.class); - @VisibleForTesting public final AbstractLockResolverClient lockResolverClient; + @VisibleForTesting + public final AbstractLockResolverClient lockResolverClient; private final TiStoreType storeType; - /** startTS -> List(locks) */ + /** + * startTS -> List(locks) + */ private final Map> resolvedLocks = new HashMap<>(); private final PDClient pdClient; @@ -236,7 +241,7 @@ public ByteString get(BackOffer backOffer, ByteString key, long version) GetRequest.newBuilder() .setContext( makeContext(getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .setKey(conf.buildRequestKey(key)) + .setKey(codec.encodeKey(key)) .setVersion(version) .build(); @@ -282,7 +287,7 @@ public List batchGet(BackOffer backOffer, List keys, long ve BatchGetRequest.newBuilder() .setContext( makeContext(getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .addAllKeys(keys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) + .addAllKeys(keys.stream().map(codec::encodeKey).collect(Collectors.toList())) .setVersion(version) .build(); KVErrorHandler handler = @@ -315,7 +320,7 @@ private List handleBatchGetResponse( for (KvPair pair : resp.getPairsList()) { if (pair.hasError()) { if (pair.getError().hasLocked()) { - Lock lock = new Lock(pair.getError().getLocked()); + Lock lock = new Lock(pair.getError().getLocked(), codec); locks.add(lock); } else { throw new KeyException(pair.getError()); @@ -330,7 +335,18 @@ private List handleBatchGetResponse( // resolveLocks already retried, just throw error to upper logic. throw new TiKVException("locks not resolved, retry"); } else { - return resp.getPairsList(); + if (conf.getApiVersion().isV1()) { + return resp.getPairsList(); + } + return resp.getPairsList() + .stream() + .map( + kvPair -> + KvPair.newBuilder() + .mergeFrom(kvPair) + .setKey(codec.decodeKey(kvPair.getKey())) + .build()) + .collect(Collectors.toList()); } } @@ -347,7 +363,7 @@ public List scan( .setContext( makeContext( getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .setStartKey(conf.buildRequestKey(startKey)) + .setStartKey(codec.encodeKey(startKey)) .setVersion(version) .setKeyOnly(keyOnly) .setLimit(getConf().getScanBatchSize()) @@ -389,12 +405,12 @@ private List doScan(ScanResponse resp) { List newKvPairs = new ArrayList<>(); for (KvPair kvPair : kvPairs) { if (kvPair.hasError()) { - Lock lock = AbstractLockResolverClient.extractLockFromKeyErr(kvPair.getError()); + Lock lock = AbstractLockResolverClient.extractLockFromKeyErr(kvPair.getError(), codec); newKvPairs.add( KvPair.newBuilder() .setError(kvPair.getError()) .setValue(kvPair.getValue()) - .setKey(conf.unwrapResponseKey(lock.getKey())) + .setKey(codec.decodeKey(lock.getKey())) .build()); } else { newKvPairs.add(kvPair); @@ -444,23 +460,23 @@ public void prewrite( () -> getIsV4() ? PrewriteRequest.newBuilder() - .setContext(makeContext(storeType, bo.getSlowLog())) - .setStartVersion(startTs) - .setPrimaryLock(conf.buildRequestKey(primaryLock)) - .addAllMutations( - mutations - .stream() - .map( - mutation -> - Mutation.newBuilder(mutation) - .setKey(conf.buildRequestKey(mutation.getKey())) - .build()) - .collect(Collectors.toList())) - .setLockTtl(ttl) - .setSkipConstraintCheck(skipConstraintCheck) - .setMinCommitTs(startTs) - .setTxnSize(16) - .build() + .setContext(makeContext(storeType, bo.getSlowLog())) + .setStartVersion(startTs) + .setPrimaryLock(codec.encodeKey(primaryLock)) + .addAllMutations( + mutations + .stream() + .map( + mutation -> + Mutation.newBuilder(mutation) + .setKey(codec.encodeKey(mutation.getKey())) + .build()) + .collect(Collectors.toList())) + .setLockTtl(ttl) + .setSkipConstraintCheck(skipConstraintCheck) + .setMinCommitTs(startTs) + .setTxnSize(16) + .build() : PrewriteRequest.newBuilder() .setContext(makeContext(storeType, bo.getSlowLog())) .setStartVersion(startTs) @@ -492,8 +508,8 @@ public void prewrite( * @param backOffer backOffer * @param resp response * @return Return true means the rpc call success. Return false means the rpc call fail, - * RegionStoreClient should retry. Throw an Exception means the rpc call fail, - * RegionStoreClient cannot handle this kind of error + * RegionStoreClient should retry. Throw an Exception means the rpc call fail, RegionStoreClient + * cannot handle this kind of error * @throws TiClientInternalException * @throws RegionException * @throws KeyException @@ -514,7 +530,7 @@ private boolean isPrewriteSuccess(BackOffer backOffer, PrewriteResponse resp, lo for (KeyError err : resp.getErrorsList()) { if (err.hasLocked()) { isSuccess = false; - Lock lock = new Lock(err.getLocked()); + Lock lock = new Lock(err.getLocked(), codec); locks.add(lock); } else { throw new KeyException(err.toString()); @@ -535,7 +551,9 @@ private boolean isPrewriteSuccess(BackOffer backOffer, PrewriteResponse resp, lo return false; } - /** TXN Heart Beat: update primary key ttl */ + /** + * TXN Heart Beat: update primary key ttl + */ public void txnHeartBeat(BackOffer bo, ByteString primaryLock, long startTs, long ttl) { boolean forWrite = false; while (true) { @@ -599,7 +617,7 @@ public void commit(BackOffer backOffer, List keys, long startTs, lon CommitRequest.newBuilder() .setStartVersion(startTs) .setCommitVersion(commitTs) - .addAllKeys(keys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) + .addAllKeys(keys.stream().map(codec::encodeKey).collect(Collectors.toList())) .setContext(makeContext(storeType, backOffer.getSlowLog())) .build(); KVErrorHandler handler = @@ -718,7 +736,7 @@ private List handleCopResponse( } if (response.hasLocked()) { - Lock lock = new Lock(response.getLocked()); + Lock lock = new Lock(response.getLocked(), codec); logger.debug(String.format("coprocessor encounters locks: %s", lock)); ResolveLockResult resolveLockResult = lockResolverClient.resolveLocks( @@ -827,7 +845,7 @@ public List splitRegion(List splitKeys) { SplitRegionRequest.newBuilder() .setContext(makeContext(storeType, SlowLogEmptyImpl.INSTANCE)) .addAllSplitKeys( - splitKeys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) + splitKeys.stream().map(codec::encodeKey).collect(Collectors.toList())) .setIsRawKv(conf.isRawKVMode()) .build(); @@ -860,7 +878,11 @@ public List splitRegion(List splitKeys) { "failed to split region %d because %s", region.getId(), resp.getRegionError())); } - return resp.getRegionsList(); + if (conf.getApiVersion().isV1()) { + return resp.getRegionsList(); + } + return resp.getRegionsList().stream().map(codec::decodeRegion).collect( + Collectors.toList()); } // APIs for Raw Scan/Put/Get/Delete @@ -874,7 +896,7 @@ public Optional rawGet(BackOffer backOffer, ByteString key) { () -> RawGetRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(conf.buildRequestKey(key)) + .setKey(codec.encodeKey(key)) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -916,7 +938,7 @@ public Optional rawGetKeyTTL(BackOffer backOffer, ByteString key) { () -> RawGetKeyTTLRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(conf.buildRequestKey(key)) + .setKey(codec.encodeKey(key)) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -958,7 +980,7 @@ public void rawDelete(BackOffer backOffer, ByteString key, boolean atomicForCAS) () -> RawDeleteRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(conf.buildRequestKey(key)) + .setKey(codec.encodeKey(key)) .setForCas(atomicForCAS) .build(); @@ -997,7 +1019,7 @@ public void rawPut( () -> RawPutRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(conf.buildRequestKey(key)) + .setKey(codec.encodeKey(key)) .setValue(value) .setTtl(ttl) .setForCas(atomicForCAS) @@ -1044,7 +1066,7 @@ public void rawCompareAndSet( () -> RawCASRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setKey(conf.buildRequestKey(key)) + .setKey(codec.encodeKey(key)) .setValue(value) .setPreviousValue(prevValue.orElse(ByteString.EMPTY)) .setPreviousNotExist(!prevValue.isPresent()) @@ -1100,7 +1122,7 @@ public List rawBatchGet(BackOffer backoffer, List keys) { () -> RawBatchGetRequest.newBuilder() .setContext(makeContext(storeType, backoffer.getSlowLog())) - .addAllKeys(keys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) + .addAllKeys(keys.stream().map(codec::encodeKey).collect(Collectors.toList())) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -1121,7 +1143,20 @@ private List handleRawBatchGet(RawBatchGetResponse resp) { if (resp.hasRegionError()) { throw new RegionException(resp.getRegionError()); } - return resp.getPairsList(); + + if (conf.getApiVersion().isV1()) { + return resp.getPairsList(); + } + + return resp.getPairsList() + .stream() + .map( + kvPair -> + KvPair.newBuilder() + .mergeFrom(kvPair) + .setKey(codec.decodeKey(kvPair.getKey())) + .build()) + .collect(Collectors.toList()); } public void rawBatchPut( @@ -1160,7 +1195,7 @@ public void rawBatchPut(BackOffer backOffer, Batch batch, long ttl, boolean atom for (int i = 0; i < batch.getKeys().size(); i++) { pairs.add( KvPair.newBuilder() - .setKey(conf.buildRequestKey(batch.getKeys().get(i))) + .setKey(codec.encodeKey(batch.getKeys().get(i))) .setValue(batch.getValues().get(i)) .build()); } @@ -1195,7 +1230,7 @@ public void rawBatchDelete(BackOffer backoffer, List keys, boolean a () -> RawBatchDeleteRequest.newBuilder() .setContext(makeContext(storeType, backoffer.getSlowLog())) - .addAllKeys(keys.stream().map(conf::buildRequestKey).collect(Collectors.toList())) + .addAllKeys(keys.stream().map(codec::encodeKey).collect(Collectors.toList())) .setForCas(atomicForCAS) .build(); RegionErrorHandler handler = @@ -1238,14 +1273,17 @@ public List rawScan(BackOffer backOffer, ByteString key, int limit, bool GRPC_RAW_REQUEST_LATENCY.labels("client_grpc_raw_scan", clusterId.toString()).startTimer(); try { Supplier factory = - () -> - RawScanRequest.newBuilder() - .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setStartKey(conf.buildRequestKey(key)) - .setEndKey(conf.getEndKey()) - .setKeyOnly(keyOnly) - .setLimit(limit) - .build(); + () -> { + org.apache.commons.lang3.tuple.Pair range = + codec.encodeRange(key, ByteString.EMPTY); + return RawScanRequest.newBuilder() + .setContext(makeContext(storeType, backOffer.getSlowLog())) + .setStartKey(range.getLeft()) + .setEndKey(range.getRight()) + .setKeyOnly(keyOnly) + .setLimit(limit) + .build(); + }; RegionErrorHandler handler = new RegionErrorHandler( @@ -1275,7 +1313,7 @@ private List rawScanHelper(RawScanResponse resp) { .map( kvPair -> KvPair.newBuilder() - .setKey(conf.unwrapResponseKey(kvPair.getKey())) + .setKey(codec.decodeKey(kvPair.getKey())) .setValue(kvPair.getValue()) .build()) .collect(Collectors.toList()); @@ -1296,12 +1334,15 @@ public void rawDeleteRange(BackOffer backOffer, ByteString startKey, ByteString .startTimer(); try { Supplier factory = - () -> - RawDeleteRangeRequest.newBuilder() - .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setStartKey(conf.buildRequestKey(startKey)) - .setEndKey(conf.buildRequestKey(endKey, true)) - .build(); + () -> { + org.apache.commons.lang3.tuple.Pair range = + codec.encodeRange(startKey, endKey); + return RawDeleteRangeRequest.newBuilder() + .setContext(makeContext(storeType, backOffer.getSlowLog())) + .setStartKey(range.getLeft()) + .setEndKey(range.getRight()) + .build(); + }; RegionErrorHandler handler = new RegionErrorHandler( diff --git a/src/main/java/org/tikv/txn/AbstractLockResolverClient.java b/src/main/java/org/tikv/txn/AbstractLockResolverClient.java index 1914fc41cee..6e535e7ab65 100644 --- a/src/main/java/org/tikv/txn/AbstractLockResolverClient.java +++ b/src/main/java/org/tikv/txn/AbstractLockResolverClient.java @@ -22,6 +22,7 @@ import org.tikv.common.StoreVersion; import org.tikv.common.TiConfiguration; import org.tikv.common.Version; +import org.tikv.common.apiversion.RequestKeyCodec; import org.tikv.common.exception.KeyException; import org.tikv.common.region.RegionManager; import org.tikv.common.region.RegionStoreClient; @@ -39,9 +40,9 @@ public interface AbstractLockResolverClient { /** transaction involves keys exceed this threshold can be treated as `big transaction`. */ long BIG_TXN_THRESHOLD = 16; - static Lock extractLockFromKeyErr(Kvrpcpb.KeyError keyError) { + static Lock extractLockFromKeyErr(Kvrpcpb.KeyError keyError, RequestKeyCodec codec) { if (keyError.hasLocked()) { - return new Lock(keyError.getLocked()); + return new Lock(keyError.getLocked(), codec); } if (keyError.hasConflict()) { diff --git a/src/main/java/org/tikv/txn/Lock.java b/src/main/java/org/tikv/txn/Lock.java index 2562d70a005..6dbb9735267 100644 --- a/src/main/java/org/tikv/txn/Lock.java +++ b/src/main/java/org/tikv/txn/Lock.java @@ -18,6 +18,7 @@ package org.tikv.txn; import com.google.protobuf.ByteString; +import org.tikv.common.apiversion.RequestKeyCodec; import org.tikv.kvproto.Kvrpcpb; public class Lock { @@ -30,10 +31,10 @@ public class Lock { private final Kvrpcpb.Op lockType; private final long lockForUpdateTs; - public Lock(Kvrpcpb.LockInfo l) { + public Lock(Kvrpcpb.LockInfo l, RequestKeyCodec codec) { txnID = l.getLockVersion(); - key = l.getKey(); - primary = l.getPrimaryLock(); + key = codec.decodeKey(l.getKey()); + primary = codec.decodeKey(l.getPrimaryLock()); ttl = l.getLockTtl() == 0 ? DEFAULT_LOCK_TTL : l.getLockTtl(); txnSize = l.getTxnSize(); lockType = l.getLockType(); diff --git a/src/main/java/org/tikv/txn/LockResolverClientV4.java b/src/main/java/org/tikv/txn/LockResolverClientV4.java index 7f428cfdb42..e433df42011 100644 --- a/src/main/java/org/tikv/txn/LockResolverClientV4.java +++ b/src/main/java/org/tikv/txn/LockResolverClientV4.java @@ -39,8 +39,12 @@ import org.tikv.common.exception.RegionException; import org.tikv.common.exception.TiClientInternalException; import org.tikv.common.operation.KVErrorHandler; -import org.tikv.common.region.*; +import org.tikv.common.region.AbstractRegionStoreClient; +import org.tikv.common.region.RegionManager; +import org.tikv.common.region.RegionStoreClient; +import org.tikv.common.region.TiRegion; import org.tikv.common.region.TiRegion.RegionVerID; +import org.tikv.common.region.TiStore; import org.tikv.common.util.BackOffer; import org.tikv.common.util.ChannelFactory; import org.tikv.common.util.TsoUtils; @@ -168,7 +172,7 @@ private void resolvePessimisticLock(BackOffer bo, Lock lock, Set cl () -> Kvrpcpb.PessimisticRollbackRequest.newBuilder() .setContext(makeContext()) - .addKeys(conf.buildRequestKey(lock.getKey())) + .addKeys(codec.encodeKey(lock.getKey())) .setStartVersion(lock.getTxnID()) .setForUpdateTs(forUpdateTS) .build(); @@ -286,7 +290,7 @@ private TxnStatus getTxnStatus( TiRegion primaryKeyRegion = regionManager.getRegionByKey(primary); return Kvrpcpb.CheckTxnStatusRequest.newBuilder() .setContext(primaryKeyRegion.getLeaderContext()) - .setPrimaryKey(conf.buildRequestKey(primary)) + .setPrimaryKey(codec.encodeKey(primary)) .setLockTs(txnID) .setCallerStartTs(callerStartTS) .setCurrentTs(currentTS) @@ -373,7 +377,7 @@ private void resolveLock( if (lock.getTxnSize() < BIG_TXN_THRESHOLD) { // Only resolve specified keys when it is a small transaction, // prevent from scanning the whole region in this case. - builder.addKeys(conf.buildRequestKey(lock.getKey())); + builder.addKeys(codec.encodeKey(lock.getKey())); } Supplier factory = builder::build; diff --git a/src/test/java/org/tikv/common/PDClientV2MockTest.java b/src/test/java/org/tikv/common/PDClientV2MockTest.java index d5ffcdcb005..8c3f156ff35 100644 --- a/src/test/java/org/tikv/common/PDClientV2MockTest.java +++ b/src/test/java/org/tikv/common/PDClientV2MockTest.java @@ -19,6 +19,7 @@ import com.google.protobuf.ByteString; import java.util.List; +import org.apache.commons.lang3.tuple.Pair; import org.junit.Assert; import org.junit.Before; import org.junit.Test; @@ -52,10 +53,15 @@ private GetRegionResponse makeGetRegionResponse(String start, String end) { } private Metapb.Region makeRegion(String start, String end) { + Pair range = + session + .getPDClient() + .getCodec() + .encodePdQueryRange(ByteString.copyFromUtf8(start), ByteString.copyFromUtf8(end)); return GrpcUtils.makeRegion( 1, - encode(session.getConf().buildRequestKey(ByteString.copyFromUtf8(start))), - encode(session.getConf().buildRequestKey(ByteString.copyFromUtf8(end), true)), + range.getLeft(), + range.getRight(), GrpcUtils.makeRegionEpoch(2, 3), GrpcUtils.makePeer(1, 10), GrpcUtils.makePeer(2, 20)); diff --git a/src/test/java/org/tikv/common/RegionErrorTest.java b/src/test/java/org/tikv/common/RegionErrorTest.java index 94f70115abc..5341cbff608 100644 --- a/src/test/java/org/tikv/common/RegionErrorTest.java +++ b/src/test/java/org/tikv/common/RegionErrorTest.java @@ -45,7 +45,7 @@ public void testOnEpochNotMatch() { ByteString key = ByteString.copyFromUtf8("key-test-epoch-not-match"); ByteString value = ByteString.copyFromUtf8("value"); - ByteString requestKey = client.getSession().getConf().buildRequestKey(key); + ByteString requestKey = client.getSession().getPDClient().getCodec().encodeKey(key); put(requestKey, value); Assert.assertEquals(Optional.of(value), client.get(key)); From a72a6cf9bee75800bb270ce0fbe703cde833066a Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Mon, 30 May 2022 20:18:49 +0800 Subject: [PATCH 21/32] ./dev/javafmt Signed-off-by: iosmanthus --- .../common/apiversion/RequestKeyCodec.java | 1 - .../common/apiversion/RequestKeyV2Codec.java | 4 +- .../tikv/common/importer/ImporterClient.java | 2 - .../tikv/common/region/RegionStoreClient.java | 56 ++++++++----------- 4 files changed, 26 insertions(+), 37 deletions(-) diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java index 3fbbea82b33..5eb7637c205 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java @@ -11,7 +11,6 @@ public interface RequestKeyCodec { Pair encodeRange(ByteString start, ByteString end); - ByteString encodePdQuery(ByteString key); default Pair encodePdQueryRange(ByteString start, ByteString end) { diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java index 22000afc87e..7388960a24c 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java @@ -8,11 +8,11 @@ public class RequestKeyV2Codec implements RequestKeyCodec { protected static ByteString RAW_KEY_PREFIX = ByteString.copyFromUtf8("r"); protected static ByteString RAW_END_KEY = - ByteString.copyFrom(new byte[]{(byte) (RAW_KEY_PREFIX.toByteArray()[0] + 1)}); + ByteString.copyFrom(new byte[] {(byte) (RAW_KEY_PREFIX.toByteArray()[0] + 1)}); protected static ByteString TXN_KEY_PREFIX = ByteString.copyFromUtf8("x"); protected static ByteString TXN_END_KEY = - ByteString.copyFrom(new byte[]{(byte) (TXN_KEY_PREFIX.toByteArray()[0] + 1)}); + ByteString.copyFrom(new byte[] {(byte) (TXN_KEY_PREFIX.toByteArray()[0] + 1)}); protected ByteString keyPrefix; protected ByteString infiniteEndKey; diff --git a/src/main/java/org/tikv/common/importer/ImporterClient.java b/src/main/java/org/tikv/common/importer/ImporterClient.java index f8ddce88242..2b89373842b 100644 --- a/src/main/java/org/tikv/common/importer/ImporterClient.java +++ b/src/main/java/org/tikv/common/importer/ImporterClient.java @@ -30,8 +30,6 @@ import org.tikv.common.TiConfiguration; import org.tikv.common.TiSession; import org.tikv.common.apiversion.RequestKeyCodec; -import org.tikv.common.codec.Codec; -import org.tikv.common.codec.CodecDataOutput; import org.tikv.common.exception.GrpcException; import org.tikv.common.exception.RegionException; import org.tikv.common.exception.TiKVException; diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index 9d0039e28ca..04f61699d5e 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -124,17 +124,12 @@ // need to be re-split across regions/stores, region info outdated, e.t.c., you // should retry it in an upper client logic (KVClient, TxnClient, e.t.c.) -/** - * Note that RegionStoreClient itself is not thread-safe - */ +/** Note that RegionStoreClient itself is not thread-safe */ public class RegionStoreClient extends AbstractRegionStoreClient { private static final Logger logger = LoggerFactory.getLogger(RegionStoreClient.class); - @VisibleForTesting - public final AbstractLockResolverClient lockResolverClient; + @VisibleForTesting public final AbstractLockResolverClient lockResolverClient; private final TiStoreType storeType; - /** - * startTS -> List(locks) - */ + /** startTS -> List(locks) */ private final Map> resolvedLocks = new HashMap<>(); private final PDClient pdClient; @@ -460,23 +455,23 @@ public void prewrite( () -> getIsV4() ? PrewriteRequest.newBuilder() - .setContext(makeContext(storeType, bo.getSlowLog())) - .setStartVersion(startTs) - .setPrimaryLock(codec.encodeKey(primaryLock)) - .addAllMutations( - mutations - .stream() - .map( - mutation -> - Mutation.newBuilder(mutation) - .setKey(codec.encodeKey(mutation.getKey())) - .build()) - .collect(Collectors.toList())) - .setLockTtl(ttl) - .setSkipConstraintCheck(skipConstraintCheck) - .setMinCommitTs(startTs) - .setTxnSize(16) - .build() + .setContext(makeContext(storeType, bo.getSlowLog())) + .setStartVersion(startTs) + .setPrimaryLock(codec.encodeKey(primaryLock)) + .addAllMutations( + mutations + .stream() + .map( + mutation -> + Mutation.newBuilder(mutation) + .setKey(codec.encodeKey(mutation.getKey())) + .build()) + .collect(Collectors.toList())) + .setLockTtl(ttl) + .setSkipConstraintCheck(skipConstraintCheck) + .setMinCommitTs(startTs) + .setTxnSize(16) + .build() : PrewriteRequest.newBuilder() .setContext(makeContext(storeType, bo.getSlowLog())) .setStartVersion(startTs) @@ -508,8 +503,8 @@ public void prewrite( * @param backOffer backOffer * @param resp response * @return Return true means the rpc call success. Return false means the rpc call fail, - * RegionStoreClient should retry. Throw an Exception means the rpc call fail, RegionStoreClient - * cannot handle this kind of error + * RegionStoreClient should retry. Throw an Exception means the rpc call fail, + * RegionStoreClient cannot handle this kind of error * @throws TiClientInternalException * @throws RegionException * @throws KeyException @@ -551,9 +546,7 @@ private boolean isPrewriteSuccess(BackOffer backOffer, PrewriteResponse resp, lo return false; } - /** - * TXN Heart Beat: update primary key ttl - */ + /** TXN Heart Beat: update primary key ttl */ public void txnHeartBeat(BackOffer bo, ByteString primaryLock, long startTs, long ttl) { boolean forWrite = false; while (true) { @@ -881,8 +874,7 @@ public List splitRegion(List splitKeys) { if (conf.getApiVersion().isV1()) { return resp.getRegionsList(); } - return resp.getRegionsList().stream().map(codec::decodeRegion).collect( - Collectors.toList()); + return resp.getRegionsList().stream().map(codec::decodeRegion).collect(Collectors.toList()); } // APIs for Raw Scan/Put/Get/Delete From de87cbbcc88a7399cf33c13f18dd1363a72cd5bc Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Mon, 30 May 2022 20:20:15 +0800 Subject: [PATCH 22/32] remove codec in TiConfiguration Signed-off-by: iosmanthus --- .../java/org/tikv/common/TiConfiguration.java | 55 ------------------- 1 file changed, 55 deletions(-) diff --git a/src/main/java/org/tikv/common/TiConfiguration.java b/src/main/java/org/tikv/common/TiConfiguration.java index be6e0314f45..9f6bb097d59 100644 --- a/src/main/java/org/tikv/common/TiConfiguration.java +++ b/src/main/java/org/tikv/common/TiConfiguration.java @@ -1260,61 +1260,6 @@ public TiConfiguration setApiVersion(ApiVersion version) { return this; } - public ByteString getKeyPrefix() { - if (apiVersion.isV2()) { - return isRawKVMode() ? API_V2_RAW_PREFIX : API_V2_TXN_PREFIX; - } - return ByteString.EMPTY; - } - - public ByteString getEndKey() { - if (apiVersion.isV2()) { - byte end = (byte) (getKeyPrefix().toByteArray()[0] + 1); - return ByteString.copyFrom(new byte[] {end}); - } - return ByteString.EMPTY; - } - - public ByteString buildRequestKey(ByteString key, boolean isEndKey) { - switch (apiVersion) { - case V1: - return key; - case V2: - if (isEndKey && key.isEmpty()) { - return getEndKey(); - } - return getKeyPrefix().concat(key); - default: - throw new IllegalArgumentException("unknown api version or kv mode"); - } - } - - public ByteString buildRequestKey(ByteString key) { - return buildRequestKey(key, false); - } - - public ByteString unwrapResponseKey(ByteString key) { - return unwrapResponseKey(key, false); - } - - public ByteString unwrapResponseKey(ByteString key, boolean isEndKey) { - switch (apiVersion) { - case V1: - return key; - case V2: - if (isEndKey && key.equals(getEndKey())) { - return ByteString.EMPTY; - } - // TODO: empty key is not always scan back!!!! - if (!key.isEmpty() && !key.startsWith(getKeyPrefix())) { - throw new IllegalArgumentException("key corrupted, wrong prefix"); - } - return key.substring(key.isEmpty() ? 0 : 1); - default: - throw new IllegalArgumentException("unknown api version or kv mode"); - } - } - public enum ApiVersion { V1, V2; From 6d43080227ebab3af5b8f20d845df91d2f3c4015 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Mon, 30 May 2022 20:22:25 +0800 Subject: [PATCH 23/32] remove decodeRegion in PDClient Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/PDClient.java | 59 --------------------- 1 file changed, 59 deletions(-) diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index c957054ca71..58756ad2e53 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -826,65 +826,6 @@ public String toString() { } } - // public Metapb.Region decodeRegion(Metapb.Region region) { - // final boolean isRawRegion = conf.isRawKVMode(); - // Metapb.Region.Builder builder = - // Metapb.Region.newBuilder() - // .setId(region.getId()) - // .setRegionEpoch(region.getRegionEpoch()) - // .addAllPeers(region.getPeersList()); - - // if (region.getStartKey().isEmpty() || (conf.getApiVersion().isV1() && isRawRegion)) { - // builder.setStartKey(region.getStartKey()); - // } else { - // try { - // ByteString decodedStartKey = - // ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getStartKey()))); - // if (conf.getApiVersion().isV2()) { - // if (ByteString.unsignedLexicographicalComparator() - // .compare(decodedStartKey, conf.getKeyPrefix()) - // < 0) { - // decodedStartKey = ByteString.EMPTY; - // } else { - // decodedStartKey = decodedStartKey.substring(conf.getKeyPrefix().size()); - // } - // } - // builder.setStartKey(decodedStartKey); - // } catch (Exception e) { - // if (!conf.isTest()) { - // throw e; - // } - // builder.setStartKey(region.getStartKey()); - // } - // } - - // if (region.getEndKey().isEmpty() || (conf.getApiVersion().isV1() && isRawRegion)) { - // builder.setEndKey(region.getEndKey()); - // } else { - // try { - // ByteString decodedEndKey = - // ByteString.copyFrom(BytesCodec.readBytes(new CodecDataInput(region.getEndKey()))); - // if (conf.getApiVersion().isV2()) { - // if (ByteString.unsignedLexicographicalComparator() - // .compare(decodedEndKey, conf.getEndKey()) - // >= 0) { - // decodedEndKey = ByteString.EMPTY; - // } else { - // decodedEndKey = decodedEndKey.substring(conf.getKeyPrefix().size()); - // } - // } - // builder.setEndKey(decodedEndKey); - // } catch (Exception e) { - // if (!conf.isTest()) { - // throw e; - // } - // builder.setEndKey(region.getEndKey()); - // } - // } - - // return builder.build(); - // } - public Long getClusterId() { return header.getClusterId(); } From ef71788fdc8a5effcbf1521f1371b981b3a88de8 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 31 May 2022 00:35:38 +0800 Subject: [PATCH 24/32] remove extra ; Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/ConfigUtils.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/org/tikv/common/ConfigUtils.java b/src/main/java/org/tikv/common/ConfigUtils.java index 4067391e104..4697db8140a 100644 --- a/src/main/java/org/tikv/common/ConfigUtils.java +++ b/src/main/java/org/tikv/common/ConfigUtils.java @@ -122,7 +122,7 @@ public class ConfigUtils { public static final String TIFLASH_ENABLE = "tiflash.enable"; public static final String TIKV_WARM_UP_ENABLE = "tikv.warm_up.enable"; - public static final String TIKV_API_VERSION = "tikv.api_version";; + public static final String TIKV_API_VERSION = "tikv.api_version"; public static final String DEF_PD_ADDRESSES = "127.0.0.1:2379"; public static final String DEF_TIMEOUT = "200ms"; From fec799c849e65b329bdba71fcc3b993c2fa31b13 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 31 May 2022 18:57:19 +0800 Subject: [PATCH 25/32] add batch encode method for RequestKeyCodec Signed-off-by: iosmanthus --- .../tikv/common/apiversion/CodecUtils.java | 3 +- .../common/apiversion/RequestKeyCodec.java | 19 +++++-- .../common/apiversion/RequestKeyV1Codec.java | 49 +++++++++++++++++++ .../apiversion/RequestKeyV1RawCodec.java | 31 +----------- .../apiversion/RequestKeyV1TxnCodec.java | 20 ++------ .../common/apiversion/RequestKeyV2Codec.java | 6 +++ .../tikv/common/region/RegionStoreClient.java | 27 +++++----- 7 files changed, 89 insertions(+), 66 deletions(-) create mode 100644 src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java diff --git a/src/main/java/org/tikv/common/apiversion/CodecUtils.java b/src/main/java/org/tikv/common/apiversion/CodecUtils.java index 803007bbc78..a3dea00bb32 100644 --- a/src/main/java/org/tikv/common/apiversion/CodecUtils.java +++ b/src/main/java/org/tikv/common/apiversion/CodecUtils.java @@ -5,7 +5,8 @@ import org.tikv.common.codec.CodecDataInput; import org.tikv.common.codec.CodecDataOutput; -public class CodecUtils { +// TODO(iosmanthus): use ByteString.wrap to avoid once more copying. +class CodecUtils { public static ByteString encode(ByteString key) { CodecDataOutput cdo = new CodecDataOutput(); BytesCodec.writeBytes(cdo, key.toByteArray()); diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java index 5eb7637c205..681cc797db9 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java @@ -1,22 +1,33 @@ package org.tikv.common.apiversion; import com.google.protobuf.ByteString; +import java.util.List; +import java.util.stream.Collectors; import org.apache.commons.lang3.tuple.Pair; +import org.tikv.kvproto.Kvrpcpb.Mutation; import org.tikv.kvproto.Metapb; public interface RequestKeyCodec { ByteString encodeKey(ByteString key); + default List encodeKeys(List keys) { + return keys.stream().map(this::encodeKey).collect(Collectors.toList()); + } + + default List encodeMutations(List mutations) { + return mutations + .stream() + .map(mut -> Mutation.newBuilder().mergeFrom(mut).setKey(encodeKey(mut.getKey())).build()) + .collect(Collectors.toList()); + } + ByteString decodeKey(ByteString key); Pair encodeRange(ByteString start, ByteString end); ByteString encodePdQuery(ByteString key); - default Pair encodePdQueryRange(ByteString start, ByteString end) { - Pair range = encodeRange(start, end); - return Pair.of(CodecUtils.encode(range.getLeft()), CodecUtils.encode(range.getRight())); - } + Pair encodePdQueryRange(ByteString start, ByteString end); Metapb.Region decodeRegion(Metapb.Region region); } diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java new file mode 100644 index 00000000000..084587ff24e --- /dev/null +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java @@ -0,0 +1,49 @@ +package org.tikv.common.apiversion; + +import com.google.protobuf.ByteString; +import java.util.List; +import org.apache.commons.lang3.tuple.Pair; +import org.tikv.kvproto.Kvrpcpb.Mutation; +import org.tikv.kvproto.Metapb.Region; + +public class RequestKeyV1Codec implements RequestKeyCodec { + @Override + public ByteString encodeKey(ByteString key) { + return key; + } + + @Override + public List encodeKeys(List keys) { + return keys; + } + + @Override + public List encodeMutations(List mutations) { + return mutations; + } + + @Override + public ByteString decodeKey(ByteString key) { + return key; + } + + @Override + public Pair encodeRange(ByteString start, ByteString end) { + return Pair.of(start, end); + } + + @Override + public ByteString encodePdQuery(ByteString key) { + return key; + } + + @Override + public Pair encodePdQueryRange(ByteString start, ByteString end) { + return Pair.of(start, end); + } + + @Override + public Region decodeRegion(Region region) { + return region; + } +} diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java index 2611f30b924..0926f39638c 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java @@ -1,34 +1,5 @@ package org.tikv.common.apiversion; -import com.google.protobuf.ByteString; -import org.apache.commons.lang3.tuple.Pair; -import org.tikv.kvproto.Metapb; - -public class RequestKeyV1RawCodec implements RequestKeyCodec { +public class RequestKeyV1RawCodec extends RequestKeyV1Codec implements RequestKeyCodec { public RequestKeyV1RawCodec() {} - - @Override - public ByteString encodeKey(ByteString key) { - return key; - } - - @Override - public ByteString decodeKey(ByteString key) { - return key; - } - - @Override - public Pair encodeRange(ByteString start, ByteString end) { - return Pair.of(start, end); - } - - @Override - public ByteString encodePdQuery(ByteString key) { - return key; - } - - @Override - public Metapb.Region decodeRegion(Metapb.Region region) { - return region; - } } diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java index 6cf3cdcb0a1..d0f71e6f8dd 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java @@ -4,27 +4,17 @@ import org.apache.commons.lang3.tuple.Pair; import org.tikv.kvproto.Metapb; -public class RequestKeyV1TxnCodec implements RequestKeyCodec { +public class RequestKeyV1TxnCodec extends RequestKeyV1Codec implements RequestKeyCodec { public RequestKeyV1TxnCodec() {} @Override - public ByteString encodeKey(ByteString key) { - return key; - } - - @Override - public ByteString decodeKey(ByteString key) { - return key; - } - - @Override - public Pair encodeRange(ByteString start, ByteString end) { - return Pair.of(start, end); + public ByteString encodePdQuery(ByteString key) { + return CodecUtils.encode(key); } @Override - public ByteString encodePdQuery(ByteString key) { - return CodecUtils.encode(key); + public Pair encodePdQueryRange(ByteString start, ByteString end) { + return Pair.of(CodecUtils.encode(start), CodecUtils.encode(end)); } @Override diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java index 7388960a24c..92cd5548531 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java @@ -49,6 +49,12 @@ public ByteString encodePdQuery(ByteString key) { return CodecUtils.encode(encodeKey(key)); } + @Override + public Pair encodePdQueryRange(ByteString start, ByteString end) { + Pair range = encodeRange(start, end); + return Pair.of(CodecUtils.encode(range.getLeft()), CodecUtils.encode(range.getRight())); + } + @Override public Region decodeRegion(Region region) { Metapb.Region.Builder builder = Metapb.Region.newBuilder().mergeFrom(region); diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index 04f61699d5e..fffb1bfdf18 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -44,6 +44,7 @@ import java.util.Set; import java.util.function.Supplier; import java.util.stream.Collectors; +import java.util.stream.StreamSupport; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.tikv.common.PDClient; @@ -282,7 +283,7 @@ public List batchGet(BackOffer backOffer, List keys, long ve BatchGetRequest.newBuilder() .setContext( makeContext(getResolvedLocks(version), this.storeType, backOffer.getSlowLog())) - .addAllKeys(keys.stream().map(codec::encodeKey).collect(Collectors.toList())) + .addAllKeys(codec.encodeKeys(keys)) .setVersion(version) .build(); KVErrorHandler handler = @@ -458,15 +459,7 @@ public void prewrite( .setContext(makeContext(storeType, bo.getSlowLog())) .setStartVersion(startTs) .setPrimaryLock(codec.encodeKey(primaryLock)) - .addAllMutations( - mutations - .stream() - .map( - mutation -> - Mutation.newBuilder(mutation) - .setKey(codec.encodeKey(mutation.getKey())) - .build()) - .collect(Collectors.toList())) + .addAllMutations(codec.encodeMutations(mutations)) .setLockTtl(ttl) .setSkipConstraintCheck(skipConstraintCheck) .setMinCommitTs(startTs) @@ -602,7 +595,7 @@ private boolean isTxnHeartBeatSuccess(TxnHeartBeatResponse resp) * @param startTs start version * @param commitTs commit version */ - public void commit(BackOffer backOffer, List keys, long startTs, long commitTs) + public void commit(BackOffer backOffer, Iterable keys, long startTs, long commitTs) throws KeyException { boolean forWrite = true; Supplier factory = @@ -610,7 +603,10 @@ public void commit(BackOffer backOffer, List keys, long startTs, lon CommitRequest.newBuilder() .setStartVersion(startTs) .setCommitVersion(commitTs) - .addAllKeys(keys.stream().map(codec::encodeKey).collect(Collectors.toList())) + .addAllKeys( + StreamSupport.stream(keys.spliterator(), false) + .map(codec::encodeKey) + .collect(Collectors.toList())) .setContext(makeContext(storeType, backOffer.getSlowLog())) .build(); KVErrorHandler handler = @@ -837,8 +833,7 @@ public List splitRegion(List splitKeys) { () -> SplitRegionRequest.newBuilder() .setContext(makeContext(storeType, SlowLogEmptyImpl.INSTANCE)) - .addAllSplitKeys( - splitKeys.stream().map(codec::encodeKey).collect(Collectors.toList())) + .addAllSplitKeys(codec.encodeKeys(splitKeys)) .setIsRawKv(conf.isRawKVMode()) .build(); @@ -1114,7 +1109,7 @@ public List rawBatchGet(BackOffer backoffer, List keys) { () -> RawBatchGetRequest.newBuilder() .setContext(makeContext(storeType, backoffer.getSlowLog())) - .addAllKeys(keys.stream().map(codec::encodeKey).collect(Collectors.toList())) + .addAllKeys(codec.encodeKeys(keys)) .build(); RegionErrorHandler handler = new RegionErrorHandler( @@ -1222,7 +1217,7 @@ public void rawBatchDelete(BackOffer backoffer, List keys, boolean a () -> RawBatchDeleteRequest.newBuilder() .setContext(makeContext(storeType, backoffer.getSlowLog())) - .addAllKeys(keys.stream().map(codec::encodeKey).collect(Collectors.toList())) + .addAllKeys(codec.encodeKeys(keys)) .setForCas(atomicForCAS) .build(); RegionErrorHandler handler = From 0a64f84218ba63f0a1b26df22403386280260624 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Tue, 31 May 2022 23:52:05 +0800 Subject: [PATCH 26/32] print stack trace for unstable tests Signed-off-by: iosmanthus --- src/test/java/org/tikv/common/PDClientMockTest.java | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/test/java/org/tikv/common/PDClientMockTest.java b/src/test/java/org/tikv/common/PDClientMockTest.java index a8074d94572..8937350afe5 100644 --- a/src/test/java/org/tikv/common/PDClientMockTest.java +++ b/src/test/java/org/tikv/common/PDClientMockTest.java @@ -82,8 +82,8 @@ public void testTso() throws Exception { @Test public void testGetRegionByKey() throws Exception { - byte[] startKey = new byte[] {1, 0, 2, 4}; - byte[] endKey = new byte[] {1, 0, 2, 5}; + byte[] startKey = new byte[]{1, 0, 2, 4}; + byte[] endKey = new byte[]{1, 0, 2, 5}; int confVer = 1026; int ver = 1027; leader.addGetRegionListener( @@ -108,13 +108,16 @@ public void testGetRegionByKey() throws Exception { assertEquals(r.getRegionEpoch().getVersion(), ver); assertEquals(1, l.getId()); assertEquals(10, l.getStoreId()); + } catch (Exception e) { + e.printStackTrace(); + throw e; } } @Test public void testGetRegionById() throws Exception { - byte[] startKey = new byte[] {1, 0, 2, 4}; - byte[] endKey = new byte[] {1, 0, 2, 5}; + byte[] startKey = new byte[]{1, 0, 2, 4}; + byte[] endKey = new byte[]{1, 0, 2, 5}; int confVer = 1026; int ver = 1027; From cedc40dd47574166ebec02db021e09b90aa8793f Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Wed, 1 Jun 2022 00:37:28 +0800 Subject: [PATCH 27/32] fix wrong base class for PDClientV2MockTest Signed-off-by: iosmanthus --- .../org/tikv/common/apiversion/CodecUtils.java | 17 +++++++++++++++++ .../tikv/common/apiversion/RequestKeyCodec.java | 17 +++++++++++++++++ .../common/apiversion/RequestKeyV1Codec.java | 17 +++++++++++++++++ .../common/apiversion/RequestKeyV1RawCodec.java | 17 +++++++++++++++++ .../common/apiversion/RequestKeyV1TxnCodec.java | 17 +++++++++++++++++ .../common/apiversion/RequestKeyV2Codec.java | 17 +++++++++++++++++ .../common/apiversion/RequestKeyV2RawCodec.java | 17 +++++++++++++++++ .../common/apiversion/RequestKeyV2TxnCodec.java | 17 +++++++++++++++++ .../java/org/tikv/common/PDClientMockTest.java | 11 ++++------- .../org/tikv/common/PDClientV2MockTest.java | 2 +- 10 files changed, 141 insertions(+), 8 deletions(-) diff --git a/src/main/java/org/tikv/common/apiversion/CodecUtils.java b/src/main/java/org/tikv/common/apiversion/CodecUtils.java index a3dea00bb32..1c6cfea9fa1 100644 --- a/src/main/java/org/tikv/common/apiversion/CodecUtils.java +++ b/src/main/java/org/tikv/common/apiversion/CodecUtils.java @@ -1,3 +1,20 @@ +/* + * 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.apiversion; import com.google.protobuf.ByteString; diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java index 681cc797db9..e63df9a6bb8 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java @@ -1,3 +1,20 @@ +/* + * 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.apiversion; import com.google.protobuf.ByteString; diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java index 084587ff24e..b3d07d0fbd4 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java @@ -1,3 +1,20 @@ +/* + * 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.apiversion; import com.google.protobuf.ByteString; diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java index 0926f39638c..a72457d0bbe 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1RawCodec.java @@ -1,3 +1,20 @@ +/* + * 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.apiversion; public class RequestKeyV1RawCodec extends RequestKeyV1Codec implements RequestKeyCodec { diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java index d0f71e6f8dd..93bfc1645eb 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java @@ -1,3 +1,20 @@ +/* + * 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.apiversion; import com.google.protobuf.ByteString; diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java index 92cd5548531..594b9b59a74 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java @@ -1,3 +1,20 @@ +/* + * 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.apiversion; import com.google.protobuf.ByteString; diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2RawCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2RawCodec.java index 3f6868a9877..255815d8403 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV2RawCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2RawCodec.java @@ -1,3 +1,20 @@ +/* + * 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.apiversion; public class RequestKeyV2RawCodec extends RequestKeyV2Codec { diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2TxnCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2TxnCodec.java index 47b7ce2427a..9fe0effccd5 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV2TxnCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2TxnCodec.java @@ -1,3 +1,20 @@ +/* + * 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.apiversion; public class RequestKeyV2TxnCodec extends RequestKeyV2Codec { diff --git a/src/test/java/org/tikv/common/PDClientMockTest.java b/src/test/java/org/tikv/common/PDClientMockTest.java index 8937350afe5..a8074d94572 100644 --- a/src/test/java/org/tikv/common/PDClientMockTest.java +++ b/src/test/java/org/tikv/common/PDClientMockTest.java @@ -82,8 +82,8 @@ public void testTso() throws Exception { @Test public void testGetRegionByKey() throws Exception { - byte[] startKey = new byte[]{1, 0, 2, 4}; - byte[] endKey = new byte[]{1, 0, 2, 5}; + byte[] startKey = new byte[] {1, 0, 2, 4}; + byte[] endKey = new byte[] {1, 0, 2, 5}; int confVer = 1026; int ver = 1027; leader.addGetRegionListener( @@ -108,16 +108,13 @@ public void testGetRegionByKey() throws Exception { assertEquals(r.getRegionEpoch().getVersion(), ver); assertEquals(1, l.getId()); assertEquals(10, l.getStoreId()); - } catch (Exception e) { - e.printStackTrace(); - throw e; } } @Test public void testGetRegionById() throws Exception { - byte[] startKey = new byte[]{1, 0, 2, 4}; - byte[] endKey = new byte[]{1, 0, 2, 5}; + byte[] startKey = new byte[] {1, 0, 2, 4}; + byte[] endKey = new byte[] {1, 0, 2, 5}; int confVer = 1026; int ver = 1027; diff --git a/src/test/java/org/tikv/common/PDClientV2MockTest.java b/src/test/java/org/tikv/common/PDClientV2MockTest.java index 8c3f156ff35..53a5968bbdd 100644 --- a/src/test/java/org/tikv/common/PDClientV2MockTest.java +++ b/src/test/java/org/tikv/common/PDClientV2MockTest.java @@ -32,7 +32,7 @@ import org.tikv.kvproto.Pdpb.Region; import org.tikv.kvproto.Pdpb.ScanRegionsResponse; -public class PDClientV2MockTest extends PDClientMockTest { +public class PDClientV2MockTest extends PDMockServerTest { @Before public void init() throws Exception { upgradeToV2Cluster(); From afd5858edeb5eb7483b6801d281001f0d421a8c1 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Wed, 1 Jun 2022 02:03:55 +0800 Subject: [PATCH 28/32] wip: add tests for RequestKeyCodec Signed-off-by: iosmanthus --- .../common/apiversion/RequestKeyCodec.java | 11 ++++ .../common/apiversion/RequestKeyV1Codec.java | 11 ++++ .../tikv/common/region/RegionStoreClient.java | 43 +++------------- .../tikv/txn/AbstractLockResolverClient.java | 2 +- .../apiversion/RequestKeyCodecTest.java | 51 +++++++++++++++++++ 5 files changed, 80 insertions(+), 38 deletions(-) create mode 100644 src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java index e63df9a6bb8..92ebba18464 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java @@ -21,6 +21,7 @@ import java.util.List; import java.util.stream.Collectors; import org.apache.commons.lang3.tuple.Pair; +import org.tikv.kvproto.Kvrpcpb.KvPair; import org.tikv.kvproto.Kvrpcpb.Mutation; import org.tikv.kvproto.Metapb; @@ -40,6 +41,16 @@ default List encodeMutations(List mutations) { ByteString decodeKey(ByteString key); + default KvPair decodeKvPair(KvPair pair) { + return KvPair.newBuilder().mergeFrom(pair).setKey(decodeKey(pair.getKey())).build(); + } + + default List decodeKvPairs(List pairs) { + return pairs.stream() + .map(this::decodeKvPair) + .collect(Collectors.toList()); + } + Pair encodeRange(ByteString start, ByteString end); ByteString encodePdQuery(ByteString key); diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java index b3d07d0fbd4..6f8379fc670 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java @@ -20,6 +20,7 @@ import com.google.protobuf.ByteString; import java.util.List; import org.apache.commons.lang3.tuple.Pair; +import org.tikv.kvproto.Kvrpcpb.KvPair; import org.tikv.kvproto.Kvrpcpb.Mutation; import org.tikv.kvproto.Metapb.Region; @@ -44,6 +45,16 @@ public ByteString decodeKey(ByteString key) { return key; } + @Override + public KvPair decodeKvPair(KvPair pair) { + return pair; + } + + @Override + public List decodeKvPairs(List pairs) { + return pairs; + } + @Override public Pair encodeRange(ByteString start, ByteString end) { return Pair.of(start, end); diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index fffb1bfdf18..b472d3c892b 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -330,20 +330,9 @@ private List handleBatchGetResponse( addResolvedLocks(version, resolveLockResult.getResolvedLocks()); // resolveLocks already retried, just throw error to upper logic. throw new TiKVException("locks not resolved, retry"); - } else { - if (conf.getApiVersion().isV1()) { - return resp.getPairsList(); - } - return resp.getPairsList() - .stream() - .map( - kvPair -> - KvPair.newBuilder() - .mergeFrom(kvPair) - .setKey(codec.decodeKey(kvPair.getKey())) - .build()) - .collect(Collectors.toList()); } + + return codec.decodeKvPairs(resp.getPairsList()); } public List scan( @@ -406,10 +395,10 @@ private List doScan(ScanResponse resp) { KvPair.newBuilder() .setError(kvPair.getError()) .setValue(kvPair.getValue()) - .setKey(codec.decodeKey(lock.getKey())) + .setKey(lock.getKey()) .build()); } else { - newKvPairs.add(kvPair); + newKvPairs.add(codec.decodeKvPair(kvPair)); } } return Collections.unmodifiableList(newKvPairs); @@ -1131,19 +1120,7 @@ private List handleRawBatchGet(RawBatchGetResponse resp) { throw new RegionException(resp.getRegionError()); } - if (conf.getApiVersion().isV1()) { - return resp.getPairsList(); - } - - return resp.getPairsList() - .stream() - .map( - kvPair -> - KvPair.newBuilder() - .mergeFrom(kvPair) - .setKey(codec.decodeKey(kvPair.getKey())) - .build()) - .collect(Collectors.toList()); + return codec.decodeKvPairs(resp.getPairsList()); } public void rawBatchPut( @@ -1295,15 +1272,7 @@ private List rawScanHelper(RawScanResponse resp) { if (resp.hasRegionError()) { throw new RegionException(resp.getRegionError()); } - return resp.getKvsList() - .stream() - .map( - kvPair -> - KvPair.newBuilder() - .setKey(codec.decodeKey(kvPair.getKey())) - .setValue(kvPair.getValue()) - .build()) - .collect(Collectors.toList()); + return codec.decodeKvPairs(resp.getKvsList()); } /** diff --git a/src/main/java/org/tikv/txn/AbstractLockResolverClient.java b/src/main/java/org/tikv/txn/AbstractLockResolverClient.java index 6e535e7ab65..2adc87a78f4 100644 --- a/src/main/java/org/tikv/txn/AbstractLockResolverClient.java +++ b/src/main/java/org/tikv/txn/AbstractLockResolverClient.java @@ -50,7 +50,7 @@ static Lock extractLockFromKeyErr(Kvrpcpb.KeyError keyError, RequestKeyCodec cod throw new KeyException( String.format( "scan meet key conflict on primary key %s at commit ts %s", - conflict.getPrimary(), conflict.getConflictTs())); + codec.decodeKey(conflict.getPrimary()), conflict.getConflictTs())); } if (!keyError.getRetryable().isEmpty()) { diff --git a/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java b/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java new file mode 100644 index 00000000000..fe08aa41f4d --- /dev/null +++ b/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java @@ -0,0 +1,51 @@ +/* + * 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.apiversion; + +import static org.junit.Assert.assertEquals; + +import com.google.protobuf.ByteString; +import org.apache.commons.lang3.tuple.Pair; +import org.junit.Test; +import org.tikv.kvproto.Metapb.Region; + +public class RequestKeyCodecTest { + @Test + public void testV1RawCodec() { + RequestKeyCodec v1 = new RequestKeyV1RawCodec(); + ByteString key = ByteString.copyFromUtf8("testV1RawCodec"); + + assertEquals(key, v1.encodeKey(key)); + assertEquals(key, v1.decodeKey(v1.encodeKey(key))); + + assertEquals(key, v1.encodePdQuery(key)); + + ByteString start = ByteString.copyFromUtf8("testV1RawCodec_start"); + ByteString end = ByteString.copyFromUtf8("testV1RawCodec_end"); + Pair range = v1.encodeRange(start, end); + assertEquals(start, range.getLeft()); + assertEquals(end, range.getRight()); + + range = v1.encodePdQueryRange(start, end); + assertEquals(start, range.getLeft()); + assertEquals(end, range.getRight()); + + Region region = Region.newBuilder().setStartKey(start).setEndKey(end).build(); + assertEquals(region, v1.decodeRegion(region)); + } +} From 7c021543322fdbbcd440cd31ce50609a81ea53a7 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Wed, 1 Jun 2022 02:04:23 +0800 Subject: [PATCH 29/32] ./dev/javafmt Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java index 92ebba18464..fdb67162563 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java @@ -46,9 +46,7 @@ default KvPair decodeKvPair(KvPair pair) { } default List decodeKvPairs(List pairs) { - return pairs.stream() - .map(this::decodeKvPair) - .collect(Collectors.toList()); + return pairs.stream().map(this::decodeKvPair).collect(Collectors.toList()); } Pair encodeRange(ByteString start, ByteString end); From 0e0bd466eda4b67858e187b4703fb2fef3a3eb6b Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Wed, 1 Jun 2022 16:47:15 +0800 Subject: [PATCH 30/32] complete RequestKeyCodecTest Signed-off-by: iosmanthus --- .../apiversion/RequestKeyV1TxnCodec.java | 10 +- .../common/apiversion/RequestKeyV2Codec.java | 8 +- .../apiversion/RequestKeyCodecTest.java | 119 ++++++++++++++++++ 3 files changed, 132 insertions(+), 5 deletions(-) diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java index 93bfc1645eb..c994b6189cd 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java @@ -31,7 +31,15 @@ public ByteString encodePdQuery(ByteString key) { @Override public Pair encodePdQueryRange(ByteString start, ByteString end) { - return Pair.of(CodecUtils.encode(start), CodecUtils.encode(end)); + if (!start.isEmpty()) { + start = CodecUtils.encode(start); + } + + if (!end.isEmpty()) { + end = CodecUtils.encode(end); + } + + return Pair.of(start, end); } @Override diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java index 594b9b59a74..231a244ed1a 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java @@ -23,12 +23,12 @@ import org.tikv.kvproto.Metapb.Region; public class RequestKeyV2Codec implements RequestKeyCodec { - protected static ByteString RAW_KEY_PREFIX = ByteString.copyFromUtf8("r"); - protected static ByteString RAW_END_KEY = + protected static final ByteString RAW_KEY_PREFIX = ByteString.copyFromUtf8("r"); + protected static final ByteString RAW_END_KEY = ByteString.copyFrom(new byte[] {(byte) (RAW_KEY_PREFIX.toByteArray()[0] + 1)}); - protected static ByteString TXN_KEY_PREFIX = ByteString.copyFromUtf8("x"); - protected static ByteString TXN_END_KEY = + protected static final ByteString TXN_KEY_PREFIX = ByteString.copyFromUtf8("x"); + protected static final ByteString TXN_END_KEY = ByteString.copyFrom(new byte[] {(byte) (TXN_KEY_PREFIX.toByteArray()[0] + 1)}); protected ByteString keyPrefix; diff --git a/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java b/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java index fe08aa41f4d..d38155c7d5a 100644 --- a/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java +++ b/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java @@ -17,6 +17,7 @@ package org.tikv.common.apiversion; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import com.google.protobuf.ByteString; @@ -48,4 +49,122 @@ public void testV1RawCodec() { Region region = Region.newBuilder().setStartKey(start).setEndKey(end).build(); assertEquals(region, v1.decodeRegion(region)); } + + @Test + public void testV1TxnCodec() { + RequestKeyCodec v1 = new RequestKeyV1TxnCodec(); + + ByteString key = ByteString.copyFromUtf8("testV1TxnCodec"); + + assertEquals(CodecUtils.encode(key), v1.encodePdQuery(key)); + + ByteString start = ByteString.copyFromUtf8("testV1TxnCodec_start"); + ByteString end = ByteString.copyFromUtf8("testV1TxnCodec_end"); + + // Test start and end are both non-empty. + Pair range = v1.encodePdQueryRange(start, end); + assertEquals(CodecUtils.encode(start), range.getLeft()); + assertEquals(CodecUtils.encode(end), range.getRight()); + + Region region = + Region.newBuilder() + .setStartKey(CodecUtils.encode(start)) + .setEndKey(CodecUtils.encode(end)) + .build(); + Region decoded = v1.decodeRegion(region); + assertEquals(start, decoded.getStartKey()); + assertEquals(end, decoded.getEndKey()); + + // Test start is empty. + start = ByteString.EMPTY; + region = + Region.newBuilder() + .setStartKey(CodecUtils.encode(start)) + .setEndKey(CodecUtils.encode(end)) + .build(); + decoded = v1.decodeRegion(region); + assertEquals(start, decoded.getStartKey()); + assertEquals(end, decoded.getEndKey()); + + range = v1.encodePdQueryRange(start, end); + assertEquals(ByteString.EMPTY, range.getLeft()); + assertEquals(CodecUtils.encode(end), range.getRight()); + + // Test end is empty. + end = ByteString.EMPTY; + region = + Region.newBuilder() + .setStartKey(CodecUtils.encode(start)) + .setEndKey(CodecUtils.encode(end)) + .build(); + decoded = v1.decodeRegion(region); + assertEquals(start, decoded.getStartKey()); + assertEquals(ByteString.EMPTY, decoded.getEndKey()); + + range = v1.encodePdQueryRange(start, end); + assertEquals(start, range.getLeft()); + assertEquals(ByteString.EMPTY, range.getRight()); + } + + @Test + public void testV2Codec() { + testV2Codec(new RequestKeyV2RawCodec()); + testV2Codec(new RequestKeyV2TxnCodec()); + } + + void testV2Codec(RequestKeyV2Codec v2) { + ByteString key = ByteString.copyFromUtf8("testV2RawCodec"); + + assertEquals(key, v2.decodeKey(v2.encodeKey(key))); + assertEquals(CodecUtils.encode(v2.encodeKey(key)), v2.encodePdQuery(key)); + + ByteString start = ByteString.copyFromUtf8("testV1TxnCodec_start"); + ByteString end = ByteString.copyFromUtf8("testV1TxnCodec_end"); + + // Test start and end are both non-empty. + Pair range = v2.encodePdQueryRange(start, end); + assertEquals(CodecUtils.encode(v2.encodeKey(start)), range.getLeft()); + assertEquals(CodecUtils.encode(v2.encodeKey(end)), range.getRight()); + + Region region = + Region.newBuilder() + .setStartKey(CodecUtils.encode(v2.encodeKey(start))) + .setEndKey(CodecUtils.encode(v2.encodeKey(end))) + .build(); + Region decoded = v2.decodeRegion(region); + assertEquals(start, decoded.getStartKey()); + assertEquals(end, decoded.getEndKey()); + + // Test start is empty. + start = ByteString.EMPTY; + region = + Region.newBuilder() + .setStartKey(CodecUtils.encode(v2.encodeKey(start))) + .setEndKey(CodecUtils.encode(v2.encodeKey(end))) + .build(); + decoded = v2.decodeRegion(region); + assertEquals(start, decoded.getStartKey()); + assertEquals(end, decoded.getEndKey()); + + range = v2.encodePdQueryRange(start, end); + assertEquals(CodecUtils.encode(v2.encodeKey(start)), range.getLeft()); + assertEquals(CodecUtils.encode(v2.encodeKey(end)), range.getRight()); + + // Test end is empty. + end = ByteString.EMPTY; + range = v2.encodeRange(start, end); + assertEquals(v2.encodeKey(start), range.getLeft()); + assertArrayEquals( + new byte[] {(byte) (v2.encodeKey(ByteString.EMPTY).byteAt(0) + 1)}, + range.getRight().toByteArray()); + + region = + Region.newBuilder() + .setStartKey(CodecUtils.encode(range.getLeft())) + .setEndKey(CodecUtils.encode(range.getRight())) + .build(); + decoded = v2.decodeRegion(region); + assertEquals(start, decoded.getStartKey()); + assertEquals(ByteString.EMPTY, decoded.getEndKey()); + } } From 1db88aaa1c882c46d3f534010624a179317eae3f Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Wed, 1 Jun 2022 23:27:49 +0800 Subject: [PATCH 31/32] remove if in scanRegions Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/PDClient.java | 13 +------------ .../org/tikv/common/apiversion/RequestKeyCodec.java | 13 +++++++++++++ .../tikv/common/apiversion/RequestKeyV1Codec.java | 6 ++++++ .../tikv/common/apiversion/RequestKeyCodecTest.java | 10 ++++++++++ 4 files changed, 30 insertions(+), 12 deletions(-) diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index 58756ad2e53..ea45f177f78 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -368,18 +368,7 @@ public List scanRegions( return null; } - if (conf.getApiVersion().isV1() && conf.isRawKVMode()) { - return resp.getRegionsList(); - } - - return resp.getRegionsList() - .stream() - .map( - r -> { - Metapb.Region decoded = codec.decodeRegion(r.getRegion()); - return Pdpb.Region.newBuilder().mergeFrom(r).setRegion(decoded).build(); - }) - .collect(Collectors.toList()); + return codec.decodePdRegions(resp.getRegionsList()); } private Supplier buildGetStoreReq(long storeId) { diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java index fdb67162563..6a24a931618 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java @@ -24,6 +24,7 @@ import org.tikv.kvproto.Kvrpcpb.KvPair; import org.tikv.kvproto.Kvrpcpb.Mutation; import org.tikv.kvproto.Metapb; +import org.tikv.kvproto.Pdpb; public interface RequestKeyCodec { ByteString encodeKey(ByteString key); @@ -56,4 +57,16 @@ default List decodeKvPairs(List pairs) { Pair encodePdQueryRange(ByteString start, ByteString end); Metapb.Region decodeRegion(Metapb.Region region); + + default List decodePdRegions(List regions) { + return regions + .stream() + .map( + r -> + Pdpb.Region.newBuilder() + .mergeFrom(r) + .setRegion(this.decodeRegion(r.getRegion())) + .build()) + .collect(Collectors.toList()); + } } diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java index 6f8379fc670..b3bf8fe00ee 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java @@ -23,6 +23,7 @@ import org.tikv.kvproto.Kvrpcpb.KvPair; import org.tikv.kvproto.Kvrpcpb.Mutation; import org.tikv.kvproto.Metapb.Region; +import org.tikv.kvproto.Pdpb; public class RequestKeyV1Codec implements RequestKeyCodec { @Override @@ -74,4 +75,9 @@ public Pair encodePdQueryRange(ByteString start, ByteStr public Region decodeRegion(Region region) { return region; } + + @Override + public List decodePdRegions(List regions) { + return regions; + } } diff --git a/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java b/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java index d38155c7d5a..441aea9c4ef 100644 --- a/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java +++ b/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java @@ -20,10 +20,13 @@ import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; +import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; +import java.util.stream.Collectors; import org.apache.commons.lang3.tuple.Pair; import org.junit.Test; import org.tikv.kvproto.Metapb.Region; +import org.tikv.kvproto.Pdpb; public class RequestKeyCodecTest { @Test @@ -48,6 +51,13 @@ public void testV1RawCodec() { Region region = Region.newBuilder().setStartKey(start).setEndKey(end).build(); assertEquals(region, v1.decodeRegion(region)); + + assertEquals( + ImmutableList.of(region), + v1.decodePdRegions(ImmutableList.of(Pdpb.Region.newBuilder().setRegion(region).build())) + .stream() + .map(Pdpb.Region::getRegion) + .collect(Collectors.toList())); } @Test From 629d4faa790c135c550e0ae379f7de75c93295a9 Mon Sep 17 00:00:00 2001 From: iosmanthus Date: Thu, 2 Jun 2022 16:08:44 +0800 Subject: [PATCH 32/32] change Pair type to org.tikv.common.util.Pair Signed-off-by: iosmanthus --- src/main/java/org/tikv/common/PDClient.java | 7 ++-- src/main/java/org/tikv/common/TiSession.java | 4 +- .../common/apiversion/RequestKeyCodec.java | 2 +- .../common/apiversion/RequestKeyV1Codec.java | 6 +-- .../apiversion/RequestKeyV1TxnCodec.java | 4 +- .../common/apiversion/RequestKeyV2Codec.java | 6 +-- .../tikv/common/importer/ImporterClient.java | 7 +--- .../tikv/common/region/RegionStoreClient.java | 14 +++---- .../org/tikv/common/PDClientV2MockTest.java | 6 +-- .../apiversion/RequestKeyCodecTest.java | 38 +++++++++---------- 10 files changed, 44 insertions(+), 50 deletions(-) diff --git a/src/main/java/org/tikv/common/PDClient.java b/src/main/java/org/tikv/common/PDClient.java index ea45f177f78..cfafbff4550 100644 --- a/src/main/java/org/tikv/common/PDClient.java +++ b/src/main/java/org/tikv/common/PDClient.java @@ -354,13 +354,12 @@ public List scanRegions( // introduce a warm-up timeout for ScanRegions requests PDGrpc.PDBlockingStub stub = getBlockingStub().withDeadlineAfter(conf.getWarmUpTimeout(), TimeUnit.MILLISECONDS); - org.apache.commons.lang3.tuple.Pair range = - codec.encodePdQueryRange(startKey, endKey); + Pair range = codec.encodePdQueryRange(startKey, endKey); Pdpb.ScanRegionsRequest request = Pdpb.ScanRegionsRequest.newBuilder() .setHeader(header) - .setStartKey(range.getLeft()) - .setEndKey(range.getRight()) + .setStartKey(range.first) + .setEndKey(range.second) .setLimit(limit) .build(); Pdpb.ScanRegionsResponse resp = stub.scanRegions(request); diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java index bf2365390f6..1cc3ec684a1 100644 --- a/src/main/java/org/tikv/common/TiSession.java +++ b/src/main/java/org/tikv/common/TiSession.java @@ -180,7 +180,7 @@ public TiSession(TiConfiguration conf) { logger.info("enable grpc forward for high available"); } if (conf.isWarmUpEnable() && conf.isRawKVMode()) { - warmup(); + warmUp(); } this.circuitBreaker = new CircuitBreakerImpl(conf, client.getClusterId()); logger.info( @@ -206,7 +206,7 @@ private static VersionInfo getVersionInfo() { } @VisibleForTesting - public synchronized void warmup() { + public synchronized void warmUp() { long warmUpStartTime = System.nanoTime(); BackOffer backOffer = ConcreteBackOffer.newRawKVBackOff(getPDClient().getClusterId()); try { diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java index 6a24a931618..b70e660cd53 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyCodec.java @@ -20,7 +20,7 @@ import com.google.protobuf.ByteString; import java.util.List; import java.util.stream.Collectors; -import org.apache.commons.lang3.tuple.Pair; +import org.tikv.common.util.Pair; import org.tikv.kvproto.Kvrpcpb.KvPair; import org.tikv.kvproto.Kvrpcpb.Mutation; import org.tikv.kvproto.Metapb; diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java index b3bf8fe00ee..9d9a92c6859 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1Codec.java @@ -19,7 +19,7 @@ import com.google.protobuf.ByteString; import java.util.List; -import org.apache.commons.lang3.tuple.Pair; +import org.tikv.common.util.Pair; import org.tikv.kvproto.Kvrpcpb.KvPair; import org.tikv.kvproto.Kvrpcpb.Mutation; import org.tikv.kvproto.Metapb.Region; @@ -58,7 +58,7 @@ public List decodeKvPairs(List pairs) { @Override public Pair encodeRange(ByteString start, ByteString end) { - return Pair.of(start, end); + return Pair.create(start, end); } @Override @@ -68,7 +68,7 @@ public ByteString encodePdQuery(ByteString key) { @Override public Pair encodePdQueryRange(ByteString start, ByteString end) { - return Pair.of(start, end); + return Pair.create(start, end); } @Override diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java index c994b6189cd..ea3949ddb0c 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV1TxnCodec.java @@ -18,7 +18,7 @@ package org.tikv.common.apiversion; import com.google.protobuf.ByteString; -import org.apache.commons.lang3.tuple.Pair; +import org.tikv.common.util.Pair; import org.tikv.kvproto.Metapb; public class RequestKeyV1TxnCodec extends RequestKeyV1Codec implements RequestKeyCodec { @@ -39,7 +39,7 @@ public Pair encodePdQueryRange(ByteString start, ByteStr end = CodecUtils.encode(end); } - return Pair.of(start, end); + return Pair.create(start, end); } @Override diff --git a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java index 231a244ed1a..2af25a1ab98 100644 --- a/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java +++ b/src/main/java/org/tikv/common/apiversion/RequestKeyV2Codec.java @@ -18,7 +18,7 @@ package org.tikv.common.apiversion; import com.google.protobuf.ByteString; -import org.apache.commons.lang3.tuple.Pair; +import org.tikv.common.util.Pair; import org.tikv.kvproto.Metapb; import org.tikv.kvproto.Metapb.Region; @@ -58,7 +58,7 @@ public Pair encodeRange(ByteString start, ByteString end end = end.isEmpty() ? infiniteEndKey : encodeKey(end); - return Pair.of(start, end); + return Pair.create(start, end); } @Override @@ -69,7 +69,7 @@ public ByteString encodePdQuery(ByteString key) { @Override public Pair encodePdQueryRange(ByteString start, ByteString end) { Pair range = encodeRange(start, end); - return Pair.of(CodecUtils.encode(range.getLeft()), CodecUtils.encode(range.getRight())); + return Pair.create(CodecUtils.encode(range.first), CodecUtils.encode(range.second)); } @Override diff --git a/src/main/java/org/tikv/common/importer/ImporterClient.java b/src/main/java/org/tikv/common/importer/ImporterClient.java index 2b89373842b..6f45a9a4d17 100644 --- a/src/main/java/org/tikv/common/importer/ImporterClient.java +++ b/src/main/java/org/tikv/common/importer/ImporterClient.java @@ -140,14 +140,11 @@ public void write(Iterator> iterator) throws TiKVEx private void init() { long regionId = region.getId(); Metapb.RegionEpoch regionEpoch = region.getRegionEpoch(); - org.apache.commons.lang3.tuple.Pair keyRange = + Pair keyRange = codec.encodePdQueryRange(minKey.toByteString(), maxKey.toByteString()); ImportSstpb.Range range = - ImportSstpb.Range.newBuilder() - .setStart(keyRange.getLeft()) - .setEnd(keyRange.getRight()) - .build(); + ImportSstpb.Range.newBuilder().setStart(keyRange.first).setEnd(keyRange.second).build(); sstMeta = ImportSstpb.SSTMeta.newBuilder() diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index b472d3c892b..42576fa7ea7 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -1238,12 +1238,11 @@ public List rawScan(BackOffer backOffer, ByteString key, int limit, bool try { Supplier factory = () -> { - org.apache.commons.lang3.tuple.Pair range = - codec.encodeRange(key, ByteString.EMPTY); + Pair range = codec.encodeRange(key, ByteString.EMPTY); return RawScanRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setStartKey(range.getLeft()) - .setEndKey(range.getRight()) + .setStartKey(range.first) + .setEndKey(range.second) .setKeyOnly(keyOnly) .setLimit(limit) .build(); @@ -1291,12 +1290,11 @@ public void rawDeleteRange(BackOffer backOffer, ByteString startKey, ByteString try { Supplier factory = () -> { - org.apache.commons.lang3.tuple.Pair range = - codec.encodeRange(startKey, endKey); + Pair range = codec.encodeRange(startKey, endKey); return RawDeleteRangeRequest.newBuilder() .setContext(makeContext(storeType, backOffer.getSlowLog())) - .setStartKey(range.getLeft()) - .setEndKey(range.getRight()) + .setStartKey(range.first) + .setEndKey(range.second) .build(); }; diff --git a/src/test/java/org/tikv/common/PDClientV2MockTest.java b/src/test/java/org/tikv/common/PDClientV2MockTest.java index 53a5968bbdd..1a5278cb6ef 100644 --- a/src/test/java/org/tikv/common/PDClientV2MockTest.java +++ b/src/test/java/org/tikv/common/PDClientV2MockTest.java @@ -19,13 +19,13 @@ import com.google.protobuf.ByteString; import java.util.List; -import org.apache.commons.lang3.tuple.Pair; import org.junit.Assert; import org.junit.Before; import org.junit.Test; import org.tikv.common.codec.Codec.BytesCodec; import org.tikv.common.codec.CodecDataOutput; import org.tikv.common.util.ConcreteBackOffer; +import org.tikv.common.util.Pair; import org.tikv.kvproto.Metapb; import org.tikv.kvproto.Pdpb; import org.tikv.kvproto.Pdpb.GetRegionResponse; @@ -60,8 +60,8 @@ private Metapb.Region makeRegion(String start, String end) { .encodePdQueryRange(ByteString.copyFromUtf8(start), ByteString.copyFromUtf8(end)); return GrpcUtils.makeRegion( 1, - range.getLeft(), - range.getRight(), + range.first, + range.second, GrpcUtils.makeRegionEpoch(2, 3), GrpcUtils.makePeer(1, 10), GrpcUtils.makePeer(2, 20)); diff --git a/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java b/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java index 441aea9c4ef..24cbf3012a2 100644 --- a/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java +++ b/src/test/java/org/tikv/common/apiversion/RequestKeyCodecTest.java @@ -23,8 +23,8 @@ import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import java.util.stream.Collectors; -import org.apache.commons.lang3.tuple.Pair; import org.junit.Test; +import org.tikv.common.util.Pair; import org.tikv.kvproto.Metapb.Region; import org.tikv.kvproto.Pdpb; @@ -42,12 +42,12 @@ public void testV1RawCodec() { ByteString start = ByteString.copyFromUtf8("testV1RawCodec_start"); ByteString end = ByteString.copyFromUtf8("testV1RawCodec_end"); Pair range = v1.encodeRange(start, end); - assertEquals(start, range.getLeft()); - assertEquals(end, range.getRight()); + assertEquals(start, range.first); + assertEquals(end, range.second); range = v1.encodePdQueryRange(start, end); - assertEquals(start, range.getLeft()); - assertEquals(end, range.getRight()); + assertEquals(start, range.first); + assertEquals(end, range.second); Region region = Region.newBuilder().setStartKey(start).setEndKey(end).build(); assertEquals(region, v1.decodeRegion(region)); @@ -73,8 +73,8 @@ public void testV1TxnCodec() { // Test start and end are both non-empty. Pair range = v1.encodePdQueryRange(start, end); - assertEquals(CodecUtils.encode(start), range.getLeft()); - assertEquals(CodecUtils.encode(end), range.getRight()); + assertEquals(CodecUtils.encode(start), range.first); + assertEquals(CodecUtils.encode(end), range.second); Region region = Region.newBuilder() @@ -97,8 +97,8 @@ public void testV1TxnCodec() { assertEquals(end, decoded.getEndKey()); range = v1.encodePdQueryRange(start, end); - assertEquals(ByteString.EMPTY, range.getLeft()); - assertEquals(CodecUtils.encode(end), range.getRight()); + assertEquals(ByteString.EMPTY, range.first); + assertEquals(CodecUtils.encode(end), range.second); // Test end is empty. end = ByteString.EMPTY; @@ -112,8 +112,8 @@ public void testV1TxnCodec() { assertEquals(ByteString.EMPTY, decoded.getEndKey()); range = v1.encodePdQueryRange(start, end); - assertEquals(start, range.getLeft()); - assertEquals(ByteString.EMPTY, range.getRight()); + assertEquals(start, range.first); + assertEquals(ByteString.EMPTY, range.second); } @Test @@ -133,8 +133,8 @@ void testV2Codec(RequestKeyV2Codec v2) { // Test start and end are both non-empty. Pair range = v2.encodePdQueryRange(start, end); - assertEquals(CodecUtils.encode(v2.encodeKey(start)), range.getLeft()); - assertEquals(CodecUtils.encode(v2.encodeKey(end)), range.getRight()); + assertEquals(CodecUtils.encode(v2.encodeKey(start)), range.first); + assertEquals(CodecUtils.encode(v2.encodeKey(end)), range.second); Region region = Region.newBuilder() @@ -157,21 +157,21 @@ void testV2Codec(RequestKeyV2Codec v2) { assertEquals(end, decoded.getEndKey()); range = v2.encodePdQueryRange(start, end); - assertEquals(CodecUtils.encode(v2.encodeKey(start)), range.getLeft()); - assertEquals(CodecUtils.encode(v2.encodeKey(end)), range.getRight()); + assertEquals(CodecUtils.encode(v2.encodeKey(start)), range.first); + assertEquals(CodecUtils.encode(v2.encodeKey(end)), range.second); // Test end is empty. end = ByteString.EMPTY; range = v2.encodeRange(start, end); - assertEquals(v2.encodeKey(start), range.getLeft()); + assertEquals(v2.encodeKey(start), range.first); assertArrayEquals( new byte[] {(byte) (v2.encodeKey(ByteString.EMPTY).byteAt(0) + 1)}, - range.getRight().toByteArray()); + range.second.toByteArray()); region = Region.newBuilder() - .setStartKey(CodecUtils.encode(range.getLeft())) - .setEndKey(CodecUtils.encode(range.getRight())) + .setStartKey(CodecUtils.encode(range.first)) + .setEndKey(CodecUtils.encode(range.second)) .build(); decoded = v2.decodeRegion(region); assertEquals(start, decoded.getStartKey());