Introduce KeySpace interface for API v2.

src/lib.rs

@@ -90,6 +90,8 @@
 //! # })}
 //! ```
 
+#![feature(specialization)]
+#![feature(explicit_generic_args_with_impl_trait)]
 #[macro_use]
 pub mod request;
 #[macro_use]

src/pd/client.rs

@@ -1,23 +1,27 @@
 // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
 
+use std::{collections::HashMap, marker::PhantomData, sync::Arc, thread};
+
+use async_trait::async_trait;
+use futures::{prelude::*, stream::BoxStream};
+use grpcio::{EnvBuilder, Environment};
+use slog::Logger;
+use tokio::sync::RwLock;
+
+use tikv_client_pd::Cluster;
+use tikv_client_proto::{kvrpcpb, metapb};
+use tikv_client_store::{KvClient, KvConnect, TikvConnect};
+
 use crate::{
     compat::stream_fn,
     kv::codec,
     pd::{retry::RetryClientTrait, RetryClient},
     region::{RegionId, RegionVerId, RegionWithLeader},
     region_cache::RegionCache,
+    request::request_codec::RequestCodec,
     store::RegionStore,
     BoundRange, Config, Key, Result, SecurityManager, Timestamp,
 };
-use async_trait::async_trait;
-use futures::{prelude::*, stream::BoxStream};
-use grpcio::{EnvBuilder, Environment};
-use slog::Logger;
-use std::{collections::HashMap, sync::Arc, thread};
-use tikv_client_pd::Cluster;
-use tikv_client_proto::{kvrpcpb, metapb};
-use tikv_client_store::{KvClient, KvConnect, TikvConnect};
-use tokio::sync::RwLock;
 
 const CQ_COUNT: usize = 1;
 const CLIENT_PREFIX: &str = "tikv-client";
@@ -42,6 +46,7 @@ const CLIENT_PREFIX: &str = "tikv-client";
 #[async_trait]
 pub trait PdClient: Send + Sync + 'static {
     type KvClient: KvClient + Send + Sync + 'static;
+    type RequestCodec: RequestCodec;
 
     /// In transactional API, `region` is decoded (keys in raw format).
     async fn map_region_to_store(self: Arc<Self>, region: RegionWithLeader) -> Result<RegionStore>;
@@ -204,22 +209,27 @@ pub trait PdClient: Send + Sync + 'static {
     async fn update_leader(&self, ver_id: RegionVerId, leader: metapb::Peer) -> Result<()>;
 
     async fn invalidate_region_cache(&self, ver_id: RegionVerId);
+
+    fn get_request_codec(&self) -> Self::RequestCodec;
 }
 
 /// This client converts requests for the logical TiKV cluster into requests
 /// for a single TiKV store using PD and internal logic.
-pub struct PdRpcClient<KvC: KvConnect + Send + Sync + 'static = TikvConnect, Cl = Cluster> {
+pub struct PdRpcClient<C, KvC: KvConnect + Send + Sync + 'static = TikvConnect, Cl = Cluster> {
     pd: Arc<RetryClient<Cl>>,
     kv_connect: KvC,
     kv_client_cache: Arc<RwLock<HashMap<String, KvC::KvClient>>>,
     enable_codec: bool,
     region_cache: RegionCache<RetryClient<Cl>>,
     logger: Logger,
+    // TODO: change to a real codec.
+    _phantom: PhantomData<C>,
 }
 
 #[async_trait]
-impl<KvC: KvConnect + Send + Sync + 'static> PdClient for PdRpcClient<KvC> {
+impl<C: RequestCodec, KvC: KvConnect + Send + Sync + 'static> PdClient for PdRpcClient<C, KvC> {
     type KvClient = KvC::KvClient;
+    type RequestCodec = C;
 
     async fn map_region_to_store(self: Arc<Self>, region: RegionWithLeader) -> Result<RegionStore> {
         let store_id = region.get_store_id()?;
@@ -260,15 +270,19 @@ impl<KvC: KvConnect + Send + Sync + 'static> PdClient for PdRpcClient<KvC> {
     async fn invalidate_region_cache(&self, ver_id: RegionVerId) {
         self.region_cache.invalidate_region_cache(ver_id).await
     }
+
+    fn get_request_codec(&self) -> Self::RequestCodec {
+        todo!()
+    }
 }
 
-impl PdRpcClient<TikvConnect, Cluster> {
+impl<C> PdRpcClient<C, TikvConnect, Cluster> {
     pub async fn connect(
         pd_endpoints: &[String],
         config: Config,
         enable_codec: bool,
         logger: Logger,
-    ) -> Result<PdRpcClient> {
+    ) -> Result<PdRpcClient<C, TikvConnect, Cluster>> {
         PdRpcClient::new(
             config.clone(),
             |env, security_mgr| TikvConnect::new(env, security_mgr, config.timeout),
@@ -291,14 +305,14 @@ fn thread_name(prefix: &str) -> String {
         .unwrap_or_else(|| prefix.to_owned())
 }
 
-impl<KvC: KvConnect + Send + Sync + 'static, Cl> PdRpcClient<KvC, Cl> {
+impl<C, KvC: KvConnect + Send + Sync + 'static, Cl> PdRpcClient<C, KvC, Cl> {
     pub async fn new<PdFut, MakeKvC, MakePd>(
         config: Config,
         kv_connect: MakeKvC,
         pd: MakePd,
         enable_codec: bool,
         logger: Logger,
-    ) -> Result<PdRpcClient<KvC, Cl>>
+    ) -> Result<PdRpcClient<C, KvC, Cl>>
     where
         PdFut: Future<Output = Result<RetryClient<Cl>>>,
         MakeKvC: FnOnce(Arc<Environment>, Arc<SecurityManager>) -> KvC,
@@ -329,6 +343,8 @@ impl<KvC: KvConnect + Send + Sync + 'static, Cl> PdRpcClient<KvC, Cl> {
             enable_codec,
             region_cache: RegionCache::new(pd),
             logger,
+            // TODO
+            _phantom: PhantomData,
         })
     }
 
@@ -352,10 +368,11 @@ impl<KvC: KvConnect + Send + Sync + 'static, Cl> PdRpcClient<KvC, Cl> {
 
 #[cfg(test)]
 pub mod test {
-    use super::*;
+    use futures::{executor, executor::block_on};
+
     use crate::mock::*;
 
-    use futures::{executor, executor::block_on};
+    use super::*;
 
     #[tokio::test]
     async fn test_kv_client_caching() {
@@ -396,7 +413,7 @@ pub mod test {
                 vec![1].into(),
                 vec![2].into(),
                 vec![3].into(),
-                vec![5, 2].into()
+                vec![5, 2].into(),
             ]
         );
         assert_eq!(
@@ -458,36 +475,36 @@ pub mod test {
             vec![
                 kvrpcpb::KeyRange {
                     start_key: k1.clone(),
-                    end_key: k2.clone()
+                    end_key: k2.clone(),
                 },
                 kvrpcpb::KeyRange {
                     start_key: k1,
-                    end_key: k_split.clone()
-                }
+                    end_key: k_split.clone(),
+                },
             ]
         );
         assert_eq!(ranges2.0, 2);
         assert_eq!(
             ranges2.1,
             vec![kvrpcpb::KeyRange {
                 start_key: k_split.clone(),
-                end_key: k3
+                end_key: k3,
             }]
         );
         assert_eq!(ranges3.0, 1);
         assert_eq!(
             ranges3.1,
             vec![kvrpcpb::KeyRange {
                 start_key: k2,
-                end_key: k_split.clone()
+                end_key: k_split.clone(),
             }]
         );
         assert_eq!(ranges4.0, 2);
         assert_eq!(
             ranges4.1,
             vec![kvrpcpb::KeyRange {
                 start_key: k_split,
-                end_key: k4
+                end_key: k4,
             }]
         );
         assert!(stream.next().is_none());

src/raw/client.rs

@@ -1,7 +1,7 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
 use core::ops::Range;
-use std::{str::FromStr, sync::Arc, u32};
+use std::{marker::PhantomData, str::FromStr, sync::Arc, u32};
 
 use slog::{Drain, Logger};
 use tikv_client_common::Error;
@@ -12,7 +12,7 @@ use crate::{
     config::Config,
     pd::{PdClient, PdRpcClient},
     raw::lowering::*,
-    request::{Collect, CollectSingle, Plan},
+    request::{request_codec::RequestCodec, Collect, CollectSingle, Plan},
     Backoff, BoundRange, ColumnFamily, Key, KvPair, Result, Value,
 };
 
@@ -26,15 +26,16 @@ const MAX_RAW_KV_SCAN_LIMIT: u32 = 10240;
 /// The returned results of raw request methods are [`Future`](std::future::Future)s that must be
 /// awaited to execute.
 #[derive(Clone)]
-pub struct Client<PdC: PdClient = PdRpcClient> {
+pub struct Client<C, PdC: PdClient = PdRpcClient<C>> {
     rpc: Arc<PdC>,
     cf: Option<ColumnFamily>,
     /// Whether to use the [`atomic mode`](Client::with_atomic_for_cas).
     atomic: bool,
     logger: Logger,
+    _phantom: std::marker::PhantomData<C>,
 }
 
-impl Client<PdRpcClient> {
+impl<C: RequestCodec> Client<C, PdRpcClient<C>> {
     /// Create a raw [`Client`] and connect to the TiKV cluster.
     ///
     /// Because TiKV is managed by a [PD](https://github.com/pingcap/pd/) cluster, the endpoints for
@@ -103,6 +104,7 @@ impl Client<PdRpcClient> {
             cf: None,
             atomic: false,
             logger,
+            _phantom: PhantomData,
         })
     }
 
@@ -137,6 +139,7 @@ impl Client<PdRpcClient> {
             cf: Some(cf),
             atomic: self.atomic,
             logger: self.logger.clone(),
+            _phantom: PhantomData,
         }
     }
 
@@ -154,11 +157,12 @@ impl Client<PdRpcClient> {
             cf: self.cf.clone(),
             atomic: true,
             logger: self.logger.clone(),
+            _phantom: PhantomData,
         }
     }
 }
 
-impl<PdC: PdClient> Client<PdC> {
+impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
     /// Create a new 'get' request.
     ///
     /// Once resolved this request will result in the fetching of the value associated with the
@@ -179,7 +183,7 @@ impl<PdC: PdClient> Client<PdC> {
     /// ```
     pub async fn get(&self, key: impl Into<Key>) -> Result<Option<Value>> {
         debug!(self.logger, "invoking raw get request");
-        let request = new_raw_get_request(key.into(), self.cf.clone());
+        let request = new_raw_get_request::<C>(key.into(), self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
             .merge(CollectSingle)
@@ -211,7 +215,8 @@ impl<PdC: PdClient> Client<PdC> {
         keys: impl IntoIterator<Item = impl Into<Key>>,
     ) -> Result<Vec<KvPair>> {
         debug!(self.logger, "invoking raw batch_get request");
-        let request = new_raw_batch_get_request(keys.into_iter().map(Into::into), self.cf.clone());
+        let request =
+            new_raw_batch_get_request::<C>(keys.into_iter().map(Into::into), self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
             .merge(Collect)
@@ -239,7 +244,8 @@ impl<PdC: PdClient> Client<PdC> {
     /// ```
     pub async fn put(&self, key: impl Into<Key>, value: impl Into<Value>) -> Result<()> {
         debug!(self.logger, "invoking raw put request");
-        let request = new_raw_put_request(key.into(), value.into(), self.cf.clone(), self.atomic);
+        let request =
+            new_raw_put_request::<C>(key.into(), value.into(), self.cf.clone(), self.atomic);
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
             .merge(CollectSingle)
@@ -271,7 +277,7 @@ impl<PdC: PdClient> Client<PdC> {
         pairs: impl IntoIterator<Item = impl Into<KvPair>>,
     ) -> Result<()> {
         debug!(self.logger, "invoking raw batch_put request");
-        let request = new_raw_batch_put_request(
+        let request = new_raw_batch_put_request::<C>(
             pairs.into_iter().map(Into::into),
             self.cf.clone(),
             self.atomic,
@@ -303,7 +309,7 @@ impl<PdC: PdClient> Client<PdC> {
     /// ```
     pub async fn delete(&self, key: impl Into<Key>) -> Result<()> {
         debug!(self.logger, "invoking raw delete request");
-        let request = new_raw_delete_request(key.into(), self.cf.clone(), self.atomic);
+        let request = new_raw_delete_request::<C>(key.into(), self.cf.clone(), self.atomic);
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
             .merge(CollectSingle)
@@ -334,7 +340,7 @@ impl<PdC: PdClient> Client<PdC> {
         debug!(self.logger, "invoking raw batch_delete request");
         self.assert_non_atomic()?;
         let request =
-            new_raw_batch_delete_request(keys.into_iter().map(Into::into), self.cf.clone());
+            new_raw_batch_delete_request::<C>(keys.into_iter().map(Into::into), self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
             .extract_error()
@@ -369,7 +375,7 @@ impl<PdC: PdClient> Client<PdC> {
     ) -> Result<()> {
         debug!(self.logger, "invoking raw delete_range request");
         self.assert_non_atomic()?;
-        let request = new_raw_delete_range_request(range.into(), self.cf.clone());
+        let request = new_raw_delete_range_request::<C>(range.into(), self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .retry_multi_region(backoff)
             .extract_error()
@@ -520,7 +526,7 @@ impl<PdC: PdClient> Client<PdC> {
     ) -> Result<(Option<Value>, bool)> {
         debug!(self.logger, "invoking raw compare_and_swap request");
         self.assert_atomic()?;
-        let req = new_cas_request(
+        let req = new_cas_request::<C>(
             key.into(),
             new_value.into(),
             previous_value.into(),
@@ -543,7 +549,7 @@ impl<PdC: PdClient> Client<PdC> {
     ) -> Result<Vec<(Vec<u8>, Vec<Range<Key>>)>> {
         let copr_version_req = copr_version_req.into();
         semver::VersionReq::from_str(&copr_version_req)?;
-        let req = new_raw_coprocessor_request(
+        let req = new_raw_coprocessor_request::<C>(
             copr_name.into(),
             copr_version_req,
             ranges.into_iter().map(Into::into),
@@ -570,7 +576,7 @@ impl<PdC: PdClient> Client<PdC> {
             });
         }
 
-        let request = new_raw_scan_request(range.into(), limit, key_only, self.cf.clone());
+        let request = new_raw_scan_request::<C>(range.into(), limit, key_only, self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
             .merge(Collect)
@@ -595,7 +601,7 @@ impl<PdC: PdClient> Client<PdC> {
             });
         }
 
-        let request = new_raw_batch_scan_request(
+        let request = new_raw_batch_scan_request::<C>(
             ranges.into_iter().map(Into::into),
             each_limit,
             key_only,