tikv
diff --git a/‎Cargo.toml
Lines changed: 2 additions & 0 deletions b/‎Cargo.toml
Lines changed: 2 additions & 0 deletions
diff --git a/‎mock-tikv/src/server.rs
Lines changed: 11 additions & 56 deletions b/‎mock-tikv/src/server.rs
Lines changed: 11 additions & 56 deletions
diff --git a/‎src/mock.rs
Lines changed: 21 additions & 10 deletions b/‎src/mock.rs
Lines changed: 21 additions & 10 deletions
diff --git a/‎src/raw/client.rs
Lines changed: 115 additions & 10 deletions b/‎src/raw/client.rs
Lines changed: 115 additions & 10 deletions
@@ -20,6 +20,7 @@ name = "tikv_client"
 [dependencies]
 async-trait = "0.1"
 derive-new = "0.5"
+either = "1.6"
 fail = "0.4"
 futures = { version = "0.3", features = ["async-await", "thread-pool"] }
 futures-timer = "3.0"
@@ -29,6 +30,7 @@ log = "0.4"
 prometheus = { version = "0.12", features = [ "push", "process" ], default-features = false } 
 rand = "0.8"
 regex = "1"
+semver = "0.11"
 serde = "1.0"
 serde_derive = "1.0"
 slog = { version = "2.3", features = ["max_level_trace", "release_max_level_debug"] }
 
@@ -5,7 +5,7 @@ use derive_new::new;
 use futures::{FutureExt, TryFutureExt};
 use grpcio::{Environment, Server, ServerBuilder};
 use std::sync::Arc;
-use tikv_client_proto::{coprocessor_v2::*, kvrpcpb::*, tikvpb::*};
+use tikv_client_proto::{kvrpcpb::*, tikvpb::*};
 
 pub const MOCK_TIKV_PORT: u16 = 50019;
 
@@ -271,60 +271,6 @@ impl Tikv for MockTikv {
         todo!()
     }
 
-    fn ver_get(
-        &mut self,
-        _ctx: grpcio::RpcContext,
-        _req: tikv_client_proto::kvrpcpb::VerGetRequest,
-        _sink: grpcio::UnarySink<tikv_client_proto::kvrpcpb::VerGetResponse>,
-    ) {
-        todo!()
-    }
-
-    fn ver_batch_get(
-        &mut self,
-        _ctx: grpcio::RpcContext,
-        _req: tikv_client_proto::kvrpcpb::VerBatchGetRequest,
-        _sink: grpcio::UnarySink<tikv_client_proto::kvrpcpb::VerBatchGetResponse>,
-    ) {
-        todo!()
-    }
-
-    fn ver_mut(
-        &mut self,
-        _ctx: grpcio::RpcContext,
-        _req: tikv_client_proto::kvrpcpb::VerMutRequest,
-        _sink: grpcio::UnarySink<tikv_client_proto::kvrpcpb::VerMutResponse>,
-    ) {
-        todo!()
-    }
-
-    fn ver_batch_mut(
-        &mut self,
-        _ctx: grpcio::RpcContext,
-        _req: tikv_client_proto::kvrpcpb::VerBatchMutRequest,
-        _sink: grpcio::UnarySink<tikv_client_proto::kvrpcpb::VerBatchMutResponse>,
-    ) {
-        todo!()
-    }
-
-    fn ver_scan(
-        &mut self,
-        _ctx: grpcio::RpcContext,
-        _req: tikv_client_proto::kvrpcpb::VerScanRequest,
-        _sink: grpcio::UnarySink<tikv_client_proto::kvrpcpb::VerScanResponse>,
-    ) {
-        todo!()
-    }
-
-    fn ver_delete_range(
-        &mut self,
-        _ctx: grpcio::RpcContext,
-        _req: tikv_client_proto::kvrpcpb::VerDeleteRangeRequest,
-        _sink: grpcio::UnarySink<tikv_client_proto::kvrpcpb::VerDeleteRangeResponse>,
-    ) {
-        todo!()
-    }
-
     fn unsafe_destroy_range(
         &mut self,
         _ctx: grpcio::RpcContext,
@@ -532,7 +478,7 @@ impl Tikv for MockTikv {
         todo!()
     }
 
-    fn coprocessor_v2(
+    fn raw_coprocessor(
         &mut self,
         _: grpcio::RpcContext,
         _: RawCoprocessorRequest,
@@ -549,4 +495,13 @@ impl Tikv for MockTikv {
     ) {
         todo!()
     }
+
+    fn get_lock_wait_info(
+        &mut self,
+        _: grpcio::RpcContext,
+        _: GetLockWaitInfoRequest,
+        _: grpcio::UnarySink<GetLockWaitInfoResponse>,
+    ) {
+        todo!()
+    }
 }
@@ -107,7 +107,7 @@ impl MockPdClient {
     pub fn region1() -> RegionWithLeader {
         let mut region = RegionWithLeader::default();
         region.region.id = 1;
-        region.region.set_start_key(vec![0]);
+        region.region.set_start_key(vec![]);
         region.region.set_end_key(vec![10]);
 
         let leader = metapb::Peer {
@@ -133,6 +133,21 @@ impl MockPdClient {
 
         region
     }
+
+    pub fn region3() -> RegionWithLeader {
+        let mut region = RegionWithLeader::default();
+        region.region.id = 3;
+        region.region.set_start_key(vec![250, 250]);
+        region.region.set_end_key(vec![]);
+
+        let leader = metapb::Peer {
+            store_id: 43,
+            ..Default::default()
+        };
+        region.leader = Some(leader);
+
+        region
+    }
 }
 
 #[async_trait]
@@ -145,10 +160,12 @@ impl PdClient for MockPdClient {
 
     async fn region_for_key(&self, key: &Key) -> Result<RegionWithLeader> {
         let bytes: &[_] = key.into();
-        let region = if bytes.is_empty() || bytes[0] < 10 {
+        let region = if bytes.is_empty() || bytes < &[10][..] {
             Self::region1()
-        } else {
+        } else if bytes >= &[10][..] && bytes < &[250, 250][..] {
             Self::region2()
+        } else {
+            Self::region3()
         };
 
         Ok(region)
@@ -158,6 +175,7 @@ impl PdClient for MockPdClient {
         match id {
             1 => Ok(Self::region1()),
             2 => Ok(Self::region2()),
+            3 => Ok(Self::region3()),
             _ => Err(Error::RegionNotFoundInResponse { region_id: id }),
         }
     }
@@ -180,10 +198,3 @@ impl PdClient for MockPdClient {
 
     async fn invalidate_region_cache(&self, _ver_id: crate::region::RegionVerId) {}
 }
-
-pub fn mock_store() -> RegionStore {
-    RegionStore {
-        region_with_leader: RegionWithLeader::default(),
-        client: Arc::new(MockKvClient::new("foo".to_owned(), None)),
-    }
-}
 
@@ -1,17 +1,20 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
+use core::ops::Range;
+use std::{str::FromStr, sync::Arc, u32};
+
+use slog::{Drain, Logger};
 use tikv_client_common::Error;
+use tikv_client_proto::metapb;
 
 use crate::{
     backoff::DEFAULT_REGION_BACKOFF,
     config::Config,
-    pd::PdRpcClient,
+    pd::{PdClient, PdRpcClient},
     raw::lowering::*,
     request::{Collect, CollectSingle, Plan},
     BoundRange, ColumnFamily, Key, KvPair, Result, Value,
 };
-use slog::{Drain, Logger};
-use std::{sync::Arc, u32};
 
 const MAX_RAW_KV_SCAN_LIMIT: u32 = 10240;
 
@@ -23,15 +26,15 @@ 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 {
-    rpc: Arc<PdRpcClient>,
+pub struct Client<PdC: PdClient = PdRpcClient> {
+    rpc: Arc<PdC>,
     cf: Option<ColumnFamily>,
     /// Whether to use the [`atomic mode`](Client::with_atomic_for_cas).
     atomic: bool,
     logger: Logger,
 }
 
-impl Client {
+impl Client<PdRpcClient> {
     /// 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
@@ -50,7 +53,7 @@ impl Client {
     pub async fn new<S: Into<String>>(
         pd_endpoints: Vec<S>,
         logger: Option<Logger>,
-    ) -> Result<Client> {
+    ) -> Result<Self> {
         Self::new_with_config(pd_endpoints, Config::default(), logger).await
     }
 
@@ -78,7 +81,7 @@ impl Client {
         pd_endpoints: Vec<S>,
         config: Config,
         optional_logger: Option<Logger>,
-    ) -> Result<Client> {
+    ) -> Result<Self> {
         let logger = optional_logger.unwrap_or_else(|| {
             let plain = slog_term::PlainSyncDecorator::new(std::io::stdout());
             Logger::root(
@@ -125,7 +128,7 @@ impl Client {
     /// let get_request = client.get("foo".to_owned());
     /// # });
     /// ```
-    pub fn with_cf(&self, cf: ColumnFamily) -> Client {
+    pub fn with_cf(&self, cf: ColumnFamily) -> Self {
         Client {
             rpc: self.rpc.clone(),
             cf: Some(cf),
@@ -141,15 +144,17 @@ impl Client {
     /// the atomicity of CAS, write operations like [`put`](Client::put) or
     /// [`delete`](Client::delete) in atomic mode are more expensive. Some
     /// operations are not supported in the mode.
-    pub fn with_atomic_for_cas(&self) -> Client {
+    pub fn with_atomic_for_cas(&self) -> Self {
         Client {
             rpc: self.rpc.clone(),
             cf: self.cf.clone(),
             atomic: true,
             logger: self.logger.clone(),
         }
     }
+}
 
+impl<PdC: PdClient> Client<PdC> {
     /// Create a new 'get' request.
     ///
     /// Once resolved this request will result in the fetching of the value associated with the
@@ -517,6 +522,29 @@ impl Client {
         plan.execute().await
     }
 
+    pub async fn coprocessor(
+        &self,
+        copr_name: impl Into<String>,
+        copr_version_req: impl Into<String>,
+        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();
+        semver::VersionReq::from_str(&copr_version_req)?;
+        let req = new_raw_coprocessor_request(
+            copr_name.into(),
+            copr_version_req,
+            ranges.into_iter().map(Into::into),
+            request_builder,
+        );
+        let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req)
+            .preserve_shard()
+            .retry_multi_region(DEFAULT_REGION_BACKOFF)
+            .post_process_default()
+            .plan();
+        plan.execute().await
+    }
+
     async fn scan_inner(
         &self,
         range: impl Into<BoundRange>,
@@ -576,3 +604,80 @@ impl Client {
         self.atomic.then(|| ()).ok_or(Error::UnsupportedMode)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        mock::{MockKvClient, MockPdClient},
+        Result,
+    };
+    use std::{any::Any, sync::Arc};
+    use tikv_client_proto::kvrpcpb;
+
+    #[tokio::test]
+    async fn test_raw_coprocessor() -> Result<()> {
+        let plain = slog_term::PlainSyncDecorator::new(std::io::stdout());
+        let logger = Logger::root(
+            slog_term::FullFormat::new(plain)
+                .build()
+                .filter_level(slog::Level::Info)
+                .fuse(),
+            o!(),
+        );
+        let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook(
+            move |req: &dyn Any| {
+                if let Some(req) = req.downcast_ref::<kvrpcpb::RawCoprocessorRequest>() {
+                    assert_eq!(req.copr_name, "example");
+                    assert_eq!(req.copr_version_req, "0.1.0");
+                    let resp = kvrpcpb::RawCoprocessorResponse {
+                        data: req.data.clone(),
+                        ..Default::default()
+                    };
+                    Ok(Box::new(resp) as Box<dyn Any>)
+                } else {
+                    unreachable!()
+                }
+            },
+        )));
+        let client = Client {
+            rpc: pd_client,
+            cf: Some(ColumnFamily::Default),
+            atomic: false,
+            logger,
+        };
+        let resps = client
+            .coprocessor(
+                "example",
+                "0.1.0",
+                vec![vec![5]..vec![15], vec![20]..vec![]],
+                |region, ranges| format!("{:?}:{:?}", region.id, ranges).into_bytes(),
+            )
+            .await?;
+        let resps: Vec<_> = resps
+            .into_iter()
+            .map(|(data, ranges)| (String::from_utf8(data).unwrap(), ranges))
+            .collect();
+        assert_eq!(
+            resps,
+            vec![
+                (
+                    "1:[Key(05)..Key(0A)]".to_string(),
+                    vec![Key::from(vec![5])..Key::from(vec![10])]
+                ),
+                (
+                    "2:[Key(0A)..Key(0F), Key(14)..Key(FAFA)]".to_string(),
+                    vec![
+                        Key::from(vec![10])..Key::from(vec![15]),
+                        Key::from(vec![20])..Key::from(vec![250, 250])
+                    ]
+                ),
+                (
+                    "3:[Key(FAFA)..Key()]".to_string(),
+                    vec![Key::from(vec![250, 250])..Key::from(vec![])]
+                )
+            ]
+        );
+        Ok(())
+    }
+}