diff --git a/Cargo.toml b/Cargo.toml index 77f97216..fad2ae0c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,6 +38,7 @@ slog-term = { version = "2.4" } thiserror = "1" tokio = { version = "1", features = [ "sync", "rt-multi-thread", "macros" ] } async-recursion = "0.3" +paste = "1.0" tikv-client-common = { version = "0.1.0", path = "tikv-client-common" } tikv-client-pd = { version = "0.1.0", path = "tikv-client-pd" } diff --git a/Makefile b/Makefile index 268dc574..c70ff9a8 100644 --- a/Makefile +++ b/Makefile @@ -34,3 +34,6 @@ tiup: tiup playground nightly --mode tikv-slim --kv 3 --without-monitor --kv.config $(shell pwd)/config/tikv.toml --pd.config $(shell pwd)/config/pd.toml & all: check doc test + +format: + cargo fmt --all \ No newline at end of file diff --git a/examples/pessimistic.rs b/examples/pessimistic.rs index bfdb9e27..565942c1 100644 --- a/examples/pessimistic.rs +++ b/examples/pessimistic.rs @@ -1,9 +1,12 @@ // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0. -mod common; +use tikv_client::{ + transaction::ApiV1, Config, Key, TransactionClient as Client, TransactionOptions, Value, +}; use crate::common::parse_args; -use tikv_client::{Config, Key, TransactionClient as Client, TransactionOptions, Value}; + +mod common; #[tokio::main] async fn main() { @@ -20,7 +23,7 @@ async fn main() { }; // init - let client = Client::new_with_config(args.pd, config, None) + let client = Client::::new_with_config(args.pd, config, None) .await .expect("Could not connect to tikv"); diff --git a/examples/raw.rs b/examples/raw.rs index de583d80..00eb59a5 100644 --- a/examples/raw.rs +++ b/examples/raw.rs @@ -5,7 +5,9 @@ mod common; use crate::common::parse_args; -use tikv_client::{Config, IntoOwnedRange, Key, KvPair, RawClient as Client, Result, Value}; +use tikv_client::{ + raw::ApiV1, Config, IntoOwnedRange, Key, KvPair, RawClient as Client, Result, Value, +}; const KEY: &str = "TiKV"; const VALUE: &str = "Rust"; @@ -26,8 +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, None).await?; - let client = client.clone(); + let client = Client::::new_with_config(args.pd, config, 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 diff --git a/examples/transaction.rs b/examples/transaction.rs index 5d58a59b..6aff1179 100644 --- a/examples/transaction.rs +++ b/examples/transaction.rs @@ -3,9 +3,12 @@ mod common; use crate::common::parse_args; -use tikv_client::{BoundRange, Config, Key, KvPair, TransactionClient as Client, Value}; +use tikv_client::{ + request::codec::TxnCodec, transaction::ApiV1, BoundRange, Config, Key, KvPair, + TransactionClient as Client, Value, +}; -async fn puts(client: &Client, pairs: impl IntoIterator>) { +async fn puts(client: &Client, pairs: impl IntoIterator>) { let mut txn = client .begin_optimistic() .await @@ -17,7 +20,7 @@ async fn puts(client: &Client, pairs: impl IntoIterator txn.commit().await.expect("Could not commit transaction"); } -async fn get(client: &Client, key: Key) -> Option { +async fn get(client: &Client, key: Key) -> Option { let mut txn = client .begin_optimistic() .await @@ -29,7 +32,7 @@ async fn get(client: &Client, key: Key) -> Option { res } -async fn key_exists(client: &Client, key: Key) -> bool { +async fn key_exists(client: &Client, key: Key) -> bool { let mut txn = client .begin_optimistic() .await @@ -44,7 +47,7 @@ async fn key_exists(client: &Client, key: Key) -> bool { res } -async fn scan(client: &Client, range: impl Into, limit: u32) { +async fn scan(client: &Client, range: impl Into, limit: u32) { let mut txn = client .begin_optimistic() .await @@ -56,7 +59,7 @@ async fn scan(client: &Client, range: impl Into, limit: u32) { txn.commit().await.expect("Could not commit transaction"); } -async fn dels(client: &Client, keys: impl IntoIterator) { +async fn dels(client: &Client, keys: impl IntoIterator) { let mut txn = client .begin_optimistic() .await @@ -81,7 +84,7 @@ async fn main() { Config::default() }; - let txn = Client::new_with_config(args.pd, config, None) + let txn = Client::::new_with_config(args.pd, config, None) .await .expect("Could not connect to tikv"); diff --git a/src/lib.rs b/src/lib.rs index e94785b0..eedc2f6f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -66,8 +66,9 @@ //! # use tikv_client::{RawClient, Result}; //! # use futures::prelude::*; //! # fn main() -> Result<()> { -//! # futures::executor::block_on(async { -//! let client = RawClient::new(vec!["127.0.0.1:2379"], None).await?; +//! # use tikv_client::raw::ApiV1; +//! futures::executor::block_on(async { +//! let client = RawClient::::new(vec!["127.0.0.1:2379"], None).await?; //! client.put("key".to_owned(), "value".to_owned()).await?; //! let value = client.get("key".to_owned()).await?; //! # Ok(()) @@ -80,8 +81,9 @@ //! # use tikv_client::{TransactionClient, Result}; //! # use futures::prelude::*; //! # fn main() -> Result<()> { -//! # futures::executor::block_on(async { -//! let txn_client = TransactionClient::new(vec!["127.0.0.1:2379"], None).await?; +//! # use tikv_client::transaction::ApiV1; +//! futures::executor::block_on(async { +//! let txn_client = TransactionClient::::new(vec!["127.0.0.1:2379"], None).await?; //! let mut txn = txn_client.begin_optimistic().await?; //! txn.put("key".to_owned(), "value".to_owned()).await?; //! let value = txn.get("key".to_owned()).await?; @@ -90,6 +92,11 @@ //! # })} //! ``` +#![feature(min_specialization)] +#![feature(core_intrinsics)] +#![feature(async_closure)] +#![feature(get_mut_unchecked)] + #[macro_use] pub mod request; #[macro_use] @@ -124,6 +131,8 @@ pub use crate::backoff::Backoff; #[doc(inline)] pub use crate::kv::{BoundRange, IntoOwnedRange, Key, KvPair, Value}; #[doc(inline)] +pub use crate::pd::PdClient; +#[doc(inline)] pub use crate::raw::{lowering as raw_lowering, Client as RawClient, ColumnFamily}; #[doc(inline)] pub use crate::request::RetryOptions; diff --git a/src/mock.rs b/src/mock.rs index 02a0bef9..cf313496 100644 --- a/src/mock.rs +++ b/src/mock.rs @@ -7,6 +7,7 @@ use crate::{ pd::{PdClient, PdRpcClient, RetryClient}, + raw::ApiV1, region::{RegionId, RegionWithLeader}, store::RegionStore, Config, Error, Key, Result, Timestamp, @@ -20,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 { +pub async fn pd_rpc_client() -> PdRpcClient { let config = Config::default(); let plain = slog_term::PlainSyncDecorator::new(std::io::stdout()); let logger = Logger::root( @@ -41,7 +42,7 @@ pub async fn pd_rpc_client() -> PdRpcClient { MockCluster, )) }, - false, + async move |_| Ok(ApiV1::default()), logger, ) .await @@ -153,6 +154,7 @@ impl MockPdClient { #[async_trait] impl PdClient for MockPdClient { type KvClient = MockKvClient; + type RequestCodec = ApiV1; async fn map_region_to_store(self: Arc, region: RegionWithLeader) -> Result { Ok(RegionStore::new(region, Arc::new(self.client.clone()))) @@ -197,4 +199,8 @@ impl PdClient for MockPdClient { } async fn invalidate_region_cache(&self, _ver_id: crate::region::RegionVerId) {} + + fn get_request_codec(&self) -> Self::RequestCodec { + ApiV1::default() + } } diff --git a/src/pd/client.rs b/src/pd/client.rs index 75610544..acdd3779 100644 --- a/src/pd/client.rs +++ b/src/pd/client.rs @@ -1,23 +1,26 @@ // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0. +use std::{collections::HashMap, 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::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 +45,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, region: RegionWithLeader) -> Result; @@ -191,33 +195,27 @@ pub trait PdClient: Send + Sync + 'static { .boxed() } - fn decode_region(mut region: RegionWithLeader, enable_codec: bool) -> Result { - if enable_codec { - codec::decode_bytes_in_place(region.region.mut_start_key(), false)?; - codec::decode_bytes_in_place(region.region.mut_end_key(), false)?; - } - Ok(region) - } - 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 { +pub struct PdRpcClient { pd: Arc>, kv_connect: KvC, kv_client_cache: Arc>>, - enable_codec: bool, - region_cache: RegionCache>, + region_cache: RegionCache>, logger: Logger, } #[async_trait] -impl PdClient for PdRpcClient { +impl PdClient for PdRpcClient { type KvClient = KvC::KvClient; + type RequestCodec = C; async fn map_region_to_store(self: Arc, region: RegionWithLeader) -> Result { let store_id = region.get_store_id()?; @@ -227,20 +225,11 @@ impl PdClient for PdRpcClient { } async fn region_for_key(&self, key: &Key) -> Result { - 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 { - 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) -> Result { @@ -258,22 +247,30 @@ impl PdClient for PdRpcClient { 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 { + self.region_cache.get_request_codec() + } } -impl PdRpcClient { - pub async fn connect( +impl PdRpcClient { + pub async fn connect( pd_endpoints: &[String], config: Config, - enable_codec: bool, + codec_factory: F, logger: Logger, - ) -> Result { + ) -> Result> + where + F: Fn(Arc) -> Fut, + Fut: Future>, + { PdRpcClient::new( config.clone(), |env, security_mgr| TikvConnect::new(env, security_mgr, config.timeout), |env, security_mgr| { RetryClient::connect(env, pd_endpoints, security_mgr, config.timeout) }, - enable_codec, + codec_factory, logger, ) .await @@ -289,18 +286,20 @@ fn thread_name(prefix: &str) -> String { .unwrap_or_else(|| prefix.to_owned()) } -impl PdRpcClient { - pub async fn new( +impl PdRpcClient { + pub async fn new( config: Config, kv_connect: MakeKvC, pd: MakePd, - enable_codec: bool, + codec: MakeCodec, logger: Logger, - ) -> Result> + ) -> Result> where PdFut: Future>>, MakeKvC: FnOnce(Arc, Arc) -> KvC, MakePd: FnOnce(Arc, Arc) -> PdFut, + MakeCodec: Fn(Arc>) -> CodecFut, + CodecFut: Future>, { let env = Arc::new( EnvBuilder::new() @@ -320,12 +319,12 @@ impl PdRpcClient { let pd = Arc::new(pd(env.clone(), security_mgr.clone()).await?); let kv_client_cache = Default::default(); + Ok(PdRpcClient { pd: pd.clone(), kv_client_cache, kv_connect: kv_connect(env, security_mgr), - enable_codec, - region_cache: RegionCache::new(pd), + region_cache: RegionCache::new(codec(pd.clone()).await?, pd), logger, }) } @@ -350,10 +349,11 @@ impl PdRpcClient { #[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() { @@ -394,7 +394,7 @@ pub mod test { vec![1].into(), vec![2].into(), vec![3].into(), - vec![5, 2].into() + vec![5, 2].into(), ] ); assert_eq!( @@ -456,12 +456,12 @@ 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.id(), 2); @@ -469,7 +469,7 @@ pub mod test { ranges2.1, vec![kvrpcpb::KeyRange { start_key: k_split.clone(), - end_key: k3 + end_key: k3, }] ); assert_eq!(ranges3.0.id(), 1); @@ -477,7 +477,7 @@ pub mod test { ranges3.1, vec![kvrpcpb::KeyRange { start_key: k2, - end_key: k_split.clone() + end_key: k_split.clone(), }] ); assert_eq!(ranges4.0.id(), 2); @@ -485,7 +485,7 @@ pub mod test { ranges4.1, vec![kvrpcpb::KeyRange { start_key: k_split, - end_key: k4 + end_key: k4, }] ); assert!(stream.next().is_none()); diff --git a/src/pd/retry.rs b/src/pd/retry.rs index f6524cad..9b743aa1 100644 --- a/src/pd/retry.rs +++ b/src/pd/retry.rs @@ -17,7 +17,7 @@ use std::{ }; use tikv_client_pd::{Cluster, Connection}; use tikv_client_proto::{ - metapb, + keyspacepb, metapb, pdpb::{self, Timestamp}, }; use tokio::sync::RwLock; @@ -43,6 +43,9 @@ pub trait RetryClientTrait { async fn get_timestamp(self: Arc) -> Result; async fn update_safepoint(self: Arc, safepoint: u64) -> Result; + + async fn load_keyspace(self: Arc, name: &str) + -> Result; } /// Client for communication with a PD cluster. Has the facility to reconnect to the cluster. pub struct RetryClient { @@ -182,6 +185,15 @@ impl RetryClientTrait for RetryClient { .map(|resp| resp.get_new_safe_point() == safepoint) }) } + + async fn load_keyspace( + self: Arc, + name: &str, + ) -> Result { + retry!(self, "load_keyspace", |cluster| async { + cluster.load_keyspace(name.to_string(), self.timeout).await + }) + } } impl fmt::Debug for RetryClient { diff --git a/src/raw/client.rs b/src/raw/client.rs index da67c56a..fbc9cdc1 100644 --- a/src/raw/client.rs +++ b/src/raw/client.rs @@ -1,18 +1,20 @@ // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. use core::ops::Range; -use std::{str::FromStr, sync::Arc, u32}; +use std::{future::Future, marker::PhantomData, str::FromStr, sync::Arc, u32}; + +use slog::Logger; -use slog::{Drain, Logger}; use tikv_client_common::Error; use tikv_client_proto::metapb; use crate::{ backoff::DEFAULT_REGION_BACKOFF, config::Config, - pd::{PdClient, PdRpcClient}, - raw::lowering::*, - request::{Collect, CollectSingle, Plan}, + pd::{PdClient, PdRpcClient, RetryClient}, + raw::{lowering::*, ApiV2}, + request::{codec::RawCodec, Collect, CollectSingle, Plan}, + util::client::ClientContext, Backoff, BoundRange, ColumnFamily, Key, KvPair, Result, Value, }; @@ -25,26 +27,81 @@ 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. -pub struct Client { +pub struct Client> +where + C: RawCodec, + PdC: PdClient, +{ rpc: Arc, cf: Option, /// Whether to use the [`atomic mode`](Client::with_atomic_for_cas). atomic: bool, logger: Logger, + _phantom: PhantomData, } -impl Clone for Client { +impl Clone for Client { fn clone(&self) -> Self { Self { rpc: self.rpc.clone(), cf: self.cf.clone(), atomic: self.atomic, logger: self.logger.clone(), + _phantom: PhantomData, } } } -impl Client { +impl Client { + /// Create a raw [`Client`] and connect to the TiKV cluster, with `keyspace` name. + /// + /// Because TiKV is managed by a [PD](https://github.com/pingcap/pd/) cluster, the endpoints for + /// PD must be provided, not the TiKV nodes. It's important to include more than one PD endpoint + /// (include all endpoints, if possible), this helps avoid having a single point of failure. + /// + /// # Examples + /// + /// ```rust,no_run + /// # use tikv_client::{Config, RawClient}; + /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV2; + /// # futures::executor::block_on(async { + /// let client = RawClient::::new_with_keyspace( + /// vec!["192.168.0.100"], + /// "space_for_raw", + /// Config::default(), + /// None, + /// ) + /// .await + /// .unwrap(); + /// # }); + /// ``` + pub async fn new_with_keyspace( + pd_endpoints: Vec, + keyspace: K, + config: Config, + optional_logger: Option, + ) -> Result + where + S: Into, + K: Into, + { + let keyspace = &keyspace.into(); + let this = Self::new_with_codec_factory( + pd_endpoints, + config, + async move |pd| ApiV2::with_keyspace(keyspace, pd).await, + optional_logger, + ) + .await?; + + info!(this.logger, "created raw client"; "keyspace" => keyspace); + + Ok(this) + } +} + +impl Client { /// 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 @@ -56,8 +113,11 @@ impl Client { /// ```rust,no_run /// # use tikv_client::RawClient; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// let client = RawClient::::new(vec!["192.168.0.100"], None) + /// .await + /// .unwrap(); /// # }); /// ``` pub async fn new>( @@ -79,8 +139,9 @@ impl Client { /// # use tikv_client::{Config, RawClient}; /// # use futures::prelude::*; /// # use std::time::Duration; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// let client = RawClient::new_with_config( + /// let client = RawClient::::new_with_config( /// vec!["192.168.0.100"], /// Config::default().with_timeout(Duration::from_secs(60)), /// None, @@ -94,26 +155,13 @@ impl Client { config: Config, optional_logger: Option, ) -> Result { - let logger = optional_logger.unwrap_or_else(|| { - let plain = slog_term::PlainSyncDecorator::new(std::io::stdout()); - Logger::root( - slog_term::FullFormat::new(plain) - .build() - .filter_level(slog::Level::Info) - .fuse(), - o!(), - ) - }); - debug!(logger, "creating new raw client"); - let pd_endpoints: Vec = pd_endpoints.into_iter().map(Into::into).collect(); - let rpc = - Arc::new(PdRpcClient::connect(&pd_endpoints, config, false, logger.clone()).await?); - Ok(Client { - rpc, - cf: None, - atomic: false, - logger, - }) + Self::new_with_codec_factory( + pd_endpoints, + config, + async move |_| Ok(C::default()), + optional_logger, + ) + .await } /// Create a new client which is a clone of `self`, but which uses an explicit column family for @@ -131,8 +179,9 @@ impl Client { /// # use tikv_client::{Config, RawClient, ColumnFamily}; /// # use futures::prelude::*; /// # use std::convert::TryInto; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// let client = RawClient::new(vec!["192.168.0.100"], None) + /// let client = RawClient::::new(vec!["192.168.0.100"], None) /// .await /// .unwrap() /// .with_cf(ColumnFamily::Write); @@ -147,6 +196,7 @@ impl Client { cf: Some(cf), atomic: self.atomic, logger: self.logger.clone(), + _phantom: PhantomData, } } @@ -164,11 +214,39 @@ impl Client { cf: self.cf.clone(), atomic: true, logger: self.logger.clone(), + _phantom: PhantomData, } } + + async fn new_with_codec_factory( + pd_endpoints: Vec, + config: Config, + codec_factory: F, + optional_logger: Option, + ) -> Result + where + S: Into, + F: Fn(Arc) -> Fut, + Fut: Future>, + { + let ClientContext { pd: rpc, logger } = + ClientContext::new(pd_endpoints, config, codec_factory, optional_logger).await?; + + Ok(Client { + rpc, + cf: None, + atomic: false, + logger, + _phantom: PhantomData, + }) + } } -impl Client { +impl Client +where + C: RawCodec, + PdC: PdClient, +{ /// Create a new 'get' request. /// /// Once resolved this request will result in the fetching of the value associated with the @@ -180,8 +258,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Value, Config, RawClient}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let key = "TiKV".to_owned(); /// let req = client.get(key); /// let result: Option = req.await.unwrap(); @@ -214,8 +293,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{KvPair, Config, RawClient}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let keys = vec!["TiKV".to_owned(), "TiDB".to_owned()]; /// let req = client.batch_get(keys); /// let result: Vec = req.await.unwrap(); @@ -253,8 +333,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Key, Value, Config, RawClient}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let key = "TiKV".to_owned(); /// let val = "TiKV".to_owned(); /// let req = client.put(key, val); @@ -291,8 +372,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Result, KvPair, Key, Value, Config, RawClient, IntoOwnedRange}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let kvpair1 = ("PD".to_owned(), "Go".to_owned()); /// let kvpair2 = ("TiKV".to_owned(), "Rust".to_owned()); /// let iterable = vec![kvpair1, kvpair2]; @@ -337,8 +419,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Key, Config, RawClient}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let key = "TiKV".to_owned(); /// let req = client.delete(key); /// let result: () = req.await.unwrap(); @@ -371,8 +454,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Config, RawClient}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let keys = vec!["TiKV".to_owned(), "TiDB".to_owned()]; /// let req = client.batch_delete(keys); /// let result: () = req.await.unwrap(); @@ -408,8 +492,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Key, Config, RawClient, IntoOwnedRange}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let inclusive_range = "TiKV"..="TiDB"; /// let req = client.delete_range(inclusive_range.into_owned()); /// let result: () = req.await.unwrap(); @@ -448,8 +533,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{KvPair, Config, RawClient, IntoOwnedRange}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let inclusive_range = "TiKV"..="TiDB"; /// let req = client.scan(inclusive_range.into_owned(), 2); /// let result: Vec = req.await.unwrap(); @@ -482,8 +568,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Key, Config, RawClient, IntoOwnedRange}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let inclusive_range = "TiKV"..="TiDB"; /// let req = client.scan_keys(inclusive_range.into_owned(), 2); /// let result: Vec = req.await.unwrap(); @@ -524,8 +611,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Key, Config, RawClient, IntoOwnedRange}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let inclusive_range1 = "TiDB"..="TiKV"; /// let inclusive_range2 = "TiKV"..="TiSpark"; /// let iterable = vec![inclusive_range1.into_owned(), inclusive_range2.into_owned()]; @@ -568,8 +656,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Key, Config, RawClient, IntoOwnedRange}; /// # use futures::prelude::*; + /// use tikv_client::raw::ApiV1; /// # futures::executor::block_on(async { - /// # let client = RawClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = RawClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let inclusive_range1 = "TiDB"..="TiKV"; /// let inclusive_range2 = "TiKV"..="TiSpark"; /// let iterable = vec![inclusive_range1.into_owned(), inclusive_range2.into_owned()]; @@ -744,23 +833,28 @@ impl Client { } fn assert_non_atomic(&self) -> Result<()> { - (!self.atomic).then(|| ()).ok_or(Error::UnsupportedMode) + (!self.atomic).then_some(()).ok_or(Error::UnsupportedMode) } fn assert_atomic(&self) -> Result<()> { - self.atomic.then(|| ()).ok_or(Error::UnsupportedMode) + self.atomic.then_some(()).ok_or(Error::UnsupportedMode) } } #[cfg(test)] mod tests { - use super::*; + use slog::Drain; + use std::{any::Any, sync::Arc}; + + use tikv_client_proto::kvrpcpb; + use crate::{ mock::{MockKvClient, MockPdClient}, + raw::ApiV1, Result, }; - use std::{any::Any, sync::Arc}; - use tikv_client_proto::kvrpcpb; + + use super::*; #[tokio::test] async fn test_raw_coprocessor() -> Result<()> { @@ -792,6 +886,7 @@ mod tests { cf: Some(ColumnFamily::Default), atomic: false, logger, + _phantom: PhantomData::, }; let resps = client .coprocessor( @@ -816,13 +911,13 @@ mod tests { "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]) + 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(()) diff --git a/src/raw/codec.rs b/src/raw/codec.rs new file mode 100644 index 00000000..cc530ee2 --- /dev/null +++ b/src/raw/codec.rs @@ -0,0 +1,4 @@ +use crate::request::codec::{self, RawMode}; + +pub type ApiV1 = codec::ApiV1; +pub type ApiV2 = codec::ApiV2; diff --git a/src/raw/mod.rs b/src/raw/mod.rs index 823e0e21..4166d5c9 100644 --- a/src/raw/mod.rs +++ b/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, ApiV2}, +}; use crate::Error; use std::{convert::TryFrom, fmt}; mod client; +mod codec; pub mod lowering; mod requests; diff --git a/src/raw/requests.rs b/src/raw/requests.rs index bd678aaf..40676b2a 100644 --- a/src/raw/requests.rs +++ b/src/raw/requests.rs @@ -5,16 +5,17 @@ use std::{any::Any, ops::Range, sync::Arc}; use async_trait::async_trait; use futures::stream::BoxStream; use grpcio::CallOption; + use tikv_client_proto::{kvrpcpb, metapb, tikvpb::TikvClient}; use tikv_client_store::Request; -use super::RawRpcRequest; use crate::{ collect_first, pd::PdClient, request::{ - plan::ResponseWithShard, Collect, CollectSingle, DefaultProcessor, KvRequest, Merge, - Process, Shardable, SingleKey, + codec::{RequestCodec, RequestCodecExt}, + plan::ResponseWithShard, + Collect, CollectSingle, DefaultProcessor, KvRequest, Merge, Process, Shardable, SingleKey, }, store::{store_stream_for_keys, store_stream_for_ranges, RegionStore}, transaction::HasLocks, @@ -22,6 +23,8 @@ use crate::{ ColumnFamily, Key, KvPair, Result, Value, }; +use super::RawRpcRequest; + pub fn new_raw_get_request(key: Vec, cf: Option) -> kvrpcpb::RawGetRequest { let mut req = kvrpcpb::RawGetRequest::default(); req.set_key(key); @@ -30,10 +33,7 @@ pub fn new_raw_get_request(key: Vec, cf: Option) -> kvrpcpb::R req } -impl KvRequest for kvrpcpb::RawGetRequest { - type Response = kvrpcpb::RawGetResponse; -} - +impl_kv_request!(kvrpcpb::RawGetRequest, key; kvrpcpb::RawGetResponse); shardable_key!(kvrpcpb::RawGetRequest); collect_first!(kvrpcpb::RawGetResponse); @@ -67,10 +67,7 @@ pub fn new_raw_batch_get_request( req } -impl KvRequest for kvrpcpb::RawBatchGetRequest { - type Response = kvrpcpb::RawBatchGetResponse; -} - +impl_kv_request!(kvrpcpb::RawBatchGetRequest, keys; kvrpcpb::RawBatchGetResponse, pairs); shardable_keys!(kvrpcpb::RawBatchGetRequest); impl Merge for Collect { @@ -99,12 +96,10 @@ pub fn new_raw_put_request( req } -impl KvRequest for kvrpcpb::RawPutRequest { - type Response = kvrpcpb::RawPutResponse; -} - +impl_kv_request!(kvrpcpb::RawPutRequest, key; kvrpcpb::RawPutResponse); shardable_key!(kvrpcpb::RawPutRequest); collect_first!(kvrpcpb::RawPutResponse); + impl SingleKey for kvrpcpb::RawPutRequest { fn key(&self) -> &Vec { &self.key @@ -124,9 +119,7 @@ pub fn new_raw_batch_put_request( req } -impl KvRequest for kvrpcpb::RawBatchPutRequest { - type Response = kvrpcpb::RawBatchPutResponse; -} +impl_kv_request!(kvrpcpb::RawBatchPutRequest, pairs; kvrpcpb::RawBatchPutResponse); impl Shardable for kvrpcpb::RawBatchPutRequest { type Shard = Vec; @@ -163,10 +156,7 @@ pub fn new_raw_delete_request( req } -impl KvRequest for kvrpcpb::RawDeleteRequest { - type Response = kvrpcpb::RawDeleteResponse; -} - +impl_kv_request!(kvrpcpb::RawDeleteRequest, key; kvrpcpb::RawDeleteResponse); shardable_key!(kvrpcpb::RawDeleteRequest); collect_first!(kvrpcpb::RawDeleteResponse); impl SingleKey for kvrpcpb::RawDeleteRequest { @@ -186,10 +176,7 @@ pub fn new_raw_batch_delete_request( req } -impl KvRequest for kvrpcpb::RawBatchDeleteRequest { - type Response = kvrpcpb::RawBatchDeleteResponse; -} - +impl_kv_request!(kvrpcpb::RawBatchDeleteRequest, keys; kvrpcpb::RawBatchDeleteResponse); shardable_keys!(kvrpcpb::RawBatchDeleteRequest); pub fn new_raw_delete_range_request( @@ -205,9 +192,10 @@ pub fn new_raw_delete_range_request( req } -impl KvRequest for kvrpcpb::RawDeleteRangeRequest { - type Response = kvrpcpb::RawDeleteRangeResponse; -} +impl_kv_request!( + kvrpcpb::RawDeleteRangeRequest; + kvrpcpb::RawDeleteRangeResponse +); shardable_range!(kvrpcpb::RawDeleteRangeRequest); @@ -228,10 +216,8 @@ pub fn new_raw_scan_request( req } -impl KvRequest for kvrpcpb::RawScanRequest { - type Response = kvrpcpb::RawScanResponse; -} - +has_reverse!(kvrpcpb::RawScanRequest); +impl_kv_request!(kvrpcpb::RawScanRequest; kvrpcpb::RawScanResponse, kvs); shardable_range!(kvrpcpb::RawScanRequest); impl Merge for Collect { @@ -260,8 +246,14 @@ pub fn new_raw_batch_scan_request( req } -impl KvRequest for kvrpcpb::RawBatchScanRequest { +impl KvRequest for kvrpcpb::RawBatchScanRequest { type Response = kvrpcpb::RawBatchScanResponse; + fn encode_request(mut self, codec: &C) -> Self { + *self.mut_ranges() = codec.encode_ranges(self.take_ranges(), self.get_reverse()); + self + } + + impl_decode_response! {kvs} } impl Shardable for kvrpcpb::RawBatchScanRequest { @@ -309,9 +301,10 @@ pub fn new_cas_request( req } -impl KvRequest for kvrpcpb::RawCasRequest { - type Response = kvrpcpb::RawCasResponse; -} +impl_kv_request!( + kvrpcpb::RawCasRequest, key; + kvrpcpb::RawCasResponse +); shardable_key!(kvrpcpb::RawCasRequest); collect_first!(kvrpcpb::RawCasResponse); @@ -376,9 +369,13 @@ impl Request for RawCoprocessorRequest { fn set_context(&mut self, context: kvrpcpb::Context) { self.inner.set_context(context); } + + fn mut_context(&mut self) -> &mut kvrpcpb::Context { + self.inner.mut_context() + } } -impl KvRequest for RawCoprocessorRequest { +impl KvRequest for RawCoprocessorRequest { type Response = kvrpcpb::RawCoprocessorResponse; } @@ -455,29 +452,43 @@ impl_raw_rpc_request!(RawDeleteRangeRequest); impl_raw_rpc_request!(RawCasRequest); impl HasLocks for kvrpcpb::RawGetResponse {} + impl HasLocks for kvrpcpb::RawBatchGetResponse {} + impl HasLocks for kvrpcpb::RawPutResponse {} + impl HasLocks for kvrpcpb::RawBatchPutResponse {} + impl HasLocks for kvrpcpb::RawDeleteResponse {} + impl HasLocks for kvrpcpb::RawBatchDeleteResponse {} + impl HasLocks for kvrpcpb::RawScanResponse {} + impl HasLocks for kvrpcpb::RawBatchScanResponse {} + impl HasLocks for kvrpcpb::RawDeleteRangeResponse {} + impl HasLocks for kvrpcpb::RawCasResponse {} + impl HasLocks for kvrpcpb::RawCoprocessorResponse {} #[cfg(test)] mod test { - use super::*; + use std::any::Any; + + use futures::executor; + + use tikv_client_proto::kvrpcpb; + use crate::{ backoff::{DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF}, mock::{MockKvClient, MockPdClient}, request::Plan, Key, }; - use futures::executor; - use std::any::Any; - use tikv_client_proto::kvrpcpb; + + use super::*; #[test] #[ignore] diff --git a/src/region_cache.rs b/src/region_cache.rs index b278ef74..c7ab43d8 100644 --- a/src/region_cache.rs +++ b/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::codec::RequestCodec, + Key, Result, +}; const MAX_RETRY_WAITING_CONCURRENT_REQUEST: usize = 4; @@ -44,23 +48,25 @@ impl RegionCacheMap { } } -pub struct RegionCache> { +pub struct RegionCache> { region_cache: RwLock, store_cache: RwLock>, inner_client: Arc, + codec: C, } -impl RegionCache { - pub fn new(inner_client: Arc) -> RegionCache { +impl RegionCache { + pub fn new(codec: C, inner_client: Arc) -> Self { RegionCache { region_cache: RwLock::new(RegionCacheMap::new()), store_cache: RwLock::new(HashMap::new()), inner_client, + codec, } } } -impl RegionCache { +impl RegionCache { // 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 { let region_cache_guard = self.region_cache.read().await; @@ -126,9 +132,14 @@ impl RegionCache { /// Force read through (query from PD) and update cache pub async fn read_through_region_by_key(&self, key: Key) -> Result { - let region = self.inner_client.clone().get_region(key.into()).await?; - self.add_region(region.clone()).await; - Ok(region) + let mut r = self + .inner_client + .clone() + .get_region(self.codec.encode_pd_query(key.into())) + .await?; + self.codec.decode_region(&mut r.region)?; + self.add_region(r.clone()).await; + Ok(r) } /// Force read through (query from PD) and update cache @@ -140,7 +151,8 @@ impl RegionCache { 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?; + self.codec.decode_region(&mut region.region)?; self.add_region(region.clone()).await; // notify others @@ -226,17 +238,14 @@ impl RegionCache { 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::{ @@ -244,10 +253,23 @@ mod test { Arc, }, }; - use tikv_client_common::Error; - use tikv_client_proto::metapb; + + use async_trait::async_trait; use tokio::sync::Mutex; + use tikv_client_common::Error; + use tikv_client_proto::{keyspacepb, metapb}; + + use crate::{ + pd::RetryClientTrait, + raw::ApiV1, + region::{RegionId, RegionWithLeader}, + request::codec::RequestCodec, + Key, Result, + }; + + use super::RegionCache; + #[derive(Default)] struct MockRetryClient { pub regions: Mutex>, @@ -304,12 +326,19 @@ mod test { async fn update_safepoint(self: Arc, _safepoint: u64) -> Result { todo!() } + + async fn load_keyspace( + self: Arc, + _name: &str, + ) -> Result { + todo!() + } } #[tokio::test] async fn cache_is_used() -> Result<()> { let retry_client = Arc::new(MockRetryClient::default()); - let cache = RegionCache::new(retry_client.clone()); + let cache = RegionCache::new(ApiV1::default(), retry_client.clone()); retry_client.regions.lock().await.insert( 1, RegionWithLeader { @@ -379,7 +408,7 @@ mod test { #[tokio::test] async fn test_add_disjoint_regions() { let retry_client = Arc::new(MockRetryClient::default()); - let cache = RegionCache::new(retry_client.clone()); + let cache = RegionCache::new(ApiV1::default(), 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![]); @@ -398,7 +427,7 @@ mod test { #[tokio::test] async fn test_add_intersecting_regions() { let retry_client = Arc::new(MockRetryClient::default()); - let cache = RegionCache::new(retry_client.clone()); + let cache = RegionCache::new(ApiV1::default(), retry_client.clone()); cache.add_region(region(1, vec![], vec![10])).await; cache.add_region(region(2, vec![10], vec![20])).await; @@ -436,7 +465,7 @@ mod test { #[tokio::test] async fn test_get_region_by_key() -> Result<()> { let retry_client = Arc::new(MockRetryClient::default()); - let cache = RegionCache::new(retry_client.clone()); + let cache = RegionCache::new(ApiV1::default(), retry_client.clone()); let region1 = region(1, vec![], vec![10]); let region2 = region(2, vec![10], vec![20]); @@ -466,8 +495,8 @@ mod test { } // a helper function to assert the cache is in expected state - async fn assert( - cache: &RegionCache, + async fn assert( + cache: &RegionCache, expected_cache: &BTreeMap, ) { let guard = cache.region_cache.read().await; diff --git a/src/request/codec.rs b/src/request/codec.rs new file mode 100644 index 00000000..8a6b3836 --- /dev/null +++ b/src/request/codec.rs @@ -0,0 +1,452 @@ +use core::intrinsics::copy; +use std::{ + borrow::Cow, + marker::PhantomData, + ops::{Deref, DerefMut}, + sync::Arc, +}; + +use tikv_client_common::Error; +use tikv_client_proto::{errorpb, keyspacepb, kvrpcpb, metapb::Region}; + +use crate::{kv::codec::decode_bytes_in_place, request::KvRequest, Key, Result}; + +use crate::pd::{RetryClient, RetryClientTrait}; +use derive_new::new; + +type Prefix = [u8; KEYSPACE_PREFIX_LEN]; + +const KEYSPACE_PREFIX_LEN: usize = 4; + +pub const DEFAULT_KEYSPACE: &str = "DEFAULT"; + +pub trait RequestCodec: Default + Sized + Clone + Sync + Send + 'static { + fn encode_request<'a, R: KvRequest>(&self, req: &'a R) -> Cow<'a, R> { + Cow::Borrowed(req) + } + + fn encode_key(&self, key: Vec) -> Vec { + key + } + + fn decode_key(&self, _key: &mut Vec) -> Result<()> { + Ok(()) + } + + fn encode_range(&self, start: Vec, end: Vec, _reverse: bool) -> (Vec, Vec) { + (start, end) + } + + fn encode_pd_query(&self, key: Vec) -> Vec { + key + } + + fn decode_region(&self, _region: &mut Region) -> Result<()> { + Ok(()) + } + + fn decode_response>( + &self, + _req: &R, + resp: R::Response, + ) -> Result { + Ok(resp) + } +} + +pub trait RequestCodecExt: RequestCodec { + fn encode_primary_lock(&self, lock: Vec) -> Vec { + self.encode_key(lock) + } + + fn encode_primary_key(&self, key: Vec) -> Vec { + self.encode_key(key) + } + + fn encode_mutations(&self, mutations: Vec) -> Vec { + mutations + .into_iter() + .map(|mut m| { + let key = m.take_key(); + m.set_key(self.encode_key(key)); + m + }) + .collect() + } + + fn encode_keys(&self, keys: Vec>) -> Vec> { + keys.into_iter().map(|key| self.encode_key(key)).collect() + } + + fn encode_secondaries(&self, secondaries: Vec>) -> Vec> { + self.encode_keys(secondaries) + } + + fn encode_pairs(&self, mut pairs: Vec) -> Vec { + for pair in pairs.iter_mut() { + *pair.mut_key() = self.encode_key(pair.take_key()); + } + + pairs + } + + fn encode_ranges( + &self, + mut ranges: Vec, + reverse: bool, + ) -> Vec { + for range in ranges.iter_mut() { + let (start, end) = + self.encode_range(range.take_start_key(), range.take_end_key(), reverse); + *range.mut_start_key() = start; + *range.mut_end_key() = end; + } + + ranges + } + + fn decode_keys(&self, keys: &mut [Vec]) -> Result<()> { + for key in keys.iter_mut() { + self.decode_key(key)?; + } + + Ok(()) + } + + fn decode_error(&self, err: &mut kvrpcpb::KeyError) -> Result<()> { + if err.has_locked() { + let locked = err.mut_locked(); + self.decode_lock_info(locked)?; + } + + if err.has_conflict() { + let conflict = err.mut_conflict(); + self.decode_key(conflict.mut_key())?; + self.decode_key(conflict.mut_primary())?; + } + + if err.has_already_exist() { + let already_exist = err.mut_already_exist(); + self.decode_key(already_exist.mut_key())?; + } + + // We do not decode key in `Deadlock` since there is no use for the key right now in client side. + // All we need is the key hash to detect deadlock. + // TODO: while we check the keys against the deadlock key hash, we need to encode the key. + + if err.has_commit_ts_expired() { + let commit_ts_expired = err.mut_commit_ts_expired(); + self.decode_key(commit_ts_expired.mut_key())?; + } + + if err.has_txn_not_found() { + let txn_not_found = err.mut_txn_not_found(); + self.decode_key(txn_not_found.mut_primary_key())?; + } + + if err.has_assertion_failed() { + let assertion_failed = err.mut_assertion_failed(); + self.decode_key(assertion_failed.mut_key())?; + } + + Ok(()) + } + + fn decode_errors(&self, errors: &mut [kvrpcpb::KeyError]) -> Result<()> { + for err in errors.iter_mut() { + self.decode_error(err)?; + } + + Ok(()) + } + + fn decode_lock_info(&self, lock: &mut kvrpcpb::LockInfo) -> Result<()> { + self.decode_key(lock.mut_primary_lock())?; + self.decode_key(lock.mut_key())?; + self.decode_keys(lock.mut_secondaries()) + } + + fn decode_locks(&self, locks: &mut [kvrpcpb::LockInfo]) -> Result<()> { + for lock in locks.iter_mut() { + self.decode_lock_info(lock)?; + } + + Ok(()) + } + + fn decode_pairs(&self, pairs: &mut [kvrpcpb::KvPair]) -> Result<()> { + for pair in pairs.iter_mut() { + self.decode_key(pair.mut_key())?; + } + + Ok(()) + } + + fn decode_kvs(&self, kvs: &mut [kvrpcpb::KvPair]) -> Result<()> { + self.decode_pairs(kvs) + } + + fn decode_regions(&self, regions: &mut [Region]) -> Result<()> { + for region in regions.iter_mut() { + self.decode_region(region)?; + } + Ok(()) + } + + fn decode_region_error(&self, err: &mut errorpb::Error) -> Result<()> { + if err.has_epoch_not_match() { + self.decode_regions(err.mut_epoch_not_match().mut_current_regions())?; + } + Ok(()) + } +} + +impl RequestCodecExt for T {} + +pub trait RawCodec: RequestCodec {} + +pub trait TxnCodec: RequestCodec {} + +#[derive(new, Clone, Copy, Default, PartialEq, Eq)] +pub struct KeySpaceId([u8; 3]); + +impl Deref for KeySpaceId { + type Target = [u8; 3]; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for KeySpaceId { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl KeySpaceId { + pub async fn from_name(name: &str, pd: Arc) -> Result { + if name == DEFAULT_KEYSPACE { + return Ok(KeySpaceId::default()); + } + + let resp = pd.load_keyspace(name).await?; + let keyspace = resp.get_keyspace(); + if keyspace.get_state() == keyspacepb::KeyspaceState::Enabled { + let id = keyspace.get_id().to_be_bytes(); + Ok(KeySpaceId::new([id[1], id[2], id[3]])) + } else { + Err(Error::KeyspaceNotEnabled { + name: name.to_string(), + }) + } + } +} + +pub trait Mode: Clone + Copy + Sync + Send + 'static { + const PREFIX: u8; + const MIN_KEY: &'static [u8] = &[Self::PREFIX, 0, 0, 0]; + const MAX_KEY: &'static [u8] = &[Self::PREFIX + 1, 0, 0, 0]; +} + +#[derive(Default, Clone, Copy)] +pub struct RawMode; + +#[derive(Default, Clone, Copy)] +pub struct TxnMode; + +impl Mode for RawMode { + const PREFIX: u8 = b'r'; +} + +impl Mode for TxnMode { + const PREFIX: u8 = b'x'; +} + +#[derive(Clone)] +pub struct ApiV1 { + _phantom: PhantomData, +} + +impl Default for ApiV1 { + fn default() -> Self { + ApiV1 { + _phantom: PhantomData, + } + } +} + +impl RequestCodec for ApiV1 {} + +impl RawCodec for ApiV1 {} + +impl RequestCodec for ApiV1 { + fn encode_pd_query(&self, key: Vec) -> Vec { + 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 TxnCodec for ApiV1 {} + +#[derive(Clone, Copy)] +pub struct KeySpace { + id: KeySpaceId, + _phantom: PhantomData, +} + +impl KeySpace { + pub fn new(id: KeySpaceId) -> Self { + KeySpace { + id, + _phantom: PhantomData, + } + } +} + +impl Default for KeySpace { + fn default() -> Self { + KeySpace { + id: KeySpaceId::default(), + _phantom: PhantomData, + } + } +} + +impl From> for Prefix { + fn from(s: KeySpace) -> Self { + [M::PREFIX, s.id[0], s.id[1], s.id[2]] + } +} + +impl KeySpace { + fn start(self) -> Prefix { + self.into() + } + + fn end(self) -> Prefix { + (u32::from_be_bytes(self.into()) + 1).to_be_bytes() + } +} + +#[derive(Clone)] +pub struct ApiV2 { + keyspace: KeySpace, +} + +impl ApiV2 { + pub async fn with_keyspace(name: &str, pd: Arc) -> Result { + let id = KeySpaceId::from_name(name, pd).await?; + Ok(ApiV2 { + keyspace: KeySpace::new(id), + }) + } +} + +impl Default for ApiV2 { + fn default() -> Self { + ApiV2 { + keyspace: KeySpace::default(), + } + } +} + +impl RequestCodec for ApiV2 { + fn encode_request<'a, R: KvRequest>(&self, req: &'a R) -> Cow<'a, R> { + let mut req = req.clone(); + req.mut_context().set_api_version(kvrpcpb::ApiVersion::V2); + Cow::Owned(req.encode_request(self)) + } + + fn encode_key(&self, mut key: Vec) -> Vec { + let mut encoded = Vec::with_capacity(key.len() + KEYSPACE_PREFIX_LEN); + let prefix: Prefix = self.keyspace.into(); + + encoded.extend_from_slice(&prefix); + encoded.append(&mut key); + encoded + } + + fn decode_key(&self, key: &mut Vec) -> Result<()> { + let prefix: Prefix = self.keyspace.into(); + + if !key.starts_with(&prefix) { + return Err(Error::CorruptedKeyspace { + expected: prefix.to_vec(), + key: key.to_vec(), + }); + } + + unsafe { + let trimmed_len = key.len() - KEYSPACE_PREFIX_LEN; + let ptr = key.as_mut_ptr(); + let trimmed = key[KEYSPACE_PREFIX_LEN..].as_mut_ptr(); + + copy(trimmed, ptr, trimmed_len); + + key.set_len(trimmed_len); + } + Ok(()) + } + + fn encode_range(&self, start: Vec, end: Vec, reverse: bool) -> (Vec, Vec) { + if reverse { + let (start, end) = self.encode_range(end, start, false); + return (end, start); + } + + let start = self.encode_key(start); + + let end = if end.is_empty() { + self.keyspace.end().into() + } else { + self.encode_key(end) + }; + + (start, end) + } + + fn encode_pd_query(&self, key: Vec) -> Vec { + Key::from(self.encode_key(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)?; + + // Map the region's start key to the keyspace start key. + if region.get_start_key() <= self.keyspace.start().as_slice() { + *region.mut_start_key() = vec![]; + } else { + self.decode_key(region.mut_start_key())?; + } + + // Map the region's end key to the keyspace end key. + if region.get_end_key().is_empty() || region.get_end_key() >= self.keyspace.end().as_slice() + { + *region.mut_end_key() = vec![]; + } else { + self.decode_key(region.mut_end_key())?; + } + + Ok(()) + } + + fn decode_response>( + &self, + req: &R, + resp: R::Response, + ) -> Result { + req.decode_response(self, resp) + } +} + +impl RawCodec for ApiV2 {} + +impl TxnCodec for ApiV2 {} diff --git a/src/request/mod.rs b/src/request/mod.rs index 959ff5e7..004e03cb 100644 --- a/src/request/mod.rs +++ b/src/request/mod.rs @@ -1,12 +1,14 @@ // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. +use async_trait::async_trait; +use derive_new::new; + +use tikv_client_store::{HasKeyErrors, HasRegionError, Request}; + use crate::{ backoff::{Backoff, DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF, PESSIMISTIC_BACKOFF}, transaction::HasLocks, }; -use async_trait::async_trait; -use derive_new::new; -use tikv_client_store::{HasKeyErrors, Request}; pub use self::{ plan::{ @@ -22,12 +24,120 @@ pub mod plan; mod plan_builder; #[macro_use] mod shard; +pub mod codec; /// Abstracts any request sent to a TiKV server. #[async_trait] -pub trait KvRequest: Request + Sized + Clone + Sync + Send + 'static { +pub trait KvRequest: Request + Sized + Clone + Sync + Send + 'static { /// The expected response to the request. - type Response: HasKeyErrors + HasLocks + Clone + Send + 'static; + type Response: HasKeyErrors + HasLocks + HasRegionError + Clone + Send + 'static; + + fn encode_request(self, _codec: &C) -> Self { + self + } + + fn decode_response(&self, _codec: &C, resp: Self::Response) -> crate::Result { + Ok(resp) + } +} + +pub(crate) trait IsDefault { + fn is_default(&self) -> bool; +} + +impl IsDefault for T { + fn is_default(&self) -> bool { + *self == T::default() + } +} + +pub(crate) trait HasReverse { + fn has_reverse(&self) -> bool; +} + +impl HasReverse for T { + default fn has_reverse(&self) -> bool { + false + } +} + +macro_rules! has_reverse { + ($t:ty) => { + impl $crate::request::HasReverse for $t { + fn has_reverse(&self) -> bool { + self.get_reverse() + } + } + }; +} + +macro_rules! impl_decode_response { + ($($o:ident)*) => { + fn decode_response(&self, codec: &C, mut resp: Self::Response) -> Result { + #[allow(unused_imports)] + use $crate::request::{IsDefault, codec::RequestCodecExt}; + + // decode errors + if resp.has_region_error() { + codec.decode_region_error(resp.mut_region_error())?; + } + + $( + paste::paste! { + if !resp.[]().is_default() { + codec.[](resp.[]())?; + } + } + )* + + Ok(resp) + } + }; +} + +macro_rules! impl_kv_request { + ($req:ty $(,$i:ident)+; $resp:ty $(,$o:ident)*) => { + impl KvRequest for $req + where C: RequestCodec + { + type Response = $resp; + + fn encode_request(mut self, codec: &C) -> Self { + #[allow(unused_imports)] + use $crate::request::codec::RequestCodecExt; + + $( + paste::paste! { + *self.[]() = codec.[](self.[]()); + } + )* + + self + } + + impl_decode_response!{$($o)*} + } + }; + + ($req:ty; $resp:ty $(,$o:ident)*) => { + impl KvRequest for $req + where C: RequestCodec, + $req: $crate::request::HasReverse + { + type Response = $resp; + + fn encode_request(mut self, codec: &C) -> Self { + use $crate::request::HasReverse; + let (start, end) = codec.encode_range(self.take_start_key(), self.take_end_key(), self.has_reverse()); + *self.mut_start_key() = start; + *self.mut_end_key() = end; + + self + } + + impl_decode_response!{$($o)*} + } + }; } #[derive(Clone, Debug, new, Eq, PartialEq)] @@ -63,22 +173,27 @@ impl RetryOptions { #[cfg(test)] mod test { - use super::*; - use crate::{ - mock::{MockKvClient, MockPdClient}, - store::store_stream_for_keys, - transaction::lowering::new_commit_request, - Error, Key, Result, - }; - use grpcio::CallOption; use std::{ any::Any, iter, sync::{atomic::AtomicUsize, Arc}, }; + + use grpcio::CallOption; + use tikv_client_proto::{kvrpcpb, pdpb::Timestamp, tikvpb::TikvClient}; use tikv_client_store::HasRegionError; + use crate::{ + mock::{MockKvClient, MockPdClient}, + request::{codec::RequestCodec, KvRequest}, + store::store_stream_for_keys, + transaction::lowering::new_commit_request, + Error, Key, Result, + }; + + use super::*; + #[tokio::test] async fn test_region_retry() { #[derive(Clone)] @@ -120,11 +235,23 @@ mod test { fn set_context(&mut self, _: kvrpcpb::Context) { unreachable!(); } + + fn mut_context(&mut self) -> &mut kvrpcpb::Context { + unreachable!() + } } #[async_trait] - impl KvRequest for MockKvRequest { + impl KvRequest for MockKvRequest { type Response = MockRpcResponse; + + fn encode_request(self, _codec: &C) -> Self { + self + } + + fn decode_response(&self, _codec: &C, resp: Self::Response) -> Result { + Ok(resp) + } } impl Shardable for MockKvRequest { @@ -184,6 +311,7 @@ mod test { region_error: None, error: Some(kvrpcpb::KeyError::default()), commit_version: 0, + ..Default::default() }) as Box) }, ))); diff --git a/src/request/plan.rs b/src/request/plan.rs index ce785262..5b51fcae 100644 --- a/src/request/plan.rs +++ b/src/request/plan.rs @@ -5,14 +5,15 @@ use std::{marker::PhantomData, sync::Arc}; use async_recursion::async_recursion; use async_trait::async_trait; use futures::{future::try_join_all, prelude::*}; +use tokio::sync::Semaphore; + use tikv_client_proto::{errorpb, errorpb::EpochNotMatch, kvrpcpb}; use tikv_client_store::{HasKeyErrors, HasRegionError, HasRegionErrors, KvClient}; -use tokio::sync::Semaphore; use crate::{ backoff::Backoff, pd::PdClient, - request::{KvRequest, Shardable}, + request::{codec::RequestCodec, KvRequest, Shardable}, stats::tikv_stats, store::RegionStore, transaction::{resolve_locks, HasLocks}, @@ -33,27 +34,44 @@ pub trait Plan: Sized + Clone + Sync + Send + 'static { /// The simplest plan which just dispatches a request to a specific kv server. #[derive(Clone)] -pub struct Dispatch { +pub struct Dispatch> { pub request: Req, pub kv_client: Option>, + codec: C, +} + +impl> Dispatch { + pub fn new(request: Req, kv_client: Option>, codec: C) -> Self { + Self { + request, + kv_client, + codec, + } + } } #[async_trait] -impl Plan for Dispatch { +impl> Plan for Dispatch { type Result = Req::Response; async fn execute(&self) -> Result { + let req = self.codec.encode_request(&self.request); + let stats = tikv_stats(self.request.label()); let result = self .kv_client .as_ref() .expect("Unreachable: kv_client has not been initialised in Dispatch") - .dispatch(&self.request) + .dispatch(req.as_ref()) .await; let result = stats.done(result); - result.map(|r| { - *r.downcast() - .expect("Downcast failed: request and response type mismatch") + + result.and_then(|r| { + let resp = *r + .downcast() + .expect("Downcast failed: request and response type mismatch"); + + self.codec.decode_response(req.as_ref(), resp) }) } } @@ -544,11 +562,14 @@ impl HasRegionError for ResponseWithShard { @@ -25,19 +29,26 @@ pub struct PlanBuilder { /// Used to ensure that a plan has a designated target or targets, a target is /// a particular TiKV server. pub trait PlanBuilderPhase {} + pub struct NoTarget; + impl PlanBuilderPhase for NoTarget {} + pub struct Targetted; + impl PlanBuilderPhase for Targetted {} -impl PlanBuilder, NoTarget> { +impl PlanBuilder, NoTarget> +where + C: RequestCodec, + Req: KvRequest, + PdC: PdClient, +{ pub fn new(pd_client: Arc, request: Req) -> Self { + let codec = pd_client.get_request_codec(); PlanBuilder { pd_client, - plan: Dispatch { - request, - kv_client: None, - }, + plan: Dispatch::new(request, None, codec), phantom: PhantomData, } } @@ -144,11 +155,13 @@ where } } -impl PlanBuilder, NoTarget> { +impl + SingleKey> + PlanBuilder, NoTarget> +{ /// Target the request at a single region. *Note*: single region plan will /// cannot automatically retry on region errors. It's only used for requests /// that target at a specific region but not keys (e.g. ResolveLockRequest). - pub async fn single_region(self) -> Result, Targetted>> { + pub async fn single_region(self) -> Result, Targetted>> { let key = self.plan.request.key(); // TODO: retry when region error occurred let store = self.pd_client.clone().store_for_key(key.into()).await?; @@ -156,12 +169,12 @@ impl PlanBuilder, NoTa } } -impl PlanBuilder, NoTarget> { +impl> PlanBuilder, NoTarget> { /// Target the request at a single region; caller supplies the store to target. pub async fn single_region_with_store( self, store: RegionStore, - ) -> Result, Targetted>> { + ) -> Result, Targetted>> { set_single_region_store(self.plan, store, self.pd_client) } } @@ -195,11 +208,11 @@ where } } -fn set_single_region_store( - mut plan: Dispatch, +fn set_single_region_store>( + mut plan: Dispatch, store: RegionStore, pd_client: Arc, -) -> Result, Targetted>> { +) -> Result, Targetted>> { plan.request .set_context(store.region_with_leader.context()?); plan.kv_client = Some(store.client); diff --git a/src/request/shard.rs b/src/request/shard.rs index cdf3f198..bbe48de2 100644 --- a/src/request/shard.rs +++ b/src/request/shard.rs @@ -38,7 +38,7 @@ pub trait Shardable { fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()>; } -impl Shardable for Dispatch { +impl + Shardable> Shardable for Dispatch { type Shard = Req::Shard; fn shards( diff --git a/src/transaction/client.rs b/src/transaction/client.rs index 69f11bce..750d859d 100644 --- a/src/transaction/client.rs +++ b/src/transaction/client.rs @@ -1,18 +1,23 @@ // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. -use super::{requests::new_scan_lock_request, resolve_locks}; +use std::{future::Future, mem, sync::Arc}; + +use slog::Logger; + +use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; + use crate::{ backoff::{DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF}, config::Config, - pd::{PdClient, PdRpcClient}, - request::Plan, + pd::{PdClient, PdRpcClient, RetryClient}, + request::{codec::TxnCodec, Plan}, timestamp::TimestampExt, - transaction::{Snapshot, Transaction, TransactionOptions}, + transaction::{ApiV2, Snapshot, Transaction, TransactionOptions}, + util::client::ClientContext, Result, }; -use slog::{Drain, Logger}; -use std::{mem, sync::Arc}; -use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; + +use super::{requests::new_scan_lock_request, resolve_locks}; // FIXME: cargo-culted value const SCAN_LOCK_BATCH_SIZE: u32 = 1024; @@ -33,12 +38,12 @@ const SCAN_LOCK_BATCH_SIZE: u32 = 1024; /// /// The returned results of transactional requests are [`Future`](std::future::Future)s that must be /// awaited to execute. -pub struct Client { - pd: Arc, +pub struct Client { + pd: Arc>, logger: Logger, } -impl Clone for Client { +impl Clone for Client { fn clone(&self) -> Self { Self { pd: self.pd.clone(), @@ -47,7 +52,59 @@ impl Clone for Client { } } -impl Client { +impl Client { + /// Create a transaction [`Client`] and connect to the TiKV cluster, with `keyspace` name. + /// + /// Because TiKV is managed by a [PD](https://github.com/pingcap/pd/) cluster, the endpoints for + /// PD must be provided, not the TiKV nodes. It's important to include more than one PD endpoint + /// (include all endpoints, if possible), this helps avoid having a single point of failure. + /// + /// # Examples + /// + /// ```rust,no_run + /// # use tikv_client::{Config, TransactionClient}; + /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV2; + /// # futures::executor::block_on(async { + /// let client = TransactionClient::::new_with_keyspace( + /// vec!["192.168.0.100"], + /// "space_for_txn", + /// Config::default(), + /// None, + /// ) + /// .await + /// .unwrap(); + /// # }); + /// ``` + pub async fn new_with_keyspace( + pd_endpoints: Vec, + keyspace: K, + config: Config, + optional_logger: Option, + ) -> Result + where + S: Into, + K: Into, + { + let keyspace = &keyspace.into(); + let this = Self::new_with_codec_factory( + pd_endpoints, + config, + async move |pd| ApiV2::with_keyspace(keyspace, pd).await, + optional_logger, + ) + .await?; + + info!(this.logger, "created transaction client"; "keyspace" => keyspace); + + Ok(this) + } +} + +impl Client +where + C: TxnCodec, +{ /// Create a transactional [`Client`] and connect to the TiKV cluster. /// /// Because TiKV is managed by a [PD](https://github.com/pingcap/pd/) cluster, the endpoints for @@ -59,8 +116,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// let client = TransactionClient::new(vec!["192.168.0.100"], None) + /// let client = TransactionClient::::new(vec!["192.168.0.100"], None) /// .await /// .unwrap(); /// # }); @@ -68,7 +126,7 @@ impl Client { pub async fn new>( pd_endpoints: Vec, logger: Option, - ) -> Result { + ) -> Result> { // debug!(self.logger, "creating transactional client"); Self::new_with_config(pd_endpoints, Config::default(), logger).await } @@ -85,8 +143,9 @@ impl Client { /// # use tikv_client::{Config, TransactionClient}; /// # use futures::prelude::*; /// # use std::time::Duration; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// let client = TransactionClient::new_with_config( + /// let client = TransactionClient::::new_with_config( /// vec!["192.168.0.100"], /// Config::default().with_timeout(Duration::from_secs(60)), /// None, @@ -99,20 +158,30 @@ impl Client { pd_endpoints: Vec, config: Config, optional_logger: Option, - ) -> Result { - let logger = optional_logger.unwrap_or_else(|| { - let plain = slog_term::PlainSyncDecorator::new(std::io::stdout()); - Logger::root( - slog_term::FullFormat::new(plain) - .build() - .filter_level(slog::Level::Info) - .fuse(), - o!(), - ) - }); - debug!(logger, "creating new transactional client"); - let pd_endpoints: Vec = pd_endpoints.into_iter().map(Into::into).collect(); - let pd = Arc::new(PdRpcClient::connect(&pd_endpoints, config, true, logger.clone()).await?); + ) -> Result> { + Self::new_with_codec_factory( + pd_endpoints, + config, + async move |_| Ok(C::default()), + optional_logger, + ) + .await + } + + pub async fn new_with_codec_factory( + pd_endpoints: Vec, + config: Config, + codec_factory: F, + optional_logger: Option, + ) -> Result> + where + S: Into, + F: Fn(Arc) -> Fut, + Fut: Future>, + { + let ClientContext { pd, logger } = + ClientContext::new(pd_endpoints, config, codec_factory, optional_logger).await?; + Ok(Client { pd, logger }) } @@ -129,8 +198,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// let client = TransactionClient::new(vec!["192.168.0.100"], None) + /// let client = TransactionClient::::new(vec!["192.168.0.100"], None) /// .await /// .unwrap(); /// let mut transaction = client.begin_optimistic().await.unwrap(); @@ -138,7 +208,7 @@ impl Client { /// transaction.commit().await.unwrap(); /// # }); /// ``` - pub async fn begin_optimistic(&self) -> Result { + pub async fn begin_optimistic(&self) -> Result> { debug!(self.logger, "creating new optimistic transaction"); let timestamp = self.current_timestamp().await?; Ok(self.new_transaction(timestamp, TransactionOptions::new_optimistic())) @@ -154,8 +224,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// let client = TransactionClient::new(vec!["192.168.0.100"], None) + /// let client = TransactionClient::::new(vec!["192.168.0.100"], None) /// .await /// .unwrap(); /// let mut transaction = client.begin_pessimistic().await.unwrap(); @@ -163,7 +234,7 @@ impl Client { /// transaction.commit().await.unwrap(); /// # }); /// ``` - pub async fn begin_pessimistic(&self) -> Result { + pub async fn begin_pessimistic(&self) -> Result> { debug!(self.logger, "creating new pessimistic transaction"); let timestamp = self.current_timestamp().await?; Ok(self.new_transaction(timestamp, TransactionOptions::new_pessimistic())) @@ -176,8 +247,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Config, TransactionClient, TransactionOptions}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// let client = TransactionClient::new(vec!["192.168.0.100"], None) + /// let client = TransactionClient::::new(vec!["192.168.0.100"], None) /// .await /// .unwrap(); /// let mut transaction = client @@ -188,14 +260,14 @@ impl Client { /// transaction.commit().await.unwrap(); /// # }); /// ``` - pub async fn begin_with_options(&self, options: TransactionOptions) -> Result { + pub async fn begin_with_options(&self, options: TransactionOptions) -> Result> { debug!(self.logger, "creating new customized transaction"); let timestamp = self.current_timestamp().await?; Ok(self.new_transaction(timestamp, options)) } /// Create a new [`Snapshot`](Snapshot) at the given [`Timestamp`](Timestamp). - pub fn snapshot(&self, timestamp: Timestamp, options: TransactionOptions) -> Snapshot { + pub fn snapshot(&self, timestamp: Timestamp, options: TransactionOptions) -> Snapshot { debug!(self.logger, "creating new snapshot"); let logger = self.logger.new(o!("child" => 1)); Snapshot::new(self.new_transaction(timestamp, options.read_only()), logger) @@ -208,8 +280,9 @@ impl Client { /// ```rust,no_run /// # use tikv_client::{Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// let client = TransactionClient::new(vec!["192.168.0.100"], None) + /// let client = TransactionClient::::new(vec!["192.168.0.100"], None) /// .await /// .unwrap(); /// let timestamp = client.current_timestamp().await.unwrap(); @@ -276,7 +349,7 @@ impl Client { Ok(res) } - fn new_transaction(&self, timestamp: Timestamp, options: TransactionOptions) -> Transaction { + fn new_transaction(&self, timestamp: Timestamp, options: TransactionOptions) -> Transaction { let logger = self.logger.new(o!("child" => 1)); Transaction::new(timestamp, self.pd.clone(), options, logger) } diff --git a/src/transaction/codec.rs b/src/transaction/codec.rs new file mode 100644 index 00000000..ad34acf1 --- /dev/null +++ b/src/transaction/codec.rs @@ -0,0 +1,4 @@ +use crate::request::codec::{self, TxnMode}; + +pub type ApiV1 = codec::ApiV1; +pub type ApiV2 = codec::ApiV2; diff --git a/src/transaction/lock.rs b/src/transaction/lock.rs index 871effc6..e70836d6 100644 --- a/src/transaction/lock.rs +++ b/src/transaction/lock.rs @@ -140,11 +140,14 @@ pub trait HasLocks { #[cfg(test)] mod tests { - use super::*; - use crate::mock::{MockKvClient, MockPdClient}; use std::any::Any; + use tikv_client_proto::errorpb; + use crate::mock::{MockKvClient, MockPdClient}; + + use super::*; + #[tokio::test] async fn test_resolve_lock_with_retry() { // Test resolve lock within retry limit diff --git a/src/transaction/mod.rs b/src/transaction/mod.rs index 2fee1aee..b8e46c63 100644 --- a/src/transaction/mod.rs +++ b/src/transaction/mod.rs @@ -9,6 +9,7 @@ //! **Warning:** It is not advisable to use both raw and transactional functionality in the same keyspace. pub use client::Client; +pub use codec::{ApiV1, ApiV2}; pub(crate) use lock::{resolve_locks, HasLocks}; pub use snapshot::Snapshot; #[doc(hidden)] @@ -20,6 +21,7 @@ mod client; pub mod lowering; #[macro_use] mod requests; +mod codec; mod lock; mod snapshot; #[allow(clippy::module_inception)] diff --git a/src/transaction/requests.rs b/src/transaction/requests.rs index 16243d83..673c7840 100644 --- a/src/transaction/requests.rs +++ b/src/transaction/requests.rs @@ -1,11 +1,21 @@ // Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0. +use std::{collections::HashMap, iter, sync::Arc}; + +use either::Either; +use futures::stream::BoxStream; + +use tikv_client_common::Error::PessimisticLockError; +use tikv_client_proto::{ + kvrpcpb::{self, TxnHeartBeatResponse}, + pdpb::Timestamp, +}; use crate::{ collect_first, pd::PdClient, request::{ - Collect, CollectSingle, CollectWithShard, DefaultProcessor, KvRequest, Merge, Process, - ResponseWithShard, Shardable, SingleKey, + codec::RequestCodec, Collect, CollectSingle, CollectWithShard, DefaultProcessor, KvRequest, + Merge, Process, ResponseWithShard, Shardable, SingleKey, }, store::{store_stream_for_keys, store_stream_for_range_by_start_key, RegionStore}, timestamp::TimestampExt, @@ -13,14 +23,6 @@ use crate::{ util::iter::FlatMapOkIterExt, Key, KvPair, Result, Value, }; -use either::Either; -use futures::stream::BoxStream; -use std::{collections::HashMap, iter, sync::Arc}; -use tikv_client_common::Error::PessimisticLockError; -use tikv_client_proto::{ - kvrpcpb::{self, TxnHeartBeatResponse}, - pdpb::Timestamp, -}; // implement HasLocks for a response type that has a `pairs` field, // where locks can be extracted from both the `pairs` and `error` fields @@ -70,10 +72,7 @@ pub fn new_get_request(key: Vec, timestamp: u64) -> kvrpcpb::GetRequest { req } -impl KvRequest for kvrpcpb::GetRequest { - type Response = kvrpcpb::GetResponse; -} - +impl_kv_request!(kvrpcpb::GetRequest, key; kvrpcpb::GetResponse, error); shardable_key!(kvrpcpb::GetRequest); collect_first!(kvrpcpb::GetResponse); impl SingleKey for kvrpcpb::GetRequest { @@ -102,10 +101,10 @@ pub fn new_batch_get_request(keys: Vec>, timestamp: u64) -> kvrpcpb::Bat req } -impl KvRequest for kvrpcpb::BatchGetRequest { - type Response = kvrpcpb::BatchGetResponse; -} - +impl_kv_request!( + kvrpcpb::BatchGetRequest, keys; + kvrpcpb::BatchGetResponse, pairs, error +); shardable_keys!(kvrpcpb::BatchGetRequest); impl Merge for Collect { @@ -142,10 +141,8 @@ pub fn new_scan_request( req } -impl KvRequest for kvrpcpb::ScanRequest { - type Response = kvrpcpb::ScanResponse; -} - +has_reverse!(kvrpcpb::ScanRequest); +impl_kv_request!(kvrpcpb::ScanRequest; kvrpcpb::ScanResponse, pairs, error); shardable_range!(kvrpcpb::ScanRequest); impl Merge for Collect { @@ -174,9 +171,8 @@ pub fn new_resolve_lock_request( // region without keys. So it's not Shardable. And we don't automatically retry // on its region errors (in the Plan level). The region error must be manually // handled (in the upper level). -impl KvRequest for kvrpcpb::ResolveLockRequest { - type Response = kvrpcpb::ResolveLockResponse; -} + +impl_kv_request!(kvrpcpb::ResolveLockRequest, keys; kvrpcpb::ResolveLockResponse, error); pub fn new_cleanup_request(key: Vec, start_version: u64) -> kvrpcpb::CleanupRequest { let mut req = kvrpcpb::CleanupRequest::default(); @@ -186,9 +182,10 @@ pub fn new_cleanup_request(key: Vec, start_version: u64) -> kvrpcpb::Cleanup req } -impl KvRequest for kvrpcpb::CleanupRequest { - type Response = kvrpcpb::CleanupResponse; -} +impl_kv_request!( + kvrpcpb::CleanupRequest, key; + kvrpcpb::CleanupResponse, error +); shardable_key!(kvrpcpb::CleanupRequest); collect_first!(kvrpcpb::CleanupResponse); @@ -237,9 +234,10 @@ pub fn new_pessimistic_prewrite_request( req } -impl KvRequest for kvrpcpb::PrewriteRequest { - type Response = kvrpcpb::PrewriteResponse; -} +impl_kv_request!( + kvrpcpb::PrewriteRequest, mutations, primary_lock, secondaries; + kvrpcpb::PrewriteResponse, errors +); impl Shardable for kvrpcpb::PrewriteRequest { type Shard = Vec; @@ -284,10 +282,7 @@ pub fn new_commit_request( req } -impl KvRequest for kvrpcpb::CommitRequest { - type Response = kvrpcpb::CommitResponse; -} - +impl_kv_request!(kvrpcpb::CommitRequest, keys; kvrpcpb::CommitResponse, error); shardable_keys!(kvrpcpb::CommitRequest); pub fn new_batch_rollback_request( @@ -301,10 +296,7 @@ pub fn new_batch_rollback_request( req } -impl KvRequest for kvrpcpb::BatchRollbackRequest { - type Response = kvrpcpb::BatchRollbackResponse; -} - +impl_kv_request!(kvrpcpb::BatchRollbackRequest, keys; kvrpcpb::BatchRollbackResponse, error); shardable_keys!(kvrpcpb::BatchRollbackRequest); pub fn new_pessimistic_rollback_request( @@ -320,9 +312,12 @@ pub fn new_pessimistic_rollback_request( req } -impl KvRequest for kvrpcpb::PessimisticRollbackRequest { - type Response = kvrpcpb::PessimisticRollbackResponse; -} +impl_kv_request!( + kvrpcpb::PessimisticRollbackRequest, + keys; + kvrpcpb::PessimisticRollbackResponse, + errors +); shardable_keys!(kvrpcpb::PessimisticRollbackRequest); @@ -351,9 +346,12 @@ pub fn new_pessimistic_lock_request( req } -impl KvRequest for kvrpcpb::PessimisticLockRequest { - type Response = kvrpcpb::PessimisticLockResponse; -} +impl_kv_request!( + kvrpcpb::PessimisticLockRequest, + mutations, primary_lock; + kvrpcpb::PessimisticLockResponse, + errors +); impl Shardable for kvrpcpb::PessimisticLockRequest { type Shard = Vec; @@ -449,9 +447,11 @@ pub fn new_scan_lock_request( req } -impl KvRequest for kvrpcpb::ScanLockRequest { - type Response = kvrpcpb::ScanLockResponse; -} +impl_kv_request!( + kvrpcpb::ScanLockRequest; + kvrpcpb::ScanLockResponse, + error, locks +); impl Shardable for kvrpcpb::ScanLockRequest { type Shard = Vec; @@ -493,9 +493,12 @@ pub fn new_heart_beat_request( req } -impl KvRequest for kvrpcpb::TxnHeartBeatRequest { - type Response = kvrpcpb::TxnHeartBeatResponse; -} +impl_kv_request!( + kvrpcpb::TxnHeartBeatRequest, + primary_lock; + kvrpcpb::TxnHeartBeatResponse, + error +); impl Shardable for kvrpcpb::TxnHeartBeatRequest { type Shard = Vec>; @@ -531,9 +534,10 @@ impl Process for DefaultProcessor { } } -impl KvRequest for kvrpcpb::CheckTxnStatusRequest { - type Response = kvrpcpb::CheckTxnStatusResponse; -} +impl_kv_request!( + kvrpcpb::CheckTxnStatusRequest, primary_key; + kvrpcpb::CheckTxnStatusResponse, lock_info, error +); impl Shardable for kvrpcpb::CheckTxnStatusRequest { type Shard = Vec>; @@ -600,9 +604,10 @@ impl From<(u64, u64, Option)> for TransactionStatusKind { } } -impl KvRequest for kvrpcpb::CheckSecondaryLocksRequest { - type Response = kvrpcpb::CheckSecondaryLocksResponse; -} +impl_kv_request!( + kvrpcpb::CheckSecondaryLocksRequest, keys; + kvrpcpb::CheckSecondaryLocksResponse, locks, error +); shardable_keys!(kvrpcpb::CheckSecondaryLocksRequest); @@ -651,6 +656,7 @@ error_locks!(kvrpcpb::CheckTxnStatusResponse); error_locks!(kvrpcpb::CheckSecondaryLocksResponse); impl HasLocks for kvrpcpb::CleanupResponse {} + impl HasLocks for kvrpcpb::ScanLockResponse {} impl HasLocks for kvrpcpb::PessimisticRollbackResponse { @@ -682,12 +688,13 @@ impl HasLocks for kvrpcpb::PrewriteResponse { #[cfg(test)] mod tests { + use tikv_client_common::Error::{PessimisticLockError, ResolveLockError}; + use tikv_client_proto::kvrpcpb; + use crate::{ request::{plan::Merge, CollectWithShard, ResponseWithShard}, KvPair, }; - use tikv_client_common::Error::{PessimisticLockError, ResolveLockError}; - use tikv_client_proto::kvrpcpb; #[tokio::test] async fn test_merge_pessimistic_lock_response() { diff --git a/src/transaction/snapshot.rs b/src/transaction/snapshot.rs index 52f85b64..286da338 100644 --- a/src/transaction/snapshot.rs +++ b/src/transaction/snapshot.rs @@ -1,8 +1,13 @@ // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. -use crate::{BoundRange, Key, KvPair, Result, Transaction, Value}; +use crate::{ + pd::{PdClient, PdRpcClient}, + request::codec::RequestCodec, + BoundRange, Key, KvPair, Result, Transaction, Value, +}; use derive_new::new; use slog::Logger; +use std::marker::PhantomData; /// A read-only transaction which reads at the given timestamp. /// @@ -12,12 +17,17 @@ use slog::Logger; /// /// See the [Transaction](struct@crate::Transaction) docs for more information on the methods. #[derive(new)] -pub struct Snapshot { - transaction: Transaction, +pub struct Snapshot = PdRpcClient> { + transaction: Transaction, logger: Logger, + _phantom: PhantomData, } -impl Snapshot { +impl Snapshot +where + C: RequestCodec, + PdC: PdClient, +{ /// Get the value associated with the given key. pub async fn get(&mut self, key: impl Into) -> Result> { debug!(self.logger, "invoking get request on snapshot"); diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index 2edc8868..627ec88e 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -1,22 +1,26 @@ // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. +use std::{iter, marker::PhantomData, sync::Arc, time::Instant}; + +use derive_new::new; +use fail::fail_point; +use futures::prelude::*; +use slog::Logger; +use tokio::{sync::RwLock, time::Duration}; + +use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; + use crate::{ backoff::{Backoff, DEFAULT_REGION_BACKOFF}, pd::{PdClient, PdRpcClient}, request::{ - Collect, CollectError, CollectSingle, CollectWithShard, Plan, PlanBuilder, RetryOptions, + codec::RequestCodec, Collect, CollectError, CollectSingle, CollectWithShard, Plan, + PlanBuilder, RetryOptions, }, timestamp::TimestampExt, transaction::{buffer::Buffer, lowering::*}, BoundRange, Error, Key, KvPair, Result, Value, }; -use derive_new::new; -use fail::fail_point; -use futures::prelude::*; -use slog::Logger; -use std::{iter, sync::Arc, time::Instant}; -use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; -use tokio::{sync::RwLock, time::Duration}; /// An undo-able set of actions on the dataset. /// @@ -48,8 +52,9 @@ use tokio::{sync::RwLock, time::Duration}; /// ```rust,no_run /// # use tikv_client::{Config, TransactionClient}; /// # use futures::prelude::*; +/// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { -/// let client = TransactionClient::new(vec!["192.168.0.100"], None) +/// let client = TransactionClient::::new(vec!["192.168.0.100"], None) /// .await /// .unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); @@ -58,7 +63,7 @@ use tokio::{sync::RwLock, time::Duration}; /// txn.commit().await.unwrap(); /// # }); /// ``` -pub struct Transaction { +pub struct Transaction = PdRpcClient> { status: Arc>, timestamp: Timestamp, buffer: Buffer, @@ -67,15 +72,20 @@ pub struct Transaction { is_heartbeat_started: bool, start_instant: Instant, logger: Logger, + _phantom: PhantomData, } -impl Transaction { +impl Transaction +where + C: RequestCodec, + PdC: PdClient, +{ pub(crate) fn new( timestamp: Timestamp, rpc: Arc, options: TransactionOptions, logger: Logger, - ) -> Transaction { + ) -> Transaction { let status = if options.read_only { TransactionStatus::ReadOnly } else { @@ -90,6 +100,7 @@ impl Transaction { is_heartbeat_started: false, start_instant: std::time::Instant::now(), logger, + _phantom: PhantomData, } } @@ -104,8 +115,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Value, Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// let key = "TiKV".to_owned(); /// let result: Option = txn.get(key).await.unwrap(); @@ -165,8 +177,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Value, Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_pessimistic().await.unwrap(); /// let key = "TiKV".to_owned(); /// let result: Value = txn.get_for_update(key).await.unwrap().unwrap(); @@ -199,8 +212,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Value, Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_pessimistic().await.unwrap(); /// let exists = txn.key_exists("k1".to_owned()).await.unwrap(); /// txn.commit().await.unwrap(); @@ -226,8 +240,9 @@ impl Transaction { /// # use tikv_client::{Key, Value, Config, TransactionClient}; /// # use futures::prelude::*; /// # use std::collections::HashMap; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// let keys = vec!["TiKV".to_owned(), "TiDB".to_owned()]; /// let result: HashMap = txn @@ -279,8 +294,9 @@ impl Transaction { /// # use tikv_client::{Key, Value, Config, TransactionClient, KvPair}; /// # use futures::prelude::*; /// # use std::collections::HashMap; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_pessimistic().await.unwrap(); /// let keys = vec!["foo".to_owned(), "bar".to_owned()]; /// let result: Vec = txn @@ -324,8 +340,9 @@ impl Transaction { /// # use tikv_client::{Key, KvPair, Value, Config, TransactionClient}; /// # use futures::prelude::*; /// # use std::collections::HashMap; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// let key1: Key = b"foo".to_vec().into(); /// let key2: Key = b"bar".to_vec().into(); @@ -360,8 +377,9 @@ impl Transaction { /// # use tikv_client::{Key, KvPair, Value, Config, TransactionClient}; /// # use futures::prelude::*; /// # use std::collections::HashMap; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// let key1: Key = b"foo".to_vec().into(); /// let key2: Key = b"bar".to_vec().into(); @@ -423,8 +441,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Key, Value, Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// let key = "foo".to_owned(); /// let val = "FOO".to_owned(); @@ -454,8 +473,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Key, Value, Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// let key = "foo".to_owned(); /// let val = "FOO".to_owned(); @@ -490,8 +510,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Key, Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100", "192.168.0.101"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// let key = "foo".to_owned(); /// txn.delete(key); @@ -525,8 +546,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Config, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// txn.lock_keys(vec!["TiKV".to_owned(), "Rust".to_owned()]); /// // ... Do some actions. @@ -561,8 +583,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Config, Timestamp, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// // ... Do some actions. /// let result: Timestamp = txn.commit().await.unwrap().unwrap(); @@ -619,8 +642,9 @@ impl Transaction { /// ```rust,no_run /// # use tikv_client::{Config, Timestamp, TransactionClient}; /// # use futures::prelude::*; + /// use tikv_client::transaction::ApiV1; /// # futures::executor::block_on(async { - /// # let client = TransactionClient::new(vec!["192.168.0.100"], None).await.unwrap(); + /// # let client = TransactionClient::::new(vec!["192.168.0.100"], None).await.unwrap(); /// let mut txn = client.begin_optimistic().await.unwrap(); /// // ... Do some actions. /// txn.rollback().await.unwrap(); @@ -914,7 +938,11 @@ impl Transaction { } } -impl Drop for Transaction { +impl Drop for Transaction +where + C: RequestCodec, + PdC: PdClient, +{ fn drop(&mut self) { debug!(self.logger, "dropping transaction"); if std::thread::panicking() { @@ -1124,7 +1152,7 @@ impl HeartbeatOption { /// The committer implements `prewrite`, `commit` and `rollback` functions. #[allow(clippy::too_many_arguments)] #[derive(new)] -struct Committer { +struct Committer { primary_key: Option, mutations: Vec, start_version: Timestamp, @@ -1353,13 +1381,6 @@ enum TransactionStatus { #[cfg(test)] mod tests { - use crate::{ - mock::{MockKvClient, MockPdClient}, - transaction::HeartbeatOption, - Transaction, TransactionOptions, - }; - use fail::FailScenario; - use slog::{Drain, Logger}; use std::{ any::Any, io, @@ -1369,8 +1390,18 @@ mod tests { }, time::Duration, }; + + use fail::FailScenario; + use slog::{Drain, Logger}; + use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; + use crate::{ + mock::{MockKvClient, MockPdClient}, + transaction::HeartbeatOption, + Transaction, TransactionOptions, + }; + #[tokio::test] async fn test_optimistic_heartbeat() -> Result<(), io::Error> { let plain = slog_term::PlainSyncDecorator::new(std::io::stdout()); diff --git a/src/util/client.rs b/src/util/client.rs new file mode 100644 index 00000000..4235b580 --- /dev/null +++ b/src/util/client.rs @@ -0,0 +1,45 @@ +use std::{future::Future, sync::Arc}; + +use slog::{Drain, Logger}; + +use crate::{ + pd::{PdRpcClient, RetryClient}, + Config, Result, +}; + +pub(crate) struct ClientContext { + pub logger: Logger, + pub pd: Arc>, +} + +impl ClientContext { + pub async fn new( + pd_endpoints: Vec, + config: Config, + codec_factory: F, + optional_logger: Option, + ) -> Result + where + S: Into, + F: Fn(Arc) -> Fut, + Fut: Future>, + { + let logger = optional_logger.unwrap_or_else(|| { + let plain = slog_term::PlainSyncDecorator::new(std::io::stdout()); + Logger::root( + slog_term::FullFormat::new(plain) + .build() + .filter_level(slog::Level::Info) + .fuse(), + o!(), + ) + }); + + let pd_endpoints: Vec = pd_endpoints.into_iter().map(Into::into).collect(); + let pd = Arc::new( + PdRpcClient::connect(&pd_endpoints, config, codec_factory, logger.clone()).await?, + ); + + Ok(ClientContext { pd, logger }) + } +} diff --git a/src/util/mod.rs b/src/util/mod.rs index 6d3bf763..3282a7ef 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,3 +1,4 @@ // Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0. +pub mod client; pub mod iter; diff --git a/tests/common/mod.rs b/tests/common/mod.rs index d666be77..39c139f0 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -3,7 +3,7 @@ mod ctl; use futures_timer::Delay; use log::{info, warn}; use std::{env, time::Duration}; -use tikv_client::{ColumnFamily, Key, RawClient, Result, TransactionClient}; +use tikv_client::{raw, transaction, ColumnFamily, Key, RawClient, Result, TransactionClient}; const ENV_PD_ADDRS: &str = "PD_ADDRS"; const ENV_ENABLE_MULIT_REGION: &str = "MULTI_REGION"; @@ -19,7 +19,10 @@ pub async fn clear_tikv() { // DEFAULT_REGION_BACKOFF is not long enough for CI environment. So set a longer backoff. let backoff = tikv_client::Backoff::no_jitter_backoff(100, 10000, 10); for cf in cfs { - let raw_client = RawClient::new(pd_addrs(), None).await.unwrap().with_cf(cf); + let raw_client = RawClient::::new(pd_addrs(), None) + .await + .unwrap() + .with_cf(cf); raw_client .delete_range_opt(vec![].., backoff.clone()) .await @@ -60,7 +63,7 @@ async fn ensure_region_split( // 1. write plenty transactional keys // 2. wait until regions split - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; let mut txn = client.begin_optimistic().await?; for key in keys.into_iter() { txn.put(key.into(), vec![0, 0, 0, 0]).await?; diff --git a/tests/failpoint_tests.rs b/tests/failpoint_tests.rs index 80519c27..68524412 100644 --- a/tests/failpoint_tests.rs +++ b/tests/failpoint_tests.rs @@ -5,7 +5,10 @@ use common::{init, pd_addrs}; use fail::FailScenario; use serial_test::serial; use std::time::Duration; -use tikv_client::{transaction::HeartbeatOption, Result, TransactionClient, TransactionOptions}; +use tikv_client::{ + transaction::{ApiV1, HeartbeatOption}, + Result, TransactionClient, TransactionOptions, +}; #[tokio::test] #[serial] @@ -16,7 +19,7 @@ async fn txn_optimistic_heartbeat() -> Result<()> { let key1 = "key1".to_owned(); let key2 = "key2".to_owned(); - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; let mut heartbeat_txn = client .begin_with_options( diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index c933ba7a..a51ced4a 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -22,8 +22,11 @@ use std::{ iter, }; use tikv_client::{ - transaction::HeartbeatOption, BoundRange, Error, Key, KvPair, RawClient, Result, Transaction, - TransactionClient, TransactionOptions, Value, + raw, + request::codec::{RawCodec, RequestCodec}, + transaction::{self, HeartbeatOption}, + BoundRange, Error, Key, KvPair, RawClient, Result, Transaction, TransactionClient, + TransactionOptions, Value, }; // Parameters used in test @@ -34,7 +37,7 @@ const NUM_TRNASFER: u32 = 100; #[serial] async fn txn_get_timestamp() -> Result<()> { const COUNT: usize = 1 << 8; // use a small number to make test fast - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; let mut versions = future::join_all((0..COUNT).map(|_| client.current_timestamp())) .await @@ -55,7 +58,7 @@ async fn txn_get_timestamp() -> Result<()> { async fn txn_crud() -> Result<()> { init().await?; - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; let mut txn = client.begin_optimistic().await?; // Get non-existent keys @@ -139,7 +142,7 @@ async fn txn_crud() -> Result<()> { async fn txn_insert_duplicate_keys() -> Result<()> { init().await?; - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; // Initialize TiKV store with {foo => bar} let mut txn = client.begin_optimistic().await?; txn.put("foo".to_owned(), "bar".to_owned()).await?; @@ -163,7 +166,7 @@ async fn txn_insert_duplicate_keys() -> Result<()> { async fn txn_pessimistic() -> Result<()> { init().await?; - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; let mut txn = client.begin_pessimistic().await?; txn.put("foo".to_owned(), "foo".to_owned()).await.unwrap(); @@ -180,7 +183,7 @@ async fn txn_pessimistic() -> Result<()> { #[serial] async fn raw_bank_transfer() -> Result<()> { init().await?; - let client = RawClient::new(pd_addrs(), None).await?; + let client = RawClient::::new(pd_addrs(), None).await?; let mut rng = thread_rng(); let people = gen_u32_keys(NUM_PEOPLE, &mut rng); @@ -232,7 +235,7 @@ async fn txn_read() -> Result<()> { let value = "large_value".repeat(10); init().await?; - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; for i in 0..2u32.pow(NUM_BITS_TXN) { let mut cur = i * 2u32.pow(32 - NUM_BITS_TXN); @@ -324,7 +327,7 @@ async fn txn_read() -> Result<()> { #[serial] async fn txn_bank_transfer() -> Result<()> { init().await?; - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; let mut rng = thread_rng(); let options = TransactionOptions::new_optimistic() .use_async_commit() @@ -377,7 +380,7 @@ async fn txn_bank_transfer() -> Result<()> { #[serial] async fn raw_req() -> Result<()> { init().await?; - let client = RawClient::new(pd_addrs(), None).await?; + let client = RawClient::::new(pd_addrs(), None).await?; // empty; get non-existent key let res = client.get("k1".to_owned()).await; @@ -507,7 +510,7 @@ async fn raw_req() -> Result<()> { #[serial] async fn txn_update_safepoint() -> Result<()> { init().await?; - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; let res = client.gc(client.current_timestamp().await?).await?; assert!(res); Ok(()) @@ -522,7 +525,7 @@ async fn raw_write_million() -> Result<()> { let interval = 2u32.pow(32 - NUM_BITS_TXN - NUM_BITS_KEY_PER_TXN); init().await?; - let client = RawClient::new(pd_addrs(), None).await?; + let client = RawClient::::new(pd_addrs(), None).await?; for i in 0..2u32.pow(NUM_BITS_TXN) { let mut cur = i * 2u32.pow(32 - NUM_BITS_TXN); @@ -568,7 +571,12 @@ async fn raw_write_million() -> Result<()> { #[serial] async fn txn_pessimistic_rollback() -> Result<()> { init().await?; - let client = TransactionClient::new_with_config(pd_addrs(), Default::default(), None).await?; + let client = TransactionClient::::new_with_config( + pd_addrs(), + Default::default(), + None, + ) + .await?; let mut preload_txn = client.begin_optimistic().await?; let key1 = vec![1]; let key2 = vec![2]; @@ -600,9 +608,12 @@ async fn txn_pessimistic_rollback() -> Result<()> { #[serial] async fn txn_pessimistic_delete() -> Result<()> { init().await?; - let client = - TransactionClient::new_with_config(vec!["127.0.0.1:2379"], Default::default(), None) - .await?; + let client = TransactionClient::::new_with_config( + vec!["127.0.0.1:2379"], + Default::default(), + None, + ) + .await?; // The transaction will lock the keys and must release the locks on commit, // even when values are not written to the DB. @@ -644,7 +655,12 @@ async fn txn_pessimistic_delete() -> Result<()> { #[serial] async fn txn_lock_keys() -> Result<()> { init().await?; - let client = TransactionClient::new_with_config(pd_addrs(), Default::default(), None).await?; + let client = TransactionClient::::new_with_config( + pd_addrs(), + Default::default(), + None, + ) + .await?; let k1 = b"key1".to_vec(); let k2 = b"key2".to_vec(); @@ -678,7 +694,12 @@ async fn txn_lock_keys() -> Result<()> { #[serial] async fn txn_lock_keys_error_handle() -> Result<()> { init().await?; - let client = TransactionClient::new_with_config(pd_addrs(), Default::default(), None).await?; + let client = TransactionClient::::new_with_config( + pd_addrs(), + Default::default(), + None, + ) + .await?; // Keys in `k` should locate in different regions. See `init()` for boundary of regions. let k: Vec = vec![ @@ -715,7 +736,12 @@ async fn txn_lock_keys_error_handle() -> Result<()> { #[serial] async fn txn_get_for_update() -> Result<()> { init().await?; - let client = TransactionClient::new_with_config(pd_addrs(), Default::default(), None).await?; + let client = TransactionClient::::new_with_config( + pd_addrs(), + Default::default(), + None, + ) + .await?; let key1 = "key".to_owned(); let key2 = "another key".to_owned(); let value1 = b"some value".to_owned(); @@ -762,7 +788,7 @@ async fn txn_pessimistic_heartbeat() -> Result<()> { let key1 = "key1".to_owned(); let key2 = "key2".to_owned(); - let client = TransactionClient::new(pd_addrs(), None).await?; + let client = TransactionClient::::new(pd_addrs(), None).await?; let mut heartbeat_txn = client .begin_with_options(TransactionOptions::new_pessimistic()) @@ -802,7 +828,7 @@ async fn txn_pessimistic_heartbeat() -> Result<()> { #[serial] async fn raw_cas() -> Result<()> { init().await?; - let client = RawClient::new(pd_addrs(), None) + let client = RawClient::::new(pd_addrs(), None) .await? .with_atomic_for_cas(); let key = "key".to_owned(); @@ -847,7 +873,7 @@ async fn raw_cas() -> Result<()> { client.batch_delete(vec![key.clone()]).await.err().unwrap(), Error::UnsupportedMode )); - let client = RawClient::new(pd_addrs(), None).await?; + let client = RawClient::::new(pd_addrs(), None).await?; assert!(matches!( client .compare_and_swap(key.clone(), None, vec![]) @@ -864,7 +890,12 @@ async fn raw_cas() -> Result<()> { #[serial] async fn txn_scan() -> Result<()> { init().await?; - let client = TransactionClient::new_with_config(pd_addrs(), Default::default(), None).await?; + let client = TransactionClient::::new_with_config( + pd_addrs(), + Default::default(), + None, + ) + .await?; let k1 = b"a".to_vec(); let v = b"b".to_vec(); @@ -887,7 +918,12 @@ async fn txn_scan() -> Result<()> { #[serial] async fn txn_scan_reverse() -> Result<()> { init().await?; - let client = TransactionClient::new_with_config(pd_addrs(), Default::default(), None).await?; + let client = TransactionClient::::new_with_config( + pd_addrs(), + Default::default(), + None, + ) + .await?; let k1 = b"a1".to_vec(); let k2 = b"a2".to_vec(); @@ -920,7 +956,7 @@ async fn txn_scan_reverse() -> Result<()> { } // helper function -async fn get_u32(client: &RawClient, key: Vec) -> Result { +async fn get_u32(client: &RawClient, key: Vec) -> Result { let x = client.get(key).await?.unwrap(); let boxed_slice = x.into_boxed_slice(); let array: Box<[u8; 4]> = boxed_slice @@ -930,7 +966,7 @@ async fn get_u32(client: &RawClient, key: Vec) -> Result { } // helper function -async fn get_txn_u32(txn: &mut Transaction, key: Vec) -> Result { +async fn get_txn_u32(txn: &mut Transaction, key: Vec) -> Result { let x = txn.get(key).await?.unwrap(); let boxed_slice = x.into_boxed_slice(); let array: Box<[u8; 4]> = boxed_slice diff --git a/tests/mock_tikv_tests.rs b/tests/mock_tikv_tests.rs index f047feb0..a855eab4 100644 --- a/tests/mock_tikv_tests.rs +++ b/tests/mock_tikv_tests.rs @@ -4,7 +4,7 @@ mod test { use log::debug; use mock_tikv::{start_mock_pd_server, start_mock_tikv_server, MOCK_PD_PORT}; use simple_logger::SimpleLogger; - use tikv_client::{KvPair, RawClient}; + use tikv_client::{raw::ApiV1, KvPair, RawClient}; #[tokio::test] #[ignore] @@ -15,7 +15,7 @@ mod test { let mut tikv_server = start_mock_tikv_server(); let _pd_server = start_mock_pd_server(); - let client = RawClient::new(vec![format!("localhost:{}", MOCK_PD_PORT)], None) + let client = RawClient::::new(vec![format!("localhost:{}", MOCK_PD_PORT)], None) .await .unwrap(); diff --git a/tikv-client-common/src/errors.rs b/tikv-client-common/src/errors.rs index db309107..97ddb6e8 100644 --- a/tikv-client-common/src/errors.rs +++ b/tikv-client-common/src/errors.rs @@ -93,6 +93,12 @@ pub enum Error { inner: Box, success_keys: Vec>, }, + #[error("Corrupted keyspace: expected: {:?}, key: {:?}", expected, key)] + CorruptedKeyspace { expected: Vec, key: Vec }, + #[error("Keyspace is unsupported in the cluster")] + KeyspaceUnsupported, + #[error("Keyspace {} is not enabled", name)] + KeyspaceNotEnabled { name: String }, } impl From for Error { diff --git a/tikv-client-pd/src/cluster.rs b/tikv-client-pd/src/cluster.rs index 063532e9..18d2c9f7 100644 --- a/tikv-client-pd/src/cluster.rs +++ b/tikv-client-pd/src/cluster.rs @@ -2,14 +2,17 @@ use crate::{timestamp::TimestampOracle, Error, Result, SecurityManager}; use async_trait::async_trait; -use grpcio::{CallOption, Environment}; +use grpcio::{CallOption, Environment, RpcStatusCode}; use std::{ collections::HashSet, sync::Arc, time::{Duration, Instant}, }; use tikv_client_common::internal_err; -use tikv_client_proto::pdpb::{self, Timestamp}; +use tikv_client_proto::{ + keyspacepb, + pdpb::{self, Timestamp}, +}; /// A PD cluster. pub struct Cluster { @@ -17,6 +20,7 @@ pub struct Cluster { client: pdpb::PdClient, members: pdpb::GetMembersResponse, tso: TimestampOracle, + keyspace_mgr: Option, } macro_rules! pd_request { @@ -75,6 +79,26 @@ impl Cluster { req.set_safe_point(safepoint); req.send(&self.client, timeout).await } + + pub async fn load_keyspace( + &self, + name: String, + timeout: Duration, + ) -> Result { + let mut req = pd_request!(self.id, keyspacepb::LoadKeyspaceRequest); + let mgr = self + .keyspace_mgr + .as_ref() + .ok_or(Error::KeyspaceUnsupported)?; + req.set_name(name); + req.send(mgr, timeout).await + } +} + +struct ConnectResponse { + client: pdpb::PdClient, + members: pdpb::GetMembersResponse, + keyspace_mgr: Option, } /// An object for connecting and reconnecting to a PD cluster. @@ -94,7 +118,11 @@ impl Connection { timeout: Duration, ) -> Result { let members = self.validate_endpoints(endpoints, timeout).await?; - let (client, members) = self.try_connect_leader(&members, timeout).await?; + let ConnectResponse { + client, + members, + keyspace_mgr, + } = self.try_connect_leader(&members, timeout).await?; let id = members.get_header().get_cluster_id(); let tso = TimestampOracle::new(id, &client)?; let cluster = Cluster { @@ -102,6 +130,7 @@ impl Connection { client, members, tso, + keyspace_mgr, }; Ok(cluster) } @@ -110,13 +139,18 @@ impl Connection { pub async fn reconnect(&self, cluster: &mut Cluster, timeout: Duration) -> Result<()> { warn!("updating pd client"); let start = Instant::now(); - let (client, members) = self.try_connect_leader(&cluster.members, timeout).await?; + let ConnectResponse { + client, + members, + keyspace_mgr, + } = self.try_connect_leader(&cluster.members, timeout).await?; let tso = TimestampOracle::new(cluster.id, &client)?; *cluster = Cluster { id: cluster.id, client, members, tso, + keyspace_mgr, }; info!("updating PD client done, spent {:?}", start.elapsed()); @@ -137,8 +171,8 @@ impl Connection { return Err(internal_err!("duplicated PD endpoint {}", ep)); } - let (_, resp) = match self.connect(ep, timeout).await { - Ok(resp) => resp, + let resp = match self.connect(ep, timeout).await { + Ok(ConnectResponse { members, .. }) => members, // Ignore failed PD node. Err(e) => { warn!("PD endpoint {} failed to respond: {:?}", ep, e); @@ -175,20 +209,39 @@ impl Connection { } } - async fn connect( - &self, - addr: &str, - timeout: Duration, - ) -> Result<(pdpb::PdClient, pdpb::GetMembersResponse)> { + async fn connect(&self, addr: &str, timeout: Duration) -> Result { let client = self .security_mgr .connect(self.env.clone(), addr, pdpb::PdClient::new)?; + + let keyspace_mgr = self.connect_keyspace_mgr(addr)?; + let option = CallOption::default().timeout(timeout); - let resp = client + let members = client .get_members_async_opt(&pdpb::GetMembersRequest::default(), option) .map_err(Error::from)? .await?; - Ok((client, resp)) + + Ok(ConnectResponse { + client, + members, + keyspace_mgr, + }) + } + + fn connect_keyspace_mgr(&self, addr: &str) -> Result> { + let result = + self.security_mgr + .connect(self.env.clone(), addr, keyspacepb::KeyspaceClient::new); + match result { + Ok(mgr) => Ok(Some(mgr)), + Err(Error::Grpc(grpcio::Error::RpcFailure(status))) + if status.code() == RpcStatusCode::UNIMPLEMENTED => + { + Ok(None) + } + Err(e) => Err(e), + } } async fn try_connect( @@ -196,10 +249,10 @@ impl Connection { addr: &str, cluster_id: u64, timeout: Duration, - ) -> Result<(pdpb::PdClient, pdpb::GetMembersResponse)> { - let (client, r) = self.connect(addr, timeout).await?; - Connection::validate_cluster_id(addr, &r, cluster_id)?; - Ok((client, r)) + ) -> Result { + let connect_resp = self.connect(addr, timeout).await?; + Connection::validate_cluster_id(addr, &connect_resp.members, cluster_id)?; + Ok(connect_resp) } fn validate_cluster_id( @@ -224,7 +277,7 @@ impl Connection { &self, previous: &pdpb::GetMembersResponse, timeout: Duration, - ) -> Result<(pdpb::PdClient, pdpb::GetMembersResponse)> { + ) -> Result { let previous_leader = previous.get_leader(); let members = previous.get_members(); let cluster_id = previous.get_header().get_cluster_id(); @@ -238,8 +291,8 @@ impl Connection { { for ep in m.get_client_urls() { match self.try_connect(ep.as_str(), cluster_id, timeout).await { - Ok((_, r)) => { - resp = Some(r); + Ok(ConnectResponse { members, .. }) => { + resp = Some(members); break 'outer; } Err(e) => { @@ -270,10 +323,11 @@ type GrpcResult = std::result::Result; #[async_trait] trait PdMessage { type Response: PdResponse; + type Client: Send + Sync; - async fn rpc(&self, client: &pdpb::PdClient, opt: CallOption) -> GrpcResult; + async fn rpc(&self, client: &Self::Client, opt: CallOption) -> GrpcResult; - async fn send(&self, client: &pdpb::PdClient, timeout: Duration) -> Result { + async fn send(&self, client: &Self::Client, timeout: Duration) -> Result { let option = CallOption::default().timeout(timeout); let response = self.rpc(client, option).await?; @@ -288,6 +342,7 @@ trait PdMessage { #[async_trait] impl PdMessage for pdpb::GetRegionRequest { type Response = pdpb::GetRegionResponse; + type Client = pdpb::PdClient; async fn rpc(&self, client: &pdpb::PdClient, opt: CallOption) -> GrpcResult { client.get_region_async_opt(self, opt)?.await @@ -297,6 +352,7 @@ impl PdMessage for pdpb::GetRegionRequest { #[async_trait] impl PdMessage for pdpb::GetRegionByIdRequest { type Response = pdpb::GetRegionResponse; + type Client = pdpb::PdClient; async fn rpc(&self, client: &pdpb::PdClient, opt: CallOption) -> GrpcResult { client.get_region_by_id_async_opt(self, opt)?.await @@ -306,6 +362,7 @@ impl PdMessage for pdpb::GetRegionByIdRequest { #[async_trait] impl PdMessage for pdpb::GetStoreRequest { type Response = pdpb::GetStoreResponse; + type Client = pdpb::PdClient; async fn rpc(&self, client: &pdpb::PdClient, opt: CallOption) -> GrpcResult { client.get_store_async_opt(self, opt)?.await @@ -315,6 +372,7 @@ impl PdMessage for pdpb::GetStoreRequest { #[async_trait] impl PdMessage for pdpb::GetAllStoresRequest { type Response = pdpb::GetAllStoresResponse; + type Client = pdpb::PdClient; async fn rpc(&self, client: &pdpb::PdClient, opt: CallOption) -> GrpcResult { client.get_all_stores_async_opt(self, opt)?.await @@ -324,12 +382,27 @@ impl PdMessage for pdpb::GetAllStoresRequest { #[async_trait] impl PdMessage for pdpb::UpdateGcSafePointRequest { type Response = pdpb::UpdateGcSafePointResponse; + type Client = pdpb::PdClient; async fn rpc(&self, client: &pdpb::PdClient, opt: CallOption) -> GrpcResult { client.update_gc_safe_point_async_opt(self, opt)?.await } } +#[async_trait] +impl PdMessage for keyspacepb::LoadKeyspaceRequest { + type Response = keyspacepb::LoadKeyspaceResponse; + type Client = keyspacepb::KeyspaceClient; + + async fn rpc( + &self, + client: &keyspacepb::KeyspaceClient, + opt: CallOption, + ) -> GrpcResult { + client.load_keyspace_async_opt(self, opt)?.await + } +} + trait PdResponse { fn header(&self) -> &pdpb::ResponseHeader; } @@ -357,3 +430,9 @@ impl PdResponse for pdpb::UpdateGcSafePointResponse { self.get_header() } } + +impl PdResponse for keyspacepb::LoadKeyspaceResponse { + fn header(&self) -> &pdpb::ResponseHeader { + self.get_header() + } +} diff --git a/tikv-client-proto/proto/brpb.proto b/tikv-client-proto/proto/brpb.proto index c25da40a..66560fdb 100644 --- a/tikv-client-proto/proto/brpb.proto +++ b/tikv-client-proto/proto/brpb.proto @@ -54,6 +54,7 @@ message BackupMeta { // Note: `raw_ranges` is deprecated, as it bloats backupmeta. It is kept for // compatibility, so new BR can restore older backups. repeated RawRange raw_ranges = 9; + // An index to files contains RawRanges. MetaFile raw_range_index = 15; @@ -69,6 +70,12 @@ message BackupMeta { // API version implies the encode of the key and value. kvrpcpb.APIVersion api_version = 18; + + // the placement policy info in backup cluster. we assume the policy won't be too much for one cluster. + repeated PlacementPolicy policies = 19; + + // new_collations_enabled specifies the config `new_collations_enabled_on_first_bootstrap` in tidb. + string new_collations_enabled = 20; } message File { @@ -106,6 +113,10 @@ message MetaFile { repeated bytes ddls = 5; } +message PlacementPolicy { + bytes info = 1; +} + message Schema { bytes db = 1; bytes table = 2; @@ -182,6 +193,47 @@ message BackupRequest { int32 compression_level = 13; // The cipher_info is Used to encrypt sst CipherInfo cipher_info = 14; + + // dst_api_version indicates the key-value encoding version used by the + // generated SST file. Accepted values: + // 1. "v1": the generated SST files are encoded with api-v1, can be restored + // to TiKV clusters whose api version is set to v1. + // 2. "v2": the generated SST files are encoded with api-v2, can be restored + // to TiKV clusters whose api version is set to v2. + kvrpcpb.APIVersion dst_api_version = 15; +} + +message StreamBackupTaskInfo { + // The storage for backup, parsed by BR. + StorageBackend storage = 1; + // The time range for backing up. + uint64 start_ts = 2; + uint64 end_ts = 3; + + // Misc meta datas. + // The name of the task, also the ID of the task. + string name = 4; + // The table filter of the task. + // Only for displaying the task info. + repeated string table_filter = 5; + + // The last timestamp of the task has been updated. + // This is a simple solution for unfrequent config changing: + // When we watched a config change(via polling or etcd watching), + // We perform a incremental scan between [last_update_ts, now), + // for filling the diff data during conf changing. + // The current implementation scan [0, now) for every task ranges newly added, + // So this field is reserved for future usage. + reserved "last_update_ts"; + + // We use '/tidb/br-stream/ranges// -> ' + // for storing all the target ranges of the task, the "ranges" field will not be saved in the taskinfo. + reserved "ranges"; + + // The status field allow us mark the task as 'paused'. + // We use '/tidb/br-stream/pause/ -> ()' for pausing a task for now. + // If we need more complexed status transforming, maybe add the status field back would be fine. + reserved "status"; } message StorageBackend { @@ -287,6 +339,9 @@ message BackupResponse { bytes start_key = 2; bytes end_key = 3; repeated File files = 4; + + // API version implies the encode of the key and value. + kvrpcpb.APIVersion api_version = 5; } service Backup { @@ -321,3 +376,66 @@ message ExternalStorageSaveRequest { message ExternalStorageSaveResponse { } +message Metadata { + repeated DataFileInfo files = 1; + int64 store_id = 2; + uint64 resolved_ts = 3; + uint64 max_ts = 4; + uint64 min_ts = 5; +} + +enum FileType { + Delete = 0; + Put = 1; +} + +message DataFileInfo { + // SHA256 of the file. + bytes sha256 = 1; + // Path of the file. + string path = 2; + int64 number_of_entries = 3; + + /// Below are extra information of the file, for better filtering files. + // The min ts of the keys in the file. + uint64 min_ts = 4; + // The max ts of the keys in the file. + uint64 max_ts = 5; + // The resolved ts of the region when saving the file. + uint64 resolved_ts = 6; + // The region of the file. + int64 region_id = 7; + // The key range of the file. + // Encoded and starts with 'z'(internal key). + bytes start_key = 8; + bytes end_key = 9; + // The column family of the file. + string cf = 10; + // The operation type of the file. + FileType type = 11; + + // Whether the data file contains meta keys(m prefixed keys) only. + bool is_meta = 12; + // The table ID of the file contains, when `is_meta` is true, would be ignored. + int64 table_id = 13; + + // The file length. + uint64 length = 14; + + // The minimal begin ts in default cf if this file is write cf. + uint64 min_begin_ts_in_default_cf = 15; + + // It may support encrypting at future. + reserved "iv"; +} + +message StreamBackupError { + // the unix epoch time (in millisecs) of the time the error reported. + uint64 happen_at = 1; + // the unified error code of the error. + string error_code = 2; + // the user-friendly error message. + string error_message = 3; + // the store id of who issues the error. + uint64 store_id = 4; +} diff --git a/tikv-client-proto/proto/cdcpb.proto b/tikv-client-proto/proto/cdcpb.proto index 5682eafd..969beff8 100644 --- a/tikv-client-proto/proto/cdcpb.proto +++ b/tikv-client-proto/proto/cdcpb.proto @@ -83,6 +83,10 @@ message Event { bytes key = 5; bytes value = 6; bytes old_value = 7; + // expire_ts_unix_secs is used for RawKV (see `ChangeDataRequest.KvApi`), + // and represents the expiration time of this row. + // Absolute time, seconds since Unix epoch. + uint64 expire_ts_unix_secs = 8; } message Entries { @@ -150,6 +154,16 @@ message ChangeDataRequest { // min_commit_ts so that the resolved_ts can be advanced. NotifyTxnStatus notify_txn_status = 10; } + + // KvAPI specifies to capture data written by different KV API. + // See more details in https://github.com/tikv/rfcs/blob/master/text/0069-api-v2.md. + enum KvAPI { + TiDB = 0; + RawKV = 1; + TxnKV = 2; + } + + KvAPI kv_api = 11; } service ChangeData { diff --git a/tikv-client-proto/proto/coprocessor.proto b/tikv-client-proto/proto/coprocessor.proto index 4078f23e..45c724aa 100644 --- a/tikv-client-proto/proto/coprocessor.proto +++ b/tikv-client-proto/proto/coprocessor.proto @@ -6,7 +6,6 @@ import "kvrpcpb.proto"; import "gogoproto/gogo.proto"; import "rustproto.proto"; import "metapb.proto"; -import "span.proto"; option (gogoproto.marshaler_all) = true; option (gogoproto.sizer_all) = true; @@ -57,7 +56,12 @@ message Response { bool is_cache_hit = 7; uint64 cache_last_version = 8; bool can_be_cached = 9; - repeated span.SpanSet spans = 10; + + reserved 10; + + // Contains the latest buckets version of the region. + // Clients should query PD to update buckets in cache if its is stale. + uint64 latest_buckets_version = 12; } message RegionInfo { @@ -66,6 +70,11 @@ message RegionInfo { repeated KeyRange ranges = 3; } +message TableRegions { + int64 physical_table_id = 1; + repeated RegionInfo regions = 2; +} + message BatchRequest { kvrpcpb.Context context = 1; int64 tp = 2; @@ -74,6 +83,8 @@ message BatchRequest { uint64 start_ts = 5; // Any schema-ful storage to validate schema correctness if necessary. int64 schema_ver = 6; + // Used for partition table scan + repeated TableRegions table_regions = 7; } message BatchResponse { diff --git a/tikv-client-proto/proto/debugpb.proto b/tikv-client-proto/proto/debugpb.proto index 400da51d..a5eab219 100644 --- a/tikv-client-proto/proto/debugpb.proto +++ b/tikv-client-proto/proto/debugpb.proto @@ -73,6 +73,10 @@ service Debug { // Get all region IDs in the store rpc GetAllRegionsInStore(GetAllRegionsInStoreRequest) returns (GetAllRegionsInStoreResponse) {} + + // Make this TiKV node return to the status on this node to certain ts. + rpc ResetToVersion(ResetToVersionRequest) returns (ResetToVersionResponse) {} + } enum DB { @@ -259,3 +263,10 @@ message GetAllRegionsInStoreRequest { message GetAllRegionsInStoreResponse { repeated uint64 regions = 1; } + +message ResetToVersionRequest { + uint64 ts = 1; +} + +message ResetToVersionResponse { +} diff --git a/tikv-client-proto/proto/encryptionpb.proto b/tikv-client-proto/proto/encryptionpb.proto index ddae704e..2328088c 100644 --- a/tikv-client-proto/proto/encryptionpb.proto +++ b/tikv-client-proto/proto/encryptionpb.proto @@ -44,6 +44,7 @@ enum EncryptionMethod { AES128_CTR = 2; AES192_CTR = 3; AES256_CTR = 4; + SM4_CTR = 5; } // The key used to encrypt the user data. diff --git a/tikv-client-proto/proto/errorpb.proto b/tikv-client-proto/proto/errorpb.proto index 8fd9c233..4a81c042 100644 --- a/tikv-client-proto/proto/errorpb.proto +++ b/tikv-client-proto/proto/errorpb.proto @@ -128,6 +128,11 @@ message DataIsNotReady { uint64 safe_ts = 3; } +message RecoveryInProgress { + // The requested region ID + uint64 region_id = 1; +} + // Error wraps all region errors, indicates an error encountered by a request. message Error { reserved "stale_epoch"; @@ -148,4 +153,6 @@ message Error { DataIsNotReady data_is_not_ready = 13; RegionNotInitialized region_not_initialized = 14; DiskFull disk_full = 15; + // Online recovery is still in performing, reject writes to avoid potential issues + RecoveryInProgress RecoveryInProgress = 16; } diff --git a/tikv-client-proto/proto/gcpb.proto b/tikv-client-proto/proto/gcpb.proto new file mode 100644 index 00000000..053872b8 --- /dev/null +++ b/tikv-client-proto/proto/gcpb.proto @@ -0,0 +1,117 @@ +syntax = "proto3"; +package gcpb; + +import "gogoproto/gogo.proto"; +import "rustproto.proto"; + +option (gogoproto.sizer_all) = true; +option (gogoproto.marshaler_all) = true; +option (gogoproto.unmarshaler_all) = true; +option (rustproto.lite_runtime_all) = true; + +option java_package = "org.tikv.kvproto"; + +service GC { + rpc ListKeySpaces(ListKeySpacesRequest) returns (ListKeySpacesResponse) {} + + rpc GetMinServiceSafePoint(GetMinServiceSafePointRequest) returns (GetMinServiceSafePointResponse) {} + + rpc UpdateGCSafePoint(UpdateGCSafePointRequest) returns (UpdateGCSafePointResponse) {} + + rpc UpdateServiceSafePoint(UpdateServiceSafePointRequest) returns (UpdateServiceSafePointResponse) {} +} + +message RequestHeader { + // cluster_id is the ID of the cluster which be sent to. + uint64 cluster_id = 1; + // sender_id is the ID of the sender server, also member ID or etcd ID. + uint64 sender_id = 2; +} + +message ResponseHeader { + // cluster_id is the ID of the cluster which sent the response. + uint64 cluster_id = 1; + Error error = 2; +} + +enum ErrorType { + OK = 0; + UNKNOWN = 1; + NOT_BOOTSTRAPPED = 2; + // revision supplied does not match the current etcd revision + REVISION_MISMATCH = 3; + // if the proposed safe point is earlier than old safe point or gc safe point + SAFEPOINT_ROLLBACK = 4; +} + +message Error { + ErrorType type = 1; + string message = 2; +} + +message KeySpace { + bytes space_id = 1; + uint64 gc_safe_point = 2; +} + +message ListKeySpacesRequest { + RequestHeader header = 1; + // set with_gc_safe_point to true to also receive gc safe point for each key space + bool with_gc_safe_point = 2; +} + +message ListKeySpacesResponse { + ResponseHeader header = 1; + repeated KeySpace key_spaces = 2; +} + +message GetMinServiceSafePointRequest { + RequestHeader header = 1; + bytes space_id = 2; +} + +message GetMinServiceSafePointResponse { + ResponseHeader header = 1; + uint64 safe_point = 2; + // revision here is to safeguard the validity of the obtained min, + // preventing cases where new services register their safe points after min is obtained by gc worker + int64 revision = 3; +} + +message UpdateGCSafePointRequest { + RequestHeader header = 1; + bytes space_id = 2; + uint64 safe_point = 3; + // here client need to provide the revision obtained from GetMinServiceSafePoint, + // so server can check if it's still valid + int64 revision = 4; +} + +message UpdateGCSafePointResponse { + ResponseHeader header = 1; + // update will be successful if revision is valid and new safepoint > old safe point + // if failed, previously obtained min might be incorrect, should retry from GetMinServiceSafePoint + bool succeeded = 2; + uint64 new_safe_point = 3; +} + +message UpdateServiceSafePointRequest { + RequestHeader header = 1; + bytes space_id = 2; + bytes service_id = 3; + // safe point will be set to expire on (PD Server time + TTL) + // pass in a ttl < 0 to remove target safe point + // pass in MAX_INT64 to set a safe point that never expire + int64 TTL = 4; + uint64 safe_point = 5; +} + +message UpdateServiceSafePointResponse { + ResponseHeader header = 1; + // update will be successful if ttl < 0 (a removal request) + // or if new safe point >= old safe point and new safe point >= gc safe point + bool succeeded = 2; + uint64 gc_safe_point = 3; + uint64 old_safe_point = 4; + uint64 new_safe_point = 5; +} diff --git a/tikv-client-proto/proto/import_sstpb.proto b/tikv-client-proto/proto/import_sstpb.proto index ec1f0461..04834210 100644 --- a/tikv-client-proto/proto/import_sstpb.proto +++ b/tikv-client-proto/proto/import_sstpb.proto @@ -5,7 +5,6 @@ package import_sstpb; import "metapb.proto"; import "errorpb.proto"; import "kvrpcpb.proto"; -import "raft_serverpb.proto"; import "gogoproto/gogo.proto"; import "rustproto.proto"; import "brpb.proto"; @@ -56,6 +55,12 @@ service ImportSST { // Collect duplicate data from TiKV. rpc DuplicateDetect(DuplicateDetectRequest) returns (stream DuplicateDetectResponse) {} + + // Apply download & apply increment kv files to TiKV. + rpc Apply(ApplyRequest) returns (ApplyResponse) {} + + // ClearFiles clear applied file after restore succeed. + rpc ClearFiles(ClearRequest) returns (ClearResponse) {} } enum SwitchMode { @@ -176,6 +181,9 @@ message DownloadRequest { // of gRPC, add more concrete types if it is necessary later. message Error { string message = 1; + + // We meet some internal errors of the store. + errorpb.Error store_error = 2; } message DownloadResponse { @@ -235,6 +243,13 @@ message WriteResponse { message RawWriteBatch { uint64 ttl = 1; repeated Pair pairs = 2; + + // To be compatible with the key encoding of API V2. + // This field should be generated from the client instead of the server, + // since the message will be send to all the replicas of a region. + // Otherwise, the underlying data generated by the server would be inconsistent which is hard to scale + // for other features like MVCC over RawKV. + uint64 ts = 3; } message RawWriteRequest { @@ -280,3 +295,80 @@ message DuplicateDetectResponse { // ] repeated KvPair pairs = 3; } + +message KVMeta { + // The file name of the KV file. + string name = 1; + + // file length for check. + uint64 length = 2; + + // tell us which cf should apply. WRITE_CF or DEFAULT_CF e.g. + string cf = 3; + + // is_delete represents whether we should delete the kv in tikv. + // it may not be too much delete file. only rollBack operation will generate delete kv file. + bool is_delete = 4; + + // the key ts space being smaller than start_ts can be filter. + uint64 start_ts = 10; + + // the key ts space large than restore_ts can be filter. + uint64 restore_ts = 5; + + bytes start_key = 6; + + bytes end_key = 7; + + // used for checksum when download kv file. + bytes sha256 = 8; + + // the key ts space less than start_snapshot_ts can be filter. + // Deprecated: this field 'start_snapshot_ts' is replaced by the field 'start_ts'. + uint64 start_snapshot_ts = 9; +} + + +message ApplyRequest { + // The meta of the KV file. + KVMeta meta = 1; + + // Performs a key prefix rewrite after downloading the file. + // All keys in the files will be rewritten as: + // + // new_key = new_key_prefix + old_key[len(old_key_prefix)..] + // + // When used for TiDB, rewriting the prefix changes the table ID. Please + // note that key-rewrite is applied on the origin keys in encoded + // representation. + // + // You need to ensure that the keys before and after rewriting are in the + // same order, otherwise the RPC request will fail. + RewriteRule rewrite_rule = 2 [(gogoproto.nullable) = false]; + + backup.StorageBackend storage_backend = 3; + + // context represents region info and it used to build raft commands. + kvrpcpb.Context context = 4; + + // cipher_info is used to decrypt kv file when download file. + backup.CipherInfo cipher_info = 11; +} + +message ApplyResponse { + // The actual key range (after rewrite) of the downloaded file. The range is + // inclusive in both ends. + Range range = 1 [(gogoproto.nullable) = false]; + + Error error = 2; +} + +message ClearRequest { + // clear files in import directory with given prefix. + string prefix = 1; +} + +message ClearResponse { + Error error = 1; +} + diff --git a/tikv-client-proto/proto/keyspacepb.proto b/tikv-client-proto/proto/keyspacepb.proto new file mode 100644 index 00000000..70a58d10 --- /dev/null +++ b/tikv-client-proto/proto/keyspacepb.proto @@ -0,0 +1,80 @@ +syntax = "proto3"; +package keyspacepb; + +import "pdpb.proto"; + +import "gogoproto/gogo.proto"; +import "rustproto.proto"; + +option (gogoproto.sizer_all) = true; +option (gogoproto.marshaler_all) = true; +option (gogoproto.unmarshaler_all) = true; +option (rustproto.lite_runtime_all) = true; + +option java_package = "org.tikv.kvproto"; + +// Keyspace provides services to manage keyspaces. +service Keyspace { + rpc UpdateKeyspaceConfig (UpdateKeyspaceConfigRequest) returns (UpdateKeyspaceConfigResponse) {} + + rpc LoadKeyspace (LoadKeyspaceRequest) returns (LoadKeyspaceResponse) {} + // WatchKeyspaces first return all current keyspaces' metadata as its first response. + // Then, it returns responses containing keyspaces that had their metadata changed. + rpc WatchKeyspaces (WatchKeyspacesRequest) returns (stream WatchKeyspacesResponse) {} +} + +message KeyspaceMeta { + uint32 id = 1; + string name = 2; + KeyspaceState state = 3; + int64 created_at = 4; + int64 state_changed_at = 5; + map config = 7; +} + +enum KeyspaceState { + ENABLED = 0; + DISABLED = 1; + ARCHIVED = 2; +} + +message UpdateKeyspaceConfigRequest { + pdpb.RequestHeader header = 1; + string name = 2; + repeated Mutation mutations = 3; +} + +message Mutation { + Op op = 1; + bytes key = 2; + bytes value = 3; +} + +enum Op { + PUT = 0; + DEL = 1; +} + +message UpdateKeyspaceConfigResponse { + pdpb.ResponseHeader header = 1; + KeyspaceMeta keyspace = 2; +} + +message LoadKeyspaceRequest { + pdpb.RequestHeader header = 1; + string name = 2; +} + +message LoadKeyspaceResponse { + pdpb.ResponseHeader header = 1; + KeyspaceMeta keyspace = 2; +} + +message WatchKeyspacesRequest { + pdpb.RequestHeader header = 1; +} + +message WatchKeyspacesResponse { + pdpb.ResponseHeader header = 1; + repeated KeyspaceMeta keyspaces = 2; +} diff --git a/tikv-client-proto/proto/kvrpcpb.proto b/tikv-client-proto/proto/kvrpcpb.proto index 34d08be4..e7e5a291 100644 --- a/tikv-client-proto/proto/kvrpcpb.proto +++ b/tikv-client-proto/proto/kvrpcpb.proto @@ -6,6 +6,7 @@ import "errorpb.proto"; import "gogoproto/gogo.proto"; import "rustproto.proto"; import "deadlock.proto"; +import "tracepb.proto"; option (gogoproto.marshaler_all) = true; option (gogoproto.sizer_all) = true; @@ -121,6 +122,8 @@ message PrewriteResponse { // the commit ts of the transaction. Otherwise, if TiKV failed to commit it with 1PC or the // transaction is not 1PC, the value will be 0. uint64 one_pc_commit_ts = 4; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 5; } // Lock a set of keys to prepare to write to them. @@ -172,6 +175,8 @@ message PessimisticLockResponse { // In legacy TiKV, this field is not used even 'force' is false. In that case, an empty value indicates // two possible situations: (1) the key does not exist. (2) the key exists but the value is empty. repeated bool not_founds = 6; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 7; } // Unlock keys locked using `PessimisticLockRequest`. @@ -185,6 +190,8 @@ message PessimisticRollbackRequest { message PessimisticRollbackResponse { errorpb.Error region_error = 1; repeated KeyError errors = 2; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 3; } // Used to update the lock_ttl of a psessimistic and/or large transaction to prevent it from been killed. @@ -203,6 +210,8 @@ message TxnHeartBeatResponse { KeyError error = 2; // The TTL actually set on the requested lock. uint64 lock_ttl = 3; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 4; } // CheckTxnStatusRequest checks the status of a transaction. @@ -246,6 +255,8 @@ message CheckTxnStatusResponse { // The action performed by TiKV (and why if the action is to rollback). Action action = 5; LockInfo lock_info = 6; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 7; } // Part of the async commit protocol, checks for locks on all supplied keys. If a lock is missing, @@ -268,6 +279,8 @@ message CheckSecondaryLocksResponse { // If any of the locks have been committed, this is the commit ts used. If no // locks have been committed, it will be zero. uint64 commit_ts = 4; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 5; } // The second phase of writing to TiKV. If there are no errors or conflicts, then this request @@ -289,6 +302,8 @@ message CommitResponse { KeyError error = 2; // If the commit ts is derived from min_commit_ts, this field should be set. uint64 commit_version = 3; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 4; } // Not yet implemented. @@ -353,6 +368,8 @@ message BatchRollbackRequest { message BatchRollbackResponse { errorpb.Error region_error = 1; KeyError error = 2; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 3; } // Scan the database for locks. Used at the start of the GC process to find all @@ -374,6 +391,8 @@ message ScanLockResponse { KeyError error = 2; // Info on all locks found by the scan. repeated LockInfo locks = 3; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 4; } // For all keys locked by the transaction identified by `start_version`, either @@ -392,6 +411,8 @@ message ResolveLockRequest { message ResolveLockResponse { errorpb.Error region_error = 1; KeyError error = 2; + // Execution details about the request processing. + ExecDetailsV2 exec_details_v2 = 3; } // Request TiKV to garbage collect all non-current data older than `safe_point`. @@ -466,11 +487,13 @@ message RawBatchPutRequest { Context context = 1; repeated KvPair pairs = 2; string cf = 3; + uint64 ttl = 4 [deprecated=true]; + bool for_cas = 5; + // The time-to-live for each keys in seconds, and if the length of `ttls` // is exactly one, the ttl will be applied to all keys. Otherwise, the length // mismatch between `ttls` and `pairs` will return an error. - repeated uint64 ttls = 4; - bool for_cas = 5; + repeated uint64 ttls = 6; } message RawBatchPutResponse { @@ -710,6 +733,13 @@ message Context { // Read request should read through locks belonging to these transactions because these // transactions are committed and theirs commit_ts <= read request's start_ts. repeated uint64 committed_locks = 22; + + // The informantion to trace a request sent to TiKV. + tracepb.TraceContext trace_context = 23; + + // The source of the request, will be used as the tag of the metrics reporting. + // This field can be set for any requests that require to report metrics with any extra labels. + string request_source = 24; } // The API version the server and the client is using. @@ -841,6 +871,7 @@ enum CommandPri { enum IsolationLevel { SI = 0; // SI = snapshot isolation RC = 1; // RC = read committed + RCCheckTS = 2; // RC read and it's needed to check if there exists more recent versions. } // Operation allowed info during each TiKV storage threshold. @@ -852,16 +883,18 @@ enum DiskFullOpt { message TimeDetail { // Off-cpu wall time elapsed in TiKV side. Usually this includes queue waiting time and - // other kind of waitings in series. - int64 wait_wall_time_ms = 1; + // other kind of waitings in series. (Wait time in the raftstore is not included.) + uint64 wait_wall_time_ms = 1; // Off-cpu and on-cpu wall time elapsed to actually process the request payload. It does not // include `wait_wall_time`. // This field is very close to the CPU time in most cases. Some wait time spend in RocksDB // cannot be excluded for now, like Mutex wait time, which is included in this field, so that // this field is called wall time instead of CPU time. - int64 process_wall_time_ms = 2; + uint64 process_wall_time_ms = 2; // KV read wall Time means the time used in key/value scan and get. - int64 kv_read_wall_time_ms = 3; + uint64 kv_read_wall_time_ms = 3; + // Total wall clock time spent on this RPC in TiKV . + uint64 total_rpc_wall_time_ns = 4; } message ScanInfo { @@ -911,6 +944,12 @@ message ScanDetailV2 { // Total number of bytes from block reads. uint64 rocksdb_block_read_byte = 7; + + // Total time used for block reads. + uint64 rocksdb_block_read_nanos = 9; + + // Time used for getting a raftstore snapshot (including leader confirmation and getting the RocksDB snapshot). + uint64 get_snapshot_nanos = 10; } message ExecDetails { @@ -931,6 +970,50 @@ message ExecDetailsV2 { // Available when ctx.record_scan_stat = true or meet slow query. ScanDetailV2 scan_detail_v2 = 2; + + // Raftstore writing durations of the request. Only available for some write requests. + WriteDetail write_detail = 3; +} + +message WriteDetail { + // Wait duration in the store loop. + uint64 store_batch_wait_nanos = 1; + + // Wait duration before sending proposal to peers. + uint64 propose_send_wait_nanos = 2; + + // Total time spent on persisting the log. + uint64 persist_log_nanos = 3; + + // Wait time until the Raft log write leader begins to write. + uint64 raft_db_write_leader_wait_nanos = 4; + + // Time spent on synchronizing the Raft log to the disk. + uint64 raft_db_sync_log_nanos = 5; + + // Time spent on writing the Raft log to the Raft memtable. + uint64 raft_db_write_memtable_nanos = 6; + + // Time waiting for peers to confirm the proposal (counting from the instant when the leader sends the proposal message). + uint64 commit_log_nanos = 7; + + // Wait duration in the apply loop. + uint64 apply_batch_wait_nanos = 8; + + // Total time spend to applying the log. + uint64 apply_log_nanos = 9; + + // Wait time until the KV RocksDB lock is acquired. + uint64 apply_mutex_lock_nanos = 10; + + // Wait time until becoming the KV RocksDB write leader. + uint64 apply_write_leader_wait_nanos = 11; + + // Time spent on writing the KV DB WAL to the disk. + uint64 apply_write_wal_nanos = 12; + + // Time spent on writing to the memtable of the KV RocksDB. + uint64 apply_write_memtable_nanos = 13; } message KvPair { @@ -1144,3 +1227,53 @@ message RawChecksumResponse { uint64 total_kvs = 4; uint64 total_bytes = 5; } + +message CompactError { + oneof error { + CompactErrorInvalidStartKey err_invalid_start_key = 1; + CompactErrorPhysicalTableNotExist err_physical_table_not_exist = 2; + CompactErrorCompactInProgress err_compact_in_progress = 3; + CompactErrorTooManyPendingTasks err_too_many_pending_tasks = 4; + } +} + +message CompactErrorInvalidStartKey {} + +message CompactErrorPhysicalTableNotExist {} + +message CompactErrorCompactInProgress {} + +message CompactErrorTooManyPendingTasks {} + +message CompactRequest { + // If specified, the compaction will start from this start key. + // If unspecified, the compaction will start from beginning. + // NOTE 1: The start key should be never manually constructed. You should always use a key + // returned in CompactResponse. + // NOTE 2: the compaction range will be always restricted by physical_table_id. + bytes start_key = 1; + + // The physical table that will be compacted. + // + // TODO: this is information that TiKV doesn't need to know. + // See https://github.com/pingcap/kvproto/issues/912 + int64 physical_table_id = 2; + + // The logical table id of the compaction. When receiving parallel requests with the same + // logical table id, err_compact_in_progress will be returned. + // + // TODO: this is information that TiKV doesn't need to know. + // See https://github.com/pingcap/kvproto/issues/912 + int64 logical_table_id = 3; +} + +message CompactResponse { + CompactError error = 1; + + // The compaction is done incrementally. If there are more data to compact, this field + // will be set. The client can request to compact more data according to the `compacted_end_key`. + bool has_remaining = 2; + + bytes compacted_start_key = 3; + bytes compacted_end_key = 4; +} diff --git a/tikv-client-proto/proto/logbackuppb.proto b/tikv-client-proto/proto/logbackuppb.proto new file mode 100644 index 00000000..c377828a --- /dev/null +++ b/tikv-client-proto/proto/logbackuppb.proto @@ -0,0 +1,43 @@ +syntax = "proto3"; +package logbackup; + +import "gogoproto/gogo.proto"; +import "rustproto.proto"; +import "errorpb.proto"; + +option (gogoproto.sizer_all) = true; +option (gogoproto.marshaler_all) = true; +option (gogoproto.unmarshaler_all) = true; +option (rustproto.lite_runtime_all) = true; + +option java_package = "org.tikv.kvproto"; + +// The minimal information for identify a region. +message RegionIdentity { + uint64 id = 1; + uint64 epoch_version = 2; + // We omitted epoch_conf_version because config change won't make range change. +} + +// The last flush ts with region information. +message RegionCheckpoint { + errorpb.Error err = 1; + RegionIdentity region = 2; + uint64 checkpoint = 3; +} + +message GetLastFlushTSOfRegionRequest { + repeated RegionIdentity regions = 1; +} + +message GetLastFlushTSOfRegionResponse { + repeated RegionCheckpoint checkpoints = 1; +} + +// The log backup service. +// Generally, most essential interfaces of log backup (say, checkpoint management, task management) are +// provided by adding some key in the embed etcd of PD. +// This interface is mainly provided for the checkpoint advancer and debug usage. +service LogBackup { + rpc GetLastFlushTSOfRegion(GetLastFlushTSOfRegionRequest) returns (GetLastFlushTSOfRegionResponse) {} +} diff --git a/tikv-client-proto/proto/metapb.proto b/tikv-client-proto/proto/metapb.proto index f59de845..1f9aacb4 100644 --- a/tikv-client-proto/proto/metapb.proto +++ b/tikv-client-proto/proto/metapb.proto @@ -26,6 +26,17 @@ enum StoreState { Tombstone = 2; } +// NodeState is going to replace StoreState to make the state concept more clear. +// "Up" is devided into "Preparing" and "Serving" stages so that we can better describe the online process. +// "Removing" is just like previous `Offline` which is more accurate. +// "Removed" has the same meaning with `Tombstone`. +enum NodeState { + Preparing = 0; + Serving = 1; + Removing = 2; + Removed = 3; +} + // Case insensitive key/value for replica constraints. message StoreLabel { string key = 1; @@ -52,6 +63,8 @@ message Store { int64 last_heartbeat = 11; // If the store is physically destroyed, which means it can never up again. bool physically_destroyed = 12; + // NodeState is used to replace StoreState which will be deprecated in the future. + NodeState node_state = 13; } message RegionEpoch { @@ -61,6 +74,42 @@ message RegionEpoch { uint64 version = 2; } +message BucketStats { + // total read in bytes of each bucket + repeated uint64 read_bytes = 1; + + // total write in bytes of each bucket + repeated uint64 write_bytes = 2; + + // total read qps of each bucket + repeated uint64 read_qps = 3; + + // total write qps of each bucket + repeated uint64 write_qps = 4; + + // total read keys of each bucket + repeated uint64 read_keys = 5; + + // total write keys of each bucket + repeated uint64 write_keys = 6; +} + +message Buckets { + uint64 region_id = 1; + + // A hint indicate if keys have changed. + uint64 version = 2; + + // keys of buckets, include start/end key of region + repeated bytes keys = 3; + + // bucket stats + BucketStats stats = 4; + + // The period in milliseconds that stats are collected with in + uint64 period_in_ms = 5; +} + message Region { uint64 id = 1; // Region key range [start_key, end_key). diff --git a/tikv-client-proto/proto/mpp.proto b/tikv-client-proto/proto/mpp.proto index 4167b49c..13b73d92 100644 --- a/tikv-client-proto/proto/mpp.proto +++ b/tikv-client-proto/proto/mpp.proto @@ -34,6 +34,8 @@ message DispatchTaskRequest { repeated coprocessor.RegionInfo regions = 4; // If this task contains table scan, we still need their region info. int64 schema_ver = 5; + // Used for partition table scan + repeated coprocessor.TableRegions table_regions = 6; } // Get response of DispatchTaskRequest. @@ -64,6 +66,7 @@ message MPPDataPacket { bytes data = 1; Error error = 2; repeated bytes chunks = 3; + repeated uint64 stream_ids = 4; } message Error { diff --git a/tikv-client-proto/proto/pdpb.proto b/tikv-client-proto/proto/pdpb.proto index 4f3d08f5..0fb827f3 100644 --- a/tikv-client-proto/proto/pdpb.proto +++ b/tikv-client-proto/proto/pdpb.proto @@ -84,6 +84,46 @@ service PD { rpc SplitAndScatterRegions(SplitAndScatterRegionsRequest) returns (SplitAndScatterRegionsResponse) {} rpc GetDCLocationInfo(GetDCLocationInfoRequest) returns (GetDCLocationInfoResponse) {} + + rpc StoreGlobalConfig(StoreGlobalConfigRequest) returns (StoreGlobalConfigResponse) {} + + rpc LoadGlobalConfig(LoadGlobalConfigRequest) returns (LoadGlobalConfigResponse) {} + + rpc WatchGlobalConfig(WatchGlobalConfigRequest) returns (stream WatchGlobalConfigResponse) {} + + rpc ReportBuckets(stream ReportBucketsRequest) returns (ReportBucketsResponse) {} + + rpc ReportMinResolvedTS(ReportMinResolvedTsRequest) returns (ReportMinResolvedTsResponse) {} +} + +message WatchGlobalConfigRequest { + +} + +message WatchGlobalConfigResponse { + repeated GlobalConfigItem changes = 1; +} + +message StoreGlobalConfigRequest { + repeated GlobalConfigItem changes = 1; +} + +message StoreGlobalConfigResponse { + Error error = 1; +} + +message LoadGlobalConfigRequest { + repeated string names = 1; +} + +message LoadGlobalConfigResponse { + repeated GlobalConfigItem items = 1; +} + +message GlobalConfigItem { + string name = 1; + string value = 2; + Error error = 3; } message RequestHeader { @@ -107,6 +147,9 @@ enum ErrorType { ALREADY_BOOTSTRAPPED = 4; INCOMPATIBLE_VERSION = 5; REGION_NOT_FOUND = 6; + GLOBAL_CONFIG_NOT_FOUND = 7; + DUPLICATED_ENTRY = 8; + ENTRY_NOT_FOUND = 9; } message Error { @@ -207,6 +250,7 @@ message GetRegionRequest { RequestHeader header = 1; bytes region_key = 2; + bool need_buckets = 3; } message GetRegionResponse { @@ -221,12 +265,15 @@ message GetRegionResponse { // Pending peers are the peers that the leader can't consider as // working followers. repeated metapb.Peer pending_peers = 6; + // buckets isn't nil if GetRegion.* requests set need_buckets. + metapb.Buckets buckets = 7; } message GetRegionByIDRequest { RequestHeader header = 1; uint64 region_id = 2; + bool need_buckets = 3; } // Use GetRegionResponse as the response of GetRegionByIDRequest. @@ -407,7 +454,9 @@ message RegionHeartbeatResponse { // Multiple change peer operations atomically. // Note: PD can use both ChangePeer and ChangePeerV2 at the same time // (not in the same RegionHeartbeatResponse). - // Now, PD use ChangePeerV2 only for replacing peers. + // Now, PD use ChangePeerV2 in following scenarios: + // 1. replacing peers + // 2. demoting voter directly ChangePeerV2 change_peer_v2 = 9; } @@ -484,6 +533,8 @@ message PeerStat { uint64 read_keys = 2; uint64 read_bytes = 3; QueryStats query_stats = 4; + uint64 written_keys = 5; + uint64 written_bytes = 6; } message StoreStats { @@ -532,37 +583,72 @@ message StoreStats { uint64 slow_score = 22; // Damaged regions on the store that need to be removed by PD. repeated uint64 damaged_regions_id = 23; + // If the apply worker is busy, namely high apply wait duration + bool is_apply_busy = 24; } message PeerReport { raft_serverpb.RaftLocalState raft_state = 1; raft_serverpb.RegionLocalState region_state = 2; + bool is_force_leader = 3; + // The peer has proposed but uncommitted commit merge. + bool has_commit_merge = 4; } message StoreReport { repeated PeerReport peer_reports = 1; + uint64 step = 2; } message StoreHeartbeatRequest { RequestHeader header = 1; StoreStats stats = 2; - // Detailed store report that is only filled up on PD's demand for online unsafe recover. + // Detailed store report that is only filled up on PD's demand for online unsafe recovery. StoreReport store_report = 3; + replication_modepb.StoreDRAutoSyncStatus dr_autosync_status = 4; +} + +message DemoteFailedVoters { + uint64 region_id = 1; + repeated metapb.Peer failed_voters = 2; +} + +message ForceLeader { + // The store ids of the failed stores, TiKV uses it to decide if a peer is alive. + repeated uint64 failed_stores = 1; + // The region ids of the peer which is to be force leader. + repeated uint64 enter_force_leaders = 2; } message RecoveryPlan { + // Create empty regions to fill the key range hole. repeated metapb.Region creates = 1; - repeated metapb.Region updates = 2; - repeated uint64 deletes = 3; + // Update the meta of the regions, including peer lists, epoch and key range. + repeated metapb.Region updates = 2 [deprecated=true]; + // Tombstone the peers on the store locally. + repeated uint64 tombstones = 3; + // Issue conf change that demote voters on failed stores to learners on the regions. + repeated DemoteFailedVoters demotes = 4; + // Make the peers to be force leaders. + ForceLeader force_leader = 5; + // Step is an increasing number to note the round of recovery, + // It should be filled in the corresponding store report. + uint64 step = 6; } message StoreHeartbeatResponse { ResponseHeader header = 1; replication_modepb.ReplicationStatus replication_status = 2; string cluster_version = 3; - bool require_detailed_report = 4; - RecoveryPlan plan = 5; + + // Used by online unsafe recovery to request store report. + // Now it's substituted by reusing recovery_plan field. PD will send a empty + // recovery plan instead to request store report. + bool require_detailed_report = 4 [deprecated=true]; + // Operations of recovery. After the plan is executed, TiKV should attach the + // store report in store heartbeat. + RecoveryPlan recovery_plan = 5; } message ScatterRegionRequest { @@ -652,6 +738,8 @@ message SyncRegionResponse{ uint64 start_index = 3; repeated RegionStat region_stats = 4; repeated metapb.Peer region_leaders = 5; + // the buckets informations without stats. + repeated metapb.Buckets buckets =6; } message GetOperatorRequest { @@ -756,3 +844,26 @@ enum QueryKind { Commit = 10; Rollback = 11; } + +message ReportBucketsRequest { + RequestHeader header = 1; + + metapb.RegionEpoch region_epoch = 2; + metapb.Buckets buckets = 3; +} + +message ReportBucketsResponse { + ResponseHeader header = 1; +} + +message ReportMinResolvedTsRequest { + RequestHeader header = 1; + + uint64 store_id = 2; + + uint64 min_resolved_ts = 3; +} + +message ReportMinResolvedTsResponse { + ResponseHeader header = 1; +} diff --git a/tikv-client-proto/proto/raft_serverpb.proto b/tikv-client-proto/proto/raft_serverpb.proto index 8b071331..6c078b14 100644 --- a/tikv-client-proto/proto/raft_serverpb.proto +++ b/tikv-client-proto/proto/raft_serverpb.proto @@ -103,6 +103,8 @@ message RegionLocalState { PeerState state = 1; metapb.Region region = 2; MergeState merge_state = 3; + // The apply index corresponding to the storage when it's initialized. + uint64 tablet_index = 4; } enum ExtraMessageType { @@ -114,6 +116,7 @@ enum ExtraMessageType { // to make sure they all agree to sleep. MsgHibernateRequest = 4; MsgHibernateResponse = 5; + MsgRejectRaftLogCausedByMemoryUsage = 6; } message ExtraMessage { diff --git a/tikv-client-proto/proto/replication_modepb.proto b/tikv-client-proto/proto/replication_modepb.proto index 9bfb8814..c60005b6 100644 --- a/tikv-client-proto/proto/replication_modepb.proto +++ b/tikv-client-proto/proto/replication_modepb.proto @@ -17,10 +17,12 @@ message ReplicationStatus { enum DRAutoSyncState { // Raft logs need to sync between different DCs SYNC = 0; + // Wait for switching to ASYNC. Stop sync raft logs between DCs. + ASYNC_WAIT = 1; // Raft logs need to sync to majority peers - ASYNC = 1; + ASYNC = 2; // Switching from ASYNC to SYNC mode - SYNC_RECOVER = 2; + SYNC_RECOVER = 3; } // The status of dr-autosync mode. @@ -32,6 +34,10 @@ message DRAutoSync { uint64 state_id = 3; // Duration to wait before switching to SYNC by force (in seconds) int32 wait_sync_timeout_hint = 4; + // Stores should only sync messages with available stores when state is ASYNC or ASYNC_WAIT. + repeated uint64 available_stores = 5; + // Stores should forbid region split. + bool pause_region_split = 6; } enum RegionReplicationState { @@ -49,3 +55,8 @@ message RegionReplicationStatus { // Unique ID of the state, it increases after each state transfer. uint64 state_id = 2; } + +message StoreDRAutoSyncStatus { + DRAutoSyncState state = 1; + uint64 state_id = 2; +} diff --git a/tikv-client-proto/proto/resource_usage_agent.proto b/tikv-client-proto/proto/resource_usage_agent.proto index e1a30932..4250ec25 100644 --- a/tikv-client-proto/proto/resource_usage_agent.proto +++ b/tikv-client-proto/proto/resource_usage_agent.proto @@ -14,54 +14,37 @@ option java_package = "org.tikv.kvproto"; // ResourceUsageAgent is the service for storing resource usage records. service ResourceUsageAgent { - // DEPRECATED: We now use `Report` to report not only CPU time. - // - // Report the CPU time records. By default, the records with the same - // resource group tag will be batched by minute. - rpc ReportCPUTime(stream CPUTimeRecord) returns (EmptyResponse) {} - // Report the resource usage records. By default, the records with the same // resource group tag will be batched by minute. rpc Report(stream ResourceUsageRecord) returns (EmptyResponse) {} } -message CPUTimeRecord { - bytes resource_group_tag = 1; +// TiKV implements ResourceMeteringPubSub service for clients to subscribe to resource metering records. +service ResourceMeteringPubSub { + // Clients subscribe to resource metering records through this RPC, and TiKV periodically (e.g. per minute) + // publishes resource metering records to clients via gRPC stream. + rpc Subscribe(ResourceMeteringRequest) returns (stream ResourceUsageRecord) {} +} - // The following 2 repeated zipped together represents a List<(UnixTimestamp, CPUTime)> +message ResourceMeteringRequest {} - // UNIX timestamp in second. - repeated uint64 record_list_timestamp_sec = 2; - // The value can be greater than 1000ms if the requests are running parallelly. - repeated uint32 record_list_cpu_time_ms = 3; +message EmptyResponse {} +message ResourceUsageRecord { + oneof record_oneof { + GroupTagRecord record = 1; + } } -message ResourceUsageRecord { +// GroupTagRecord is a set of resource usage data grouped by resource_group_tag. +message GroupTagRecord { bytes resource_group_tag = 1; - - // The following repeated zipped together represents a List<(UnixTimestamp, Record)> - - // UNIX timestamp in second. - repeated uint64 record_list_timestamp_sec = 2; - - // The value can be greater than 1000ms if the requests are running parallelly. - repeated uint32 record_list_cpu_time_ms = 3; - - // The number of reads of keys associated with resource_group_tag. - repeated uint32 record_list_read_keys = 4; - - // The number of writes of keys associated with resource_group_tag. - repeated uint32 record_list_write_keys = 5; + repeated GroupTagRecordItem items = 2; } -message EmptyResponse {} - -// TiKV implements ResourceMeteringPubSub service for clients to subscribe to resource metering records. -service ResourceMeteringPubSub { - // Clients subscribe to resource metering records through this RPC, and TiKV periodically (e.g. per minute) - // publishes resource metering records to clients via gRPC stream. - rpc Subscribe(ResourceMeteringRequest) returns (stream ResourceUsageRecord) {} +message GroupTagRecordItem { + uint64 timestamp_sec = 1; + uint32 cpu_time_ms = 2; + uint32 read_keys = 3; + uint32 write_keys = 4; } - -message ResourceMeteringRequest {} diff --git a/tikv-client-proto/proto/tikvpb.proto b/tikv-client-proto/proto/tikvpb.proto index c3fd6ad3..dd37fcfa 100644 --- a/tikv-client-proto/proto/tikvpb.proto +++ b/tikv-client-proto/proto/tikvpb.proto @@ -104,6 +104,12 @@ service Tikv { /// Get the information about lock waiting from TiKV. rpc GetLockWaitInfo(kvrpcpb.GetLockWaitInfoRequest) returns (kvrpcpb.GetLockWaitInfoResponse); + + /// Compact a specified key range. This request is not restricted to raft leaders and will not be replicated. + /// It only compacts data on this node. + /// TODO: Currently this RPC is designed to be only compatible with TiFlash. + /// Shall be move out in https://github.com/pingcap/kvproto/issues/912 + rpc Compact(kvrpcpb.CompactRequest) returns (kvrpcpb.CompactResponse); } message BatchCommandsRequest { diff --git a/tikv-client-proto/proto/tracepb.proto b/tikv-client-proto/proto/tracepb.proto new file mode 100644 index 00000000..c229454c --- /dev/null +++ b/tikv-client-proto/proto/tracepb.proto @@ -0,0 +1,60 @@ +syntax = "proto3"; +package tracepb; + +option java_package = "org.tikv.kvproto"; + +service TraceRecordPubSub { + // Subscribe the Trace records generated on this service. The service will periodically (e.g. per minute) + // publishes Trace records to clients via gRPC stream. + rpc Subscribe(TraceRecordRequest) returns (stream TraceRecord) {} +} + +message TraceRecordRequest {} + +message TraceRecord { + oneof record_oneof { + Report report = 1; + NotifyCollect notify_collect = 2; + } +} + +message RemoteParentSpan { + // A unique id to identify the request. It's usually a UUID. + uint64 trace_id = 1; + // The span of remote caller that is awaiting the request. + uint64 span_id = 2; +} + +// The context of the request to be traced. +message TraceContext { + repeated RemoteParentSpan remote_parent_spans = 1; + // Report the trace records only if the duration of handling the request exceeds the threshold. + uint32 duration_threshold_ms = 2; +} + +// Report the spans collected when handling a request on a service. +message Report { + repeated RemoteParentSpan remote_parent_spans = 1; + repeated Span spans = 2; +} + +// Notify the subscriber to persis the spans of the trace. +message NotifyCollect { + uint64 trace_id = 1; +} + +message Span { + // The unique span id within the spans with the same `trace_id`. + // The most significant 32 bits should be random number generated by each service instance. + uint64 span_id = 1; + uint64 parent_id = 2; + uint64 begin_unix_ns = 3; + uint64 duration_ns = 4; + string event = 5; + repeated Property properties = 6; +} + +message Property { + string key = 1; + string value = 2; +} diff --git a/tikv-client-proto/src/lib.rs b/tikv-client-proto/src/lib.rs index 0a9a84af..a3bcc3c1 100644 --- a/tikv-client-proto/src/lib.rs +++ b/tikv-client-proto/src/lib.rs @@ -1,7 +1,9 @@ // Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0. use protos::*; -pub use protos::{coprocessor, errorpb, kvrpcpb, metapb, mpp, pdpb, raft_serverpb, tikvpb}; +pub use protos::{ + coprocessor, errorpb, keyspacepb, kvrpcpb, metapb, mpp, pdpb, raft_serverpb, tikvpb, +}; #[allow(dead_code)] #[allow(clippy::all)] diff --git a/tikv-client-store/src/errors.rs b/tikv-client-store/src/errors.rs index c1e915ac..7c57ed7a 100644 --- a/tikv-client-store/src/errors.rs +++ b/tikv-client-store/src/errors.rs @@ -227,6 +227,7 @@ mod test { region_error: None, error: None, commit_version: 0, + ..Default::default() }); assert!(resp.key_errors().is_none()); @@ -234,6 +235,7 @@ mod test { region_error: None, error: Some(kvrpcpb::KeyError::default()), commit_version: 0, + ..Default::default() }); assert!(resp.key_errors().is_some()); diff --git a/tikv-client-store/src/request.rs b/tikv-client-store/src/request.rs index 290f142a..7c14e2e9 100644 --- a/tikv-client-store/src/request.rs +++ b/tikv-client-store/src/request.rs @@ -12,6 +12,7 @@ pub trait Request: Any + Sync + Send + 'static { fn label(&self) -> &'static str; fn as_any(&self) -> &dyn Any; fn set_context(&mut self, context: kvrpcpb::Context); + fn mut_context(&mut self) -> &mut kvrpcpb::Context; } macro_rules! impl_request { @@ -41,6 +42,10 @@ macro_rules! impl_request { fn set_context(&mut self, context: kvrpcpb::Context) { kvrpcpb::$name::set_context(self, context) } + + fn mut_context(&mut self) -> &mut kvrpcpb::Context { + kvrpcpb::$name::mut_context(self) + } } }; }