address comment

Signed-off-by: Andy Lok <andylokandy@hotmail.com>

Author: andylokandy

examples/pessimistic.rs

@@ -25,6 +25,7 @@ async fn main() {
     } else {
         Config::default()
     }
+    // This example uses the default keyspace, so api-v2 must be enabled on the server.
     .with_default_keyspace();
 
     // init

examples/raw.rs

@@ -32,6 +32,7 @@ async fn main() -> Result<()> {
     } else {
         Config::default()
     }
+    // This example uses the default keyspace, so api-v2 must be enabled on the server.
     .with_default_keyspace();
 
     // When we first create a client we receive a `Connect` structure which must be resolved before
@@ -137,6 +138,7 @@ async fn main() -> Result<()> {
     );
     println!("Scanning batch scan from {batch_scan_keys:?} gives: {vals:?}");
 
+    // Delete all keys in the whole range.
     client.delete_range("".to_owned().."".to_owned()).await?;
 
     Ok(())

examples/transaction.rs

@@ -88,6 +88,7 @@ async fn main() {
     } else {
         Config::default()
     }
+    // This example uses the default keyspace, so api-v2 must be enabled on the server.
     .with_default_keyspace();
 
     let txn = Client::new_with_config(args.pd, config)

src/pd/retry.rs

@@ -201,7 +201,7 @@ impl RetryClientTrait for RetryClient<Cluster> {
     }
 
     async fn get_keyspace_id(&self, keyspace: &str) -> Result<u32> {
-        retry_mut!(self, "get_keyspace_id", |cluster| async {
+        retry!(self, "get_keyspace_id", |cluster| async {
             cluster.get_keyspace_id(keyspace).await
         })
     }

src/raw/client.rs

@@ -692,7 +692,7 @@ impl<PdC: PdClient> Client<PdC> {
         // limit is a soft limit, so we need check the number of results
         result.truncate(limit as usize);
 
-        // truncate the version of keys
+        // truncate the prefix of keys
         let result = result.truncate_keyspace(self.keyspace);
 
         Ok(result)

src/request/keyspace.rs

@@ -149,6 +149,15 @@ impl TruncateKeyspace for Vec<KvPair> {
     }
 }
 
+impl TruncateKeyspace for Vec<crate::proto::kvrpcpb::LockInfo> {
+    fn truncate_keyspace(mut self, keyspace: Keyspace) -> Self {
+        for lock in self.locks.iter_mut() {
+            lock.key.truncate_keyspace(self.keyspace);
+        }
+        self
+    }
+}
+
 fn keyspace_prefix(keyspace_id: u32, key_mode: KeyMode) -> [u8; KEYSPACE_PREFIX_LEN] {
     let mut prefix = keyspace_id.to_be_bytes();
     prefix[0] = match key_mode {

src/transaction/client.rs

@@ -284,7 +284,7 @@ impl Client {
             .retry_multi_region(DEFAULT_REGION_BACKOFF)
             .merge(crate::request::Collect)
             .plan();
-        plan.execute().await
+        Ok(plan.execute().await?.truncate_keyspace(self.keyspace))
     }
 
     /// Cleans up all keys in a range and quickly reclaim disk space.

src/transaction/lock.rs

@@ -121,8 +121,8 @@ async fn resolve_lock_with_retry(
         let store = pd_client.clone().store_for_key(key.into()).await?;
         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(), keyspace, request)
+        let encoded_req = EncodedRequest::new(request, pd_client.get_codec());
+        let plan = crate::request::PlanBuilder::new(pd_client.clone(), encoded_req)
             .single_region_with_store(store)
             .await?
             .resolve_lock(Backoff::no_backoff(), keyspace)

tests/common/mod.rs

@@ -25,7 +25,7 @@ const REGION_SPLIT_TIME_LIMIT: Duration = Duration::from_secs(15);
 // Delete all entries in TiKV to leave a clean space for following tests.
 pub async fn clear_tikv() {
     // DEFAULT_REGION_BACKOFF is not long enough for CI environment. So set a longer backoff.
-    // let backoff = tikv_client::Backoff::no_jitter_backoff(100, 30000, 20);
+    let backoff = tikv_client::Backoff::no_jitter_backoff(100, 30000, 20);
     let raw_client =
         RawClient::new_with_config(pd_addrs(), Config::default().with_default_keyspace())
             .await