fix api version context

Signed-off-by: iosmanthus <[email protected]>

Author: iosmanthus

src/raw/requests.rs

@@ -66,7 +66,7 @@ pub fn new_raw_batch_get_request(
     req
 }
 
-impl_kv_request!(kvrpcpb::RawBatchGetRequest, keys; kvrpcpb::RawBatchGetResponse);
+impl_kv_request!(kvrpcpb::RawBatchGetRequest, keys; kvrpcpb::RawBatchGetResponse, pairs);
 shardable_keys!(kvrpcpb::RawBatchGetRequest);
 
 impl Merge<kvrpcpb::RawBatchGetResponse> for Collect {
@@ -362,6 +362,10 @@ impl Request for RawCoprocessorRequest {
     fn set_context(&mut self, context: kvrpcpb::Context) {
         self.inner.set_context(context);
     }
+
+    fn mut_context(&mut self) -> &mut kvrpcpb::Context {
+        self.inner.mut_context()
+    }
 }
 
 impl<C: RequestCodec> KvRequest<C> for RawCoprocessorRequest {

src/request/codec.rs

@@ -40,8 +40,8 @@ pub trait RequestCodec: Sized + Clone + Sync + Send + 'static {
         Ok(())
     }
 
-    fn is_plain(&self) -> bool {
-        true
+    fn version(&self) -> kvrpcpb::ApiVersion {
+        kvrpcpb::ApiVersion::V1
     }
 }
 
@@ -331,7 +331,8 @@ impl RequestCodec for KeySpaceCodec {
         }
 
         // Map the region's end key to the keyspace's end key.
-        if region.get_end_key() > self.mode.max_key().as_slice() {
+        if region.get_end_key().is_empty() || region.get_end_key() > self.mode.max_key().as_slice()
+        {
             *region.mut_end_key() = vec![];
         } else {
             self.decode_key(region.mut_end_key())?;
@@ -340,8 +341,8 @@ impl RequestCodec for KeySpaceCodec {
         Ok(())
     }
 
-    fn is_plain(&self) -> bool {
-        false
+    fn version(&self) -> kvrpcpb::ApiVersion {
+        kvrpcpb::ApiVersion::V2
     }
 }
 
@@ -357,7 +358,31 @@ impl RawKeyspaceCodec {
     }
 }
 
-impl RequestCodec for RawKeyspaceCodec {}
+impl RequestCodec for RawKeyspaceCodec {
+    fn encode_key(&self, key: Vec<u8>) -> Vec<u8> {
+        self.0.encode_key(key)
+    }
+
+    fn decode_key(&self, key: &mut Vec<u8>) -> Result<()> {
+        self.0.decode_key(key)
+    }
+
+    fn encode_range(&self, start: Vec<u8>, end: Vec<u8>) -> (Vec<u8>, Vec<u8>) {
+        self.0.encode_range(start, end)
+    }
+
+    fn encode_pd_query(&self, key: Vec<u8>) -> Vec<u8> {
+        self.0.encode_pd_query(key)
+    }
+
+    fn decode_region(&self, region: &mut Region) -> Result<()> {
+        self.0.decode_region(region)
+    }
+
+    fn version(&self) -> kvrpcpb::ApiVersion {
+        self.0.version()
+    }
+}
 
 impl RawCodec for RawKeyspaceCodec {}
 
@@ -373,6 +398,30 @@ impl TxnKeyspaceCodec {
     }
 }
 
-impl RequestCodec for TxnKeyspaceCodec {}
+impl RequestCodec for TxnKeyspaceCodec {
+    fn encode_key(&self, key: Vec<u8>) -> Vec<u8> {
+        self.0.encode_key(key)
+    }
+
+    fn decode_key(&self, key: &mut Vec<u8>) -> Result<()> {
+        self.0.decode_key(key)
+    }
+
+    fn encode_range(&self, start: Vec<u8>, end: Vec<u8>) -> (Vec<u8>, Vec<u8>) {
+        self.0.encode_range(start, end)
+    }
+
+    fn encode_pd_query(&self, key: Vec<u8>) -> Vec<u8> {
+        self.0.encode_pd_query(key)
+    }
+
+    fn decode_region(&self, region: &mut Region) -> Result<()> {
+        self.0.decode_region(region)
+    }
+
+    fn version(&self) -> kvrpcpb::ApiVersion {
+        self.0.version()
+    }
+}
 
 impl TxnCodec for TxnKeyspaceCodec {}

src/request/mod.rs

@@ -119,88 +119,6 @@ macro_rules! impl_kv_request {
     };
 }
 
-#[macro_export]
-macro_rules! impl_kv_request_for_single_key_op {
-    ($req: ty, $resp: ty) => {
-        impl<C> KvRequest<C> for $req
-        where
-            C: RequestCodec,
-        {
-            type Response = $resp;
-
-            fn encode_request(mut self, codec: &C) -> Self {
-                *self.mut_key() = codec.encode_key(self.take_key());
-
-                self
-            }
-        }
-    };
-}
-
-#[macro_export]
-macro_rules! impl_kv_request_for_batch_get {
-    ($req: ty, $resp: ty) => {
-        impl<C> KvRequest<C> for $req
-        where
-            C: RequestCodec,
-        {
-            type Response = $resp;
-
-            fn encode_request(mut self, codec: &C) -> Self {
-                *self.mut_keys() = codec.encode_keys(self.take_keys());
-
-                self
-            }
-
-            fn decode_response(
-                &self,
-                codec: &C,
-                mut resp: Self::Response,
-            ) -> $crate::Result<Self::Response> {
-                codec.decode_pairs(resp.mut_pairs())?;
-
-                Ok(resp)
-            }
-        }
-    };
-}
-
-#[macro_export]
-macro_rules! impl_kv_request_for_scan_op {
-    ($req: ty, $resp: ty, $pairs: ident) => {
-        impl<C> KvRequest<C> for $req
-        where
-            C: RequestCodec,
-        {
-            type Response = $resp;
-
-            fn encode_request(mut self, codec: &C) -> Self {
-                let (start, end) =
-                    codec.encode_range(self.take_start_key().into(), self.take_end_key().into());
-
-                self.set_start_key(start);
-                self.set_end_key(end);
-
-                self
-            }
-
-            fn decode_response(
-                &self,
-                codec: &C,
-                mut resp: Self::Response,
-            ) -> $crate::Result<Self::Response> {
-                paste::paste! {
-                    let pairs = resp.[<mut_ $pairs>]();
-
-                    codec.decode_pairs(pairs)?;
-
-                    Ok(resp)
-                }
-            }
-        }
-    };
-}
-
 #[derive(Clone, Debug, new, Eq, PartialEq)]
 pub struct RetryOptions {
     /// How to retry when there is a region error and we need to resolve regions with PD.
@@ -296,6 +214,10 @@ mod test {
             fn set_context(&mut self, _: kvrpcpb::Context) {
                 unreachable!();
             }
+
+            fn mut_context(&mut self) -> &mut kvrpcpb::Context {
+                unreachable!()
+            }
         }
 
         #[async_trait]

src/request/plan.rs

@@ -55,10 +55,13 @@ impl<C: RequestCodec, Req: KvRequest<C>> Plan for Dispatch<C, Req> {
     type Result = Req::Response;
 
     async fn execute(&self) -> Result<Self::Result> {
-        let req = if self.codec.is_plain() {
-            Cow::Borrowed(&self.request)
-        } else {
-            Cow::Owned(self.request.clone().encode_request(&self.codec))
+        let req = match self.codec.version() {
+            kvrpcpb::ApiVersion::V2 => {
+                let mut req = self.request.clone();
+                req.mut_context().set_api_version(self.codec.version());
+                Cow::Owned(req.encode_request(&self.codec))
+            }
+            _ => Cow::Borrowed(&self.request),
         };
 
         let stats = tikv_stats(self.request.label());
@@ -75,10 +78,9 @@ impl<C: RequestCodec, Req: KvRequest<C>> Plan for Dispatch<C, Req> {
                 .downcast()
                 .expect("Downcast failed: request and response type mismatch");
 
-            if self.codec.is_plain() {
-                Ok(resp)
-            } else {
-                req.decode_response(&self.codec, resp)
+            match self.codec.version() {
+                kvrpcpb::ApiVersion::V2 => req.decode_response(&self.codec, resp),
+                _ => Ok(resp),
             }
         })
     }

tikv-client-store/src/request.rs

@@ -12,6 +12,7 @@ pub trait Request: Any + Sync + Send + 'static {
     fn label(&self) -> &'static str;
     fn as_any(&self) -> &dyn Any;
     fn set_context(&mut self, context: kvrpcpb::Context);
+    fn mut_context(&mut self) -> &mut kvrpcpb::Context;
 }
 
 macro_rules! impl_request {
@@ -41,6 +42,10 @@ macro_rules! impl_request {
             fn set_context(&mut self, context: kvrpcpb::Context) {
                 kvrpcpb::$name::set_context(self, context)
             }
+
+            fn mut_context(&mut self) -> &mut kvrpcpb::Context {
+                kvrpcpb::$name::mut_context(self)
+            }
         }
     };
 }