address review comments and fix tests

Signed-off-by: haojinming <jinming.hao@pingcap.com>

Author: haojinming

src/request/plan.rs

@@ -574,13 +574,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.meet_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/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,7 +9,8 @@ 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::sync::Arc;
@@ -242,7 +242,7 @@ impl Client {
             batch_size: SCAN_LOCK_BATCH_SIZE,
             ..Default::default()
         };
-        self.cleanup_locks(&safepoint, options).await?;
+        self.cleanup_locks(&safepoint, vec![].., options).await?;
 
         // update safepoint to PD
         let res: bool = self
@@ -259,18 +259,14 @@ impl Client {
     pub async fn cleanup_locks(
         &self,
         safepoint: &Timestamp,
+        range: impl Into<BoundRange>,
         options: ResolveLocksOptions,
     ) -> Result<CleanupLocksResult> {
         debug!(self.logger, "invoking cleanup async commit locks");
         // scan all locks with ts <= safepoint
         let ctx = ResolveLocksContext::default();
         let backoff = Backoff::equal_jitter_backoff(100, 10000, 50);
-        let req = new_scan_lock_request(
-            options.start_key.clone(),
-            options.end_key.clone(),
-            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)
@@ -286,16 +282,10 @@ impl Client {
     pub async fn scan_locks(
         &self,
         safepoint: &Timestamp,
-        mut start_key: Vec<u8>,
-        mut end_key: Vec<u8>,
+        range: impl Into<BoundRange>,
         batch_size: u32,
     ) -> Result<Vec<tikv_client_proto::kvrpcpb::LockInfo>> {
-        let req = new_scan_lock_request(
-            std::mem::take(&mut start_key),
-            std::mem::take(&mut end_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/lock.rs

@@ -152,17 +152,13 @@ pub struct ResolveLocksContext {
 #[derive(Clone, Debug)]
 pub struct ResolveLocksOptions {
     pub async_commit_only: bool,
-    pub start_key: Vec<u8>,
-    pub end_key: Vec<u8>,
     pub batch_size: u32,
 }
 
 impl Default for ResolveLocksOptions {
     fn default() -> Self {
         Self {
             async_commit_only: false,
-            start_key: vec![],
-            end_key: vec![],
             batch_size: 1024,
         }
     }

src/transaction/lowering.rs

@@ -161,12 +161,17 @@ pub fn new_pessimistic_lock_request(
 }
 
 pub fn new_scan_lock_request(
-    start_key: Key,
-    end_key: Key,
-    safepoint: Timestamp,
+    range: BoundRange,
+    safepoint: &Timestamp,
     limit: u32,
 ) -> kvrpcpb::ScanLockRequest {
-    requests::new_scan_lock_request(start_key.into(), end_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(

tests/failpoint_tests.rs

@@ -83,6 +83,19 @@ async fn txn_cleanup_locks_batch_size() -> Result<()> {
 
     init().await?;
     let scenario = FailScenario::setup();
+    let full_range = vec![]..;
+
+    let client = TransactionClient::new(pd_addrs(), Some(logger.clone())).await?;
+    // Clean all locks at the beginning to avoid other cases' side effect.
+    let safepoint = client.current_timestamp().await?;
+    let options = ResolveLocksOptions {
+        async_commit_only: false,
+        ..Default::default()
+    };
+    client
+        .cleanup_locks(&safepoint, full_range.clone(), options)
+        .await
+        .unwrap();
 
     fail::cfg("after-prewrite", "return").unwrap();
     fail::cfg("before-cleanup-locks", "return").unwrap();
@@ -91,7 +104,6 @@ async fn txn_cleanup_locks_batch_size() -> Result<()> {
         fail::cfg("before-cleanup-locks", "off").unwrap();
     }}
 
-    let client = TransactionClient::new(pd_addrs(), Some(logger.clone())).await?;
     let keys = write_data(&client, true, true).await?;
     assert_eq!(count_locks(&client).await?, keys.len());
 
@@ -101,7 +113,9 @@ async fn txn_cleanup_locks_batch_size() -> Result<()> {
         batch_size: 4,
         ..Default::default()
     };
-    let res = client.cleanup_locks(&safepoint, options).await?;
+    let res = client
+        .cleanup_locks(&safepoint, full_range, options)
+        .await?;
 
     assert_eq!(res.meet_locks, keys.len());
     assert_eq!(count_locks(&client).await?, keys.len());
@@ -117,6 +131,7 @@ async fn txn_cleanup_async_commit_locks() -> Result<()> {
 
     init().await?;
     let scenario = FailScenario::setup();
+    let full_range = vec![]..;
 
     // no commit
     {
@@ -135,7 +150,9 @@ async fn txn_cleanup_async_commit_locks() -> Result<()> {
             async_commit_only: true,
             ..Default::default()
         };
-        client.cleanup_locks(&safepoint, options).await?;
+        client
+            .cleanup_locks(&safepoint, full_range.clone(), options)
+            .await?;
 
         must_committed(&client, keys).await;
         assert_eq!(count_locks(&client).await?, 0);
@@ -160,7 +177,9 @@ async fn txn_cleanup_async_commit_locks() -> Result<()> {
             async_commit_only: true,
             ..Default::default()
         };
-        client.cleanup_locks(&safepoint, options).await?;
+        client
+            .cleanup_locks(&safepoint, full_range.clone(), options)
+            .await?;
 
         must_committed(&client, keys).await;
         assert_eq!(count_locks(&client).await?, 0);
@@ -177,7 +196,9 @@ async fn txn_cleanup_async_commit_locks() -> Result<()> {
             async_commit_only: true,
             ..Default::default()
         };
-        client.cleanup_locks(&safepoint, options).await?;
+        client
+            .cleanup_locks(&safepoint, full_range, options)
+            .await?;
 
         must_committed(&client, keys).await;
         assert_eq!(count_locks(&client).await?, 0);
@@ -217,11 +238,11 @@ async fn txn_cleanup_range_async_commit_locks() -> Result<()> {
     let safepoint = client.current_timestamp().await?;
     let options = ResolveLocksOptions {
         async_commit_only: true,
-        start_key,
-        end_key,
         ..Default::default()
     };
-    let res = client.cleanup_locks(&safepoint, options).await?;
+    let res = client
+        .cleanup_locks(&safepoint, start_key..end_key, options)
+        .await?;
 
     assert_eq!(res.meet_locks, keys.len() - 3);
 
@@ -230,7 +251,7 @@ async fn txn_cleanup_range_async_commit_locks() -> Result<()> {
         async_commit_only: false,
         ..Default::default()
     };
-    client.cleanup_locks(&safepoint, options).await?;
+    client.cleanup_locks(&safepoint, vec![].., options).await?;
     must_committed(&client, keys).await;
     assert_eq!(count_locks(&client).await?, 0);
 
@@ -245,6 +266,7 @@ async fn txn_cleanup_2pc_locks() -> Result<()> {
 
     init().await?;
     let scenario = FailScenario::setup();
+    let full_range = vec![]..;
 
     // no commit
     {
@@ -264,14 +286,18 @@ async fn txn_cleanup_2pc_locks() -> Result<()> {
                 async_commit_only: true, // Skip 2pc locks.
                 ..Default::default()
             };
-            client.cleanup_locks(&safepoint, options).await?;
+            client
+                .cleanup_locks(&safepoint, full_range.clone(), options)
+                .await?;
             assert_eq!(count_locks(&client).await?, keys.len());
         }
         let options = ResolveLocksOptions {
             async_commit_only: false,
             ..Default::default()
         };
-        client.cleanup_locks(&safepoint, options).await?;
+        client
+            .cleanup_locks(&safepoint, full_range.clone(), options)
+            .await?;
 
         must_rollbacked(&client, keys).await;
         assert_eq!(count_locks(&client).await?, 0);
@@ -289,7 +315,9 @@ async fn txn_cleanup_2pc_locks() -> Result<()> {
             async_commit_only: false,
             ..Default::default()
         };
-        client.cleanup_locks(&safepoint, options).await?;
+        client
+            .cleanup_locks(&safepoint, full_range, options)
+            .await?;
 
         must_committed(&client, keys).await;
         assert_eq!(count_locks(&client).await?, 0);
@@ -319,7 +347,7 @@ async fn must_rollbacked(client: &TransactionClient, keys: HashSet<Vec<u8>>) {
 
 async fn count_locks(client: &TransactionClient) -> Result<usize> {
     let ts = client.current_timestamp().await.unwrap();
-    let locks = client.scan_locks(&ts, vec![], vec![], 1024).await?;
+    let locks = client.scan_locks(&ts, vec![].., 1024).await?;
     // De-duplicated as `scan_locks` will return duplicated locks due to retry on region changes.
     let locks_set: HashSet<Vec<u8>> =
         HashSet::from_iter(locks.into_iter().map(|mut l| l.take_key()));