introduce region cache codec interface

Author: iosmanthus

src/lib.rs

@@ -90,6 +90,7 @@
 //! # })}
 //! ```
 
+#![allow(incomplete_features)]
 #![feature(specialization)]
 #![feature(explicit_generic_args_with_impl_trait)]
 #[macro_use]

src/pd/client.rs

@@ -1,6 +1,6 @@
 // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
 
-use std::{collections::HashMap, marker::PhantomData, sync::Arc, thread};
+use std::{collections::HashMap, sync::Arc, thread};
 
 use async_trait::async_trait;
 use futures::{prelude::*, stream::BoxStream};
@@ -219,9 +219,8 @@ pub struct PdRpcClient<C, KvC: KvConnect + Send + Sync + 'static = TikvConnect,
     pd: Arc<RetryClient<Cl>>,
     kv_connect: KvC,
     kv_client_cache: Arc<RwLock<HashMap<String, KvC::KvClient>>>,
-    region_cache: RegionCache<RetryClient<Cl>>,
+    region_cache: RegionCache<C, RetryClient<Cl>>,
     logger: Logger,
-    codec: C,
 }
 
 #[async_trait]
@@ -237,20 +236,11 @@ impl<C: RequestCodec, KvC: KvConnect + Send + Sync + 'static> PdClient for PdRpc
     }
 
     async fn region_for_key(&self, key: &Key) -> Result<RegionWithLeader> {
-        let enable_codec = self.enable_codec;
-        let key = if enable_codec {
-            key.to_encoded()
-        } else {
-            key.clone()
-        };
-
-        let region = self.region_cache.get_region_by_key(&key).await?;
-        Self::decode_region(region, enable_codec)
+        self.region_cache.get_region_by_key(&key).await
     }
 
     async fn region_for_id(&self, id: RegionId) -> Result<RegionWithLeader> {
-        let region = self.region_cache.get_region_by_id(id).await?;
-        Self::decode_region(region, self.enable_codec)
+        self.region_cache.get_region_by_id(id).await
     }
 
     async fn get_timestamp(self: Arc<Self>) -> Result<Timestamp> {
@@ -270,7 +260,7 @@ impl<C: RequestCodec, KvC: KvConnect + Send + Sync + 'static> PdClient for PdRpc
     }
 
     fn get_request_codec(&self) -> Self::RequestCodec {
-        todo!()
+        self.region_cache.get_request_codec()
     }
 }
 
@@ -338,9 +328,8 @@ impl<C, KvC: KvConnect + Send + Sync + 'static, Cl> PdRpcClient<C, KvC, Cl> {
             pd: pd.clone(),
             kv_client_cache,
             kv_connect: kv_connect(env, security_mgr),
-            region_cache: RegionCache::new(pd),
+            region_cache: RegionCache::new(codec, pd),
             logger,
-            codec
         })
     }
 

src/raw/client.rs

@@ -9,12 +9,12 @@ use tikv_client_common::Error;
 use tikv_client_proto::metapb;
 
 use crate::{
-    Backoff,
     backoff::DEFAULT_REGION_BACKOFF,
-    BoundRange,
-    ColumnFamily,
     config::Config,
-    Key, KvPair, pd::{PdClient, PdRpcClient}, raw::lowering::*, request::{Collect, CollectSingle, Plan, request_codec::RequestCodec}, Result, Value,
+    pd::{PdClient, PdRpcClient},
+    raw::lowering::*,
+    request::{request_codec::RequestCodec, Collect, CollectSingle, Plan},
+    Backoff, BoundRange, ColumnFamily, Key, KvPair, Result, Value,
 };
 
 const MAX_RAW_KV_SCAN_LIMIT: u32 = 10240;
@@ -215,7 +215,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
     /// ```
     pub async fn batch_get(
         &self,
-        keys: impl IntoIterator<Item=impl Into<Key>>,
+        keys: impl IntoIterator<Item = impl Into<Key>>,
     ) -> Result<Vec<KvPair>> {
         debug!(self.logger, "invoking raw batch_get request");
         let request =
@@ -277,7 +277,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
     /// ```
     pub async fn batch_put(
         &self,
-        pairs: impl IntoIterator<Item=impl Into<KvPair>>,
+        pairs: impl IntoIterator<Item = impl Into<KvPair>>,
     ) -> Result<()> {
         debug!(self.logger, "invoking raw batch_put request");
         let request = new_raw_batch_put_request::<C>(
@@ -339,7 +339,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
     /// let result: () = req.await.unwrap();
     /// # });
     /// ```
-    pub async fn batch_delete(&self, keys: impl IntoIterator<Item=impl Into<Key>>) -> Result<()> {
+    pub async fn batch_delete(&self, keys: impl IntoIterator<Item = impl Into<Key>>) -> Result<()> {
         debug!(self.logger, "invoking raw batch_delete request");
         self.assert_non_atomic()?;
         let request =
@@ -465,7 +465,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
     /// ```
     pub async fn batch_scan(
         &self,
-        ranges: impl IntoIterator<Item=impl Into<BoundRange>>,
+        ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
         each_limit: u32,
     ) -> Result<Vec<KvPair>> {
         debug!(self.logger, "invoking raw batch_scan request");
@@ -497,7 +497,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
     /// ```
     pub async fn batch_scan_keys(
         &self,
-        ranges: impl IntoIterator<Item=impl Into<BoundRange>>,
+        ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
         each_limit: u32,
     ) -> Result<Vec<Key>> {
         debug!(self.logger, "invoking raw batch_scan_keys request");
@@ -547,7 +547,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
         &self,
         copr_name: impl Into<String>,
         copr_version_req: impl Into<String>,
-        ranges: impl IntoIterator<Item=impl Into<BoundRange>>,
+        ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
         request_builder: impl Fn(metapb::Region, Vec<Range<Key>>) -> Vec<u8> + Send + Sync + 'static,
     ) -> Result<Vec<(Vec<u8>, Vec<Range<Key>>)>> {
         let copr_version_req = copr_version_req.into();
@@ -593,7 +593,7 @@ impl<C: RequestCodec, PdC: PdClient> Client<C, PdC> {
 
     async fn batch_scan_inner(
         &self,
-        ranges: impl IntoIterator<Item=impl Into<BoundRange>>,
+        ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
         each_limit: u32,
         key_only: bool,
     ) -> Result<Vec<KvPair>> {

src/raw/lowering.rs

@@ -4,14 +4,13 @@
 //! types (i.e., the types from the client crate) and converts these to the types used in the
 //! generated protobuf code, then calls the low-level ctor functions in the requests module.
 
-use std::{iter::Iterator, marker::PhantomData, ops::Range, sync::Arc};
+use std::{iter::Iterator, ops::Range, sync::Arc};
 
 use tikv_client_proto::{kvrpcpb, metapb};
 
 use crate::{
-    raw::requests,
-    request::{request_codec::RequestCodec, KvRequest},
-    BoundRange, ColumnFamily, Key, KvPair, Value,
+    raw::requests, request::request_codec::RequestCodec, BoundRange, ColumnFamily, Key, KvPair,
+    Value,
 };
 
 pub fn new_raw_get_request<C: RequestCodec>(

src/raw/requests.rs

@@ -13,9 +13,8 @@ use crate::{
     collect_first,
     pd::PdClient,
     request::{
-        plan::ResponseWithShard,
-        request_codec::{RawApiV1, RequestCodec},
-        Collect, CollectSingle, DefaultProcessor, KvRequest, Merge, Process, Shardable, SingleKey,
+        plan::ResponseWithShard, request_codec::RequestCodec, Collect, CollectSingle,
+        DefaultProcessor, KvRequest, Merge, Process, Shardable, SingleKey,
     },
     store::{store_stream_for_keys, store_stream_for_ranges, RegionStore},
     transaction::HasLocks,

src/region_cache.rs

@@ -1,18 +1,22 @@
 // Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
 
-use crate::{
-    pd::{RetryClient, RetryClientTrait},
-    region::{RegionId, RegionVerId, RegionWithLeader, StoreId},
-    Key, Result,
-};
 use std::{
     collections::{BTreeMap, HashMap, HashSet},
     sync::Arc,
 };
+
+use tokio::sync::{Notify, RwLock};
+
 use tikv_client_common::Error;
 use tikv_client_pd::Cluster;
 use tikv_client_proto::metapb::{self, Store};
-use tokio::sync::{Notify, RwLock};
+
+use crate::{
+    pd::{RetryClient, RetryClientTrait},
+    region::{RegionId, RegionVerId, RegionWithLeader, StoreId},
+    request::request_codec::RequestCodec,
+    Key, Result,
+};
 
 const MAX_RETRY_WAITING_CONCURRENT_REQUEST: usize = 4;
 
@@ -44,23 +48,25 @@ impl RegionCacheMap {
     }
 }
 
-pub struct RegionCache<Client = RetryClient<Cluster>> {
+pub struct RegionCache<C, Client = RetryClient<Cluster>> {
     region_cache: RwLock<RegionCacheMap>,
     store_cache: RwLock<HashMap<StoreId, Store>>,
     inner_client: Arc<Client>,
+    codec: C,
 }
 
-impl<Client> RegionCache<Client> {
-    pub fn new(inner_client: Arc<Client>) -> RegionCache<Client> {
+impl<C, Client> RegionCache<C, Client> {
+    pub fn new(codec: C, inner_client: Arc<Client>) -> Self {
         RegionCache {
             region_cache: RwLock::new(RegionCacheMap::new()),
             store_cache: RwLock::new(HashMap::new()),
             inner_client,
+            codec,
         }
     }
 }
 
-impl<C: RetryClientTrait> RegionCache<C> {
+impl<C: RequestCodec, R: RetryClientTrait> RegionCache<C, R> {
     // Retrieve cache entry by key. If there's no entry, query PD and update cache.
     pub async fn get_region_by_key(&self, key: &Key) -> Result<RegionWithLeader> {
         let region_cache_guard = self.region_cache.read().await;
@@ -126,7 +132,12 @@ impl<C: RetryClientTrait> RegionCache<C> {
 
     /// Force read through (query from PD) and update cache
     pub async fn read_through_region_by_key(&self, key: Key) -> Result<RegionWithLeader> {
-        let region = self.inner_client.clone().get_region(key.into()).await?;
+        let mut region = self
+            .inner_client
+            .clone()
+            .get_region(self.codec.encode_pd_query(key).into())
+            .await?;
+        region.region = self.codec.decode_region(region.region)?;
         self.add_region(region.clone()).await;
         Ok(region)
     }
@@ -140,7 +151,8 @@ impl<C: RetryClientTrait> RegionCache<C> {
             region_cache_guard.on_my_way_id.insert(id, notify.clone());
         }
 
-        let region = self.inner_client.clone().get_region_by_id(id).await?;
+        let mut region = self.inner_client.clone().get_region_by_id(id).await?;
+        region.region = self.codec.decode_region(region.region)?;
         self.add_region(region.clone()).await;
 
         // notify others
@@ -226,27 +238,35 @@ impl<C: RetryClientTrait> RegionCache<C> {
             cache.key_to_ver_id.remove(&start_key);
         }
     }
+
+    pub fn get_request_codec(&self) -> C {
+        self.codec.clone()
+    }
 }
 
 #[cfg(test)]
 mod test {
-    use super::RegionCache;
-    use crate::{
-        pd::RetryClientTrait,
-        region::{RegionId, RegionWithLeader},
-        Key, Result,
-    };
-    use async_trait::async_trait;
     use std::{
         collections::{BTreeMap, HashMap, HashSet},
         sync::{
             atomic::{AtomicU64, Ordering::SeqCst},
             Arc,
         },
     };
+
+    use async_trait::async_trait;
+    use tokio::sync::Mutex;
+
     use tikv_client_common::Error;
     use tikv_client_proto::metapb;
-    use tokio::sync::Mutex;
+
+    use crate::{
+        pd::RetryClientTrait,
+        region::{RegionId, RegionWithLeader},
+        Key, Result,
+    };
+
+    use super::RegionCache;
 
     #[derive(Default)]
     struct MockRetryClient {

src/request/mod.rs

@@ -9,7 +9,6 @@ use tikv_client_store::{HasKeyErrors, HasRegionError, Request};
 
 use crate::{
     backoff::{Backoff, DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF, PESSIMISTIC_BACKOFF},
-    pd::PdClient,
     transaction::HasLocks,
 };
 

src/request/plan_builder.rs

@@ -45,7 +45,7 @@ where
     PdC: PdClient<RequestCodec = C>,
 {
     pub fn new(pd_client: Arc<PdC>, request: Req) -> Self {
-        let codec = pd_client.get_request_codec().clone();
+        let codec = pd_client.get_request_codec();
         PlanBuilder {
             pd_client,
             plan: Dispatch::new(request, None, codec),

src/request/request_codec.rs

@@ -1,18 +1,18 @@
 use tikv_client_proto::metapb::Region;
 
-use crate::Result;
+use crate::{Key, Result};
 
 pub trait RequestCodec: Sized + Clone + Sync + Send + 'static {
-    fn encode_key(&self, key: Vec<u8>) -> Vec<u8> {
+    fn encode_key(&self, key: Key) -> Key {
         key
     }
-    fn decode_key(&self, key: Vec<u8>) -> Result<Vec<u8>> {
+    fn decode_key(&self, key: Key) -> Result<Key> {
         Ok(key)
     }
-    fn encode_range(&self, start: Vec<u8>, end: Vec<u8>) -> (Vec<u8>, Vec<u8>) {
+    fn encode_range(&self, start: Key, end: Key) -> (Key, Key) {
         (start, end)
     }
-    fn encode_pd_query(&self, key: Vec<u8>) -> Vec<u8> {
+    fn encode_pd_query(&self, key: Key) -> Key {
         key
     }
     fn decode_region(&self, region: Region) -> Result<Region> {

src/transaction/client.rs

@@ -10,10 +10,10 @@ use crate::{
     backoff::{DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF},
     config::Config,
     pd::{PdClient, PdRpcClient},
-    request::{Plan, request_codec::RequestCodec},
-    Result,
+    request::{request_codec::RequestCodec, Plan},
     timestamp::TimestampExt,
     transaction::{Snapshot, Transaction, TransactionOptions},
+    Result,
 };
 
 use super::{requests::new_scan_lock_request, resolve_locks};
@@ -44,8 +44,8 @@ pub struct Client<C> {
 }
 
 impl<C> Client<C>
-    where
-        C: RequestCodec,
+where
+    C: RequestCodec,
 {
     /// Create a transactional [`Client`] and connect to the TiKV cluster.
     ///
@@ -113,7 +113,8 @@ impl<C> Client<C>
         });
         debug!(logger, "creating new transactional client");
         let pd_endpoints: Vec<String> = pd_endpoints.into_iter().map(Into::into).collect();
-        let pd = Arc::new(PdRpcClient::connect(&pd_endpoints, config, codec, logger.clone()).await?);
+        let pd =
+            Arc::new(PdRpcClient::connect(&pd_endpoints, config, codec, logger.clone()).await?);
         Ok(Client {
             pd,
             logger,