style: resolve comments

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

Author: ekexium

src/kv/key.rs

@@ -128,10 +128,6 @@ impl Key {
     }
 }
 
-pub trait HasKeys {
-    fn get_keys(&self) -> Vec<Key>;
-}
-
 impl From<Vec<u8>> for Key {
     fn from(v: Vec<u8>) -> Self {
         Key(v)

src/kv/mod.rs

@@ -8,7 +8,7 @@ mod kvpair;
 mod value;
 
 pub use bound_range::{BoundRange, IntoOwnedRange};
-pub use key::{HasKeys, Key};
+pub use key::Key;
 pub use kvpair::KvPair;
 pub use value::Value;
 

src/request/mod.rs

@@ -10,8 +10,9 @@ use tikv_client_store::{HasError, Request};
 
 pub use self::{
     plan::{
-        Collect, CollectAndMatchKey, CollectError, DefaultProcessor, Dispatch, ExtractError, Merge,
-        MergeResponse, MultiRegion, Plan, Process, ProcessResponse, ResolveLock, RetryRegion,
+        Collect, CollectAndMatchKey, CollectError, DefaultProcessor, Dispatch, ExtractError,
+        HasKeys, Merge, MergeResponse, MultiRegion, Plan, PreserveKey, Process, ProcessResponse,
+        ResolveLock, RetryRegion,
     },
     plan_builder::{PlanBuilder, SingleKey},
     shard::Shardable,

src/request/plan.rs

@@ -2,7 +2,6 @@
 
 use crate::{
     backoff::Backoff,
-    kv::HasKeys,
     pd::PdClient,
     request::{KvRequest, Shardable},
     stats::tikv_stats,
@@ -340,67 +339,12 @@ where
     async fn execute(&self) -> Result<Self::Result> {
         let keys = self.inner.get_keys();
         let res = self.inner.execute().await?;
-        // TODO: should we check they have the same length?
         Ok(ResponseAndKeys(res, keys))
     }
 }
 
-#[cfg(test)]
-mod test {
-    use super::*;
-    use crate::mock::{mock_store, MockPdClient};
-    use futures::stream::BoxStream;
-    use tikv_client_proto::kvrpcpb::BatchGetResponse;
-
-    #[derive(Clone)]
-    struct ErrPlan;
-
-    #[async_trait]
-    impl Plan for ErrPlan {
-        type Result = BatchGetResponse;
-
-        async fn execute(&self) -> Result<Self::Result> {
-            Err(Error::Unimplemented)
-        }
-    }
-
-    impl Shardable for ErrPlan {
-        type Shard = u8;
-
-        fn shards(
-            &self,
-            _: &Arc<impl crate::pd::PdClient>,
-        ) -> BoxStream<'static, crate::Result<(Self::Shard, crate::store::Store)>> {
-            Box::pin(stream::iter(1..=3).map(|_| Err(Error::Unimplemented)))
-                .map_ok(|_: u8| (42, mock_store()))
-                .boxed()
-        }
-
-        fn apply_shard(&mut self, _: Self::Shard, _: &crate::store::Store) -> Result<()> {
-            Ok(())
-        }
-    }
-
-    #[tokio::test]
-    async fn test_err() {
-        let plan = RetryRegion {
-            inner: MultiRegion {
-                inner: ResolveLock {
-                    inner: ErrPlan,
-                    backoff: Backoff::no_backoff(),
-                    pd_client: Arc::new(MockPdClient::default()),
-                },
-                pd_client: Arc::new(MockPdClient::default()),
-            },
-            backoff: Backoff::no_backoff(),
-            pd_client: Arc::new(MockPdClient::default()),
-        };
-        plan.execute()
-            .await
-            .unwrap()
-            .iter()
-            .for_each(|r| assert!(r.is_err()));
-    }
+pub trait HasKeys {
+    fn get_keys(&self) -> Vec<Key>;
 }
 
 // contains a response and the corresponding keys
@@ -459,3 +403,61 @@ impl Merge<ResponseAndKeys<kvrpcpb::PessimisticLockResponse>> for CollectAndMatc
             .collect()
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::mock::{mock_store, MockPdClient};
+    use futures::stream::BoxStream;
+    use tikv_client_proto::kvrpcpb::BatchGetResponse;
+
+    #[derive(Clone)]
+    struct ErrPlan;
+
+    #[async_trait]
+    impl Plan for ErrPlan {
+        type Result = BatchGetResponse;
+
+        async fn execute(&self) -> Result<Self::Result> {
+            Err(Error::Unimplemented)
+        }
+    }
+
+    impl Shardable for ErrPlan {
+        type Shard = u8;
+
+        fn shards(
+            &self,
+            _: &Arc<impl crate::pd::PdClient>,
+        ) -> BoxStream<'static, crate::Result<(Self::Shard, crate::store::Store)>> {
+            Box::pin(stream::iter(1..=3).map(|_| Err(Error::Unimplemented)))
+                .map_ok(|_: u8| (42, mock_store()))
+                .boxed()
+        }
+
+        fn apply_shard(&mut self, _: Self::Shard, _: &crate::store::Store) -> Result<()> {
+            Ok(())
+        }
+    }
+
+    #[tokio::test]
+    async fn test_err() {
+        let plan = RetryRegion {
+            inner: MultiRegion {
+                inner: ResolveLock {
+                    inner: ErrPlan,
+                    backoff: Backoff::no_backoff(),
+                    pd_client: Arc::new(MockPdClient::default()),
+                },
+                pd_client: Arc::new(MockPdClient::default()),
+            },
+            backoff: Backoff::no_backoff(),
+            pd_client: Arc::new(MockPdClient::default()),
+        };
+        plan.execute()
+            .await
+            .unwrap()
+            .iter()
+            .for_each(|r| assert!(r.is_err()));
+    }
+}

src/request/plan_builder.rs

@@ -1,12 +1,12 @@
 // Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
 
+use super::PreserveKey;
 use crate::{
     backoff::Backoff,
-    kv::HasKeys,
     pd::PdClient,
     request::{
-        DefaultProcessor, Dispatch, ExtractError, KvRequest, Merge, MergeResponse, MultiRegion,
-        Plan, Process, ProcessResponse, ResolveLock, RetryRegion, Shardable,
+        DefaultProcessor, Dispatch, ExtractError, HasKeys, KvRequest, Merge, MergeResponse,
+        MultiRegion, Plan, Process, ProcessResponse, ResolveLock, RetryRegion, Shardable,
     },
     store::Store,
     transaction::HasLocks,
@@ -15,8 +15,6 @@ use crate::{
 use std::{marker::PhantomData, sync::Arc};
 use tikv_client_store::HasError;
 
-use super::plan::PreserveKey;
-
 /// Builder type for plans (see that module for more).
 pub struct PlanBuilder<PdC: PdClient, P: Plan, Ph: PlanBuilderPhase> {
     pd_client: Arc<PdC>,

src/request/shard.rs

@@ -1,16 +1,30 @@
 // Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
 
 use crate::{
-    kv::HasKeys,
     pd::PdClient,
-    request::{Dispatch, KvRequest, Plan, ResolveLock, RetryRegion},
+    request::{Dispatch, HasKeys, KvRequest, Plan, PreserveKey, ResolveLock, RetryRegion},
     store::Store,
     Result,
 };
 use futures::stream::BoxStream;
 use std::sync::Arc;
 
-use super::plan::PreserveKey;
+macro_rules! impl_inner_shardable {
+    () => {
+        type Shard = P::Shard;
+
+        fn shards(
+            &self,
+            pd_client: &Arc<impl PdClient>,
+        ) -> BoxStream<'static, Result<(Self::Shard, Store)>> {
+            self.inner.shards(pd_client)
+        }
+
+        fn apply_shard(&mut self, shard: Self::Shard, store: &Store) -> Result<()> {
+            self.inner.apply_shard(shard, store)
+        }
+    };
+}
 
 pub trait Shardable {
     type Shard: Send;
@@ -40,48 +54,15 @@ impl<Req: KvRequest + Shardable> Shardable for Dispatch<Req> {
 }
 
 impl<P: Plan + Shardable, PdC: PdClient> Shardable for ResolveLock<P, PdC> {
-    type Shard = P::Shard;
-
-    fn shards(
-        &self,
-        pd_client: &Arc<impl PdClient>,
-    ) -> BoxStream<'static, Result<(Self::Shard, Store)>> {
-        self.inner.shards(pd_client)
-    }
-
-    fn apply_shard(&mut self, shard: Self::Shard, store: &Store) -> Result<()> {
-        self.inner.apply_shard(shard, store)
-    }
+    impl_inner_shardable!();
 }
 
 impl<P: Plan + HasKeys + Shardable> Shardable for PreserveKey<P> {
-    type Shard = P::Shard;
-
-    fn shards(
-        &self,
-        pd_client: &Arc<impl PdClient>,
-    ) -> BoxStream<'static, Result<(Self::Shard, Store)>> {
-        self.inner.shards(pd_client)
-    }
-
-    fn apply_shard(&mut self, shard: Self::Shard, store: &Store) -> Result<()> {
-        self.inner.apply_shard(shard, store)
-    }
+    impl_inner_shardable!();
 }
 
 impl<P: Plan + Shardable, PdC: PdClient> Shardable for RetryRegion<P, PdC> {
-    type Shard = P::Shard;
-
-    fn shards(
-        &self,
-        pd_client: &Arc<impl PdClient>,
-    ) -> BoxStream<'static, Result<(Self::Shard, Store)>> {
-        self.inner.shards(pd_client)
-    }
-
-    fn apply_shard(&mut self, shard: Self::Shard, store: &Store) -> Result<()> {
-        self.inner.apply_shard(shard, store)
-    }
+    impl_inner_shardable!();
 }
 
 #[macro_export]

src/transaction/requests.rs

@@ -1,9 +1,10 @@
 // Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
 
 use crate::{
-    kv::HasKeys,
     pd::PdClient,
-    request::{Collect, DefaultProcessor, KvRequest, Merge, Process, Shardable, SingleKey},
+    request::{
+        Collect, DefaultProcessor, HasKeys, KvRequest, Merge, Process, Shardable, SingleKey,
+    },
     store::{store_stream_for_keys, store_stream_for_range_by_start_key, Store},
     timestamp::TimestampExt,
     transaction::HasLocks,

tests/integration_tests.rs

@@ -290,7 +290,7 @@ async fn txn_write_million() -> Result<()> {
     assert_eq!(sum, 2usize.pow(NUM_BITS_KEY_PER_TXN + NUM_BITS_TXN));
 
     // test batch_get and batch_get_for_update
-    const SKIP_BITS: u32 = 6; // do not retrive all because there's a limit of message size
+    const SKIP_BITS: u32 = 7; // do not retrive all because there's a limit of message size
     let mut cur = 0u32;
     let keys = iter::repeat_with(|| {
         let v = cur;