Merge pull request #386 from haojinming/resolve_lock_in_range

Resolve lock in range

Author: pingyu

.github/workflows/ci.yml

@@ -88,8 +88,9 @@ jobs:
         run: curl --proto '=https' --tlsv1.2 -sSf https://tiup-mirrors.pingcap.com/install.sh | sh
       - name: start tiup playground
         run: |
-          ~/.tiup/bin/tiup install tikv:nightly pd:nightly
-          ~/.tiup/bin/tiup playground nightly --mode tikv-slim --kv 3 --without-monitor --kv.config config/tikv.toml --pd.config config/pd.toml &
+          # use latest stable version
+          ~/.tiup/bin/tiup install tikv pd
+          ~/.tiup/bin/tiup playground --mode tikv-slim --kv 3 --without-monitor --kv.config config/tikv.toml --pd.config config/pd.toml &
           while :; do
             echo "waiting cluster to be ready"
             [[ "$(curl -I http://127.0.0.1:2379/pd/api/v1/regions 2>/dev/null | head -n 1 | cut -d$' ' -f2)" -ne "405" ]] || break

src/request/plan.rs

@@ -446,14 +446,13 @@ where
 pub struct CleanupLocksResult {
     pub region_error: Option<errorpb::Error>,
     pub key_error: Option<Vec<Error>>,
-    pub meet_locks: usize,
-    // TODO: pub resolved_locks: usize,
+    pub resolved_locks: usize,
 }
 
 impl Clone for CleanupLocksResult {
     fn clone(&self) -> Self {
         Self {
-            meet_locks: self.meet_locks,
+            resolved_locks: self.resolved_locks,
             ..Default::default() // Ignore errors, which should be extracted by `extract_error()`.
         }
     }
@@ -479,7 +478,7 @@ impl Merge<CleanupLocksResult> for Collect {
             .into_iter()
             .fold(Ok(CleanupLocksResult::default()), |acc, x| {
                 Ok(CleanupLocksResult {
-                    meet_locks: acc.unwrap().meet_locks + x?.meet_locks,
+                    resolved_locks: acc?.resolved_locks + x?.resolved_locks,
                     ..Default::default()
                 })
             })
@@ -574,13 +573,15 @@ where
                 "CleanupLocks::execute, meet locks:{}",
                 locks.len()
             );
-            result.meet_locks += locks.len();
 
+            let lock_size = locks.len();
             match lock_resolver
                 .cleanup_locks(self.store.clone().unwrap(), locks, self.pd_client.clone())
                 .await
             {
-                Ok(()) => {}
+                Ok(()) => {
+                    result.resolved_locks += lock_size;
+                }
                 Err(Error::ExtractedErrors(mut errors)) => {
                     // Propagate errors to `retry_multi_region` for retry.
                     if let Error::RegionError(e) = errors.pop().unwrap() {

src/store.rs

@@ -53,7 +53,11 @@ pub fn store_stream_for_range<PdC: PdClient>(
     range: (Vec<u8>, Vec<u8>),
     pd_client: Arc<PdC>,
 ) -> BoxStream<'static, Result<((Vec<u8>, Vec<u8>), RegionStore)>> {
-    let bnd_range = BoundRange::from(range.clone());
+    let bnd_range = if range.1.is_empty() {
+        BoundRange::range_from(range.0.clone().into())
+    } else {
+        BoundRange::from(range.clone())
+    };
     pd_client
         .stores_for_range(bnd_range)
         .map_ok(move |store| {
@@ -67,25 +71,6 @@ pub fn store_stream_for_range<PdC: PdClient>(
         .boxed()
 }
 
-pub fn store_stream_for_range_by_start_key<PdC: PdClient>(
-    start_key: Key,
-    pd_client: Arc<PdC>,
-) -> BoxStream<'static, Result<(Vec<u8>, RegionStore)>> {
-    let bnd_range = BoundRange::range_from(start_key.clone());
-    pd_client
-        .stores_for_range(bnd_range)
-        .map_ok(move |store| {
-            let region_range = store.region_with_leader.range();
-            (
-                range_intersection(region_range, (start_key.clone(), vec![].into()))
-                    .0
-                    .into(),
-                store,
-            )
-        })
-        .boxed()
-}
-
 /// The range used for request should be the intersection of `region_range` and `range`.
 fn range_intersection(region_range: (Key, Key), range: (Key, Key)) -> (Key, Key) {
     let (lower, upper) = region_range;

src/transaction/client.rs

@@ -1,6 +1,5 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
-use super::requests::new_scan_lock_request;
 use crate::{
     backoff::DEFAULT_REGION_BACKOFF,
     config::Config,
@@ -10,10 +9,11 @@ use crate::{
     transaction::{
         lock::ResolveLocksOptions, ResolveLocksContext, Snapshot, Transaction, TransactionOptions,
     },
-    Backoff, Result,
+    transaction_lowering::new_scan_lock_request,
+    Backoff, BoundRange, Result,
 };
 use slog::{Drain, Logger};
-use std::{mem, sync::Arc};
+use std::sync::Arc;
 use tikv_client_proto::pdpb::Timestamp;
 
 // FIXME: cargo-culted value
@@ -242,7 +242,7 @@ impl Client {
             batch_size: SCAN_LOCK_BATCH_SIZE,
             ..Default::default()
         };
-        self.cleanup_locks(&safepoint, options).await?;
+        self.cleanup_locks(.., &safepoint, options).await?;
 
         // update safepoint to PD
         let res: bool = self
@@ -258,24 +258,20 @@ impl Client {
 
     pub async fn cleanup_locks(
         &self,
+        range: impl Into<BoundRange>,
         safepoint: &Timestamp,
         options: ResolveLocksOptions,
     ) -> Result<CleanupLocksResult> {
         debug!(self.logger, "invoking cleanup async commit locks");
         // scan all locks with ts <= safepoint
-        let mut start_key = vec![];
         let ctx = ResolveLocksContext::default();
         let backoff = Backoff::equal_jitter_backoff(100, 10000, 50);
-        let req = new_scan_lock_request(
-            mem::take(&mut start_key),
-            safepoint.version(),
-            options.batch_size,
-        );
+        let req = new_scan_lock_request(range.into(), safepoint, options.batch_size);
         let plan = crate::request::PlanBuilder::new(self.pd.clone(), req)
             .cleanup_locks(self.logger.clone(), ctx.clone(), options, backoff)
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
-            .merge(crate::request::Collect)
             .extract_error()
+            .merge(crate::request::Collect)
             .plan();
         plan.execute().await
     }
@@ -286,10 +282,10 @@ impl Client {
     pub async fn scan_locks(
         &self,
         safepoint: &Timestamp,
-        mut start_key: Vec<u8>,
+        range: impl Into<BoundRange>,
         batch_size: u32,
     ) -> Result<Vec<tikv_client_proto::kvrpcpb::LockInfo>> {
-        let req = new_scan_lock_request(mem::take(&mut start_key), safepoint.version(), batch_size);
+        let req = new_scan_lock_request(range.into(), safepoint, batch_size);
         let plan = crate::request::PlanBuilder::new(self.pd.clone(), req)
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
             .merge(crate::request::Collect)

src/transaction/lowering.rs

@@ -161,11 +161,17 @@ pub fn new_pessimistic_lock_request(
 }
 
 pub fn new_scan_lock_request(
-    start_key: Key,
-    safepoint: Timestamp,
+    range: BoundRange,
+    safepoint: &Timestamp,
     limit: u32,
 ) -> kvrpcpb::ScanLockRequest {
-    requests::new_scan_lock_request(start_key.into(), safepoint.version(), limit)
+    let (start_key, end_key) = range.into_keys();
+    requests::new_scan_lock_request(
+        start_key.into(),
+        end_key.unwrap_or_default().into(),
+        safepoint.version(),
+        limit,
+    )
 }
 
 pub fn new_heart_beat_request(

src/transaction/requests.rs

@@ -7,7 +7,7 @@ use crate::{
         Collect, CollectSingle, CollectWithShard, DefaultProcessor, HasNextBatch, KvRequest, Merge,
         NextBatch, Process, ResponseWithShard, Shardable, SingleKey,
     },
-    store::{store_stream_for_keys, store_stream_for_range_by_start_key, RegionStore},
+    store::{store_stream_for_keys, store_stream_for_range, RegionStore},
     timestamp::TimestampExt,
     transaction::HasLocks,
     util::iter::FlatMapOkIterExt,
@@ -445,11 +445,13 @@ impl Merge<ResponseWithShard<kvrpcpb::PessimisticLockResponse, Vec<kvrpcpb::Muta
 
 pub fn new_scan_lock_request(
     start_key: Vec<u8>,
+    end_key: Vec<u8>,
     safepoint: u64,
     limit: u32,
 ) -> kvrpcpb::ScanLockRequest {
     let mut req = kvrpcpb::ScanLockRequest::default();
     req.set_start_key(start_key);
+    req.set_end_key(end_key);
     req.set_max_version(safepoint);
     req.set_limit(limit);
     req
@@ -460,25 +462,29 @@ impl KvRequest for kvrpcpb::ScanLockRequest {
 }
 
 impl Shardable for kvrpcpb::ScanLockRequest {
-    type Shard = Vec<u8>;
+    type Shard = (Vec<u8>, Vec<u8>);
 
     fn shards(
         &self,
         pd_client: &Arc<impl PdClient>,
     ) -> BoxStream<'static, Result<(Self::Shard, RegionStore)>> {
-        store_stream_for_range_by_start_key(self.start_key.clone().into(), pd_client.clone())
+        store_stream_for_range(
+            (self.start_key.clone(), self.end_key.clone()),
+            pd_client.clone(),
+        )
     }
 
     fn apply_shard(&mut self, shard: Self::Shard, store: &RegionStore) -> Result<()> {
         self.set_context(store.region_with_leader.context()?);
-        self.set_start_key(shard);
+        self.set_start_key(shard.0);
         Ok(())
     }
 }
 
 impl HasNextBatch for kvrpcpb::ScanLockResponse {
     fn has_next_batch(&self) -> Option<(Vec<u8>, Vec<u8>)> {
         self.get_locks().last().map(|lock| {
+            // TODO: if last key is larger or equal than ScanLockRequest.end_key, return None.
             let mut start_key: Vec<u8> = lock.get_key().to_vec();
             start_key.push(0);
             (start_key, vec![])

tests/common/mod.rs

@@ -46,10 +46,14 @@ pub async fn init() -> Result<()> {
             .take(count as usize - 1)
             .map(|x| x.to_be_bytes().to_vec());
 
-        ensure_region_split(keys_1.chain(keys_2), 80).await?;
+        // about 43 regions with above keys.
+        ensure_region_split(keys_1.chain(keys_2), 40).await?;
     }
 
     clear_tikv().await;
+    let region_cnt = ctl::get_region_count().await?;
+    // print log for debug convenience
+    println!("init finish with {region_cnt} regions");
     Ok(())
 }