From 71a4931b481e0b41b2f21781ba4c5293e1ed62b8 Mon Sep 17 00:00:00 2001 From: yongman Date: Sat, 6 Aug 2022 18:36:23 +0800 Subject: [PATCH 01/14] WIP add super batch for request Signed-off-by: yongman --- src/config.rs | 39 ++++++- src/raw/client.rs | 104 +++++++++-------- src/raw/requests.rs | 3 +- src/request/batch.rs | 185 +++++++++++++++++++++++++++++++ src/request/mod.rs | 35 +++--- src/request/plan.rs | 59 +++++++--- src/request/plan_builder.rs | 4 +- src/transaction/client.rs | 35 ++++-- src/transaction/lock.rs | 28 +++-- src/transaction/transaction.rs | 38 +++++-- tikv-client-store/src/request.rs | 159 ++++++++++++++++++++------ 11 files changed, 548 insertions(+), 141 deletions(-) create mode 100644 src/request/batch.rs diff --git a/src/config.rs b/src/config.rs index a0ee3abb..e42ff2fb 100644 --- a/src/config.rs +++ b/src/config.rs @@ -3,6 +3,14 @@ use serde_derive::{Deserialize, Serialize}; use std::{path::PathBuf, time::Duration}; +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2); +const DEFAULT_GRPC_KEEPALIVE_TIME: Duration = Duration::from_secs(10); +const DEFAULT_GRPC_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(3); +const DEFAULT_GRPC_COMPLETION_QUEUE_SIZE: usize = 1; +const DEFAULT_MAX_BATCH_WAIT_TIME: Duration = Duration::from_millis(0); +const DEFAULT_MAX_BATCH_SIZE: usize = 8; +const DEFAULT_OVERLOAD_THRESHOLD: usize = 200; + /// The configuration for either a [`RawClient`](crate::RawClient) or a /// [`TransactionClient`](crate::TransactionClient). /// @@ -16,9 +24,35 @@ pub struct Config { pub cert_path: Option, pub key_path: Option, pub timeout: Duration, + pub kv_client_config: KVClientConfig, } -const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2); +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default)] +#[serde(rename_all = "kebab-case")] +pub struct KVClientConfig { + pub completion_queue_size: usize, + pub grpc_keepalive_time: Duration, + pub grpc_keepalive_timeout: Duration, + pub allow_batch: bool, + pub overload_threshold: usize, + pub max_batch_wait_time: Duration, + pub max_batch_size: usize, +} + +impl Default for KVClientConfig { + fn default() -> Self { + Self { + completion_queue_size: DEFAULT_GRPC_COMPLETION_QUEUE_SIZE, + grpc_keepalive_time: DEFAULT_GRPC_KEEPALIVE_TIME, + grpc_keepalive_timeout: DEFAULT_GRPC_KEEPALIVE_TIMEOUT, + allow_batch: true, + overload_threshold: DEFAULT_OVERLOAD_THRESHOLD, + max_batch_wait_time: DEFAULT_MAX_BATCH_WAIT_TIME, + max_batch_size: DEFAULT_MAX_BATCH_SIZE, + } + } +} impl Default for Config { fn default() -> Self { @@ -27,6 +61,7 @@ impl Default for Config { cert_path: None, key_path: None, timeout: DEFAULT_REQUEST_TIMEOUT, + kv_client_config: KVClientConfig::default(), } } } @@ -80,4 +115,6 @@ impl Config { self.timeout = timeout; self } + + // TODO: add more config options for tivk client config } diff --git a/src/raw/client.rs b/src/raw/client.rs index b9ac740d..3801510c 100644 --- a/src/raw/client.rs +++ b/src/raw/client.rs @@ -9,7 +9,7 @@ use tikv_client_proto::metapb; use crate::{ backoff::DEFAULT_REGION_BACKOFF, - config::Config, + config::{Config, KVClientConfig}, pd::{PdClient, PdRpcClient}, raw::lowering::*, request::{Collect, CollectSingle, Plan}, @@ -31,6 +31,7 @@ pub struct Client { /// Whether to use the [`atomic mode`](Client::with_atomic_for_cas). atomic: bool, logger: Logger, + kv_config: KVClientConfig, } impl Clone for Client { @@ -40,6 +41,7 @@ impl Clone for Client { cf: self.cf.clone(), atomic: self.atomic, logger: self.logger.clone(), + kv_config: self.kv_config.clone(), } } } @@ -106,13 +108,15 @@ impl Client { }); 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?); + let rpc = Arc::new( + PdRpcClient::connect(&pd_endpoints, config.clone(), false, logger.clone()).await?, + ); Ok(Client { rpc, cf: None, atomic: false, logger, + kv_config: config.kv_client_config, }) } @@ -147,6 +151,7 @@ impl Client { cf: Some(cf), atomic: self.atomic, logger: self.logger.clone(), + kv_config: self.kv_config.clone(), } } @@ -164,6 +169,7 @@ impl Client { cf: self.cf.clone(), atomic: true, logger: self.logger.clone(), + kv_config: self.kv_config.clone(), } } } @@ -195,11 +201,12 @@ impl Client { pub async fn get_opt(&self, key: impl Into, backoff: Backoff) -> Result> { debug!(self.logger, "invoking raw get request"); let request = new_raw_get_request(key.into(), self.cf.clone()); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .merge(CollectSingle) - .post_process_default() - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .merge(CollectSingle) + .post_process_default() + .plan(); plan.execute().await } @@ -236,10 +243,11 @@ impl Client { ) -> Result> { debug!(self.logger, "invoking raw batch_get request"); let request = new_raw_batch_get_request(keys.into_iter().map(Into::into), self.cf.clone()); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .merge(Collect) - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .merge(Collect) + .plan(); plan.execute() .await .map(|r| r.into_iter().map(Into::into).collect()) @@ -274,11 +282,12 @@ impl Client { ) -> Result<()> { debug!(self.logger, "invoking raw put request"); let request = new_raw_put_request(key.into(), value.into(), self.cf.clone(), self.atomic); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .merge(CollectSingle) - .extract_error() - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .merge(CollectSingle) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -319,10 +328,11 @@ impl Client { self.cf.clone(), self.atomic, ); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .extract_error() - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -352,11 +362,12 @@ impl Client { pub async fn delete_opt(&self, key: impl Into, backoff: Backoff) -> Result<()> { debug!(self.logger, "invoking raw delete request"); let request = new_raw_delete_request(key.into(), self.cf.clone(), self.atomic); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .merge(CollectSingle) - .extract_error() - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .merge(CollectSingle) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -392,10 +403,11 @@ impl Client { self.assert_non_atomic()?; let request = new_raw_batch_delete_request(keys.into_iter().map(Into::into), self.cf.clone()); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .extract_error() - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -428,10 +440,11 @@ impl Client { debug!(self.logger, "invoking raw delete_range request"); self.assert_non_atomic()?; let request = new_raw_delete_range_request(range.into(), self.cf.clone()); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .extract_error() - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -640,7 +653,7 @@ impl Client { previous_value.into(), self.cf.clone(), ); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req) + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req, self.kv_config.clone()) .retry_multi_region(backoff) .merge(CollectSingle) .post_process_default() @@ -682,7 +695,7 @@ impl Client { ranges.into_iter().map(Into::into), request_builder, ); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req) + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req, self.kv_config.clone()) .preserve_shard() .retry_multi_region(backoff) .post_process_default() @@ -705,10 +718,11 @@ impl Client { } let request = new_raw_scan_request(range.into(), limit, key_only, self.cf.clone()); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .merge(Collect) - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .merge(Collect) + .plan(); let res = plan.execute().await; res.map(|mut s| { s.truncate(limit as usize); @@ -736,10 +750,11 @@ impl Client { key_only, self.cf.clone(), ); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) - .retry_multi_region(backoff) - .merge(Collect) - .plan(); + let plan = + crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + .retry_multi_region(backoff) + .merge(Collect) + .plan(); plan.execute().await } @@ -792,6 +807,7 @@ mod tests { cf: Some(ColumnFamily::Default), atomic: false, logger, + kv_config: KVClientConfig::default(), }; let resps = client .coprocessor( diff --git a/src/raw/requests.rs b/src/raw/requests.rs index bd678aaf..a62bf1c9 100644 --- a/src/raw/requests.rs +++ b/src/raw/requests.rs @@ -471,6 +471,7 @@ mod test { use super::*; use crate::{ backoff::{DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF}, + config::KVClientConfig, mock::{MockKvClient, MockPdClient}, request::Plan, Key, @@ -510,7 +511,7 @@ mod test { key_only: true, ..Default::default() }; - let plan = crate::request::PlanBuilder::new(client, scan) + let plan = crate::request::PlanBuilder::new(client, scan, KVClientConfig::default()) .resolve_lock(OPTIMISTIC_BACKOFF) .retry_multi_region(DEFAULT_REGION_BACKOFF) .merge(Collect) diff --git a/src/request/batch.rs b/src/request/batch.rs new file mode 100644 index 00000000..2acf4ab6 --- /dev/null +++ b/src/request/batch.rs @@ -0,0 +1,185 @@ +use crate::{Error, Result}; +use futures::{ + channel::{mpsc, oneshot}, + executor::block_on, + join, pin_mut, + prelude::*, + task::{AtomicWaker, Context, Poll}, +}; +use grpcio::WriteFlags; +use log::debug; +use std::{cell::RefCell, collections::VecDeque, pin::Pin, rc::Rc, thread}; +use tikv_client_common::internal_err; +use tikv_client_proto::pdpb::*; + +pub struct BatchCommandRequestEntry { + cmd: BatchCommandsRequest, + tx: oneshot::Sender, + // TODO + id: u64, +} + +impl BatchCommandRequestEntry { + pub fn new(cmd: BatchCommandsRequest, tx: oneshot::Sender) -> Self { + Self { cmd, tx, id: 0 } + } +} + +/// BatchWorker provides request in batch and return the result in batch. +pub struct BatchWorker { + request_tx: mpsc::Sender, +} + +impl BatchWorker { + pub fn new( + kv_client: &KvClient, + max_batch_size: usize, + max_inflight_requests: usize, + max_deplay_duration: std::time::Duration, + options: CallOption, + ) -> Result { + let (request_tx, request_rx) = mpsc::channel(max_batch_size); + + // Create rpc sender and receiver + let (rpc_sender, rpc_receiver) = kv_client.batch_commands_opt(options)?; + + // Start a background thread to handle batch requests and responses + thread::spawn(move || { + block_on(run_batch_worker( + rpc_sender.sink_err_into(), + rpc_receiver.err_into(), + request_rx, + max_batch_size, + max_inflight_requests, + )) + }); + + Ok(BatchWorker { request_tx }) + } + + pub async fn dispatch( + mut self, + request: BatchCommandsRequest, + ) -> Result { + let (tx, rx) = oneshot::channel(); + // Generate BatchCommandRequestEntry + let entry = BatchCommandRequestEntry::new(request, tx); + // Send request entry to the background thread to handle the request, response will be + // received in rx channel. + self.request_tx + .send(entry) + .await + .map_err(|_| Error::Internal("Failed to send request to batch worker".to_owned()))?; + Ok(rx.await?) + } +} + +async fn run_batch_worker( + mut tx: impl Sink<(BatchCommandsRequest, WriteFlags), Error = Error> + Unpin, + mut rx: impl Stream> + Unpin, + request_rx: mpsc::Receiver, + max_batch_size: usize, + max_inflight_requests: usize, + max_deplay_duration: std::time::Duration, +) { + // Inflight requests which are waiting for the response from rpc server + let mut inflight_requests = + Rc::new(RefCell::new(VecDeque::with_capacity(max_inflight_requests))); + + let waker = Rc::new(AtomicWaker::new()); + + // TODO + pin_mut!(request_rx); + + let mut request_stream = BatchCommandsRequestStream { + request_rx, + inflight_requests: inflight_requests.clone(), + self_waker: waker.clone(), + max_batch_size, + max_inflight_requests, + max_deplay_duration, + } + .map(Ok); + + let send_requests = tx.send_all(&mut request_stream); + + let recv_handle_response = async move { + while let Some(Ok(resp)) = rx.next().await { + let mut inflight_requests = inflight_requests.borrow_mut(); + + if inflight_requests.len == max_inflight_requests { + waker.wake(); + } + + // TODO more handle to response + + let request_id = resp.get_header().get_request_id(); + let request = inflight_requests.pop_front().unwrap(); + debug!( + "Received response for request_id {}: {:?}", + request_id, resp + ); + request.tx.send(resp).unwrap(); + } + }; + + let (tx_res, rx_res) = join!(send_requests, recv_handle_response); + debug!("Batch sender finished: {:?}", tx_res); + debug!("Batch receiver finished: {:?}", rx_res); +} + +struct BatchCommandsRequestStream { + request_rx: mpsc::Receiver, + inflight_requests: Rc>>>, + self_waker: Rc, + max_batch_size: usize, + max_inflight_requests: usize, + max_deplay_duration: std::time::Duration, +} + +impl Stream for BatchCommandsRequestStream { + type Item = Result<(BatchCommandsRequest, WriteFlags)>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut inflight_requests = self.inflight_requests.clone().borrow_mut(); + if inflight_requests.len() == 0 { + self.self_waker.register(cx.waker()); + } + + // Collect user requests + let requests = vec![]; + let latency_timer = Instant::now(); + while requests.len() < self.max_batch_size + && inflight_requests.len() < self.max_inflight_requests + { + // We can not wait longger than max_deplay_duration + if latency_timer.elapsed() > self.max_deplay_duration { + break; + } + + match self.request_rx.poll_next(cx) { + Poll::Ready(Some(Ok(entry))) => { + inflight_requests.push_back(entry.tx); + requests.push(entry.cmd); + } + Poll::Ready(Some(Err(e))) => { + return Poll::Ready(Some(Err(e))); + } + Poll::Ready(None) => { + return Poll::Ready(None); + } + Poll::Pending => { + break; + } + } + } + + if !requests.is_empty() { + let write_flags = WriteFlags::default().buffer_hint(false); + Poll::Ready(Some((requests, write_flags))) + } else { + self.self_waker.register(cx.waker()); + Poll::Pending + } + } +} diff --git a/src/request/mod.rs b/src/request/mod.rs index 959ff5e7..eb6d24ef 100644 --- a/src/request/mod.rs +++ b/src/request/mod.rs @@ -65,6 +65,7 @@ impl RetryOptions { mod test { use super::*; use crate::{ + config::KVClientConfig, mock::{MockKvClient, MockPdClient}, store::store_stream_for_keys, transaction::lowering::new_commit_request, @@ -165,11 +166,12 @@ mod test { |_: &dyn Any| Ok(Box::new(MockRpcResponse) as Box), ))); - let plan = crate::request::PlanBuilder::new(pd_client.clone(), request) - .resolve_lock(Backoff::no_jitter_backoff(1, 1, 3)) - .retry_multi_region(Backoff::no_jitter_backoff(1, 1, 3)) - .extract_error() - .plan(); + let plan = + crate::request::PlanBuilder::new(pd_client.clone(), request, KVClientConfig::default()) + .resolve_lock(Backoff::no_jitter_backoff(1, 1, 3)) + .retry_multi_region(Backoff::no_jitter_backoff(1, 1, 3)) + .extract_error() + .plan(); let _ = plan.execute().await; // Original call plus the 3 retries @@ -192,18 +194,23 @@ mod test { let req = new_commit_request(iter::once(key), Timestamp::default(), Timestamp::default()); // does not extract error - let plan = crate::request::PlanBuilder::new(pd_client.clone(), req.clone()) - .resolve_lock(OPTIMISTIC_BACKOFF) - .retry_multi_region(OPTIMISTIC_BACKOFF) - .plan(); + let plan = crate::request::PlanBuilder::new( + pd_client.clone(), + req.clone(), + KVClientConfig::default(), + ) + .resolve_lock(OPTIMISTIC_BACKOFF) + .retry_multi_region(OPTIMISTIC_BACKOFF) + .plan(); assert!(plan.execute().await.is_ok()); // extract error - let plan = crate::request::PlanBuilder::new(pd_client.clone(), req) - .resolve_lock(OPTIMISTIC_BACKOFF) - .retry_multi_region(OPTIMISTIC_BACKOFF) - .extract_error() - .plan(); + let plan = + crate::request::PlanBuilder::new(pd_client.clone(), req, KVClientConfig::default()) + .resolve_lock(OPTIMISTIC_BACKOFF) + .retry_multi_region(OPTIMISTIC_BACKOFF) + .extract_error() + .plan(); assert!(plan.execute().await.is_err()); } } diff --git a/src/request/plan.rs b/src/request/plan.rs index ce785262..091331c9 100644 --- a/src/request/plan.rs +++ b/src/request/plan.rs @@ -11,6 +11,7 @@ use tokio::sync::Semaphore; use crate::{ backoff::Backoff, + config::KVClientConfig, pd::PdClient, request::{KvRequest, Shardable}, stats::tikv_stats, @@ -29,6 +30,10 @@ pub trait Plan: Sized + Clone + Sync + Send + 'static { /// Execute the plan. async fn execute(&self) -> Result; + + fn kv_config(&self) -> KVClientConfig { + KVClientConfig::default() + } } /// The simplest plan which just dispatches a request to a specific kv server. @@ -36,6 +41,7 @@ pub trait Plan: Sized + Clone + Sync + Send + 'static { pub struct Dispatch { pub request: Req, pub kv_client: Option>, + pub kv_config: KVClientConfig, } #[async_trait] @@ -44,17 +50,39 @@ impl Plan for Dispatch { async fn execute(&self) -> Result { 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) - .await; - let result = stats.done(result); - result.map(|r| { - *r.downcast() - .expect("Downcast failed: request and response type mismatch") - }) + + let c = &self.kv_config; + + // Do a super batch request if possible if batching is enabled. + if c.allow_batch && c.max_batch_size > 1 { + // Generate a entry for batch request and response + + // Send the request to the batch stream channel and wait for response + let result = self + .kv_client + .as_ref() + .expect("Unreachable: kv_client has not been initialised in Dispatch") + .dispatch(&self.request) + .await?; + + Ok(Result) + } else { + let result = self + .kv_client + .as_ref() + .expect("Unreachable: kv_client has not been initialised in Dispatch") + .dispatch(&self.request) + .await; + let result = stats.done(result); + result.map(|r| { + *r.downcast() + .expect("Downcast failed: request and response type mismatch") + }) + } + } + + fn kv_config(&self) -> KVClientConfig { + self.kv_config.clone() } } @@ -85,7 +113,7 @@ where preserve_region_results: bool, ) -> Result<::Result> { let shards = current_plan.shards(&pd_client).collect::>().await; - let mut handles = Vec::new(); + let mut handles = Vec::with_capacity(shards.len()); for shard in shards { let (shard, region_store) = shard?; let mut clone = current_plan.clone(); @@ -416,6 +444,7 @@ where async fn execute(&self) -> Result { let mut result = self.inner.execute().await?; let mut clone = self.clone(); + let kv_config = self.kv_config(); loop { let locks = result.take_locks(); if locks.is_empty() { @@ -427,7 +456,7 @@ where } let pd_client = self.pd_client.clone(); - if resolve_locks(locks, pd_client.clone()).await? { + if resolve_locks(locks, pd_client.clone(), kv_config.clone()).await? { result = self.inner.execute().await?; } else { match clone.backoff.next_delay_duration() { @@ -440,6 +469,10 @@ where } } } + + fn kv_config(&self) -> KVClientConfig { + self.inner.kv_config() + } } /// When executed, the plan extracts errors from its inner plan, and returns an diff --git a/src/request/plan_builder.rs b/src/request/plan_builder.rs index ce269fac..0da1fb01 100644 --- a/src/request/plan_builder.rs +++ b/src/request/plan_builder.rs @@ -3,6 +3,7 @@ use super::plan::PreserveShard; use crate::{ backoff::Backoff, + config::KVClientConfig, pd::PdClient, request::{ DefaultProcessor, Dispatch, ExtractError, KvRequest, Merge, MergeResponse, Plan, Process, @@ -31,11 +32,12 @@ pub struct Targetted; impl PlanBuilderPhase for Targetted {} impl PlanBuilder, NoTarget> { - pub fn new(pd_client: Arc, request: Req) -> Self { + pub fn new(pd_client: Arc, request: Req, kv_config: KVClientConfig) -> Self { PlanBuilder { pd_client, plan: Dispatch { request, + kv_config, kv_client: None, }, phantom: PhantomData, diff --git a/src/transaction/client.rs b/src/transaction/client.rs index 69f11bce..37000dfa 100644 --- a/src/transaction/client.rs +++ b/src/transaction/client.rs @@ -3,7 +3,7 @@ use super::{requests::new_scan_lock_request, resolve_locks}; use crate::{ backoff::{DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF}, - config::Config, + config::{Config, KVClientConfig}, pd::{PdClient, PdRpcClient}, request::Plan, timestamp::TimestampExt, @@ -36,6 +36,7 @@ const SCAN_LOCK_BATCH_SIZE: u32 = 1024; pub struct Client { pd: Arc, logger: Logger, + kv_config: KVClientConfig, } impl Clone for Client { @@ -43,6 +44,7 @@ impl Clone for Client { Self { pd: self.pd.clone(), logger: self.logger.clone(), + kv_config: self.kv_config.clone(), } } } @@ -112,8 +114,14 @@ impl Client { }); 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?); - Ok(Client { pd, logger }) + let pd = Arc::new( + PdRpcClient::connect(&pd_endpoints, config.clone(), true, logger.clone()).await?, + ); + Ok(Client { + pd, + logger, + kv_config: config.kv_client_config, + }) } /// Creates a new optimistic [`Transaction`]. @@ -245,11 +253,12 @@ impl Client { SCAN_LOCK_BATCH_SIZE, ); - let plan = crate::request::PlanBuilder::new(self.pd.clone(), req) - .resolve_lock(OPTIMISTIC_BACKOFF) - .retry_multi_region(DEFAULT_REGION_BACKOFF) - .merge(crate::request::Collect) - .plan(); + let plan = + crate::request::PlanBuilder::new(self.pd.clone(), req, self.kv_config.clone()) + .resolve_lock(OPTIMISTIC_BACKOFF) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .merge(crate::request::Collect) + .plan(); let res: Vec = plan.execute().await?; if res.is_empty() { @@ -262,7 +271,7 @@ impl Client { // resolve locks // FIXME: (1) this is inefficient (2) when region error occurred - resolve_locks(locks, self.pd.clone()).await?; + resolve_locks(locks, self.pd.clone(), self.kv_config.clone()).await?; // update safepoint to PD let res: bool = self @@ -278,6 +287,12 @@ impl Client { 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) + Transaction::new( + timestamp, + self.pd.clone(), + options, + self.kv_config.clone(), + logger, + ) } } diff --git a/src/transaction/lock.rs b/src/transaction/lock.rs index 871effc6..af2256e7 100644 --- a/src/transaction/lock.rs +++ b/src/transaction/lock.rs @@ -2,6 +2,7 @@ use crate::{ backoff::{Backoff, DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF}, + config::KVClientConfig, pd::PdClient, region::RegionVerId, request::{CollectSingle, Plan}, @@ -28,6 +29,7 @@ const RESOLVE_LOCK_RETRY_LIMIT: usize = 10; pub async fn resolve_locks( locks: Vec, pd_client: Arc, + kv_config: KVClientConfig, ) -> Result { debug!("resolving locks"); let ts = pd_client.clone().get_timestamp().await?; @@ -62,12 +64,13 @@ pub async fn resolve_locks( Some(&commit_version) => commit_version, None => { let request = requests::new_cleanup_request(lock.primary_lock, lock.lock_version); - let plan = crate::request::PlanBuilder::new(pd_client.clone(), request) - .resolve_lock(OPTIMISTIC_BACKOFF) - .retry_multi_region(DEFAULT_REGION_BACKOFF) - .merge(CollectSingle) - .post_process_default() - .plan(); + let plan = + crate::request::PlanBuilder::new(pd_client.clone(), request, kv_config.clone()) + .resolve_lock(OPTIMISTIC_BACKOFF) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .merge(CollectSingle) + .post_process_default() + .plan(); let commit_version = plan.execute().await?; commit_versions.insert(lock.lock_version, commit_version); commit_version @@ -79,6 +82,7 @@ pub async fn resolve_locks( lock.lock_version, commit_version, pd_client.clone(), + kv_config.clone(), ) .await?; clean_regions @@ -94,6 +98,7 @@ async fn resolve_lock_with_retry( start_version: u64, commit_version: u64, pd_client: Arc, + kv_config: KVClientConfig, ) -> Result { debug!("resolving locks with retry"); // FIXME: Add backoff @@ -104,7 +109,7 @@ async fn resolve_lock_with_retry( let ver_id = store.region_with_leader.ver_id(); let request = requests::new_resolve_lock_request(start_version, commit_version); // The only place where single-region is used - let plan = crate::request::PlanBuilder::new(pd_client.clone(), request) + let plan = crate::request::PlanBuilder::new(pd_client.clone(), request, kv_config.clone()) .single_region_with_store(store) .await? .resolve_lock(Backoff::no_backoff()) @@ -165,15 +170,16 @@ mod tests { let key = vec![1]; let region1 = MockPdClient::region1(); - let resolved_region = resolve_lock_with_retry(&key, 1, 2, client.clone()) - .await - .unwrap(); + let resolved_region = + resolve_lock_with_retry(&key, 1, 2, client.clone(), KVClientConfig::default()) + .await + .unwrap(); assert_eq!(region1.ver_id(), resolved_region); // Test resolve lock over retry limit fail::cfg("region-error", "10*return").unwrap(); let key = vec![100]; - resolve_lock_with_retry(&key, 3, 4, client) + resolve_lock_with_retry(&key, 3, 4, client, KVClientConfig::default()) .await .expect_err("should return error"); } diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index 30f9602a..38c8fa4c 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -2,6 +2,7 @@ use crate::{ backoff::{Backoff, DEFAULT_REGION_BACKOFF}, + config::KVClientConfig, pd::{PdClient, PdRpcClient}, request::{ Collect, CollectError, CollectSingle, CollectWithShard, Plan, PlanBuilder, RetryOptions, @@ -64,6 +65,7 @@ pub struct Transaction { buffer: Buffer, rpc: Arc, options: TransactionOptions, + kv_config: KVClientConfig, is_heartbeat_started: bool, start_instant: Instant, logger: Logger, @@ -74,6 +76,7 @@ impl Transaction { timestamp: Timestamp, rpc: Arc, options: TransactionOptions, + kv_config: KVClientConfig, logger: Logger, ) -> Transaction { let status = if options.read_only { @@ -87,6 +90,7 @@ impl Transaction { buffer: Buffer::new(options.is_pessimistic()), rpc, options, + kv_config, is_heartbeat_started: false, start_instant: std::time::Instant::now(), logger, @@ -118,11 +122,12 @@ impl Transaction { let rpc = self.rpc.clone(); let key = key.into(); let retry_options = self.options.retry_options.clone(); + let kv_config = self.kv_config.clone(); self.buffer .get_or_else(key, |key| async move { let request = new_get_request(key, timestamp); - let plan = PlanBuilder::new(rpc, request) + let plan = PlanBuilder::new(rpc, request, kv_config) .resolve_lock(retry_options.lock_backoff) .retry_multi_region(DEFAULT_REGION_BACKOFF) .merge(CollectSingle) @@ -248,11 +253,12 @@ impl Transaction { let timestamp = self.timestamp.clone(); let rpc = self.rpc.clone(); let retry_options = self.options.retry_options.clone(); + let kv_config = self.kv_config.clone(); self.buffer .batch_get_or_else(keys.into_iter().map(|k| k.into()), move |keys| async move { let request = new_batch_get_request(keys, timestamp); - let plan = PlanBuilder::new(rpc, request) + let plan = PlanBuilder::new(rpc, request, kv_config) .resolve_lock(retry_options.lock_backoff) .retry_multi_region(retry_options.region_backoff) .merge(Collect) @@ -595,6 +601,7 @@ impl Transaction { self.timestamp.clone(), self.rpc.clone(), self.options.clone(), + self.kv_config.clone(), self.buffer.get_write_size() as u64, self.start_instant, self.logger.new(o!("child" => 1)), @@ -652,6 +659,7 @@ impl Transaction { self.timestamp.clone(), self.rpc.clone(), self.options.clone(), + self.kv_config.clone(), self.buffer.get_write_size() as u64, self.start_instant, self.logger.new(o!("child" => 1)), @@ -687,7 +695,7 @@ impl Transaction { primary_key, self.start_instant.elapsed().as_millis() as u64 + MAX_TTL, ); - let plan = PlanBuilder::new(self.rpc.clone(), request) + let plan = PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .merge(CollectSingle) @@ -707,6 +715,7 @@ impl Transaction { let timestamp = self.timestamp.clone(); let rpc = self.rpc.clone(); let retry_options = self.options.retry_options.clone(); + let kv_config = self.kv_config.clone(); self.buffer .scan_and_fetch( @@ -717,7 +726,7 @@ impl Transaction { move |new_range, new_limit| async move { let request = new_scan_request(new_range, timestamp, new_limit, key_only, reverse); - let plan = PlanBuilder::new(rpc, request) + let plan = PlanBuilder::new(rpc, request, kv_config) .resolve_lock(retry_options.lock_backoff) .retry_multi_region(retry_options.region_backoff) .merge(Collect) @@ -772,7 +781,7 @@ impl Transaction { for_update_ts.clone(), need_value, ); - let plan = PlanBuilder::new(self.rpc.clone(), request) + let plan = PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .preserve_shard() .retry_multi_region_preserve_results(self.options.retry_options.region_backoff.clone()) @@ -826,7 +835,7 @@ impl Transaction { start_version, for_update_ts, ); - let plan = PlanBuilder::new(self.rpc.clone(), req) + let plan = PlanBuilder::new(self.rpc.clone(), req, self.kv_config.clone()) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .extract_error() @@ -876,6 +885,7 @@ impl Transaction { HeartbeatOption::FixedTime(heartbeat_interval) => heartbeat_interval, }; let start_instant = self.start_instant; + let kv_config = self.kv_config.clone(); let heartbeat_task = async move { loop { @@ -896,7 +906,7 @@ impl Transaction { primary_key.clone(), start_instant.elapsed().as_millis() as u64 + MAX_TTL, ); - let plan = PlanBuilder::new(rpc.clone(), request) + let plan = PlanBuilder::new(rpc.clone(), request, kv_config.clone()) .retry_multi_region(region_backoff.clone()) .merge(CollectSingle) .plan(); @@ -1129,6 +1139,7 @@ struct Committer { start_version: Timestamp, rpc: Arc, options: TransactionOptions, + kv_config: KVClientConfig, #[new(default)] undetermined: bool, write_size: u64, @@ -1203,7 +1214,7 @@ impl Committer { .collect(); // FIXME set max_commit_ts and min_commit_ts - let plan = PlanBuilder::new(self.rpc.clone(), request) + let plan = PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .merge(CollectError) @@ -1243,7 +1254,7 @@ impl Committer { self.start_version.clone(), commit_version.clone(), ); - let plan = PlanBuilder::new(self.rpc.clone(), req) + let plan = PlanBuilder::new(self.rpc.clone(), req, self.kv_config.clone()) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .extract_error() @@ -1280,7 +1291,7 @@ impl Committer { .filter(|key| &primary_key != key); new_commit_request(keys, self.start_version, commit_version) }; - let plan = PlanBuilder::new(self.rpc, req) + let plan = PlanBuilder::new(self.rpc, req, self.kv_config.clone()) .resolve_lock(self.options.retry_options.lock_backoff) .retry_multi_region(self.options.retry_options.region_backoff) .extract_error() @@ -1301,7 +1312,7 @@ impl Committer { match self.options.kind { TransactionKind::Optimistic => { let req = new_batch_rollback_request(keys, self.start_version); - let plan = PlanBuilder::new(self.rpc, req) + let plan = PlanBuilder::new(self.rpc, req, self.kv_config.clone()) .resolve_lock(self.options.retry_options.lock_backoff) .retry_multi_region(self.options.retry_options.region_backoff) .extract_error() @@ -1310,7 +1321,7 @@ impl Committer { } TransactionKind::Pessimistic(for_update_ts) => { let req = new_pessimistic_rollback_request(keys, self.start_version, for_update_ts); - let plan = PlanBuilder::new(self.rpc, req) + let plan = PlanBuilder::new(self.rpc, req, self.kv_config.clone()) .resolve_lock(self.options.retry_options.lock_backoff) .retry_multi_region(self.options.retry_options.region_backoff) .extract_error() @@ -1353,6 +1364,7 @@ enum TransactionStatus { #[cfg(test)] mod tests { use crate::{ + config::KVClientConfig, mock::{MockKvClient, MockPdClient}, transaction::HeartbeatOption, Transaction, TransactionOptions, @@ -1402,6 +1414,7 @@ mod tests { pd_client, TransactionOptions::new_optimistic() .heartbeat_option(HeartbeatOption::FixedTime(Duration::from_secs(1))), + KVClientConfig::default(), logger.new(o!("child" => 1)), ); heartbeat_txn.put(key1.clone(), "foo").await.unwrap(); @@ -1451,6 +1464,7 @@ mod tests { pd_client, TransactionOptions::new_pessimistic() .heartbeat_option(HeartbeatOption::FixedTime(Duration::from_secs(1))), + KVClientConfig::default(), logger.new(o!("child" => 1)), ); heartbeat_txn.put(key1.clone(), "foo").await.unwrap(); diff --git a/tikv-client-store/src/request.rs b/tikv-client-store/src/request.rs index 290f142a..0a3a7442 100644 --- a/tikv-client-store/src/request.rs +++ b/tikv-client-store/src/request.rs @@ -2,9 +2,16 @@ use crate::{Error, Result}; use async_trait::async_trait; -use grpcio::CallOption; +use futures::{SinkExt, TryStreamExt}; +use grpcio::{CallOption, WriteFlags}; use std::any::Any; -use tikv_client_proto::{kvrpcpb, tikvpb::TikvClient}; +use tikv_client_proto::{ + kvrpcpb, + tikvpb::{ + batch_commands_request::{self, request::Cmd::*}, + BatchCommandsRequest, TikvClient, + }, +}; #[async_trait] pub trait Request: Any + Sync + Send + 'static { @@ -12,10 +19,11 @@ 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 into_batch(&self) -> BatchCommandsRequest; } macro_rules! impl_request { - ($name: ident, $fun: ident, $label: literal) => { + ($name: ident, $fun: ident, $label: literal, $cmd: ident) => { #[async_trait] impl Request for kvrpcpb::$name { async fn dispatch( @@ -41,87 +49,170 @@ macro_rules! impl_request { fn set_context(&mut self, context: kvrpcpb::Context) { kvrpcpb::$name::set_context(self, context) } + + fn into_batch(&self) -> BatchCommandsRequest { + let req = batch_commands_request::Request { + cmd: Some($cmd(self.clone())), + }; + BatchCommandsRequest { + requests: vec![req], + request_ids: vec![0], + } + } } }; } -impl_request!(RawGetRequest, raw_get_async_opt, "raw_get"); -impl_request!(RawBatchGetRequest, raw_batch_get_async_opt, "raw_batch_get"); -impl_request!(RawPutRequest, raw_put_async_opt, "raw_put"); -impl_request!(RawBatchPutRequest, raw_batch_put_async_opt, "raw_batch_put"); -impl_request!(RawDeleteRequest, raw_delete_async_opt, "raw_delete"); +impl_request!(RawGetRequest, raw_get_async_opt, "raw_get", RawGet); +impl_request!( + RawBatchGetRequest, + raw_batch_get_async_opt, + "raw_batch_get", + RawBatchGet +); +impl_request!(RawPutRequest, raw_put_async_opt, "raw_put", RawPut); +impl_request!( + RawBatchPutRequest, + raw_batch_put_async_opt, + "raw_batch_put", + RawBatchPut +); +impl_request!( + RawDeleteRequest, + raw_delete_async_opt, + "raw_delete", + RawDelete +); impl_request!( RawBatchDeleteRequest, raw_batch_delete_async_opt, - "raw_batch_delete" + "raw_batch_delete", + RawBatchDelete ); -impl_request!(RawScanRequest, raw_scan_async_opt, "raw_scan"); +impl_request!(RawScanRequest, raw_scan_async_opt, "raw_scan", RawScan); impl_request!( RawBatchScanRequest, raw_batch_scan_async_opt, - "raw_batch_scan" + "raw_batch_scan", + RawBatchScan ); impl_request!( RawDeleteRangeRequest, raw_delete_range_async_opt, - "raw_delete_range" -); -impl_request!( - RawCasRequest, - raw_compare_and_swap_async_opt, - "raw_compare_and_swap" + "raw_delete_range", + RawDeleteRange ); + +// TODO +// impl_request!( +// RawCasRequest, +// raw_compare_and_swap_async_opt, +// "raw_compare_and_swap", +// RawCas +// ); + impl_request!( RawCoprocessorRequest, raw_coprocessor_async_opt, - "raw_coprocessor" + "raw_coprocessor", + RawCoprocessor ); -impl_request!(GetRequest, kv_get_async_opt, "kv_get"); -impl_request!(ScanRequest, kv_scan_async_opt, "kv_scan"); -impl_request!(PrewriteRequest, kv_prewrite_async_opt, "kv_prewrite"); -impl_request!(CommitRequest, kv_commit_async_opt, "kv_commit"); -impl_request!(CleanupRequest, kv_cleanup_async_opt, "kv_cleanup"); -impl_request!(BatchGetRequest, kv_batch_get_async_opt, "kv_batch_get"); +impl_request!(GetRequest, kv_get_async_opt, "kv_get", Get); +impl_request!(ScanRequest, kv_scan_async_opt, "kv_scan", Scan); +impl_request!( + PrewriteRequest, + kv_prewrite_async_opt, + "kv_prewrite", + Prewrite +); +impl_request!(CommitRequest, kv_commit_async_opt, "kv_commit", Commit); +impl_request!(CleanupRequest, kv_cleanup_async_opt, "kv_cleanup", Cleanup); +impl_request!( + BatchGetRequest, + kv_batch_get_async_opt, + "kv_batch_get", + BatchGet +); impl_request!( BatchRollbackRequest, kv_batch_rollback_async_opt, - "kv_batch_rollback" + "kv_batch_rollback", + BatchRollback ); impl_request!( PessimisticRollbackRequest, kv_pessimistic_rollback_async_opt, - "kv_pessimistic_rollback" + "kv_pessimistic_rollback", + PessimisticRollback ); impl_request!( ResolveLockRequest, kv_resolve_lock_async_opt, - "kv_resolve_lock" + "kv_resolve_lock", + ResolveLock +); +impl_request!( + ScanLockRequest, + kv_scan_lock_async_opt, + "kv_scan_lock", + ScanLock ); -impl_request!(ScanLockRequest, kv_scan_lock_async_opt, "kv_scan_lock"); impl_request!( PessimisticLockRequest, kv_pessimistic_lock_async_opt, - "kv_pessimistic_lock" + "kv_pessimistic_lock", + PessimisticLock ); impl_request!( TxnHeartBeatRequest, kv_txn_heart_beat_async_opt, - "kv_txn_heart_beat" + "kv_txn_heart_beat", + TxnHeartBeat ); impl_request!( CheckTxnStatusRequest, kv_check_txn_status_async_opt, - "kv_check_txn_status" + "kv_check_txn_status", + CheckTxnStatus ); impl_request!( CheckSecondaryLocksRequest, kv_check_secondary_locks_async_opt, - "kv_check_secondary_locks_request" + "kv_check_secondary_locks_request", + CheckSecondaryLocks ); -impl_request!(GcRequest, kv_gc_async_opt, "kv_gc"); +impl_request!(GcRequest, kv_gc_async_opt, "kv_gc", Gc); impl_request!( DeleteRangeRequest, kv_delete_range_async_opt, - "kv_delete_range" + "kv_delete_range", + DeleteRange ); + +#[async_trait] +impl Request for BatchCommandsRequest { + async fn dispatch(&self, client: &TikvClient, options: CallOption) -> Result> { + match client.batch_commands_opt(options) { + Ok((mut tx, mut rx)) => { + tx.send((self.clone(), WriteFlags::default())).await?; + tx.close().await?; + let resp = rx.try_next().await?.unwrap(); + Ok(Box::new(resp) as Box) + } + Err(e) => Err(Error::Grpc(e)), + } + } + fn label(&self) -> &'static str { + "kv_batch_commands" + } + fn as_any(&self) -> &dyn Any { + self + } + fn set_context(&mut self, _: tikv_client_proto::kvrpcpb::Context) { + todo!() + } + fn into_batch(&self) -> BatchCommandsRequest { + self.clone() + } +} From d30632936d33889ed0c9b86ff5500648f4b79aef Mon Sep 17 00:00:00 2001 From: yongman Date: Mon, 8 Aug 2022 21:58:17 +0800 Subject: [PATCH 02/14] Add batch worker to tikv-client-store Signed-off-by: yongman --- src/config.rs | 27 ---- src/pd/client.rs | 6 +- src/request/mod.rs | 1 + src/request/plan.rs | 40 ++--- tikv-client-common/src/security.rs | 6 +- tikv-client-pd/src/cluster.rs | 6 +- tikv-client-store/Cargo.toml | 2 + .../src}/batch.rs | 138 +++++++++++------- tikv-client-store/src/client.rs | 93 ++++++++++-- tikv-client-store/src/lib.rs | 2 + tikv-client-store/src/request.rs | 13 +- 11 files changed, 195 insertions(+), 139 deletions(-) rename {src/request => tikv-client-store/src}/batch.rs (52%) diff --git a/src/config.rs b/src/config.rs index e42ff2fb..02ca2883 100644 --- a/src/config.rs +++ b/src/config.rs @@ -27,33 +27,6 @@ pub struct Config { pub kv_client_config: KVClientConfig, } -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -#[serde(default)] -#[serde(rename_all = "kebab-case")] -pub struct KVClientConfig { - pub completion_queue_size: usize, - pub grpc_keepalive_time: Duration, - pub grpc_keepalive_timeout: Duration, - pub allow_batch: bool, - pub overload_threshold: usize, - pub max_batch_wait_time: Duration, - pub max_batch_size: usize, -} - -impl Default for KVClientConfig { - fn default() -> Self { - Self { - completion_queue_size: DEFAULT_GRPC_COMPLETION_QUEUE_SIZE, - grpc_keepalive_time: DEFAULT_GRPC_KEEPALIVE_TIME, - grpc_keepalive_timeout: DEFAULT_GRPC_KEEPALIVE_TIMEOUT, - allow_batch: true, - overload_threshold: DEFAULT_OVERLOAD_THRESHOLD, - max_batch_wait_time: DEFAULT_MAX_BATCH_WAIT_TIME, - max_batch_size: DEFAULT_MAX_BATCH_SIZE, - } - } -} - impl Default for Config { fn default() -> Self { Config { diff --git a/src/pd/client.rs b/src/pd/client.rs index 75610544..5cc19664 100644 --- a/src/pd/client.rs +++ b/src/pd/client.rs @@ -2,6 +2,7 @@ use crate::{ compat::stream_fn, + config::KVClientConfig, kv::codec, pd::{retry::RetryClientTrait, RetryClient}, region::{RegionId, RegionVerId, RegionWithLeader}, @@ -210,6 +211,7 @@ pub struct PdRpcClient>, kv_connect: KvC, kv_client_cache: Arc>>, + kv_config: KVClientConfig, enable_codec: bool, region_cache: RegionCache>, logger: Logger, @@ -293,6 +295,7 @@ impl PdRpcClient { pub async fn new( config: Config, kv_connect: MakeKvC, + kv_config: KVClientConfig, pd: MakePd, enable_codec: bool, logger: Logger, @@ -324,6 +327,7 @@ impl PdRpcClient { pd: pd.clone(), kv_client_cache, kv_connect: kv_connect(env, security_mgr), + kv_config, enable_codec, region_cache: RegionCache::new(pd), logger, @@ -335,7 +339,7 @@ impl PdRpcClient { return Ok(client.clone()); }; info!(self.logger, "connect to tikv endpoint: {:?}", address); - match self.kv_connect.connect(address) { + match self.kv_connect.connect(address, kv_config) { Ok(client) => { self.kv_client_cache .write() diff --git a/src/request/mod.rs b/src/request/mod.rs index eb6d24ef..ef2ba4d5 100644 --- a/src/request/mod.rs +++ b/src/request/mod.rs @@ -22,6 +22,7 @@ pub mod plan; mod plan_builder; #[macro_use] mod shard; +pub mod batch; /// Abstracts any request sent to a TiKV server. #[async_trait] diff --git a/src/request/plan.rs b/src/request/plan.rs index 091331c9..67a5bba9 100644 --- a/src/request/plan.rs +++ b/src/request/plan.rs @@ -50,35 +50,17 @@ impl Plan for Dispatch { async fn execute(&self) -> Result { let stats = tikv_stats(self.request.label()); - - let c = &self.kv_config; - - // Do a super batch request if possible if batching is enabled. - if c.allow_batch && c.max_batch_size > 1 { - // Generate a entry for batch request and response - - // Send the request to the batch stream channel and wait for response - let result = self - .kv_client - .as_ref() - .expect("Unreachable: kv_client has not been initialised in Dispatch") - .dispatch(&self.request) - .await?; - - Ok(Result) - } else { - let result = self - .kv_client - .as_ref() - .expect("Unreachable: kv_client has not been initialised in Dispatch") - .dispatch(&self.request) - .await; - let result = stats.done(result); - result.map(|r| { - *r.downcast() - .expect("Downcast failed: request and response type mismatch") - }) - } + let result = self + .kv_client + .as_ref() + .expect("Unreachable: kv_client has not been initialised in Dispatch") + .dispatch(Box::new(self.request)) + .await; + let result = stats.done(result); + result.map(|r| { + *r.downcast() + .expect("Downcast failed: request and response type mismatch") + }) } fn kv_config(&self) -> KVClientConfig { diff --git a/tikv-client-common/src/security.rs b/tikv-client-common/src/security.rs index a89b2fca..37aa480e 100644 --- a/tikv-client-common/src/security.rs +++ b/tikv-client-common/src/security.rs @@ -67,6 +67,8 @@ impl SecurityManager { &self, env: Arc, addr: &str, + keepalive: u64, + keepalive_timeout: u64, factory: Factory, ) -> Result where @@ -77,8 +79,8 @@ impl SecurityManager { let addr = SCHEME_REG.replace(addr, ""); let cb = ChannelBuilder::new(env) - .keepalive_time(Duration::from_secs(10)) - .keepalive_timeout(Duration::from_secs(3)) + .keepalive_time(Duration::from_millis(keepalive)) + .keepalive_timeout(Duration::from_millis(keepalive_timeout)) .use_local_subchannel_pool(true); let channel = if self.ca.is_empty() { diff --git a/tikv-client-pd/src/cluster.rs b/tikv-client-pd/src/cluster.rs index 063532e9..a4e49eb2 100644 --- a/tikv-client-pd/src/cluster.rs +++ b/tikv-client-pd/src/cluster.rs @@ -180,9 +180,9 @@ impl Connection { addr: &str, timeout: Duration, ) -> Result<(pdpb::PdClient, pdpb::GetMembersResponse)> { - let client = self - .security_mgr - .connect(self.env.clone(), addr, pdpb::PdClient::new)?; + let client = + self.security_mgr + .connect(self.env.clone(), addr, 10000, 2000, pdpb::PdClient::new)?; let option = CallOption::default().timeout(timeout); let resp = client .get_members_async_opt(&pdpb::GetMembersRequest::default(), option) diff --git a/tikv-client-store/Cargo.toml b/tikv-client-store/Cargo.toml index efdc18c5..1b1883d2 100644 --- a/tikv-client-store/Cargo.toml +++ b/tikv-client-store/Cargo.toml @@ -13,5 +13,7 @@ derive-new = "0.5" futures = { version = "0.3", features = ["compat", "async-await", "thread-pool"] } grpcio = { version = "0.10", features = [ "prost-codec" ], default-features = false } log = "0.4" +serde = "1.0" +serde_derive = "1.0" tikv-client-common = { version = "0.1.0", path = "../tikv-client-common" } tikv-client-proto = { version = "0.1.0", path = "../tikv-client-proto" } diff --git a/src/request/batch.rs b/tikv-client-store/src/batch.rs similarity index 52% rename from src/request/batch.rs rename to tikv-client-store/src/batch.rs index 2acf4ab6..8bb2782f 100644 --- a/src/request/batch.rs +++ b/tikv-client-store/src/batch.rs @@ -1,4 +1,5 @@ -use crate::{Error, Result}; +use crate::{Error, Request, Result}; +use core::any::Any; use futures::{ channel::{mpsc, oneshot}, executor::block_on, @@ -6,36 +7,54 @@ use futures::{ prelude::*, task::{AtomicWaker, Context, Poll}, }; -use grpcio::WriteFlags; +use grpcio::{CallOption, WriteFlags}; use log::debug; -use std::{cell::RefCell, collections::VecDeque, pin::Pin, rc::Rc, thread}; +use std::{ + cell::RefCell, + collections::HashMap, + pin::Pin, + rc::Rc, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + thread, + time::{Duration, Instant}, +}; use tikv_client_common::internal_err; -use tikv_client_proto::pdpb::*; +use tikv_client_proto::tikvpb::{BatchCommandsRequest, BatchCommandsResponse, TikvClient}; + +static ID_ALLOC: AtomicU64 = AtomicU64::new(0); -pub struct BatchCommandRequestEntry { - cmd: BatchCommandsRequest, - tx: oneshot::Sender, - // TODO +type Response = oneshot::Sender>>; +pub struct RequestEntry { + cmd: Box, + tx: Response, id: u64, } -impl BatchCommandRequestEntry { - pub fn new(cmd: BatchCommandsRequest, tx: oneshot::Sender) -> Self { - Self { cmd, tx, id: 0 } +impl RequestEntry { + pub fn new(cmd: Box, tx: Response) -> Self { + Self { + cmd, + tx, + id: ID_ALLOC.fetch_add(1, Ordering::Relaxed), + } } } /// BatchWorker provides request in batch and return the result in batch. +#[derive(Clone)] pub struct BatchWorker { - request_tx: mpsc::Sender, + request_tx: mpsc::Sender, } impl BatchWorker { pub fn new( - kv_client: &KvClient, + kv_client: Arc, max_batch_size: usize, max_inflight_requests: usize, - max_deplay_duration: std::time::Duration, + max_delay_duration: u64, options: CallOption, ) -> Result { let (request_tx, request_rx) = mpsc::channel(max_batch_size); @@ -51,75 +70,72 @@ impl BatchWorker { request_rx, max_batch_size, max_inflight_requests, + max_delay_duration, )) }); Ok(BatchWorker { request_tx }) } - pub async fn dispatch( - mut self, - request: BatchCommandsRequest, - ) -> Result { + pub async fn dispatch(mut self, request: Box) -> Result> { let (tx, rx) = oneshot::channel(); // Generate BatchCommandRequestEntry - let entry = BatchCommandRequestEntry::new(request, tx); + let entry = RequestEntry::new(request, tx); // Send request entry to the background thread to handle the request, response will be // received in rx channel. self.request_tx .send(entry) .await - .map_err(|_| Error::Internal("Failed to send request to batch worker".to_owned()))?; - Ok(rx.await?) + .map_err(|_| internal_err!("Failed to send request to batch worker".to_owned()))?; + Ok(Box::new(rx.await?)) } } async fn run_batch_worker( mut tx: impl Sink<(BatchCommandsRequest, WriteFlags), Error = Error> + Unpin, mut rx: impl Stream> + Unpin, - request_rx: mpsc::Receiver, + request_rx: mpsc::Receiver, max_batch_size: usize, max_inflight_requests: usize, - max_deplay_duration: std::time::Duration, + max_delay_duration: u64, ) { // Inflight requests which are waiting for the response from rpc server - let mut inflight_requests = - Rc::new(RefCell::new(VecDeque::with_capacity(max_inflight_requests))); + let inflight_requests = Rc::new(RefCell::new(HashMap::new())); let waker = Rc::new(AtomicWaker::new()); - // TODO pin_mut!(request_rx); - let mut request_stream = BatchCommandsRequestStream { request_rx, inflight_requests: inflight_requests.clone(), self_waker: waker.clone(), max_batch_size, max_inflight_requests, - max_deplay_duration, - } - .map(Ok); + max_delay_duration, + }; let send_requests = tx.send_all(&mut request_stream); let recv_handle_response = async move { - while let Some(Ok(resp)) = rx.next().await { + while let Some(Ok(mut batch_resp)) = rx.next().await { + // TODO resp is BatchCommandsResponse, split it into responses let mut inflight_requests = inflight_requests.borrow_mut(); - if inflight_requests.len == max_inflight_requests { + if inflight_requests.len() == max_inflight_requests { waker.wake(); } // TODO more handle to response - - let request_id = resp.get_header().get_request_id(); - let request = inflight_requests.pop_front().unwrap(); - debug!( - "Received response for request_id {}: {:?}", - request_id, resp - ); - request.tx.send(resp).unwrap(); + for (id, resp) in batch_resp + .take_request_ids() + .into_iter() + .zip(batch_resp.take_responses()) + { + if let Some(tx) = inflight_requests.remove(&id) { + debug!("Received response for request_id {}: {:?}", id, resp.cmd); + tx.send(Ok(Box::new(resp.cmd))).unwrap(); + } + } } }; @@ -128,42 +144,42 @@ async fn run_batch_worker( debug!("Batch receiver finished: {:?}", rx_res); } -struct BatchCommandsRequestStream { - request_rx: mpsc::Receiver, - inflight_requests: Rc>>>, +struct BatchCommandsRequestStream<'a> { + request_rx: Pin<&'a mut mpsc::Receiver>, + inflight_requests: Rc>>, self_waker: Rc, max_batch_size: usize, max_inflight_requests: usize, - max_deplay_duration: std::time::Duration, + max_delay_duration: u64, } -impl Stream for BatchCommandsRequestStream { +impl Stream for BatchCommandsRequestStream<'_> { type Item = Result<(BatchCommandsRequest, WriteFlags)>; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let mut inflight_requests = self.inflight_requests.clone().borrow_mut(); + let inflight_requests = self.inflight_requests.clone(); + let mut inflight_requests = inflight_requests.borrow_mut(); if inflight_requests.len() == 0 { self.self_waker.register(cx.waker()); } // Collect user requests - let requests = vec![]; + let mut requests = vec![]; + let mut request_ids = vec![]; let latency_timer = Instant::now(); while requests.len() < self.max_batch_size && inflight_requests.len() < self.max_inflight_requests { // We can not wait longger than max_deplay_duration - if latency_timer.elapsed() > self.max_deplay_duration { + if latency_timer.elapsed() > Duration::from_millis(self.max_delay_duration) { break; } - match self.request_rx.poll_next(cx) { - Poll::Ready(Some(Ok(entry))) => { - inflight_requests.push_back(entry.tx); - requests.push(entry.cmd); - } - Poll::Ready(Some(Err(e))) => { - return Poll::Ready(Some(Err(e))); + match self.request_rx.as_mut().poll_next(cx) { + Poll::Ready(Some(entry)) => { + inflight_requests.insert(entry.id, entry.tx); + requests.push(entry.cmd.to_batch_request()); + request_ids.push(entry.id); } Poll::Ready(None) => { return Poll::Ready(None); @@ -174,12 +190,22 @@ impl Stream for BatchCommandsRequestStream { } } + // The requests is the commands will be convert to a batch request if !requests.is_empty() { + // TODO generate BatchCommandsRequest from requests + let mut batch_request = BatchCommandsRequest::new_(); + batch_request.set_requests(requests); + // TODO request id + batch_request.set_request_ids(request_ids); let write_flags = WriteFlags::default().buffer_hint(false); - Poll::Ready(Some((requests, write_flags))) + Poll::Ready(Some(Ok((batch_request, write_flags)))) } else { self.self_waker.register(cx.waker()); Poll::Pending } } + + fn size_hint(&self) -> (usize, Option) { + (0, None) + } } diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index 5e3534c5..43ff008c 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -1,17 +1,54 @@ // Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0. -use crate::{request::Request, Result, SecurityManager}; +use crate::{batch::BatchWorker, request::Request, Result, SecurityManager}; use async_trait::async_trait; use derive_new::new; use grpcio::{CallOption, Environment}; +use serde_derive::{Deserialize, Serialize}; use std::{any::Any, sync::Arc, time::Duration}; use tikv_client_proto::tikvpb::TikvClient; +const DEFAULT_REQUEST_TIMEOUT: u64 = 2000; +const DEFAULT_GRPC_KEEPALIVE_TIME: u64 = 10000; +const DEFAULT_GRPC_KEEPALIVE_TIMEOUT: u64 = 3000; +const DEFAULT_GRPC_COMPLETION_QUEUE_SIZE: usize = 1; +const DEFAULT_MAX_BATCH_WAIT_TIME: u64 = 100; +const DEFAULT_MAX_BATCH_SIZE: usize = 8; +const DEFAULT_OVERLOAD_THRESHOLD: usize = 200; /// A trait for connecting to TiKV stores. pub trait KvConnect: Sized + Send + Sync + 'static { type KvClient: KvClient + Clone + Send + Sync + 'static; - fn connect(&self, address: &str) -> Result; + fn connect(&self, address: &str, kv_config: KVClientConfig) -> Result; +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default)] +#[serde(rename_all = "kebab-case")] +pub struct KVClientConfig { + pub request_timeout: u64, + pub completion_queue_size: usize, + pub grpc_keepalive_time: u64, + pub grpc_keepalive_timeout: u64, + pub allow_batch: bool, + pub overload_threshold: usize, + pub max_batch_wait_time: u64, + pub max_batch_size: usize, +} + +impl Default for KVClientConfig { + fn default() -> Self { + Self { + request_timeout: DEFAULT_REQUEST_TIMEOUT, + completion_queue_size: DEFAULT_GRPC_COMPLETION_QUEUE_SIZE, + grpc_keepalive_time: DEFAULT_GRPC_KEEPALIVE_TIME, + grpc_keepalive_timeout: DEFAULT_GRPC_KEEPALIVE_TIMEOUT, + allow_batch: true, + overload_threshold: DEFAULT_OVERLOAD_THRESHOLD, + max_batch_wait_time: DEFAULT_MAX_BATCH_WAIT_TIME, + max_batch_size: DEFAULT_MAX_BATCH_SIZE, + } + } } #[derive(new, Clone)] @@ -24,16 +61,40 @@ pub struct TikvConnect { impl KvConnect for TikvConnect { type KvClient = KvRpcClient; - fn connect(&self, address: &str) -> Result { + fn connect(&self, address: &str, kv_config: KVClientConfig) -> Result { self.security_mgr - .connect(self.env.clone(), address, TikvClient::new) - .map(|c| KvRpcClient::new(Arc::new(c), self.timeout)) + .connect( + self.env.clone(), + address, + kv_config.grpc_keepalive_time, + kv_config.grpc_keepalive_timeout, + TikvClient::new, + ) + .map(|c| { + // Create batch worker if needed + let c = Arc::new(c); + let batch_worker = if kv_config.allow_batch { + Some( + BatchWorker::new( + c.clone(), + kv_config.max_batch_size, + kv_config.max_batch_size, + kv_config.max_batch_wait_time, + CallOption::default(), + ) + .unwrap(), + ) + } else { + None + }; + KvRpcClient::new(c.clone(), self.timeout, batch_worker) + }) } } #[async_trait] pub trait KvClient { - async fn dispatch(&self, req: &dyn Request) -> Result>; + async fn dispatch(&self, req: Box) -> Result>; } /// This client handles requests for a single TiKV node. It converts the data @@ -42,16 +103,22 @@ pub trait KvClient { pub struct KvRpcClient { rpc_client: Arc, timeout: Duration, + batch_worker: Option, } #[async_trait] impl KvClient for KvRpcClient { - async fn dispatch(&self, request: &dyn Request) -> Result> { - request - .dispatch( - &self.rpc_client, - CallOption::default().timeout(self.timeout), - ) - .await + async fn dispatch(&self, request: Box) -> Result> { + if let Some(batch_worker) = self.batch_worker.clone() { + batch_worker.dispatch(request).await + } else { + // Batch no needed if not batch enabled + request + .dispatch( + &self.rpc_client, + CallOption::default().timeout(self.timeout), + ) + .await + } } } diff --git a/tikv-client-store/src/lib.rs b/tikv-client-store/src/lib.rs index 5df938ff..15b26c28 100644 --- a/tikv-client-store/src/lib.rs +++ b/tikv-client-store/src/lib.rs @@ -1,11 +1,13 @@ // Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0. +mod batch; mod client; mod errors; mod request; #[doc(inline)] pub use crate::{ + batch::{BatchWorker, RequestEntry}, client::{KvClient, KvConnect, TikvConnect}, errors::{HasKeyErrors, HasRegionError, HasRegionErrors}, request::Request, diff --git a/tikv-client-store/src/request.rs b/tikv-client-store/src/request.rs index 0a3a7442..e6ed08e8 100644 --- a/tikv-client-store/src/request.rs +++ b/tikv-client-store/src/request.rs @@ -19,7 +19,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 into_batch(&self) -> BatchCommandsRequest; + fn to_batch_request(&self) -> batch_commands_request::Request; } macro_rules! impl_request { @@ -50,14 +50,11 @@ macro_rules! impl_request { kvrpcpb::$name::set_context(self, context) } - fn into_batch(&self) -> BatchCommandsRequest { + fn to_batch_request(&self) -> batch_commands_request::Request { let req = batch_commands_request::Request { cmd: Some($cmd(self.clone())), }; - BatchCommandsRequest { - requests: vec![req], - request_ids: vec![0], - } + req } } }; @@ -212,7 +209,7 @@ impl Request for BatchCommandsRequest { fn set_context(&mut self, _: tikv_client_proto::kvrpcpb::Context) { todo!() } - fn into_batch(&self) -> BatchCommandsRequest { - self.clone() + fn to_batch_request(&self) -> batch_commands_request::Request { + batch_commands_request::Request { cmd: None } } } From be230e59c4a8f657037494b7d367942b9f2f18a7 Mon Sep 17 00:00:00 2001 From: yongman Date: Tue, 9 Aug 2022 16:37:14 +0800 Subject: [PATCH 03/14] Add super batch request worked Signed-off-by: yongman --- src/config.rs | 12 +-- src/mock.rs | 6 +- src/pd/client.rs | 8 +- src/raw/client.rs | 96 ++++++++++----------- src/raw/requests.rs | 13 ++- src/request/mod.rs | 42 +++++---- src/request/plan.rs | 19 +---- src/request/plan_builder.rs | 4 +- src/store.rs | 18 ++-- src/transaction/client.rs | 26 +++--- src/transaction/lock.rs | 28 +++--- src/transaction/transaction.rs | 39 +++------ tikv-client-store/src/batch.rs | 29 +++---- tikv-client-store/src/client.rs | 4 +- tikv-client-store/src/lib.rs | 2 +- tikv-client-store/src/request.rs | 142 +++++++++++++++++++++++++++++-- 16 files changed, 279 insertions(+), 209 deletions(-) diff --git a/src/config.rs b/src/config.rs index 02ca2883..23663508 100644 --- a/src/config.rs +++ b/src/config.rs @@ -2,15 +2,9 @@ use serde_derive::{Deserialize, Serialize}; use std::{path::PathBuf, time::Duration}; +use tikv_client_store::KVClientConfig; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2); -const DEFAULT_GRPC_KEEPALIVE_TIME: Duration = Duration::from_secs(10); -const DEFAULT_GRPC_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(3); -const DEFAULT_GRPC_COMPLETION_QUEUE_SIZE: usize = 1; -const DEFAULT_MAX_BATCH_WAIT_TIME: Duration = Duration::from_millis(0); -const DEFAULT_MAX_BATCH_SIZE: usize = 8; -const DEFAULT_OVERLOAD_THRESHOLD: usize = 200; - /// The configuration for either a [`RawClient`](crate::RawClient) or a /// [`TransactionClient`](crate::TransactionClient). /// @@ -24,7 +18,7 @@ pub struct Config { pub cert_path: Option, pub key_path: Option, pub timeout: Duration, - pub kv_client_config: KVClientConfig, + pub kv_config: KVClientConfig, } impl Default for Config { @@ -34,7 +28,7 @@ impl Default for Config { cert_path: None, key_path: None, timeout: DEFAULT_REQUEST_TIMEOUT, - kv_client_config: KVClientConfig::default(), + kv_config: KVClientConfig::default(), } } } diff --git a/src/mock.rs b/src/mock.rs index 02a0bef9..94013b08 100644 --- a/src/mock.rs +++ b/src/mock.rs @@ -16,7 +16,7 @@ use derive_new::new; use slog::{Drain, Logger}; use std::{any::Any, sync::Arc}; use tikv_client_proto::metapb; -use tikv_client_store::{KvClient, KvConnect, Request}; +use tikv_client_store::{KVClientConfig, 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. @@ -78,7 +78,7 @@ pub struct MockPdClient { #[async_trait] impl KvClient for MockKvClient { - async fn dispatch(&self, req: &dyn Request) -> Result> { + async fn dispatch(&self, req: Box) -> Result> { match &self.dispatch { Some(f) => f(req.as_any()), None => panic!("no dispatch hook set"), @@ -89,7 +89,7 @@ impl KvClient for MockKvClient { impl KvConnect for MockKvConnect { type KvClient = MockKvClient; - fn connect(&self, address: &str) -> Result { + fn connect(&self, address: &str, kv_config: KVClientConfig) -> Result { Ok(MockKvClient { addr: address.to_owned(), dispatch: None, diff --git a/src/pd/client.rs b/src/pd/client.rs index 5cc19664..94f3c017 100644 --- a/src/pd/client.rs +++ b/src/pd/client.rs @@ -2,7 +2,6 @@ use crate::{ compat::stream_fn, - config::KVClientConfig, kv::codec, pd::{retry::RetryClientTrait, RetryClient}, region::{RegionId, RegionVerId, RegionWithLeader}, @@ -17,7 +16,7 @@ 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 tikv_client_store::{KVClientConfig, KvClient, KvConnect, TikvConnect}; use tokio::sync::RwLock; const CQ_COUNT: usize = 1; @@ -295,7 +294,6 @@ impl PdRpcClient { pub async fn new( config: Config, kv_connect: MakeKvC, - kv_config: KVClientConfig, pd: MakePd, enable_codec: bool, logger: Logger, @@ -327,7 +325,7 @@ impl PdRpcClient { pd: pd.clone(), kv_client_cache, kv_connect: kv_connect(env, security_mgr), - kv_config, + kv_config: config.kv_config, enable_codec, region_cache: RegionCache::new(pd), logger, @@ -339,7 +337,7 @@ impl PdRpcClient { return Ok(client.clone()); }; info!(self.logger, "connect to tikv endpoint: {:?}", address); - match self.kv_connect.connect(address, kv_config) { + match self.kv_connect.connect(address, self.kv_config.clone()) { Ok(client) => { self.kv_client_cache .write() diff --git a/src/raw/client.rs b/src/raw/client.rs index 3801510c..b0fa2f6d 100644 --- a/src/raw/client.rs +++ b/src/raw/client.rs @@ -6,10 +6,11 @@ use std::{str::FromStr, sync::Arc, u32}; use slog::{Drain, Logger}; use tikv_client_common::Error; use tikv_client_proto::metapb; +use tikv_client_store::KVClientConfig; use crate::{ backoff::DEFAULT_REGION_BACKOFF, - config::{Config, KVClientConfig}, + config::Config, pd::{PdClient, PdRpcClient}, raw::lowering::*, request::{Collect, CollectSingle, Plan}, @@ -116,7 +117,7 @@ impl Client { cf: None, atomic: false, logger, - kv_config: config.kv_client_config, + kv_config: config.kv_config, }) } @@ -201,12 +202,11 @@ impl Client { pub async fn get_opt(&self, key: impl Into, backoff: Backoff) -> Result> { debug!(self.logger, "invoking raw get request"); let request = new_raw_get_request(key.into(), self.cf.clone()); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .merge(CollectSingle) - .post_process_default() - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .merge(CollectSingle) + .post_process_default() + .plan(); plan.execute().await } @@ -243,11 +243,10 @@ impl Client { ) -> Result> { debug!(self.logger, "invoking raw batch_get request"); let request = new_raw_batch_get_request(keys.into_iter().map(Into::into), self.cf.clone()); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .merge(Collect) - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .merge(Collect) + .plan(); plan.execute() .await .map(|r| r.into_iter().map(Into::into).collect()) @@ -282,12 +281,11 @@ impl Client { ) -> Result<()> { debug!(self.logger, "invoking raw put request"); let request = new_raw_put_request(key.into(), value.into(), self.cf.clone(), self.atomic); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .merge(CollectSingle) - .extract_error() - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .merge(CollectSingle) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -328,11 +326,10 @@ impl Client { self.cf.clone(), self.atomic, ); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .extract_error() - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -362,12 +359,11 @@ impl Client { pub async fn delete_opt(&self, key: impl Into, backoff: Backoff) -> Result<()> { debug!(self.logger, "invoking raw delete request"); let request = new_raw_delete_request(key.into(), self.cf.clone(), self.atomic); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .merge(CollectSingle) - .extract_error() - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .merge(CollectSingle) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -403,11 +399,10 @@ impl Client { self.assert_non_atomic()?; let request = new_raw_batch_delete_request(keys.into_iter().map(Into::into), self.cf.clone()); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .extract_error() - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -440,11 +435,10 @@ impl Client { debug!(self.logger, "invoking raw delete_range request"); self.assert_non_atomic()?; let request = new_raw_delete_range_request(range.into(), self.cf.clone()); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .extract_error() - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .extract_error() + .plan(); plan.execute().await?; Ok(()) } @@ -653,7 +647,7 @@ impl Client { previous_value.into(), self.cf.clone(), ); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req, self.kv_config.clone()) + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req) .retry_multi_region(backoff) .merge(CollectSingle) .post_process_default() @@ -695,7 +689,7 @@ impl Client { ranges.into_iter().map(Into::into), request_builder, ); - let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req, self.kv_config.clone()) + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), req) .preserve_shard() .retry_multi_region(backoff) .post_process_default() @@ -718,11 +712,10 @@ impl Client { } let request = new_raw_scan_request(range.into(), limit, key_only, self.cf.clone()); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .merge(Collect) - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .merge(Collect) + .plan(); let res = plan.execute().await; res.map(|mut s| { s.truncate(limit as usize); @@ -750,11 +743,10 @@ impl Client { key_only, self.cf.clone(), ); - let plan = - crate::request::PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) - .retry_multi_region(backoff) - .merge(Collect) - .plan(); + let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request) + .retry_multi_region(backoff) + .merge(Collect) + .plan(); plan.execute().await } diff --git a/src/raw/requests.rs b/src/raw/requests.rs index a62bf1c9..fc289089 100644 --- a/src/raw/requests.rs +++ b/src/raw/requests.rs @@ -361,7 +361,11 @@ pub struct RawCoprocessorRequest { #[async_trait] impl Request for RawCoprocessorRequest { - async fn dispatch(&self, client: &TikvClient, options: CallOption) -> Result> { + async fn dispatch( + &self, + client: &TikvClient, + options: CallOption, + ) -> Result> { self.inner.dispatch(client, options).await } @@ -376,6 +380,10 @@ impl Request for RawCoprocessorRequest { fn set_context(&mut self, context: kvrpcpb::Context) { self.inner.set_context(context); } + + fn to_batch_request(&self) -> tikv_client_proto::tikvpb::batch_commands_request::Request { + todo!() + } } impl KvRequest for RawCoprocessorRequest { @@ -471,7 +479,6 @@ mod test { use super::*; use crate::{ backoff::{DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF}, - config::KVClientConfig, mock::{MockKvClient, MockPdClient}, request::Plan, Key, @@ -511,7 +518,7 @@ mod test { key_only: true, ..Default::default() }; - let plan = crate::request::PlanBuilder::new(client, scan, KVClientConfig::default()) + let plan = crate::request::PlanBuilder::new(client, scan) .resolve_lock(OPTIMISTIC_BACKOFF) .retry_multi_region(DEFAULT_REGION_BACKOFF) .merge(Collect) diff --git a/src/request/mod.rs b/src/request/mod.rs index ef2ba4d5..c1e8be24 100644 --- a/src/request/mod.rs +++ b/src/request/mod.rs @@ -22,7 +22,6 @@ pub mod plan; mod plan_builder; #[macro_use] mod shard; -pub mod batch; /// Abstracts any request sent to a TiKV server. #[async_trait] @@ -66,7 +65,6 @@ impl RetryOptions { mod test { use super::*; use crate::{ - config::KVClientConfig, mock::{MockKvClient, MockPdClient}, store::store_stream_for_keys, transaction::lowering::new_commit_request, @@ -122,6 +120,12 @@ mod test { fn set_context(&mut self, _: kvrpcpb::Context) { unreachable!(); } + + fn to_batch_request( + &self, + ) -> tikv_client_proto::tikvpb::batch_commands_request::Request { + todo!() + } } #[async_trait] @@ -167,12 +171,11 @@ mod test { |_: &dyn Any| Ok(Box::new(MockRpcResponse) as Box), ))); - let plan = - crate::request::PlanBuilder::new(pd_client.clone(), request, KVClientConfig::default()) - .resolve_lock(Backoff::no_jitter_backoff(1, 1, 3)) - .retry_multi_region(Backoff::no_jitter_backoff(1, 1, 3)) - .extract_error() - .plan(); + let plan = crate::request::PlanBuilder::new(pd_client.clone(), request) + .resolve_lock(Backoff::no_jitter_backoff(1, 1, 3)) + .retry_multi_region(Backoff::no_jitter_backoff(1, 1, 3)) + .extract_error() + .plan(); let _ = plan.execute().await; // Original call plus the 3 retries @@ -195,23 +198,18 @@ mod test { let req = new_commit_request(iter::once(key), Timestamp::default(), Timestamp::default()); // does not extract error - let plan = crate::request::PlanBuilder::new( - pd_client.clone(), - req.clone(), - KVClientConfig::default(), - ) - .resolve_lock(OPTIMISTIC_BACKOFF) - .retry_multi_region(OPTIMISTIC_BACKOFF) - .plan(); + let plan = crate::request::PlanBuilder::new(pd_client.clone(), req.clone()) + .resolve_lock(OPTIMISTIC_BACKOFF) + .retry_multi_region(OPTIMISTIC_BACKOFF) + .plan(); assert!(plan.execute().await.is_ok()); // extract error - let plan = - crate::request::PlanBuilder::new(pd_client.clone(), req, KVClientConfig::default()) - .resolve_lock(OPTIMISTIC_BACKOFF) - .retry_multi_region(OPTIMISTIC_BACKOFF) - .extract_error() - .plan(); + let plan = crate::request::PlanBuilder::new(pd_client.clone(), req) + .resolve_lock(OPTIMISTIC_BACKOFF) + .retry_multi_region(OPTIMISTIC_BACKOFF) + .extract_error() + .plan(); assert!(plan.execute().await.is_err()); } } diff --git a/src/request/plan.rs b/src/request/plan.rs index 67a5bba9..31e97cc8 100644 --- a/src/request/plan.rs +++ b/src/request/plan.rs @@ -11,7 +11,6 @@ use tokio::sync::Semaphore; use crate::{ backoff::Backoff, - config::KVClientConfig, pd::PdClient, request::{KvRequest, Shardable}, stats::tikv_stats, @@ -30,10 +29,6 @@ pub trait Plan: Sized + Clone + Sync + Send + 'static { /// Execute the plan. async fn execute(&self) -> Result; - - fn kv_config(&self) -> KVClientConfig { - KVClientConfig::default() - } } /// The simplest plan which just dispatches a request to a specific kv server. @@ -41,7 +36,6 @@ pub trait Plan: Sized + Clone + Sync + Send + 'static { pub struct Dispatch { pub request: Req, pub kv_client: Option>, - pub kv_config: KVClientConfig, } #[async_trait] @@ -54,7 +48,7 @@ impl Plan for Dispatch { .kv_client .as_ref() .expect("Unreachable: kv_client has not been initialised in Dispatch") - .dispatch(Box::new(self.request)) + .dispatch(Box::new(self.request.clone())) .await; let result = stats.done(result); result.map(|r| { @@ -62,10 +56,6 @@ impl Plan for Dispatch { .expect("Downcast failed: request and response type mismatch") }) } - - fn kv_config(&self) -> KVClientConfig { - self.kv_config.clone() - } } const MULTI_REGION_CONCURRENCY: usize = 16; @@ -426,7 +416,6 @@ where async fn execute(&self) -> Result { let mut result = self.inner.execute().await?; let mut clone = self.clone(); - let kv_config = self.kv_config(); loop { let locks = result.take_locks(); if locks.is_empty() { @@ -438,7 +427,7 @@ where } let pd_client = self.pd_client.clone(); - if resolve_locks(locks, pd_client.clone(), kv_config.clone()).await? { + if resolve_locks(locks, pd_client.clone()).await? { result = self.inner.execute().await?; } else { match clone.backoff.next_delay_duration() { @@ -451,10 +440,6 @@ where } } } - - fn kv_config(&self) -> KVClientConfig { - self.inner.kv_config() - } } /// When executed, the plan extracts errors from its inner plan, and returns an diff --git a/src/request/plan_builder.rs b/src/request/plan_builder.rs index 0da1fb01..ce269fac 100644 --- a/src/request/plan_builder.rs +++ b/src/request/plan_builder.rs @@ -3,7 +3,6 @@ use super::plan::PreserveShard; use crate::{ backoff::Backoff, - config::KVClientConfig, pd::PdClient, request::{ DefaultProcessor, Dispatch, ExtractError, KvRequest, Merge, MergeResponse, Plan, Process, @@ -32,12 +31,11 @@ pub struct Targetted; impl PlanBuilderPhase for Targetted {} impl PlanBuilder, NoTarget> { - pub fn new(pd_client: Arc, request: Req, kv_config: KVClientConfig) -> Self { + pub fn new(pd_client: Arc, request: Req) -> Self { PlanBuilder { pd_client, plan: Dispatch { request, - kv_config, kv_client: None, }, phantom: PhantomData, diff --git a/src/store.rs b/src/store.rs index 44c56a04..11696582 100644 --- a/src/store.rs +++ b/src/store.rs @@ -8,7 +8,7 @@ use std::{ sync::Arc, }; use tikv_client_proto::kvrpcpb; -use tikv_client_store::{KvClient, KvConnect, TikvConnect}; +use tikv_client_store::KvClient; #[derive(new, Clone)] pub struct RegionStore { @@ -16,15 +16,15 @@ pub struct RegionStore { pub client: Arc, } -pub trait KvConnectStore: KvConnect { - fn connect_to_store(&self, region: RegionWithLeader, address: String) -> Result { - log::info!("connect to tikv endpoint: {:?}", &address); - let client = self.connect(address.as_str())?; - Ok(RegionStore::new(region, Arc::new(client))) - } -} +// pub trait KvConnectStore: KvConnect { +// fn connect_to_store(&self, region: RegionWithLeader, address: String) -> Result { +// log::info!("connect to tikv endpoint: {:?}", &address); +// let client = self.connect(address.as_str())?; +// Ok(RegionStore::new(region, Arc::new(client))) +// } +// } -impl KvConnectStore for TikvConnect {} +// impl KvConnectStore for TikvConnect {} /// Maps keys to a stream of stores. `key_data` must be sorted in increasing order pub fn store_stream_for_keys( diff --git a/src/transaction/client.rs b/src/transaction/client.rs index 37000dfa..a4427ff0 100644 --- a/src/transaction/client.rs +++ b/src/transaction/client.rs @@ -3,7 +3,7 @@ use super::{requests::new_scan_lock_request, resolve_locks}; use crate::{ backoff::{DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF}, - config::{Config, KVClientConfig}, + config::Config, pd::{PdClient, PdRpcClient}, request::Plan, timestamp::TimestampExt, @@ -13,6 +13,7 @@ use crate::{ use slog::{Drain, Logger}; use std::{mem, sync::Arc}; use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; +use tikv_client_store::KVClientConfig; // FIXME: cargo-culted value const SCAN_LOCK_BATCH_SIZE: u32 = 1024; @@ -120,7 +121,7 @@ impl Client { Ok(Client { pd, logger, - kv_config: config.kv_client_config, + kv_config: config.kv_config, }) } @@ -253,12 +254,11 @@ impl Client { SCAN_LOCK_BATCH_SIZE, ); - let plan = - crate::request::PlanBuilder::new(self.pd.clone(), req, self.kv_config.clone()) - .resolve_lock(OPTIMISTIC_BACKOFF) - .retry_multi_region(DEFAULT_REGION_BACKOFF) - .merge(crate::request::Collect) - .plan(); + let plan = crate::request::PlanBuilder::new(self.pd.clone(), req) + .resolve_lock(OPTIMISTIC_BACKOFF) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .merge(crate::request::Collect) + .plan(); let res: Vec = plan.execute().await?; if res.is_empty() { @@ -271,7 +271,7 @@ impl Client { // resolve locks // FIXME: (1) this is inefficient (2) when region error occurred - resolve_locks(locks, self.pd.clone(), self.kv_config.clone()).await?; + resolve_locks(locks, self.pd.clone()).await?; // update safepoint to PD let res: bool = self @@ -287,12 +287,6 @@ impl Client { fn new_transaction(&self, timestamp: Timestamp, options: TransactionOptions) -> Transaction { let logger = self.logger.new(o!("child" => 1)); - Transaction::new( - timestamp, - self.pd.clone(), - options, - self.kv_config.clone(), - logger, - ) + Transaction::new(timestamp, self.pd.clone(), options, logger) } } diff --git a/src/transaction/lock.rs b/src/transaction/lock.rs index af2256e7..871effc6 100644 --- a/src/transaction/lock.rs +++ b/src/transaction/lock.rs @@ -2,7 +2,6 @@ use crate::{ backoff::{Backoff, DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF}, - config::KVClientConfig, pd::PdClient, region::RegionVerId, request::{CollectSingle, Plan}, @@ -29,7 +28,6 @@ const RESOLVE_LOCK_RETRY_LIMIT: usize = 10; pub async fn resolve_locks( locks: Vec, pd_client: Arc, - kv_config: KVClientConfig, ) -> Result { debug!("resolving locks"); let ts = pd_client.clone().get_timestamp().await?; @@ -64,13 +62,12 @@ pub async fn resolve_locks( Some(&commit_version) => commit_version, None => { let request = requests::new_cleanup_request(lock.primary_lock, lock.lock_version); - let plan = - crate::request::PlanBuilder::new(pd_client.clone(), request, kv_config.clone()) - .resolve_lock(OPTIMISTIC_BACKOFF) - .retry_multi_region(DEFAULT_REGION_BACKOFF) - .merge(CollectSingle) - .post_process_default() - .plan(); + let plan = crate::request::PlanBuilder::new(pd_client.clone(), request) + .resolve_lock(OPTIMISTIC_BACKOFF) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .merge(CollectSingle) + .post_process_default() + .plan(); let commit_version = plan.execute().await?; commit_versions.insert(lock.lock_version, commit_version); commit_version @@ -82,7 +79,6 @@ pub async fn resolve_locks( lock.lock_version, commit_version, pd_client.clone(), - kv_config.clone(), ) .await?; clean_regions @@ -98,7 +94,6 @@ async fn resolve_lock_with_retry( start_version: u64, commit_version: u64, pd_client: Arc, - kv_config: KVClientConfig, ) -> Result { debug!("resolving locks with retry"); // FIXME: Add backoff @@ -109,7 +104,7 @@ async fn resolve_lock_with_retry( let ver_id = store.region_with_leader.ver_id(); let request = requests::new_resolve_lock_request(start_version, commit_version); // The only place where single-region is used - let plan = crate::request::PlanBuilder::new(pd_client.clone(), request, kv_config.clone()) + let plan = crate::request::PlanBuilder::new(pd_client.clone(), request) .single_region_with_store(store) .await? .resolve_lock(Backoff::no_backoff()) @@ -170,16 +165,15 @@ mod tests { let key = vec![1]; let region1 = MockPdClient::region1(); - let resolved_region = - resolve_lock_with_retry(&key, 1, 2, client.clone(), KVClientConfig::default()) - .await - .unwrap(); + let resolved_region = resolve_lock_with_retry(&key, 1, 2, client.clone()) + .await + .unwrap(); assert_eq!(region1.ver_id(), resolved_region); // Test resolve lock over retry limit fail::cfg("region-error", "10*return").unwrap(); let key = vec![100]; - resolve_lock_with_retry(&key, 3, 4, client, KVClientConfig::default()) + resolve_lock_with_retry(&key, 3, 4, client) .await .expect_err("should return error"); } diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index 38c8fa4c..b600c70e 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -2,7 +2,6 @@ use crate::{ backoff::{Backoff, DEFAULT_REGION_BACKOFF}, - config::KVClientConfig, pd::{PdClient, PdRpcClient}, request::{ Collect, CollectError, CollectSingle, CollectWithShard, Plan, PlanBuilder, RetryOptions, @@ -65,7 +64,6 @@ pub struct Transaction { buffer: Buffer, rpc: Arc, options: TransactionOptions, - kv_config: KVClientConfig, is_heartbeat_started: bool, start_instant: Instant, logger: Logger, @@ -76,7 +74,6 @@ impl Transaction { timestamp: Timestamp, rpc: Arc, options: TransactionOptions, - kv_config: KVClientConfig, logger: Logger, ) -> Transaction { let status = if options.read_only { @@ -90,7 +87,6 @@ impl Transaction { buffer: Buffer::new(options.is_pessimistic()), rpc, options, - kv_config, is_heartbeat_started: false, start_instant: std::time::Instant::now(), logger, @@ -122,12 +118,11 @@ impl Transaction { let rpc = self.rpc.clone(); let key = key.into(); let retry_options = self.options.retry_options.clone(); - let kv_config = self.kv_config.clone(); self.buffer .get_or_else(key, |key| async move { let request = new_get_request(key, timestamp); - let plan = PlanBuilder::new(rpc, request, kv_config) + let plan = PlanBuilder::new(rpc, request) .resolve_lock(retry_options.lock_backoff) .retry_multi_region(DEFAULT_REGION_BACKOFF) .merge(CollectSingle) @@ -253,12 +248,11 @@ impl Transaction { let timestamp = self.timestamp.clone(); let rpc = self.rpc.clone(); let retry_options = self.options.retry_options.clone(); - let kv_config = self.kv_config.clone(); self.buffer .batch_get_or_else(keys.into_iter().map(|k| k.into()), move |keys| async move { let request = new_batch_get_request(keys, timestamp); - let plan = PlanBuilder::new(rpc, request, kv_config) + let plan = PlanBuilder::new(rpc, request) .resolve_lock(retry_options.lock_backoff) .retry_multi_region(retry_options.region_backoff) .merge(Collect) @@ -601,7 +595,6 @@ impl Transaction { self.timestamp.clone(), self.rpc.clone(), self.options.clone(), - self.kv_config.clone(), self.buffer.get_write_size() as u64, self.start_instant, self.logger.new(o!("child" => 1)), @@ -659,7 +652,6 @@ impl Transaction { self.timestamp.clone(), self.rpc.clone(), self.options.clone(), - self.kv_config.clone(), self.buffer.get_write_size() as u64, self.start_instant, self.logger.new(o!("child" => 1)), @@ -695,7 +687,7 @@ impl Transaction { primary_key, self.start_instant.elapsed().as_millis() as u64 + MAX_TTL, ); - let plan = PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + let plan = PlanBuilder::new(self.rpc.clone(), request) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .merge(CollectSingle) @@ -715,7 +707,6 @@ impl Transaction { let timestamp = self.timestamp.clone(); let rpc = self.rpc.clone(); let retry_options = self.options.retry_options.clone(); - let kv_config = self.kv_config.clone(); self.buffer .scan_and_fetch( @@ -726,7 +717,7 @@ impl Transaction { move |new_range, new_limit| async move { let request = new_scan_request(new_range, timestamp, new_limit, key_only, reverse); - let plan = PlanBuilder::new(rpc, request, kv_config) + let plan = PlanBuilder::new(rpc, request) .resolve_lock(retry_options.lock_backoff) .retry_multi_region(retry_options.region_backoff) .merge(Collect) @@ -781,7 +772,7 @@ impl Transaction { for_update_ts.clone(), need_value, ); - let plan = PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + let plan = PlanBuilder::new(self.rpc.clone(), request) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .preserve_shard() .retry_multi_region_preserve_results(self.options.retry_options.region_backoff.clone()) @@ -835,7 +826,7 @@ impl Transaction { start_version, for_update_ts, ); - let plan = PlanBuilder::new(self.rpc.clone(), req, self.kv_config.clone()) + let plan = PlanBuilder::new(self.rpc.clone(), req) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .extract_error() @@ -885,7 +876,6 @@ impl Transaction { HeartbeatOption::FixedTime(heartbeat_interval) => heartbeat_interval, }; let start_instant = self.start_instant; - let kv_config = self.kv_config.clone(); let heartbeat_task = async move { loop { @@ -906,7 +896,7 @@ impl Transaction { primary_key.clone(), start_instant.elapsed().as_millis() as u64 + MAX_TTL, ); - let plan = PlanBuilder::new(rpc.clone(), request, kv_config.clone()) + let plan = PlanBuilder::new(rpc.clone(), request) .retry_multi_region(region_backoff.clone()) .merge(CollectSingle) .plan(); @@ -1139,7 +1129,6 @@ struct Committer { start_version: Timestamp, rpc: Arc, options: TransactionOptions, - kv_config: KVClientConfig, #[new(default)] undetermined: bool, write_size: u64, @@ -1214,7 +1203,7 @@ impl Committer { .collect(); // FIXME set max_commit_ts and min_commit_ts - let plan = PlanBuilder::new(self.rpc.clone(), request, self.kv_config.clone()) + let plan = PlanBuilder::new(self.rpc.clone(), request) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .merge(CollectError) @@ -1254,7 +1243,7 @@ impl Committer { self.start_version.clone(), commit_version.clone(), ); - let plan = PlanBuilder::new(self.rpc.clone(), req, self.kv_config.clone()) + let plan = PlanBuilder::new(self.rpc.clone(), req) .resolve_lock(self.options.retry_options.lock_backoff.clone()) .retry_multi_region(self.options.retry_options.region_backoff.clone()) .extract_error() @@ -1291,7 +1280,7 @@ impl Committer { .filter(|key| &primary_key != key); new_commit_request(keys, self.start_version, commit_version) }; - let plan = PlanBuilder::new(self.rpc, req, self.kv_config.clone()) + let plan = PlanBuilder::new(self.rpc, req) .resolve_lock(self.options.retry_options.lock_backoff) .retry_multi_region(self.options.retry_options.region_backoff) .extract_error() @@ -1312,7 +1301,7 @@ impl Committer { match self.options.kind { TransactionKind::Optimistic => { let req = new_batch_rollback_request(keys, self.start_version); - let plan = PlanBuilder::new(self.rpc, req, self.kv_config.clone()) + let plan = PlanBuilder::new(self.rpc, req) .resolve_lock(self.options.retry_options.lock_backoff) .retry_multi_region(self.options.retry_options.region_backoff) .extract_error() @@ -1321,7 +1310,7 @@ impl Committer { } TransactionKind::Pessimistic(for_update_ts) => { let req = new_pessimistic_rollback_request(keys, self.start_version, for_update_ts); - let plan = PlanBuilder::new(self.rpc, req, self.kv_config.clone()) + let plan = PlanBuilder::new(self.rpc, req) .resolve_lock(self.options.retry_options.lock_backoff) .retry_multi_region(self.options.retry_options.region_backoff) .extract_error() @@ -1364,7 +1353,6 @@ enum TransactionStatus { #[cfg(test)] mod tests { use crate::{ - config::KVClientConfig, mock::{MockKvClient, MockPdClient}, transaction::HeartbeatOption, Transaction, TransactionOptions, @@ -1381,6 +1369,7 @@ mod tests { time::Duration, }; use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; + use tikv_client_store::KVClientConfig; #[tokio::test] async fn test_optimistic_heartbeat() -> Result<(), io::Error> { @@ -1414,7 +1403,6 @@ mod tests { pd_client, TransactionOptions::new_optimistic() .heartbeat_option(HeartbeatOption::FixedTime(Duration::from_secs(1))), - KVClientConfig::default(), logger.new(o!("child" => 1)), ); heartbeat_txn.put(key1.clone(), "foo").await.unwrap(); @@ -1464,7 +1452,6 @@ mod tests { pd_client, TransactionOptions::new_pessimistic() .heartbeat_option(HeartbeatOption::FixedTime(Duration::from_secs(1))), - KVClientConfig::default(), logger.new(o!("child" => 1)), ); heartbeat_txn.put(key1.clone(), "foo").await.unwrap(); diff --git a/tikv-client-store/src/batch.rs b/tikv-client-store/src/batch.rs index 8bb2782f..a17a7e7b 100644 --- a/tikv-client-store/src/batch.rs +++ b/tikv-client-store/src/batch.rs @@ -1,4 +1,4 @@ -use crate::{Error, Request, Result}; +use crate::{request::from_batch_commands_resp, Error, Request, Result}; use core::any::Any; use futures::{ channel::{mpsc, oneshot}, @@ -26,7 +26,7 @@ use tikv_client_proto::tikvpb::{BatchCommandsRequest, BatchCommandsResponse, Tik static ID_ALLOC: AtomicU64 = AtomicU64::new(0); -type Response = oneshot::Sender>>; +type Response = oneshot::Sender>; pub struct RequestEntry { cmd: Box, tx: Response, @@ -77,7 +77,7 @@ impl BatchWorker { Ok(BatchWorker { request_tx }) } - pub async fn dispatch(mut self, request: Box) -> Result> { + pub async fn dispatch(mut self, request: Box) -> Result> { let (tx, rx) = oneshot::channel(); // Generate BatchCommandRequestEntry let entry = RequestEntry::new(request, tx); @@ -87,7 +87,8 @@ impl BatchWorker { .send(entry) .await .map_err(|_| internal_err!("Failed to send request to batch worker".to_owned()))?; - Ok(Box::new(rx.await?)) + rx.await + .map_err(|_| internal_err!("Failed to receive response from batch worker".to_owned())) } } @@ -112,28 +113,28 @@ async fn run_batch_worker( max_batch_size, max_inflight_requests, max_delay_duration, - }; + } + .map(Ok); let send_requests = tx.send_all(&mut request_stream); let recv_handle_response = async move { while let Some(Ok(mut batch_resp)) = rx.next().await { - // TODO resp is BatchCommandsResponse, split it into responses let mut inflight_requests = inflight_requests.borrow_mut(); if inflight_requests.len() == max_inflight_requests { waker.wake(); } - // TODO more handle to response for (id, resp) in batch_resp .take_request_ids() .into_iter() .zip(batch_resp.take_responses()) { if let Some(tx) = inflight_requests.remove(&id) { - debug!("Received response for request_id {}: {:?}", id, resp.cmd); - tx.send(Ok(Box::new(resp.cmd))).unwrap(); + let inner_resp = from_batch_commands_resp(resp); + debug!("Received response for request_id {}", id); + tx.send(inner_resp.unwrap()).unwrap(); } } } @@ -154,7 +155,7 @@ struct BatchCommandsRequestStream<'a> { } impl Stream for BatchCommandsRequestStream<'_> { - type Item = Result<(BatchCommandsRequest, WriteFlags)>; + type Item = (BatchCommandsRequest, WriteFlags); fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let inflight_requests = self.inflight_requests.clone(); @@ -192,20 +193,14 @@ impl Stream for BatchCommandsRequestStream<'_> { // The requests is the commands will be convert to a batch request if !requests.is_empty() { - // TODO generate BatchCommandsRequest from requests let mut batch_request = BatchCommandsRequest::new_(); batch_request.set_requests(requests); - // TODO request id batch_request.set_request_ids(request_ids); let write_flags = WriteFlags::default().buffer_hint(false); - Poll::Ready(Some(Ok((batch_request, write_flags)))) + Poll::Ready(Some((batch_request, write_flags))) } else { self.self_waker.register(cx.waker()); Poll::Pending } } - - fn size_hint(&self) -> (usize, Option) { - (0, None) - } } diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index 43ff008c..6485440d 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -94,7 +94,7 @@ impl KvConnect for TikvConnect { #[async_trait] pub trait KvClient { - async fn dispatch(&self, req: Box) -> Result>; + async fn dispatch(&self, req: Box) -> Result>; } /// This client handles requests for a single TiKV node. It converts the data @@ -108,7 +108,7 @@ pub struct KvRpcClient { #[async_trait] impl KvClient for KvRpcClient { - async fn dispatch(&self, request: Box) -> Result> { + async fn dispatch(&self, request: Box) -> Result> { if let Some(batch_worker) = self.batch_worker.clone() { batch_worker.dispatch(request).await } else { diff --git a/tikv-client-store/src/lib.rs b/tikv-client-store/src/lib.rs index 15b26c28..bb59d406 100644 --- a/tikv-client-store/src/lib.rs +++ b/tikv-client-store/src/lib.rs @@ -8,7 +8,7 @@ mod request; #[doc(inline)] pub use crate::{ batch::{BatchWorker, RequestEntry}, - client::{KvClient, KvConnect, TikvConnect}, + client::{KVClientConfig, KvClient, KvConnect, TikvConnect}, errors::{HasKeyErrors, HasRegionError, HasRegionErrors}, request::Request, }; diff --git a/tikv-client-store/src/request.rs b/tikv-client-store/src/request.rs index e6ed08e8..1b96fda5 100644 --- a/tikv-client-store/src/request.rs +++ b/tikv-client-store/src/request.rs @@ -5,17 +5,22 @@ use async_trait::async_trait; use futures::{SinkExt, TryStreamExt}; use grpcio::{CallOption, WriteFlags}; use std::any::Any; +use tikv_client_common::internal_err; use tikv_client_proto::{ kvrpcpb, tikvpb::{ batch_commands_request::{self, request::Cmd::*}, - BatchCommandsRequest, TikvClient, + batch_commands_response, BatchCommandsRequest, TikvClient, }, }; #[async_trait] pub trait Request: Any + Sync + Send + 'static { - async fn dispatch(&self, client: &TikvClient, options: CallOption) -> Result>; + async fn dispatch( + &self, + client: &TikvClient, + options: CallOption, + ) -> Result>; fn label(&self) -> &'static str; fn as_any(&self) -> &dyn Any; fn set_context(&mut self, context: kvrpcpb::Context); @@ -30,11 +35,11 @@ macro_rules! impl_request { &self, client: &TikvClient, options: CallOption, - ) -> Result> { + ) -> Result> { client .$fun(self, options)? .await - .map(|r| Box::new(r) as Box) + .map(|r| Box::new(r) as Box) .map_err(Error::Grpc) } @@ -100,13 +105,40 @@ impl_request!( RawDeleteRange ); -// TODO +// TODO implement batchcommands support for rawcasrequest // impl_request!( // RawCasRequest, // raw_compare_and_swap_async_opt, // "raw_compare_and_swap", // RawCas // ); +#[async_trait] +impl Request for kvrpcpb::RawCasRequest { + async fn dispatch( + &self, + client: &TikvClient, + options: CallOption, + ) -> Result> { + client + .raw_compare_and_swap_async_opt(self, options)? + .await + .map(|r| Box::new(r) as Box) + .map_err(Error::Grpc) + } + fn label(&self) -> &'static str { + "raw_compare_and_swap" + } + fn as_any(&self) -> &dyn Any { + self + } + fn set_context(&mut self, _: tikv_client_proto::kvrpcpb::Context) { + todo!() + } + // TODO + fn to_batch_request(&self) -> batch_commands_request::Request { + batch_commands_request::Request { cmd: None } + } +} impl_request!( RawCoprocessorRequest, @@ -189,13 +221,17 @@ impl_request!( #[async_trait] impl Request for BatchCommandsRequest { - async fn dispatch(&self, client: &TikvClient, options: CallOption) -> Result> { + async fn dispatch( + &self, + client: &TikvClient, + options: CallOption, + ) -> Result> { match client.batch_commands_opt(options) { Ok((mut tx, mut rx)) => { tx.send((self.clone(), WriteFlags::default())).await?; tx.close().await?; let resp = rx.try_next().await?.unwrap(); - Ok(Box::new(resp) as Box) + Ok(Box::new(resp) as Box) } Err(e) => Err(Error::Grpc(e)), } @@ -213,3 +249,95 @@ impl Request for BatchCommandsRequest { batch_commands_request::Request { cmd: None } } } + +pub fn from_batch_commands_resp( + resp: batch_commands_response::Response, +) -> Result> { + match resp.cmd { + Some(batch_commands_response::response::Cmd::Get(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::Scan(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::Prewrite(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::Commit(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::Import(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::Cleanup(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::BatchGet(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::BatchRollback(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::ScanLock(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::ResolveLock(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::Gc(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::DeleteRange(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawGet(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawBatchGet(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawPut(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawBatchPut(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawDelete(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawBatchDelete(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawScan(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawDeleteRange(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawBatchScan(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::Coprocessor(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::PessimisticLock(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::PessimisticRollback(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::CheckTxnStatus(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::TxnHeartBeat(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::CheckSecondaryLocks(cmd)) => { + Ok(Box::new(cmd) as Box) + } + Some(batch_commands_response::response::Cmd::RawCoprocessor(cmd)) => { + Ok(Box::new(cmd) as Box) + } + _ => Err(internal_err!("batch_commands_resp.cmd is None".to_owned())), + } +} From b9af67613a8aadbd358b189efa052e9728b583fd Mon Sep 17 00:00:00 2001 From: yongman Date: Tue, 9 Aug 2022 18:11:33 +0800 Subject: [PATCH 04/14] Update kv client default config Signed-off-by: yongman --- src/config.rs | 39 +++++++++++++++++++++++++++++++++ src/pd/client.rs | 3 +-- tikv-client-store/src/client.rs | 6 ++--- 3 files changed, 43 insertions(+), 5 deletions(-) diff --git a/src/config.rs b/src/config.rs index 23663508..a0b9eebe 100644 --- a/src/config.rs +++ b/src/config.rs @@ -84,4 +84,43 @@ impl Config { } // TODO: add more config options for tivk client config + pub fn with_kv_timeout(mut self, timeout: u64) -> Self { + self.kv_config.request_timeout = timeout; + self + } + + pub fn with_kv_completion_queue_size(mut self, size: usize) -> Self { + self.kv_config.completion_queue_size = size; + self + } + + pub fn with_kv_grpc_keepalive_time(mut self, time: u64) -> Self { + self.kv_config.grpc_keepalive_time = time; + self + } + + pub fn with_kv_grpc_keepalive_timeout(mut self, timeout: u64) -> Self { + self.kv_config.grpc_keepalive_timeout = timeout; + self + } + + pub fn with_kv_allow_batch(mut self, allow_batch: bool) -> Self { + self.kv_config.allow_batch = allow_batch; + self + } + + pub fn with_kv_overload_threshold(mut self, threshold: usize) -> Self { + self.kv_config.overload_threshold = threshold; + self + } + + pub fn with_kv_max_batch_wait_time(mut self, wait: u64) -> Self { + self.kv_config.max_batch_wait_time = wait; + self + } + + pub fn with_kv_max_batch_size(mut self, size: usize) -> Self { + self.kv_config.max_batch_size = size; + self + } } diff --git a/src/pd/client.rs b/src/pd/client.rs index 94f3c017..0f1be0a7 100644 --- a/src/pd/client.rs +++ b/src/pd/client.rs @@ -19,7 +19,6 @@ use tikv_client_proto::{kvrpcpb, metapb}; use tikv_client_store::{KVClientConfig, KvClient, KvConnect, TikvConnect}; use tokio::sync::RwLock; -const CQ_COUNT: usize = 1; const CLIENT_PREFIX: &str = "tikv-client"; /// The PdClient handles all the encoding stuff. @@ -305,7 +304,7 @@ impl PdRpcClient { { let env = Arc::new( EnvBuilder::new() - .cq_count(CQ_COUNT) + .cq_count(config.kv_config.completion_queue_size) .name_prefix(thread_name(CLIENT_PREFIX)) .build(), ); diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index 6485440d..3a2cba71 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -12,8 +12,8 @@ const DEFAULT_REQUEST_TIMEOUT: u64 = 2000; const DEFAULT_GRPC_KEEPALIVE_TIME: u64 = 10000; const DEFAULT_GRPC_KEEPALIVE_TIMEOUT: u64 = 3000; const DEFAULT_GRPC_COMPLETION_QUEUE_SIZE: usize = 1; -const DEFAULT_MAX_BATCH_WAIT_TIME: u64 = 100; -const DEFAULT_MAX_BATCH_SIZE: usize = 8; +const DEFAULT_MAX_BATCH_WAIT_TIME: u64 = 10; +const DEFAULT_MAX_BATCH_SIZE: usize = 100; const DEFAULT_OVERLOAD_THRESHOLD: usize = 200; /// A trait for connecting to TiKV stores. pub trait KvConnect: Sized + Send + Sync + 'static { @@ -43,7 +43,7 @@ impl Default for KVClientConfig { completion_queue_size: DEFAULT_GRPC_COMPLETION_QUEUE_SIZE, grpc_keepalive_time: DEFAULT_GRPC_KEEPALIVE_TIME, grpc_keepalive_timeout: DEFAULT_GRPC_KEEPALIVE_TIMEOUT, - allow_batch: true, + allow_batch: false, overload_threshold: DEFAULT_OVERLOAD_THRESHOLD, max_batch_wait_time: DEFAULT_MAX_BATCH_WAIT_TIME, max_batch_size: DEFAULT_MAX_BATCH_SIZE, From 4a80ecbf7aea30db053dabb010e3a3e09c91358a Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 10 Aug 2022 21:02:13 +0800 Subject: [PATCH 05/14] Rename KvClientConfig Signed-off-by: yongman --- src/config.rs | 6 +++--- src/mock.rs | 4 ++-- src/pd/client.rs | 4 ++-- src/raw/client.rs | 6 +++--- src/transaction/client.rs | 4 ++-- src/transaction/transaction.rs | 2 +- tikv-client-store/src/client.rs | 8 ++++---- tikv-client-store/src/lib.rs | 2 +- 8 files changed, 18 insertions(+), 18 deletions(-) diff --git a/src/config.rs b/src/config.rs index a0b9eebe..aee7abdf 100644 --- a/src/config.rs +++ b/src/config.rs @@ -2,7 +2,7 @@ use serde_derive::{Deserialize, Serialize}; use std::{path::PathBuf, time::Duration}; -use tikv_client_store::KVClientConfig; +use tikv_client_store::KvClientConfig; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2); /// The configuration for either a [`RawClient`](crate::RawClient) or a @@ -18,7 +18,7 @@ pub struct Config { pub cert_path: Option, pub key_path: Option, pub timeout: Duration, - pub kv_config: KVClientConfig, + pub kv_config: KvClientConfig, } impl Default for Config { @@ -28,7 +28,7 @@ impl Default for Config { cert_path: None, key_path: None, timeout: DEFAULT_REQUEST_TIMEOUT, - kv_config: KVClientConfig::default(), + kv_config: KvClientConfig::default(), } } } diff --git a/src/mock.rs b/src/mock.rs index 94013b08..1fe97915 100644 --- a/src/mock.rs +++ b/src/mock.rs @@ -16,7 +16,7 @@ use derive_new::new; use slog::{Drain, Logger}; use std::{any::Any, sync::Arc}; use tikv_client_proto::metapb; -use tikv_client_store::{KVClientConfig, KvClient, KvConnect, Request}; +use tikv_client_store::{KvClient, KvClientConfig, KvConnect, Request}; /// Create a `PdRpcClient` with it's internals replaced with mocks so that the /// client can be tested without doing any RPC calls. @@ -89,7 +89,7 @@ impl KvClient for MockKvClient { impl KvConnect for MockKvConnect { type KvClient = MockKvClient; - fn connect(&self, address: &str, kv_config: KVClientConfig) -> Result { + fn connect(&self, address: &str, kv_config: KvClientConfig) -> Result { Ok(MockKvClient { addr: address.to_owned(), dispatch: None, diff --git a/src/pd/client.rs b/src/pd/client.rs index 0f1be0a7..9c586506 100644 --- a/src/pd/client.rs +++ b/src/pd/client.rs @@ -16,7 +16,7 @@ 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::{KVClientConfig, KvClient, KvConnect, TikvConnect}; +use tikv_client_store::{KvClientConfig, KvClient, KvConnect, TikvConnect}; use tokio::sync::RwLock; const CLIENT_PREFIX: &str = "tikv-client"; @@ -209,7 +209,7 @@ pub struct PdRpcClient>, kv_connect: KvC, kv_client_cache: Arc>>, - kv_config: KVClientConfig, + kv_config: KvClientConfig, enable_codec: bool, region_cache: RegionCache>, logger: Logger, diff --git a/src/raw/client.rs b/src/raw/client.rs index b0fa2f6d..eb980e39 100644 --- a/src/raw/client.rs +++ b/src/raw/client.rs @@ -6,7 +6,7 @@ use std::{str::FromStr, sync::Arc, u32}; use slog::{Drain, Logger}; use tikv_client_common::Error; use tikv_client_proto::metapb; -use tikv_client_store::KVClientConfig; +use tikv_client_store::KvClientConfig; use crate::{ backoff::DEFAULT_REGION_BACKOFF, @@ -32,7 +32,7 @@ pub struct Client { /// Whether to use the [`atomic mode`](Client::with_atomic_for_cas). atomic: bool, logger: Logger, - kv_config: KVClientConfig, + kv_config: KvClientConfig, } impl Clone for Client { @@ -799,7 +799,7 @@ mod tests { cf: Some(ColumnFamily::Default), atomic: false, logger, - kv_config: KVClientConfig::default(), + kv_config: KvClientConfig::default(), }; let resps = client .coprocessor( diff --git a/src/transaction/client.rs b/src/transaction/client.rs index a4427ff0..eac4eb45 100644 --- a/src/transaction/client.rs +++ b/src/transaction/client.rs @@ -13,7 +13,7 @@ use crate::{ use slog::{Drain, Logger}; use std::{mem, sync::Arc}; use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; -use tikv_client_store::KVClientConfig; +use tikv_client_store::KvClientConfig; // FIXME: cargo-culted value const SCAN_LOCK_BATCH_SIZE: u32 = 1024; @@ -37,7 +37,7 @@ const SCAN_LOCK_BATCH_SIZE: u32 = 1024; pub struct Client { pd: Arc, logger: Logger, - kv_config: KVClientConfig, + kv_config: KvClientConfig, } impl Clone for Client { diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index b600c70e..d8b0a611 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -1369,7 +1369,7 @@ mod tests { time::Duration, }; use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; - use tikv_client_store::KVClientConfig; + use tikv_client_store::KvClientConfig; #[tokio::test] async fn test_optimistic_heartbeat() -> Result<(), io::Error> { diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index 3a2cba71..1871c6c0 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -19,13 +19,13 @@ const DEFAULT_OVERLOAD_THRESHOLD: usize = 200; pub trait KvConnect: Sized + Send + Sync + 'static { type KvClient: KvClient + Clone + Send + Sync + 'static; - fn connect(&self, address: &str, kv_config: KVClientConfig) -> Result; + fn connect(&self, address: &str, kv_config: KvClientConfig) -> Result; } #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] #[serde(rename_all = "kebab-case")] -pub struct KVClientConfig { +pub struct KvClientConfig { pub request_timeout: u64, pub completion_queue_size: usize, pub grpc_keepalive_time: u64, @@ -36,7 +36,7 @@ pub struct KVClientConfig { pub max_batch_size: usize, } -impl Default for KVClientConfig { +impl Default for KvClientConfig { fn default() -> Self { Self { request_timeout: DEFAULT_REQUEST_TIMEOUT, @@ -61,7 +61,7 @@ pub struct TikvConnect { impl KvConnect for TikvConnect { type KvClient = KvRpcClient; - fn connect(&self, address: &str, kv_config: KVClientConfig) -> Result { + fn connect(&self, address: &str, kv_config: KvClientConfig) -> Result { self.security_mgr .connect( self.env.clone(), diff --git a/tikv-client-store/src/lib.rs b/tikv-client-store/src/lib.rs index bb59d406..514eb1d8 100644 --- a/tikv-client-store/src/lib.rs +++ b/tikv-client-store/src/lib.rs @@ -8,7 +8,7 @@ mod request; #[doc(inline)] pub use crate::{ batch::{BatchWorker, RequestEntry}, - client::{KVClientConfig, KvClient, KvConnect, TikvConnect}, + client::{KvClientConfig, KvClient, KvConnect, TikvConnect}, errors::{HasKeyErrors, HasRegionError, HasRegionErrors}, request::Request, }; From aab9f06e5c637cf0012270adaebe5454cd015349 Mon Sep 17 00:00:00 2001 From: yongman Date: Thu, 11 Aug 2022 10:48:26 +0800 Subject: [PATCH 06/14] Fix add send trait to tests Signed-off-by: yongman --- src/mock.rs | 8 +++++--- src/raw/client.rs | 4 ++-- src/raw/requests.rs | 6 +++--- src/request/mod.rs | 10 +++++----- src/transaction/lock.rs | 6 +++--- src/transaction/transaction.rs | 20 ++++++++++---------- tikv-client-store/src/request.rs | 8 ++++---- 7 files changed, 32 insertions(+), 30 deletions(-) diff --git a/src/mock.rs b/src/mock.rs index 1fe97915..6d5870ee 100644 --- a/src/mock.rs +++ b/src/mock.rs @@ -52,13 +52,15 @@ pub async fn pd_rpc_client() -> PdRpcClient { #[derive(new, Default, Clone)] pub struct MockKvClient { pub addr: String, - dispatch: Option Result> + Send + Sync + 'static>>, + dispatch: Option< + Arc Result> + Send + Sync + 'static>, + >, } impl MockKvClient { pub fn with_dispatch_hook(dispatch: F) -> MockKvClient where - F: Fn(&dyn Any) -> Result> + Send + Sync + 'static, + F: Fn(&(dyn Any + Send)) -> Result> + Send + Sync + 'static, { MockKvClient { addr: String::new(), @@ -78,7 +80,7 @@ pub struct MockPdClient { #[async_trait] impl KvClient for MockKvClient { - async fn dispatch(&self, req: Box) -> Result> { + async fn dispatch(&self, req: Box) -> Result> { match &self.dispatch { Some(f) => f(req.as_any()), None => panic!("no dispatch hook set"), diff --git a/src/raw/client.rs b/src/raw/client.rs index eb980e39..dcbca4d8 100644 --- a/src/raw/client.rs +++ b/src/raw/client.rs @@ -780,7 +780,7 @@ mod tests { o!(), ); let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( - move |req: &dyn Any| { + move |req: &(dyn Any + Send)| { if let Some(req) = req.downcast_ref::() { assert_eq!(req.copr_name, "example"); assert_eq!(req.copr_version_req, "0.1.0"); @@ -788,7 +788,7 @@ mod tests { data: req.data.clone(), ..Default::default() }; - Ok(Box::new(resp) as Box) + Ok(Box::new(resp) as Box) } else { unreachable!() } diff --git a/src/raw/requests.rs b/src/raw/requests.rs index fc289089..b8f0c4a8 100644 --- a/src/raw/requests.rs +++ b/src/raw/requests.rs @@ -373,7 +373,7 @@ impl Request for RawCoprocessorRequest { self.inner.label() } - fn as_any(&self) -> &dyn Any { + fn as_any(&self) -> &(dyn Any + Send) { self.inner.as_any() } @@ -491,7 +491,7 @@ mod test { #[ignore] fn test_raw_scan() { let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( - |req: &dyn Any| { + |req: &(dyn Any + Send)| { let req: &kvrpcpb::RawScanRequest = req.downcast_ref().unwrap(); assert!(req.key_only); assert_eq!(req.limit, 10); @@ -505,7 +505,7 @@ mod test { resp.kvs.push(kv); } - Ok(Box::new(resp) as Box) + Ok(Box::new(resp) as Box) }, ))); diff --git a/src/request/mod.rs b/src/request/mod.rs index c1e8be24..e429f907 100644 --- a/src/request/mod.rs +++ b/src/request/mod.rs @@ -105,7 +105,7 @@ mod test { #[async_trait] impl Request for MockKvRequest { - async fn dispatch(&self, _: &TikvClient, _: CallOption) -> Result> { + async fn dispatch(&self, _: &TikvClient, _: CallOption) -> Result> { Ok(Box::new(MockRpcResponse {})) } @@ -113,7 +113,7 @@ mod test { "mock" } - fn as_any(&self) -> &dyn Any { + fn as_any(&self) -> &(dyn Any + Send) { self } @@ -168,7 +168,7 @@ mod test { }; let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( - |_: &dyn Any| Ok(Box::new(MockRpcResponse) as Box), + |_: &(dyn Any + Send)| Ok(Box::new(MockRpcResponse) as Box), ))); let plan = crate::request::PlanBuilder::new(pd_client.clone(), request) @@ -185,12 +185,12 @@ mod test { #[tokio::test] async fn test_extract_error() { let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( - |_: &dyn Any| { + |_: &(dyn Any + Send)| { Ok(Box::new(kvrpcpb::CommitResponse { region_error: None, error: Some(kvrpcpb::KeyError::default()), commit_version: 0, - }) as Box) + }) as Box) }, ))); diff --git a/src/transaction/lock.rs b/src/transaction/lock.rs index 871effc6..05900284 100644 --- a/src/transaction/lock.rs +++ b/src/transaction/lock.rs @@ -151,15 +151,15 @@ mod tests { fail::cfg("region-error", "9*return").unwrap(); let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( - |_: &dyn Any| { + |_: &(dyn Any + Send)| { fail::fail_point!("region-error", |_| { let resp = kvrpcpb::ResolveLockResponse { region_error: Some(errorpb::Error::default()), ..Default::default() }; - Ok(Box::new(resp) as Box) + Ok(Box::new(resp) as Box) }); - Ok(Box::new(kvrpcpb::ResolveLockResponse::default()) as Box) + Ok(Box::new(kvrpcpb::ResolveLockResponse::default()) as Box) }, ))); diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index d8b0a611..7ee5049f 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -1369,7 +1369,6 @@ mod tests { time::Duration, }; use tikv_client_proto::{kvrpcpb, pdpb::Timestamp}; - use tikv_client_store::KvClientConfig; #[tokio::test] async fn test_optimistic_heartbeat() -> Result<(), io::Error> { @@ -1386,14 +1385,14 @@ mod tests { let heartbeats = Arc::new(AtomicUsize::new(0)); let heartbeats_cloned = heartbeats.clone(); let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( - move |req: &dyn Any| { + move |req: &(dyn Any + Send)| { if req.downcast_ref::().is_some() { heartbeats_cloned.fetch_add(1, Ordering::SeqCst); - Ok(Box::new(kvrpcpb::TxnHeartBeatResponse::default()) as Box) + Ok(Box::new(kvrpcpb::TxnHeartBeatResponse::default()) as Box) } else if req.downcast_ref::().is_some() { - Ok(Box::new(kvrpcpb::PrewriteResponse::default()) as Box) + Ok(Box::new(kvrpcpb::PrewriteResponse::default()) as Box) } else { - Ok(Box::new(kvrpcpb::CommitResponse::default()) as Box) + Ok(Box::new(kvrpcpb::CommitResponse::default()) as Box) } }, ))); @@ -1430,19 +1429,20 @@ mod tests { let heartbeats = Arc::new(AtomicUsize::new(0)); let heartbeats_cloned = heartbeats.clone(); let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( - move |req: &dyn Any| { + move |req: &(dyn Any + Send)| { if req.downcast_ref::().is_some() { heartbeats_cloned.fetch_add(1, Ordering::SeqCst); - Ok(Box::new(kvrpcpb::TxnHeartBeatResponse::default()) as Box) + Ok(Box::new(kvrpcpb::TxnHeartBeatResponse::default()) as Box) } else if req.downcast_ref::().is_some() { - Ok(Box::new(kvrpcpb::PrewriteResponse::default()) as Box) + Ok(Box::new(kvrpcpb::PrewriteResponse::default()) as Box) } else if req .downcast_ref::() .is_some() { - Ok(Box::new(kvrpcpb::PessimisticLockResponse::default()) as Box) + Ok(Box::new(kvrpcpb::PessimisticLockResponse::default()) + as Box) } else { - Ok(Box::new(kvrpcpb::CommitResponse::default()) as Box) + Ok(Box::new(kvrpcpb::CommitResponse::default()) as Box) } }, ))); diff --git a/tikv-client-store/src/request.rs b/tikv-client-store/src/request.rs index 1b96fda5..ea957310 100644 --- a/tikv-client-store/src/request.rs +++ b/tikv-client-store/src/request.rs @@ -22,7 +22,7 @@ pub trait Request: Any + Sync + Send + 'static { options: CallOption, ) -> Result>; fn label(&self) -> &'static str; - fn as_any(&self) -> &dyn Any; + fn as_any(&self) -> &(dyn Any + Send); fn set_context(&mut self, context: kvrpcpb::Context); fn to_batch_request(&self) -> batch_commands_request::Request; } @@ -47,7 +47,7 @@ macro_rules! impl_request { $label } - fn as_any(&self) -> &dyn Any { + fn as_any(&self) -> &(dyn Any + Send) { self } @@ -128,7 +128,7 @@ impl Request for kvrpcpb::RawCasRequest { fn label(&self) -> &'static str { "raw_compare_and_swap" } - fn as_any(&self) -> &dyn Any { + fn as_any(&self) -> &(dyn Any + Send) { self } fn set_context(&mut self, _: tikv_client_proto::kvrpcpb::Context) { @@ -239,7 +239,7 @@ impl Request for BatchCommandsRequest { fn label(&self) -> &'static str { "kv_batch_commands" } - fn as_any(&self) -> &dyn Any { + fn as_any(&self) -> &(dyn Any + Send) { self } fn set_context(&mut self, _: tikv_client_proto::kvrpcpb::Context) { From a60b6180eef8873a2c5b5e16a40c5c634c3e6be3 Mon Sep 17 00:00:00 2001 From: yongman Date: Thu, 11 Aug 2022 10:53:18 +0800 Subject: [PATCH 07/14] Add thread name for batch worker Signed-off-by: yongman --- src/mock.rs | 2 +- tikv-client-store/src/batch.rs | 24 ++++++++++++++---------- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/src/mock.rs b/src/mock.rs index 6d5870ee..107a31b7 100644 --- a/src/mock.rs +++ b/src/mock.rs @@ -91,7 +91,7 @@ impl KvClient for MockKvClient { impl KvConnect for MockKvConnect { type KvClient = MockKvClient; - fn connect(&self, address: &str, kv_config: KvClientConfig) -> Result { + fn connect(&self, address: &str, _: KvClientConfig) -> Result { Ok(MockKvClient { addr: address.to_owned(), dispatch: None, diff --git a/tikv-client-store/src/batch.rs b/tikv-client-store/src/batch.rs index a17a7e7b..58ba31b7 100644 --- a/tikv-client-store/src/batch.rs +++ b/tikv-client-store/src/batch.rs @@ -25,6 +25,7 @@ use tikv_client_common::internal_err; use tikv_client_proto::tikvpb::{BatchCommandsRequest, BatchCommandsResponse, TikvClient}; static ID_ALLOC: AtomicU64 = AtomicU64::new(0); +const BATCH_WORKER_NAME: &str = "batch-worker"; type Response = oneshot::Sender>; pub struct RequestEntry { @@ -63,16 +64,19 @@ impl BatchWorker { let (rpc_sender, rpc_receiver) = kv_client.batch_commands_opt(options)?; // Start a background thread to handle batch requests and responses - thread::spawn(move || { - block_on(run_batch_worker( - rpc_sender.sink_err_into(), - rpc_receiver.err_into(), - request_rx, - max_batch_size, - max_inflight_requests, - max_delay_duration, - )) - }); + thread::Builder::new() + .name(BATCH_WORKER_NAME.to_owned()) + .spawn(move || { + block_on(run_batch_worker( + rpc_sender.sink_err_into(), + rpc_receiver.err_into(), + request_rx, + max_batch_size, + max_inflight_requests, + max_delay_duration, + )) + }) + .unwrap(); Ok(BatchWorker { request_tx }) } From 0ebd740f94e9919fdba8641408276c825e970974 Mon Sep 17 00:00:00 2001 From: yongman Date: Thu, 11 Aug 2022 17:11:05 +0800 Subject: [PATCH 08/14] Add transport layer load threshold in batch Signed-off-by: yongman --- src/config.rs | 7 +++++- tikv-client-store/src/batch.rs | 40 ++++++++++++++++++++++++++++++--- tikv-client-store/src/client.rs | 12 ++++++---- 3 files changed, 51 insertions(+), 8 deletions(-) diff --git a/src/config.rs b/src/config.rs index aee7abdf..7f58a305 100644 --- a/src/config.rs +++ b/src/config.rs @@ -109,7 +109,7 @@ impl Config { self } - pub fn with_kv_overload_threshold(mut self, threshold: usize) -> Self { + pub fn with_kv_overload_threshold(mut self, threshold: u64) -> Self { self.kv_config.overload_threshold = threshold; self } @@ -123,4 +123,9 @@ impl Config { self.kv_config.max_batch_size = size; self } + + pub fn with_kv_max_inflight_requests(mut self, requests: usize) -> Self { + self.kv_config.max_inflight_requests = requests; + self + } } diff --git a/tikv-client-store/src/batch.rs b/tikv-client-store/src/batch.rs index 58ba31b7..62a9ec4e 100644 --- a/tikv-client-store/src/batch.rs +++ b/tikv-client-store/src/batch.rs @@ -31,14 +31,16 @@ type Response = oneshot::Sender>; pub struct RequestEntry { cmd: Box, tx: Response, + transport_layer_load: u64, id: u64, } impl RequestEntry { - pub fn new(cmd: Box, tx: Response) -> Self { + pub fn new(cmd: Box, tx: Response, transport_layer_load: u64) -> Self { Self { cmd, tx, + transport_layer_load, id: ID_ALLOC.fetch_add(1, Ordering::Relaxed), } } @@ -48,6 +50,7 @@ impl RequestEntry { #[derive(Clone)] pub struct BatchWorker { request_tx: mpsc::Sender, + last_transport_layer_load_report: Arc, } impl BatchWorker { @@ -56,6 +59,7 @@ impl BatchWorker { max_batch_size: usize, max_inflight_requests: usize, max_delay_duration: u64, + overload_threshold: u64, options: CallOption, ) -> Result { let (request_tx, request_rx) = mpsc::channel(max_batch_size); @@ -63,7 +67,10 @@ impl BatchWorker { // Create rpc sender and receiver let (rpc_sender, rpc_receiver) = kv_client.batch_commands_opt(options)?; + let last_transport_layer_load_report = Arc::new(AtomicU64::new(0)); + // Start a background thread to handle batch requests and responses + let last_transport_layer_load_report_clone = last_transport_layer_load_report.clone(); thread::Builder::new() .name(BATCH_WORKER_NAME.to_owned()) .spawn(move || { @@ -74,17 +81,27 @@ impl BatchWorker { max_batch_size, max_inflight_requests, max_delay_duration, + overload_threshold, + last_transport_layer_load_report_clone, )) }) .unwrap(); - Ok(BatchWorker { request_tx }) + Ok(BatchWorker { + request_tx, + last_transport_layer_load_report, + }) } pub async fn dispatch(mut self, request: Box) -> Result> { let (tx, rx) = oneshot::channel(); // Generate BatchCommandRequestEntry - let entry = RequestEntry::new(request, tx); + let last_transport_layer_load = self + .last_transport_layer_load_report + .load(Ordering::Relaxed); + + // Save the load of transport layer in RequestEntry + let entry = RequestEntry::new(request, tx, last_transport_layer_load); // Send request entry to the background thread to handle the request, response will be // received in rx channel. self.request_tx @@ -103,6 +120,8 @@ async fn run_batch_worker( max_batch_size: usize, max_inflight_requests: usize, max_delay_duration: u64, + overload_threshold: u64, + last_transport_layer_load_report: Arc, ) { // Inflight requests which are waiting for the response from rpc server let inflight_requests = Rc::new(RefCell::new(HashMap::new())); @@ -117,6 +136,7 @@ async fn run_batch_worker( max_batch_size, max_inflight_requests, max_delay_duration, + overload_threshold, } .map(Ok); @@ -130,6 +150,10 @@ async fn run_batch_worker( waker.wake(); } + let trasport_layer_load = batch_resp.get_transport_layer_load(); + // Store the load of transport layer + last_transport_layer_load_report.store(trasport_layer_load, Ordering::Relaxed); + for (id, resp) in batch_resp .take_request_ids() .into_iter() @@ -156,6 +180,7 @@ struct BatchCommandsRequestStream<'a> { max_batch_size: usize, max_inflight_requests: usize, max_delay_duration: u64, + overload_threshold: u64, } impl Stream for BatchCommandsRequestStream<'_> { @@ -185,6 +210,15 @@ impl Stream for BatchCommandsRequestStream<'_> { inflight_requests.insert(entry.id, entry.tx); requests.push(entry.cmd.to_batch_request()); request_ids.push(entry.id); + + // Check the transport layer load received in RequestEntry + let load_reported = entry.transport_layer_load; + if load_reported > 0 + && self.overload_threshold > 0 + && load_reported > self.overload_threshold + { + break; + } } Poll::Ready(None) => { return Poll::Ready(None); diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index 1871c6c0..a48ffcb5 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -13,8 +13,9 @@ const DEFAULT_GRPC_KEEPALIVE_TIME: u64 = 10000; const DEFAULT_GRPC_KEEPALIVE_TIMEOUT: u64 = 3000; const DEFAULT_GRPC_COMPLETION_QUEUE_SIZE: usize = 1; const DEFAULT_MAX_BATCH_WAIT_TIME: u64 = 10; -const DEFAULT_MAX_BATCH_SIZE: usize = 100; -const DEFAULT_OVERLOAD_THRESHOLD: usize = 200; +const DEFAULT_MAX_BATCH_SIZE: usize = 10; +const DEFAULT_MAX_INFLIGHT_REQUESTS: usize = 100; +const DEFAULT_OVERLOAD_THRESHOLD: u64 = 1000; /// A trait for connecting to TiKV stores. pub trait KvConnect: Sized + Send + Sync + 'static { type KvClient: KvClient + Clone + Send + Sync + 'static; @@ -31,9 +32,10 @@ pub struct KvClientConfig { pub grpc_keepalive_time: u64, pub grpc_keepalive_timeout: u64, pub allow_batch: bool, - pub overload_threshold: usize, + pub overload_threshold: u64, pub max_batch_wait_time: u64, pub max_batch_size: usize, + pub max_inflight_requests: usize, } impl Default for KvClientConfig { @@ -47,6 +49,7 @@ impl Default for KvClientConfig { overload_threshold: DEFAULT_OVERLOAD_THRESHOLD, max_batch_wait_time: DEFAULT_MAX_BATCH_WAIT_TIME, max_batch_size: DEFAULT_MAX_BATCH_SIZE, + max_inflight_requests: DEFAULT_MAX_INFLIGHT_REQUESTS, } } } @@ -78,8 +81,9 @@ impl KvConnect for TikvConnect { BatchWorker::new( c.clone(), kv_config.max_batch_size, - kv_config.max_batch_size, + kv_config.max_inflight_requests, kv_config.max_batch_wait_time, + kv_config.overload_threshold, CallOption::default(), ) .unwrap(), From 795772cc4b94cd3212b7c54253b44a81b3857e30 Mon Sep 17 00:00:00 2001 From: yongman Date: Thu, 11 Aug 2022 20:10:43 +0800 Subject: [PATCH 09/14] Add default max_inflight_requests config Signed-off-by: yongman --- tikv-client-store/src/batch.rs | 2 +- tikv-client-store/src/client.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tikv-client-store/src/batch.rs b/tikv-client-store/src/batch.rs index 62a9ec4e..5f687394 100644 --- a/tikv-client-store/src/batch.rs +++ b/tikv-client-store/src/batch.rs @@ -62,7 +62,7 @@ impl BatchWorker { overload_threshold: u64, options: CallOption, ) -> Result { - let (request_tx, request_rx) = mpsc::channel(max_batch_size); + let (request_tx, request_rx) = mpsc::channel(max_inflight_requests); // Create rpc sender and receiver let (rpc_sender, rpc_receiver) = kv_client.batch_commands_opt(options)?; diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index a48ffcb5..b6627a3d 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -14,7 +14,7 @@ const DEFAULT_GRPC_KEEPALIVE_TIMEOUT: u64 = 3000; const DEFAULT_GRPC_COMPLETION_QUEUE_SIZE: usize = 1; const DEFAULT_MAX_BATCH_WAIT_TIME: u64 = 10; const DEFAULT_MAX_BATCH_SIZE: usize = 10; -const DEFAULT_MAX_INFLIGHT_REQUESTS: usize = 100; +const DEFAULT_MAX_INFLIGHT_REQUESTS: usize = 10000; const DEFAULT_OVERLOAD_THRESHOLD: u64 = 1000; /// A trait for connecting to TiKV stores. pub trait KvConnect: Sized + Send + Sync + 'static { From df5c8f9eb0e9708c02b4f73388a28f3555c8ebba Mon Sep 17 00:00:00 2001 From: yongman Date: Fri, 12 Aug 2022 11:09:39 +0800 Subject: [PATCH 10/14] Fix cargo fmt Signed-off-by: yongman --- src/pd/client.rs | 2 +- tikv-client-store/src/lib.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pd/client.rs b/src/pd/client.rs index 9c586506..2e4db808 100644 --- a/src/pd/client.rs +++ b/src/pd/client.rs @@ -16,7 +16,7 @@ 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::{KvClientConfig, KvClient, KvConnect, TikvConnect}; +use tikv_client_store::{KvClient, KvClientConfig, KvConnect, TikvConnect}; use tokio::sync::RwLock; const CLIENT_PREFIX: &str = "tikv-client"; diff --git a/tikv-client-store/src/lib.rs b/tikv-client-store/src/lib.rs index 514eb1d8..994818db 100644 --- a/tikv-client-store/src/lib.rs +++ b/tikv-client-store/src/lib.rs @@ -8,7 +8,7 @@ mod request; #[doc(inline)] pub use crate::{ batch::{BatchWorker, RequestEntry}, - client::{KvClientConfig, KvClient, KvConnect, TikvConnect}, + client::{KvClient, KvClientConfig, KvConnect, TikvConnect}, errors::{HasKeyErrors, HasRegionError, HasRegionErrors}, request::Request, }; From 4b7e1fde52fb5ba7f7ac15f0d31892fdba64788d Mon Sep 17 00:00:00 2001 From: yongman Date: Fri, 12 Aug 2022 11:27:30 +0800 Subject: [PATCH 11/14] Fix cargo clippy Signed-off-by: yongman --- tikv-client-store/src/batch.rs | 1 + tikv-client-store/src/client.rs | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tikv-client-store/src/batch.rs b/tikv-client-store/src/batch.rs index 5f687394..16995f8b 100644 --- a/tikv-client-store/src/batch.rs +++ b/tikv-client-store/src/batch.rs @@ -113,6 +113,7 @@ impl BatchWorker { } } +#[allow(clippy::too_many_arguments)] async fn run_batch_worker( mut tx: impl Sink<(BatchCommandsRequest, WriteFlags), Error = Error> + Unpin, mut rx: impl Stream> + Unpin, diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index b6627a3d..56aabf72 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -91,7 +91,7 @@ impl KvConnect for TikvConnect { } else { None }; - KvRpcClient::new(c.clone(), self.timeout, batch_worker) + KvRpcClient::new(c, self.timeout, batch_worker) }) } } From 54315a06fcc3393371825be73aef0a09a574efe5 Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 17 Aug 2022 10:33:04 +0800 Subject: [PATCH 12/14] Fix batch worker running after gRPC stream down Signed-off-by: yongman --- tikv-client-store/src/batch.rs | 47 +++++++++++++++++++++++++++++++-- tikv-client-store/src/client.rs | 25 ++++++++++++++---- 2 files changed, 65 insertions(+), 7 deletions(-) diff --git a/tikv-client-store/src/batch.rs b/tikv-client-store/src/batch.rs index 16995f8b..81977fd0 100644 --- a/tikv-client-store/src/batch.rs +++ b/tikv-client-store/src/batch.rs @@ -15,7 +15,7 @@ use std::{ pin::Pin, rc::Rc, sync::{ - atomic::{AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, Arc, }, thread, @@ -51,6 +51,12 @@ impl RequestEntry { pub struct BatchWorker { request_tx: mpsc::Sender, last_transport_layer_load_report: Arc, + is_running: Arc, + max_batch_size: usize, + max_inflight_requests: usize, + max_delay_duration: u64, + overload_threshold: u64, + options: CallOption, } impl BatchWorker { @@ -65,9 +71,11 @@ impl BatchWorker { let (request_tx, request_rx) = mpsc::channel(max_inflight_requests); // Create rpc sender and receiver - let (rpc_sender, rpc_receiver) = kv_client.batch_commands_opt(options)?; + let (rpc_sender, rpc_receiver) = kv_client.batch_commands_opt(options.clone())?; let last_transport_layer_load_report = Arc::new(AtomicU64::new(0)); + let is_running_status = Arc::new(AtomicBool::new(true)); + let is_running_status_cloned = is_running_status.clone(); // Start a background thread to handle batch requests and responses let last_transport_layer_load_report_clone = last_transport_layer_load_report.clone(); @@ -77,6 +85,7 @@ impl BatchWorker { block_on(run_batch_worker( rpc_sender.sink_err_into(), rpc_receiver.err_into(), + is_running_status_cloned, request_rx, max_batch_size, max_inflight_requests, @@ -90,9 +99,39 @@ impl BatchWorker { Ok(BatchWorker { request_tx, last_transport_layer_load_report, + is_running: is_running_status, + max_batch_size, + max_inflight_requests, + max_delay_duration, + overload_threshold, + options, }) } + pub fn is_running(&self) -> bool { + self.is_running.load(Ordering::Relaxed) + } + + pub fn max_batch_size(&self) -> usize { + self.max_batch_size + } + + pub fn max_inflight_requests(&self) -> usize { + self.max_inflight_requests + } + + pub fn max_delay_duration(&self) -> u64 { + self.max_delay_duration + } + + pub fn overload_threshold(&self) -> u64 { + self.overload_threshold + } + + pub fn options(&self) -> CallOption { + self.options.clone() + } + pub async fn dispatch(mut self, request: Box) -> Result> { let (tx, rx) = oneshot::channel(); // Generate BatchCommandRequestEntry @@ -117,6 +156,7 @@ impl BatchWorker { async fn run_batch_worker( mut tx: impl Sink<(BatchCommandsRequest, WriteFlags), Error = Error> + Unpin, mut rx: impl Stream> + Unpin, + is_running: Arc, request_rx: mpsc::Receiver, max_batch_size: usize, max_inflight_requests: usize, @@ -170,6 +210,9 @@ async fn run_batch_worker( }; let (tx_res, rx_res) = join!(send_requests, recv_handle_response); + + is_running.store(false, Ordering::Relaxed); + debug!("Batch sender finished: {:?}", tx_res); debug!("Batch receiver finished: {:?}", rx_res); } diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index 56aabf72..dee2a511 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -3,6 +3,7 @@ use crate::{batch::BatchWorker, request::Request, Result, SecurityManager}; use async_trait::async_trait; use derive_new::new; +use futures::lock::Mutex; use grpcio::{CallOption, Environment}; use serde_derive::{Deserialize, Serialize}; use std::{any::Any, sync::Arc, time::Duration}; @@ -77,7 +78,7 @@ impl KvConnect for TikvConnect { // Create batch worker if needed let c = Arc::new(c); let batch_worker = if kv_config.allow_batch { - Some( + Some(Arc::new(Mutex::new( BatchWorker::new( c.clone(), kv_config.max_batch_size, @@ -87,7 +88,7 @@ impl KvConnect for TikvConnect { CallOption::default(), ) .unwrap(), - ) + ))) } else { None }; @@ -107,14 +108,28 @@ pub trait KvClient { pub struct KvRpcClient { rpc_client: Arc, timeout: Duration, - batch_worker: Option, + batch_worker: Option>>, } #[async_trait] impl KvClient for KvRpcClient { async fn dispatch(&self, request: Box) -> Result> { - if let Some(batch_worker) = self.batch_worker.clone() { - batch_worker.dispatch(request).await + if let Some(batch_worker_arc) = self.batch_worker.clone() { + let mut batch_worker = batch_worker_arc.lock().await; + if batch_worker.is_running() { + return batch_worker.clone().dispatch(request).await; + } + // batch worker is not running, because of gRPC channel is broken, create a new one + *batch_worker = BatchWorker::new( + self.rpc_client.clone(), + batch_worker.max_batch_size(), + batch_worker.max_inflight_requests(), + batch_worker.max_delay_duration(), + batch_worker.overload_threshold(), + batch_worker.options(), + ) + .unwrap(); + batch_worker.clone().dispatch(request).await } else { // Batch no needed if not batch enabled request From 5c62fc49ce13e78a1758c79f049e33e59053f28c Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 17 Aug 2022 11:14:55 +0800 Subject: [PATCH 13/14] Use RwLock replace mutex for batch worker Signed-off-by: yongman --- tikv-client-store/Cargo.toml | 1 + tikv-client-store/src/client.rs | 11 +++++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/tikv-client-store/Cargo.toml b/tikv-client-store/Cargo.toml index 1b1883d2..08891bc1 100644 --- a/tikv-client-store/Cargo.toml +++ b/tikv-client-store/Cargo.toml @@ -15,5 +15,6 @@ grpcio = { version = "0.10", features = [ "prost-codec" ], default-features = fa log = "0.4" serde = "1.0" serde_derive = "1.0" +tokio = { version = "1", features = [ "sync", "rt-multi-thread", "macros" ] } tikv-client-common = { version = "0.1.0", path = "../tikv-client-common" } tikv-client-proto = { version = "0.1.0", path = "../tikv-client-proto" } diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index dee2a511..98307b6c 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -3,11 +3,11 @@ use crate::{batch::BatchWorker, request::Request, Result, SecurityManager}; use async_trait::async_trait; use derive_new::new; -use futures::lock::Mutex; use grpcio::{CallOption, Environment}; use serde_derive::{Deserialize, Serialize}; use std::{any::Any, sync::Arc, time::Duration}; use tikv_client_proto::tikvpb::TikvClient; +use tokio::sync::RwLock; const DEFAULT_REQUEST_TIMEOUT: u64 = 2000; const DEFAULT_GRPC_KEEPALIVE_TIME: u64 = 10000; @@ -78,7 +78,7 @@ impl KvConnect for TikvConnect { // Create batch worker if needed let c = Arc::new(c); let batch_worker = if kv_config.allow_batch { - Some(Arc::new(Mutex::new( + Some(Arc::new(RwLock::new( BatchWorker::new( c.clone(), kv_config.max_batch_size, @@ -108,17 +108,20 @@ pub trait KvClient { pub struct KvRpcClient { rpc_client: Arc, timeout: Duration, - batch_worker: Option>>, + batch_worker: Option>>, } #[async_trait] impl KvClient for KvRpcClient { async fn dispatch(&self, request: Box) -> Result> { if let Some(batch_worker_arc) = self.batch_worker.clone() { - let mut batch_worker = batch_worker_arc.lock().await; + let batch_worker = batch_worker_arc.read().await; if batch_worker.is_running() { return batch_worker.clone().dispatch(request).await; } + drop(batch_worker); + + let mut batch_worker = batch_worker_arc.write().await; // batch worker is not running, because of gRPC channel is broken, create a new one *batch_worker = BatchWorker::new( self.rpc_client.clone(), From 603a2f8a46d84347b37bf0d681ae0f4773bf79df Mon Sep 17 00:00:00 2001 From: yongman Date: Wed, 17 Aug 2022 16:58:33 +0800 Subject: [PATCH 14/14] Fix check request support batch before dispatch Signed-off-by: yongman --- tikv-client-store/src/client.rs | 2 +- tikv-client-store/src/request.rs | 69 +++++++++----------------------- 2 files changed, 20 insertions(+), 51 deletions(-) diff --git a/tikv-client-store/src/client.rs b/tikv-client-store/src/client.rs index 98307b6c..cdf0573c 100644 --- a/tikv-client-store/src/client.rs +++ b/tikv-client-store/src/client.rs @@ -114,7 +114,7 @@ pub struct KvRpcClient { #[async_trait] impl KvClient for KvRpcClient { async fn dispatch(&self, request: Box) -> Result> { - if let Some(batch_worker_arc) = self.batch_worker.clone() { + if let Some(batch_worker_arc) = self.batch_worker.clone() && request.support_batch(){ let batch_worker = batch_worker_arc.read().await; if batch_worker.is_running() { return batch_worker.clone().dispatch(request).await; diff --git a/tikv-client-store/src/request.rs b/tikv-client-store/src/request.rs index ea957310..e5531d90 100644 --- a/tikv-client-store/src/request.rs +++ b/tikv-client-store/src/request.rs @@ -2,15 +2,14 @@ use crate::{Error, Result}; use async_trait::async_trait; -use futures::{SinkExt, TryStreamExt}; -use grpcio::{CallOption, WriteFlags}; +use grpcio::CallOption; use std::any::Any; use tikv_client_common::internal_err; use tikv_client_proto::{ kvrpcpb, tikvpb::{ batch_commands_request::{self, request::Cmd::*}, - batch_commands_response, BatchCommandsRequest, TikvClient, + batch_commands_response, TikvClient, }, }; @@ -24,7 +23,12 @@ pub trait Request: Any + Sync + Send + 'static { fn label(&self) -> &'static str; fn as_any(&self) -> &(dyn Any + Send); fn set_context(&mut self, context: kvrpcpb::Context); - fn to_batch_request(&self) -> batch_commands_request::Request; + fn to_batch_request(&self) -> batch_commands_request::Request { + batch_commands_request::Request { cmd: None } + } + fn support_batch(&self) -> bool { + false + } } macro_rules! impl_request { @@ -61,6 +65,10 @@ macro_rules! impl_request { }; req } + + fn support_batch(&self) -> bool { + true + } } }; } @@ -105,41 +113,6 @@ impl_request!( RawDeleteRange ); -// TODO implement batchcommands support for rawcasrequest -// impl_request!( -// RawCasRequest, -// raw_compare_and_swap_async_opt, -// "raw_compare_and_swap", -// RawCas -// ); -#[async_trait] -impl Request for kvrpcpb::RawCasRequest { - async fn dispatch( - &self, - client: &TikvClient, - options: CallOption, - ) -> Result> { - client - .raw_compare_and_swap_async_opt(self, options)? - .await - .map(|r| Box::new(r) as Box) - .map_err(Error::Grpc) - } - fn label(&self) -> &'static str { - "raw_compare_and_swap" - } - fn as_any(&self) -> &(dyn Any + Send) { - self - } - fn set_context(&mut self, _: tikv_client_proto::kvrpcpb::Context) { - todo!() - } - // TODO - fn to_batch_request(&self) -> batch_commands_request::Request { - batch_commands_request::Request { cmd: None } - } -} - impl_request!( RawCoprocessorRequest, raw_coprocessor_async_opt, @@ -220,24 +193,20 @@ impl_request!( ); #[async_trait] -impl Request for BatchCommandsRequest { +impl Request for kvrpcpb::RawCasRequest { async fn dispatch( &self, client: &TikvClient, options: CallOption, ) -> Result> { - match client.batch_commands_opt(options) { - Ok((mut tx, mut rx)) => { - tx.send((self.clone(), WriteFlags::default())).await?; - tx.close().await?; - let resp = rx.try_next().await?.unwrap(); - Ok(Box::new(resp) as Box) - } - Err(e) => Err(Error::Grpc(e)), - } + client + .raw_compare_and_swap_async_opt(self, options)? + .await + .map(|r| Box::new(r) as Box) + .map_err(Error::Grpc) } fn label(&self) -> &'static str { - "kv_batch_commands" + "raw_compare_and_swap" } fn as_any(&self) -> &(dyn Any + Send) { self