diff --git a/examples/transaction.rs b/examples/transaction.rs index 5d58a59b572b7831c6ae95343f1ac0945a2a80c3..6ec8c35edccf19635412f6d48789a5fc14ad0aa6 100644 --- a/examples/transaction.rs +++ b/examples/transaction.rs @@ -44,7 +44,19 @@ async fn key_exists(client: &Client, key: Key) -> bool { res } +async fn delete_range(client: &Client, range: impl Into) { + let mut txn = client + .begin_optimistic() + .await + .expect("Could not begin a transaction"); + txn.delete_range(range) + .await + .expect("Could not delete range"); + txn.commit().await.expect("Could not commit transaction"); +} + async fn scan(client: &Client, range: impl Into, limit: u32) { + println!("start scan..."); let mut txn = client .begin_optimistic() .await @@ -86,16 +98,18 @@ async fn main() { .expect("Could not connect to tikv"); // set + let key0: Key = b"aaa".to_vec().into(); + let value0: Value = b"valuea".to_vec(); let key1: Key = b"key1".to_vec().into(); let value1: Value = b"value1".to_vec(); let key2: Key = b"key2".to_vec().into(); let value2: Value = b"value2".to_vec(); - puts(&txn, vec![(key1, value1), (key2, value2)]).await; + puts(&txn, vec![(key0, value0), (key1, value1), (key2, value2)]).await; // get let key1: Key = b"key1".to_vec().into(); let value1 = get(&txn, key1.clone()).await; - println!("{:?}", (key1, value1)); + println!("get {:?}", (key1, value1)); // check key exists let key1: Key = b"key1".to_vec().into(); @@ -111,8 +125,34 @@ async fn main() { let key1: Key = b"key1".to_vec().into(); scan(&txn, key1.., 10).await; + // delete range + // insert keys + let count = 20; + // use j as it's less than k to test delete range only delete what in range + let prefix = "jey"; + let mut pairs = Vec::with_capacity(count); + format!("{}", prefix); + for i in 1..count { + let key_str = format!("{}{:?}", prefix, i.to_le_bytes()); + let key: Key = key_str.as_bytes().to_vec().into(); + let value: Value = format!("{:?}", key_str).as_bytes().to_vec().into(); + pairs.push((key, value)); + } + println!("insert {} keys.", pairs.len()); + puts(&txn, pairs).await; + // check existence + let start: Key = format!("{}{:?}", prefix, (0_usize).to_le_bytes()).as_bytes().to_vec().into(); + let end: Key = b"jez".to_vec().into(); + let scan_start: Key = b"aaa".to_vec().into(); + scan(&txn, scan_start.clone().., 1000).await; + // delete + delete_range(&txn, start.clone()..=end).await; + // check existence + scan(&txn, scan_start.., 1000).await; + // delete + let key0: Key = b"aaa".to_vec().into(); let key1: Key = b"key1".to_vec().into(); let key2: Key = b"key2".to_vec().into(); - dels(&txn, vec![key1, key2]).await; + dels(&txn, vec![key0, key1, key2]).await; } diff --git a/src/request/plan.rs b/src/request/plan.rs index 42cd948a2473f423b3d85063c4889808f12f519a..5df1750547517ad9f8b0a16afb38404720648952 100644 --- a/src/request/plan.rs +++ b/src/request/plan.rs @@ -64,7 +64,7 @@ pub struct RetryableMultiRegion { pub(super) inner: P, pub pd_client: Arc, pub backoff: Backoff, - + pub serial_scan: bool, /// Preserve all regions' results for other downstream plans to handle. /// If true, return Ok and preserve all regions' results, even if some of them are Err. /// Otherwise, return the first Err if there is any. @@ -121,6 +121,56 @@ where } } + #[async_recursion] + async fn serial_plan_handler( + pd_client: Arc, + current_plan: P, + backoff: Backoff, + permits: Arc, + preserve_region_results: bool, + ) -> Result<::Result> + { + let shards = current_plan.shards(&pd_client).collect::>().await; + let mut rets = Vec::new(); + let mut regions = 0; + for shard in shards { + let (shard, region_store) = shard?; + let mut clone = current_plan.clone(); + clone.apply_shard(shard, ®ion_store)?; + let result = Self::single_shard_handler( + pd_client.clone(), + clone, + region_store, + backoff.clone(), + permits.clone(), + preserve_region_results, + ).await; + rets.push(result); + regions+=1; + + if regions>=2 { + break; + } + } + if preserve_region_results { + Ok(rets + .into_iter() + .flat_map_ok(|x| x) + .map(|x| match x { + Ok(r) => r, + Err(e) => Err(e), + }) + .collect()) + } else { + Ok(rets + .into_iter() + .collect::>>()? + .into_iter() + .flatten() + .collect()) + } + } + #[async_recursion] async fn single_shard_handler( pd_client: Arc, @@ -286,6 +336,7 @@ impl Clone for RetryableMultiRegion { RetryableMultiRegion { inner: self.inner.clone(), pd_client: self.pd_client.clone(), + serial_scan:self.serial_scan, backoff: self.backoff.clone(), preserve_region_results: self.preserve_region_results, } @@ -304,14 +355,25 @@ where // too many concurrent requests, TiKV is more likely to return a "TiKV // is busy" error let concurrency_permits = Arc::new(Semaphore::new(MULTI_REGION_CONCURRENCY)); - Self::single_plan_handler( - self.pd_client.clone(), - self.inner.clone(), - self.backoff.clone(), - concurrency_permits.clone(), - self.preserve_region_results, - ) - .await + if self.serial_scan { + Self::serial_plan_handler( + self.pd_client.clone(), + self.inner.clone(), + self.backoff.clone(), + concurrency_permits.clone(), + self.preserve_region_results, + ) + .await + }else { + Self::single_plan_handler( + self.pd_client.clone(), + self.inner.clone(), + self.backoff.clone(), + concurrency_permits.clone(), + self.preserve_region_results, + ) + .await + } } } @@ -602,6 +664,7 @@ mod test { pd_client: Arc::new(MockPdClient::default()), }, pd_client: Arc::new(MockPdClient::default()), + serial_scan:false, backoff: Backoff::no_backoff(), preserve_region_results: false, }; diff --git a/src/request/plan_builder.rs b/src/request/plan_builder.rs index ce269fac43d09b4c55f2cdc21fe34c81449434aa..04bebf1ec68595ec6781207178227d4102d5dc22 100644 --- a/src/request/plan_builder.rs +++ b/src/request/plan_builder.rs @@ -126,6 +126,25 @@ where self.make_retry_multi_region(backoff, true) } + /// Split the request into shards sending a request to the region of each shard. + pub fn retry_multi_region_serial( + self, + backoff: Backoff, + ) -> PlanBuilder, Targetted> { + self.make_retry_multi_region_serial(backoff, false) + } + + /// Preserve all results, even some of them are Err. + /// To pass all responses to merge, and handle partial successful results correctly. + pub fn retry_multi_region_preserve_results_serial( + self, + backoff: Backoff, + ) -> PlanBuilder, Targetted> { + self.make_retry_multi_region_serial(backoff, true) + } + + + fn make_retry_multi_region( self, backoff: Backoff, @@ -136,6 +155,25 @@ where plan: RetryableMultiRegion { inner: self.plan, pd_client: self.pd_client, + serial_scan:false, + backoff, + preserve_region_results, + }, + phantom: PhantomData, + } + } + + fn make_retry_multi_region_serial( + self, + backoff: Backoff, + preserve_region_results: bool, + ) -> PlanBuilder, Targetted> { + PlanBuilder { + pd_client: self.pd_client.clone(), + plan: RetryableMultiRegion { + inner: self.plan, + pd_client: self.pd_client, + serial_scan:true, backoff, preserve_region_results, }, diff --git a/src/transaction/lowering.rs b/src/transaction/lowering.rs index c931e02cee4a37acb4deff3eaf0f06febfa3e563..b25e17ebb22b339de7fc2689f674a67795b2cc71 100644 --- a/src/transaction/lowering.rs +++ b/src/transaction/lowering.rs @@ -36,6 +36,16 @@ pub fn new_scan_request( ) } +pub fn new_delete_range_request( + range: BoundRange, +) -> kvrpcpb::DeleteRangeRequest { + let (start_key, end_key) = range.into_keys(); + requests::new_delete_range_request( + start_key.into(), + end_key.unwrap_or_default().into(), + ) +} + pub fn new_resolve_lock_request( start_version: Timestamp, commit_version: Timestamp, diff --git a/src/transaction/requests.rs b/src/transaction/requests.rs index 16243d83217fa31d95f889741963e7ac55f939cc..0e0d42191941ea0791a9082987ac4afe37d83c25 100644 --- a/src/transaction/requests.rs +++ b/src/transaction/requests.rs @@ -159,6 +159,22 @@ impl Merge for Collect { } } +pub fn new_delete_range_request( + start_key: Vec, + end_key: Vec, +) -> kvrpcpb::DeleteRangeRequest { + let mut req = kvrpcpb::DeleteRangeRequest::default(); + req.set_start_key(start_key); + req.set_end_key(end_key); + req +} + +impl KvRequest for kvrpcpb::DeleteRangeRequest { + type Response = kvrpcpb::DeleteRangeResponse; +} + +shardable_range!(kvrpcpb::DeleteRangeRequest); + pub fn new_resolve_lock_request( start_version: u64, commit_version: u64, @@ -652,6 +668,7 @@ error_locks!(kvrpcpb::CheckSecondaryLocksResponse); impl HasLocks for kvrpcpb::CleanupResponse {} impl HasLocks for kvrpcpb::ScanLockResponse {} +impl HasLocks for kvrpcpb::DeleteRangeResponse {} impl HasLocks for kvrpcpb::PessimisticRollbackResponse { fn take_locks(&mut self) -> Vec { diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index 30f9602ae7941c729a44f6e32dd2c2fc1fb35146..084ebeff46fa381e6dea755219bee3a77adc8772 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -309,6 +309,15 @@ impl Transaction { } } + /// Create a new 'delete_range' request. + pub async fn delete_range( + &mut self, + range: impl Into, + ) -> Result<()> { + debug!(self.logger, "invoking transactional delete_range request"); + self.delete_range_inner(range).await + } + /// Create a new 'scan' request. /// /// Once resolved this request will result in a `Vec` of all key-value pairs that lie in the @@ -696,6 +705,24 @@ impl Transaction { plan.execute().await } + async fn delete_range_inner( + &mut self, + range: impl Into, + ) -> Result<()> { + self.check_allow_operation().await?; + let rpc = self.rpc.clone(); + let retry_options = self.options.retry_options.clone(); + + let request = new_delete_range_request(range.into()); + let plan = PlanBuilder::new(rpc, request) + .resolve_lock(retry_options.lock_backoff) + .retry_multi_region(retry_options.region_backoff) + .extract_error() + .plan(); + plan.execute().await?; + Ok(()) + } + async fn scan_inner( &mut self, range: impl Into, @@ -708,26 +735,44 @@ impl Transaction { let rpc = self.rpc.clone(); let retry_options = self.options.retry_options.clone(); - self.buffer - .scan_and_fetch( - range.into(), - limit, - !key_only, - reverse, - 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) - .resolve_lock(retry_options.lock_backoff) - .retry_multi_region(retry_options.region_backoff) - .merge(Collect) - .plan(); - plan.execute() - .await - .map(|r| r.into_iter().map(Into::into).collect()) - }, - ) - .await + let mut redundant_limit = limit; + + let mut result = vec![]; + + let mut last_size:u32 = 1000000000; + let mut redundant_range = range.into().into_keys(); + + while redundant_limit> 0 && last_size != 0 { + //info!(self.logger,"new scan request rang {:?} limit {:?}",redundant_range,redundant_limit); + let request = + new_scan_request(BoundRange::from(redundant_range.clone()), timestamp.clone(), redundant_limit, key_only, reverse); + let plan = PlanBuilder::new(rpc.clone(), request) + .resolve_lock(retry_options.lock_backoff.clone()) + .retry_multi_region_serial(retry_options.region_backoff.clone()) + .merge(Collect) + .plan(); + let res = plan.execute() + .await?; + + last_size = res.len() as u32; + redundant_limit -= if redundant_limit>=last_size { last_size } else{ redundant_limit }; + + if redundant_limit>0 && last_size>0{ + let mut last_key:Vec = res.last().unwrap().key().clone().into(); + last_key.push(0); + redundant_range.0=last_key.into(); + } + result.extend(res); + } + + + // TODO: use `BTreeMap` instead of `HashMap` to avoid sorting. + if reverse { + result.sort_unstable_by(|a, b| b.key().cmp(a.key())); + } + + + Ok(result.into_iter().take(limit as usize)) } /// Pessimistically lock the keys, and optionally retrieve corresponding values.

AltStyle によって変換されたページ (->オリジナル) /