[to #555] enable api v2

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]})

src/main/java/org/tikv/common/ConfigUtils.java

@@ -121,6 +121,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";
@@ -204,4 +207,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;
 }

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<PDBlockingStub, PDFutureStub>
   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<PDBlockingStub, PDFutureStub>
           .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,22 +316,19 @@ public Pair<Metapb.Region, Metapb.Peer> getRegionByKey(BackOffer backOffer, Byte
     Histogram.Timer requestTimer =
         PD_GET_REGION_BY_KEY_REQUEST_LATENCY.labels(getClusterId().toString()).startTimer();
     try {
-      if (conf.isTxnKVMode()) {
-        CodecDataOutput cdo = new CodecDataOutput();
-        BytesCodec.writeBytes(cdo, key.toByteArray());
-        key = cdo.toByteString();
-      }
-      ByteString queryKey = key;
-
       Supplier<GetRegionRequest> request =
-          () -> GetRegionRequest.newBuilder().setHeader(header).setRegionKey(queryKey).build();
+          () ->
+              GetRegionRequest.newBuilder()
+                  .setHeader(header)
+                  .setRegionKey(codec.encodePdQuery(key))
+                  .build();
 
       PDErrorHandler<GetRegionResponse> handler =
           new PDErrorHandler<>(getRegionResponseErrorExtractor, this);
 
       GetRegionResponse resp =
           callWithRetry(backOffer, PDGrpc.getGetRegionMethod(), request, handler);
-      return new Pair<Metapb.Region, Metapb.Peer>(decodeRegion(resp.getRegion()), resp.getLeader());
+      return new Pair<>(codec.decodeRegion(resp.getRegion()), resp.getLeader());
     } finally {
       requestTimer.observeDuration();
     }
@@ -343,7 +343,8 @@ public Pair<Metapb.Region, Metapb.Peer> getRegionByID(BackOffer backOffer, long
 
     GetRegionResponse resp =
         callWithRetry(backOffer, PDGrpc.getGetRegionByIDMethod(), request, handler);
-    return new Pair<Metapb.Region, Metapb.Peer>(decodeRegion(resp.getRegion()), resp.getLeader());
+    return new Pair<Metapb.Region, Metapb.Peer>(
+        codec.decodeRegion(resp.getRegion()), resp.getLeader());
   }
 
   @Override
@@ -353,18 +354,20 @@ public List<Pdpb.Region> scanRegions(
     // introduce a warm-up timeout for ScanRegions requests
     PDGrpc.PDBlockingStub stub =
         getBlockingStub().withDeadlineAfter(conf.getWarmUpTimeout(), TimeUnit.MILLISECONDS);
+    Pair<ByteString, ByteString> range = codec.encodePdQueryRange(startKey, endKey);
     Pdpb.ScanRegionsRequest request =
         Pdpb.ScanRegionsRequest.newBuilder()
             .setHeader(header)
-            .setStartKey(startKey)
-            .setEndKey(endKey)
+            .setStartKey(range.first)
+            .setEndKey(range.second)
             .setLimit(limit)
             .build();
     Pdpb.ScanRegionsResponse resp = stub.scanRegions(request);
     if (resp == null) {
       return null;
     }
-    return resp.getRegionsList();
+
+    return codec.decodePdRegions(resp.getRegionsList());
   }
 
   private Supplier<GetStoreRequest> buildGetStoreReq(long storeId) {
@@ -811,54 +814,15 @@ public String toString() {
     }
   }
 
-  private 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() || isRawRegion) {
-      builder.setStartKey(region.getStartKey());
-    } else {
-      if (!conf.isTest()) {
-        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());
-        }
-      }
-    }
-
-    if (region.getEndKey().isEmpty() || isRawRegion) {
-      builder.setEndKey(region.getEndKey());
-    } else {
-      if (!conf.isTest()) {
-        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());
-        }
-      }
-    }
-
-    return builder.build();
-  }
-
   public Long getClusterId() {
     return header.getClusterId();
   }
 
   public List<URI> getPdAddrs() {
     return pdAddrs;
   }
+
+  public RequestKeyCodec getCodec() {
+    return codec;
+  }
 }

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;
@@ -69,4 +70,6 @@ List<Pdpb.Region> scanRegions(
   TiConfiguration.ReplicaRead getReplicaRead();
 
   Long getClusterId();
+
+  RequestKeyCodec getCodec();
 }

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;

src/main/java/org/tikv/common/TiConfiguration.java

@@ -40,6 +40,7 @@
 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_CONN_RECYCLE_TIME;
 import static org.tikv.common.ConfigUtils.DEF_TIKV_ENABLE_ATOMIC_FOR_CAS;
@@ -82,6 +83,7 @@
 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;
@@ -154,6 +156,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;
@@ -173,17 +176,19 @@
 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;
 
 public class TiConfiguration implements Serializable {
-
   private static final Logger logger = LoggerFactory.getLogger(TiConfiguration.class);
   private static final ConcurrentHashMap<String, String> settings = new ConcurrentHashMap<>();
   public static final Metadata.Key<String> FORWARD_META_DATA_KEY =
       Metadata.Key.of("tikv-forwarded-host", Metadata.ASCII_STRING_MARSHALLER);
   public static final Metadata.Key<String> PD_FORWARD_META_DATA_KEY =
       Metadata.Key.of("pd-forwarded-host", Metadata.ASCII_STRING_MARSHALLER);
+  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
@@ -296,6 +301,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() {
@@ -551,6 +558,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
@@ -1241,4 +1250,48 @@ 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;
+  }
+
+  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);
+      }
+    }
+  }
 }