move request codec into raw/txn owned module

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

Author: iosmanthus

examples/pessimistic.rs

@@ -1,12 +1,13 @@
 // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
 
-mod common;
-
-use crate::common::parse_args;
 use tikv_client::{
-    request::codec::TxnApiV1, Config, Key, TransactionClient as Client, TransactionOptions, Value,
+    transaction::ApiV1, Config, Key, TransactionClient as Client, TransactionOptions, Value,
 };
 
+use crate::common::parse_args;
+
+mod common;
+
 #[tokio::main]
 async fn main() {
     // You can try running this example by passing your pd endpoints
@@ -22,7 +23,7 @@ async fn main() {
     };
 
     // init
-    let client = Client::new_with_config(args.pd, config, TxnApiV1, None)
+    let client = Client::new_with_config(args.pd, config, ApiV1, None)
         .await
         .expect("Could not connect to tikv");
 

examples/raw.rs

@@ -6,8 +6,7 @@ mod common;
 
 use crate::common::parse_args;
 use tikv_client::{
-    request::codec::RawApiV1, Config, IntoOwnedRange, Key, KvPair, RawClient as Client, Result,
-    Value,
+    raw::ApiV1, Config, IntoOwnedRange, Key, KvPair, RawClient as Client, Result, Value,
 };
 
 const KEY: &str = "TiKV";
@@ -29,7 +28,7 @@ async fn main() -> Result<()> {
 
     // When we first create a client we receive a `Connect` structure which must be resolved before
     // the client is actually connected and usable.
-    let client = Client::new_with_config(args.pd, config, RawApiV1, None).await?;
+    let client = Client::new_with_config(args.pd, config, ApiV1, None).await?;
 
     // Requests are created from the connected client. These calls return structures which
     // implement `Future`. This means the `Future` must be resolved before the action ever takes

examples/transaction.rs

@@ -4,8 +4,8 @@ mod common;
 
 use crate::common::parse_args;
 use tikv_client::{
-    request::codec::{TxnApiV1, TxnCodec},
-    BoundRange, Config, Key, KvPair, TransactionClient as Client, Value,
+    request::codec::TxnCodec, transaction::ApiV1, BoundRange, Config, Key, KvPair,
+    TransactionClient as Client, Value,
 };
 
 async fn puts<C: TxnCodec>(client: &Client<C>, pairs: impl IntoIterator<Item = impl Into<KvPair>>) {
@@ -84,7 +84,7 @@ async fn main() {
         Config::default()
     };
 
-    let txn = Client::new_with_config(args.pd, config, TxnApiV1, None)
+    let txn = Client::new_with_config(args.pd, config, ApiV1, None)
         .await
         .expect("Could not connect to tikv");
 

src/mock.rs

@@ -7,8 +7,8 @@
 
 use crate::{
     pd::{PdClient, PdRpcClient, RetryClient},
+    raw::ApiV1,
     region::{RegionId, RegionWithLeader},
-    request::codec::RawApiV1,
     store::RegionStore,
     Config, Error, Key, Result, Timestamp,
 };
@@ -21,7 +21,7 @@ use tikv_client_store::{KvClient, KvConnect, Request};
 
 /// Create a `PdRpcClient` with it's internals replaced with mocks so that the
 /// client can be tested without doing any RPC calls.
-pub async fn pd_rpc_client() -> PdRpcClient<RawApiV1, MockKvConnect, MockCluster> {
+pub async fn pd_rpc_client() -> PdRpcClient<ApiV1, MockKvConnect, MockCluster> {
     let config = Config::default();
     let plain = slog_term::PlainSyncDecorator::new(std::io::stdout());
     let logger = Logger::root(
@@ -42,7 +42,7 @@ pub async fn pd_rpc_client() -> PdRpcClient<RawApiV1, MockKvConnect, MockCluster
                 MockCluster,
             ))
         },
-        RawApiV1,
+        ApiV1,
         logger,
     )
     .await
@@ -154,7 +154,7 @@ impl MockPdClient {
 #[async_trait]
 impl PdClient for MockPdClient {
     type KvClient = MockKvClient;
-    type RequestCodec = RawApiV1;
+    type RequestCodec = ApiV1;
 
     async fn map_region_to_store(self: Arc<Self>, region: RegionWithLeader) -> Result<RegionStore> {
         Ok(RegionStore::new(region, Arc::new(self.client.clone())))
@@ -201,6 +201,6 @@ impl PdClient for MockPdClient {
     async fn invalidate_region_cache(&self, _ver_id: crate::region::RegionVerId) {}
 
     fn get_request_codec(&self) -> Self::RequestCodec {
-        RawApiV1
+        ApiV1
     }
 }

src/raw/client.rs

@@ -792,7 +792,7 @@ mod tests {
 
     use crate::{
         mock::{MockKvClient, MockPdClient},
-        request::codec::RawApiV1,
+        raw::ApiV1,
         Result,
     };
 
@@ -828,7 +828,7 @@ mod tests {
             cf: Some(ColumnFamily::Default),
             atomic: false,
             logger,
-            _phantom: PhantomData::<RawApiV1>,
+            _phantom: PhantomData::<ApiV1>,
         };
         let resps = client
             .coprocessor(

src/raw/codec.rs

@@ -0,0 +1,27 @@
+use tikv_client_proto::{kvrpcpb, metapb::Region};
+
+use crate::{
+    impl_request_codec_for_new_type,
+    request::codec::{KeyMode, KeySpaceCodec, KeySpaceId, RawCodec, RequestCodec, TxnCodec},
+    Result,
+};
+
+#[derive(Clone)]
+pub struct ApiV1;
+
+impl RequestCodec for ApiV1 {}
+
+impl RawCodec for ApiV1 {}
+
+#[derive(Clone)]
+pub struct Keyspace(KeySpaceCodec);
+
+impl Keyspace {
+    pub fn new(id: KeySpaceId) -> Self {
+        Keyspace(KeySpaceCodec::new(KeyMode::Raw, id))
+    }
+}
+
+impl_request_codec_for_new_type!(Keyspace);
+
+impl TxnCodec for Keyspace {}

src/raw/mod.rs

@@ -9,11 +9,15 @@
 //!
 //! **Warning:** It is not advisable to use both raw and transactional functionality in the same keyspace.
 
-pub use self::client::Client;
+pub use self::{
+    client::Client,
+    codec::{ApiV1, Keyspace},
+};
 use crate::Error;
 use std::{convert::TryFrom, fmt};
 
 mod client;
+mod codec;
 pub mod lowering;
 mod requests;
 

src/region_cache.rs

@@ -262,8 +262,9 @@ mod test {
 
     use crate::{
         pd::RetryClientTrait,
+        raw::ApiV1,
         region::{RegionId, RegionWithLeader},
-        request::codec::{RawApiV1, RequestCodec},
+        request::codec::RequestCodec,
         Key, Result,
     };
 
@@ -330,7 +331,7 @@ mod test {
     #[tokio::test]
     async fn cache_is_used() -> Result<()> {
         let retry_client = Arc::new(MockRetryClient::default());
-        let cache = RegionCache::new(RawApiV1, retry_client.clone());
+        let cache = RegionCache::new(ApiV1, retry_client.clone());
         retry_client.regions.lock().await.insert(
             1,
             RegionWithLeader {
@@ -400,7 +401,7 @@ mod test {
     #[tokio::test]
     async fn test_add_disjoint_regions() {
         let retry_client = Arc::new(MockRetryClient::default());
-        let cache = RegionCache::new(RawApiV1, retry_client.clone());
+        let cache = RegionCache::new(ApiV1, retry_client.clone());
         let region1 = region(1, vec![], vec![10]);
         let region2 = region(2, vec![10], vec![20]);
         let region3 = region(3, vec![30], vec![]);
@@ -419,7 +420,7 @@ mod test {
     #[tokio::test]
     async fn test_add_intersecting_regions() {
         let retry_client = Arc::new(MockRetryClient::default());
-        let cache = RegionCache::new(RawApiV1, retry_client.clone());
+        let cache = RegionCache::new(ApiV1, retry_client.clone());
 
         cache.add_region(region(1, vec![], vec![10])).await;
         cache.add_region(region(2, vec![10], vec![20])).await;
@@ -457,7 +458,7 @@ mod test {
     #[tokio::test]
     async fn test_get_region_by_key() -> Result<()> {
         let retry_client = Arc::new(MockRetryClient::default());
-        let cache = RegionCache::new(RawApiV1, retry_client.clone());
+        let cache = RegionCache::new(ApiV1, retry_client.clone());
 
         let region1 = region(1, vec![], vec![10]);
         let region2 = region(2, vec![10], vec![20]);

src/request/codec.rs

@@ -1,6 +1,8 @@
 use core::intrinsics::copy;
 use std::ops::{Deref, DerefMut};
 
+use derive_new::new;
+
 use tikv_client_common::Error;
 use tikv_client_proto::{errorpb, kvrpcpb, metapb::Region};
 
@@ -185,36 +187,12 @@ pub trait RequestCodecExt: RequestCodec {
 
 impl<T: RequestCodec> RequestCodecExt for T {}
 
-#[derive(Clone)]
-pub struct TxnApiV1;
-
-#[derive(Clone)]
-pub struct RawApiV1;
-
 pub trait RawCodec: RequestCodec {}
 
 pub trait TxnCodec: RequestCodec {}
 
-impl RequestCodec for RawApiV1 {}
-
-impl RequestCodec for TxnApiV1 {
-    fn encode_pd_query(&self, key: Vec<u8>) -> Vec<u8> {
-        Key::from(key).to_encoded().into()
-    }
-
-    fn decode_region(&self, region: &mut Region) -> Result<()> {
-        decode_bytes_in_place(region.mut_start_key(), false)?;
-        decode_bytes_in_place(region.mut_end_key(), false)?;
-        Ok(())
-    }
-}
-
-impl RawCodec for RawApiV1 {}
-
-impl TxnCodec for TxnApiV1 {}
-
 #[derive(Copy, Clone)]
-enum KeyMode {
+pub(crate) enum KeyMode {
     Raw,
     Txn,
 }
@@ -263,8 +241,8 @@ impl DerefMut for KeySpaceId {
     }
 }
 
-#[derive(Copy, Clone)]
-struct KeySpaceCodec {
+#[derive(new, Copy, Clone)]
+pub(crate) struct KeySpaceCodec {
     mode: KeyMode,
     id: KeySpaceId,
 }
@@ -346,82 +324,33 @@ impl RequestCodec for KeySpaceCodec {
     }
 }
 
-#[derive(Clone)]
-pub struct RawKeyspaceCodec(KeySpaceCodec);
-
-impl RawKeyspaceCodec {
-    pub fn new(id: KeySpaceId) -> Self {
-        RawKeyspaceCodec(KeySpaceCodec {
-            mode: KeyMode::Raw,
-            id,
-        })
-    }
-}
-
-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 {}
-
-#[derive(Clone)]
-pub struct TxnKeyspaceCodec(KeySpaceCodec);
-
-impl TxnKeyspaceCodec {
-    pub fn new(id: KeySpaceId) -> Self {
-        TxnKeyspaceCodec(KeySpaceCodec {
-            mode: KeyMode::Txn,
-            id,
-        })
-    }
-}
-
-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()
-    }
+#[macro_export]
+macro_rules! impl_request_codec_for_new_type {
+    ($t:ty) => {
+        impl RequestCodec for $t {
+            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 {}