fix

tikv · Dec 16, 2023 · 8534111 · 8534111
1 parent 99e366e
commit 8534111
Show file tree

Hide file tree

Showing 11 changed files with 88 additions and 71 deletions.
diff --git a/src/common/errors.rs b/src/common/errors.rs
@@ -113,6 +113,10 @@ pub enum Error {
         inner: Box<Error>,
         success_keys: Vec<Vec<u8>>,
     },
+    #[error("Keyspace is not enabled on the server, please enable it by setting `storage.api-version = 2`")]
+    ServerKeyspaceNotEnabled,
+    #[error("Keyspace is not enabled on the client, please enable it by `Config::with_default_keyspace()` or `Config::with_keyspace()`")]
+    ClientKeyspaceNotEnabled,
 }
 
 impl From<crate::proto::errorpb::Error> for Error {

diff --git a/src/config.rs b/src/config.rs
@@ -88,15 +88,15 @@ impl Config {
 
     /// Set to use default keyspace.
     ///
-    /// Server should enable `api-version = 2` to use this feature.
+    /// Server should enable `storage.api-version = 2` to use this feature.
     #[must_use]
     pub fn with_default_keyspace(self) -> Self {
         self.with_keyspace("DEFAULT")
     }
 
     /// Set the use keyspace for the client.
     ///
-    /// Server should enable `api-version = 2` to use this feature.
+    /// Server should enable `storage.api-version = 2` to use this feature.
     #[must_use]
     pub fn with_keyspace(mut self, keyspace: &str) -> Self {
         self.keyspace = Some(keyspace.to_owned());

diff --git a/src/raw/client.rs b/src/raw/client.rs
@@ -233,6 +233,7 @@ impl<PdC: PdClient> Client<PdC> {
         let request = new_raw_get_request(key, self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), self.api_version, request)
             .retry_multi_region(self.backoff.clone())
+            .extract_error()
             .merge(CollectSingle)
             .post_process_default()
             .plan();
@@ -268,6 +269,7 @@ impl<PdC: PdClient> Client<PdC> {
         let request = new_raw_batch_get_request(keys, self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), self.api_version, request)
             .retry_multi_region(self.backoff.clone())
+            .extract_error()
             .merge(Collect)
             .plan();
         plan.execute().await.map(|r| {
@@ -581,6 +583,7 @@ impl<PdC: PdClient> Client<PdC> {
         );
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), self.api_version, req)
             .retry_multi_region(self.backoff.clone())
+            .extract_error()
             .merge(CollectSingle)
             .post_process_default()
             .plan();
@@ -660,6 +663,7 @@ impl<PdC: PdClient> Client<PdC> {
                 crate::request::PlanBuilder::new(self.rpc.clone(), self.api_version, request)
                     .single_region_with_store(region_store.clone())
                     .await?
+                    .extract_error()
                     .plan()
                     .execute()
                     .await?;
@@ -719,6 +723,7 @@ impl<PdC: PdClient> Client<PdC> {
         let request = new_raw_batch_scan_request(ranges, each_limit, key_only, self.cf.clone());
         let plan = crate::request::PlanBuilder::new(self.rpc.clone(), self.api_version, request)
             .retry_multi_region(self.backoff.clone())
+            .extract_error()
             .merge(Collect)
             .plan();
         plan.execute().await.map(|r| {

diff --git a/src/request/api_version.rs b/src/request/api_version.rs
@@ -163,7 +163,7 @@ fn keyspace_prefix(keyspace_id: u32, key_mode: KeyMode) -> [u8; KEYSPACE_PREFIX_
 fn prepend_bytes<const N: usize>(vec: &mut Vec<u8>, prefix: &[u8; N]) {
     unsafe {
         vec.reserve_exact(N);
-        std::ptr::copy(vec.as_ptr(), vec.as_mut_ptr().offset(N as isize), vec.len());
+        std::ptr::copy(vec.as_ptr(), vec.as_mut_ptr().add(N), vec.len());
         std::ptr::copy_nonoverlapping(prefix.as_ptr(), vec.as_mut_ptr(), N);
         vec.set_len(vec.len() + N);
     }
@@ -172,11 +172,7 @@ fn prepend_bytes<const N: usize>(vec: &mut Vec<u8>, prefix: &[u8; N]) {
 fn pretruncate_bytes<const N: usize>(vec: &mut Vec<u8>) {
     assert!(vec.len() >= N);
     unsafe {
-        std::ptr::copy(
-            vec.as_ptr().offset(N as isize),
-            vec.as_mut_ptr(),
-            vec.len() - N,
-        );
+        std::ptr::copy(vec.as_ptr().add(N), vec.as_mut_ptr(), vec.len() - N);
         vec.set_len(vec.len() - N);
     }
 }

diff --git a/src/request/mod.rs b/src/request/mod.rs
@@ -112,7 +112,7 @@ mod test {
 
     #[tokio::test]
     async fn test_region_retry() {
-        #[derive(Clone)]
+        #[derive(Debug, Clone)]
         struct MockRpcResponse;
 
         impl HasKeyErrors for MockRpcResponse {

diff --git a/src/request/plan.rs b/src/request/plan.rs
@@ -599,7 +599,7 @@ where
     }
 }
 
-#[derive(Default)]
+#[derive(Debug, Default)]
 pub struct CleanupLocksResult {
     pub region_error: Option<errorpb::Error>,
     pub key_error: Option<Vec<Error>>,
@@ -784,18 +784,54 @@ where
     type Result = P::Result;
 
     async fn execute(&self) -> Result<Self::Result> {
-        let mut result = self.inner.execute().await?;
-        if let Some(errors) = result.key_errors() {
-            Err(Error::ExtractedErrors(errors))
-        } else if let Some(errors) = result.region_errors() {
-            Err(Error::ExtractedErrors(
-                errors
-                    .into_iter()
-                    .map(|e| Error::RegionError(Box::new(e)))
-                    .collect(),
-            ))
-        } else {
-            Ok(result)
+        let result = self.inner.execute().await;
+        let mut errors = Vec::new();
+        match result {
+            Ok(mut resp) => {
+                if let Some(e) = resp.key_errors() {
+                    errors.extend(e);
+                }
+                if let Some(e) = resp.region_errors() {
+                    errors.extend(e.into_iter().map(|e| Error::RegionError(Box::new(e))));
+                }
+
+                if errors.is_empty() {
+                    return Ok(resp);
+                }
+            }
+            Err(Error::MultipleKeyErrors(e)) => errors.extend(e),
+            Err(e) => errors.push(e),
+        };
+        for error in &mut errors {
+            match error {
+                Error::KvError { message }
+                    if message
+                        .contains("Api version in request does not match with TiKV storage")
+                        && message.contains("storage: V1, request: V2") =>
+                {
+                    *error = Error::ServerKeyspaceNotEnabled
+                }
+                Error::KvError { message }
+                    if message
+                        .contains("Api version in request does not match with TiKV storage")
+                        && message.contains("storage: V1ttl, request: V2") =>
+                {
+                    *error = Error::ServerKeyspaceNotEnabled
+                }
+                Error::KvError { message }
+                    if message
+                        .contains("Api version in request does not match with TiKV storage")
+                        && message.contains("storage: V2, request: V1") =>
+                {
+                    *error = Error::ClientKeyspaceNotEnabled
+                }
+                _ => (),
+            }
+        }
+        match errors.len() {
+            0 => unreachable!(),
+            1 => Err(errors.pop().unwrap()),
+            _ => Err(Error::ExtractedErrors(errors)),
         }
     }
 }

diff --git a/src/store/errors.rs b/src/store/errors.rs
@@ -1,7 +1,5 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
-use std::fmt::Display;
-
 use crate::proto::kvrpcpb;
 use crate::Error;
 
@@ -162,10 +160,11 @@ impl HasKeyErrors for kvrpcpb::PessimisticRollbackResponse {
     }
 }
 
-impl<T: HasKeyErrors, E: Display> HasKeyErrors for Result<T, E> {
+impl<T: HasKeyErrors> HasKeyErrors for Result<T, Error> {
     fn key_errors(&mut self) -> Option<Vec<Error>> {
         match self {
             Ok(x) => x.key_errors(),
+            Err(Error::MultipleKeyErrors(e)) => Some(std::mem::take(e)),
             Err(e) => Some(vec![Error::StringError(e.to_string())]),
         }
     }

diff --git a/src/transaction/client.rs b/src/transaction/client.rs
@@ -299,6 +299,7 @@ impl Client {
         let req = new_unsafe_destroy_range_request(range);
         let plan = crate::request::PlanBuilder::new(self.pd.clone(), self.api_version, req)
             .all_stores(DEFAULT_STORE_BACKOFF)
+            .extract_error()
             .merge(crate::request::Collect)
             .plan();
         plan.execute().await

diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs
@@ -146,6 +146,7 @@ impl<PdC: PdClient> Transaction<PdC> {
                 let plan = PlanBuilder::new(rpc, api_version, request)
                     .resolve_lock(retry_options.lock_backoff, api_version)
                     .retry_multi_region(DEFAULT_REGION_BACKOFF)
+                    .extract_error()
                     .merge(CollectSingle)
                     .post_process_default()
                     .plan();
@@ -281,6 +282,7 @@ impl<PdC: PdClient> Transaction<PdC> {
                 let plan = PlanBuilder::new(rpc, api_version, request)
                     .resolve_lock(retry_options.lock_backoff, api_version)
                     .retry_multi_region(retry_options.region_backoff)
+                    .extract_error()
                     .merge(Collect)
                     .plan();
                 plan.execute()
@@ -333,7 +335,11 @@ impl<PdC: PdClient> Transaction<PdC> {
             let keys = keys
                 .into_iter()
                 .map(move |k| k.into().encode_version(api_version, KeyMode::Txn));
-            self.pessimistic_lock(keys, true).await
+            let pairs = self
+                .pessimistic_lock(keys, true)
+                .await?
+                .truncate_version(api_version);
+            Ok(pairs)
         }
     }
 
@@ -541,23 +547,14 @@ impl<PdC: PdClient> Transaction<PdC> {
     /// # Examples
     ///
     /// ```rust,no_run
-    /// # use tikv_client::{Key, Config, TransactionClient, proto::kvrpcpb};
+    /// # use tikv_client::{Key, Config, TransactionClient, transaction::Mutation};
     /// # use futures::prelude::*;
     /// # futures::executor::block_on(async {
     /// # let client = TransactionClient::new(vec!["192.168.0.100", "192.168.0.101"]).await.unwrap();
     /// let mut txn = client.begin_optimistic().await.unwrap();
     /// let mutations = vec![
-    ///     kvrpcpb::Mutation {
-    ///         op: kvrpcpb::Op::Del.into(),
-    ///         key: b"k0".to_vec(),
-    ///         ..Default::default()
-    ///     },
-    ///     kvrpcpb::Mutation {
-    ///         op: kvrpcpb::Op::Put.into(),
-    ///         key: b"k1".to_vec(),
-    ///         value: b"v1".to_vec(),
-    ///         ..Default::default()
-    ///     },
+    ///     Mutation::Delete("k0".to_owned().into()),
+    ///     Mutation::Put("k1".to_owned().into(), b"v1".to_vec()),
     /// ];
     /// txn.batch_mutate(mutations).await.unwrap();
     /// txn.commit().await.unwrap();
@@ -574,7 +571,7 @@ impl<PdC: PdClient> Transaction<PdC> {
             .map(|mutation| mutation.encode_version(self.api_version, KeyMode::Txn))
             .collect();
         if self.is_pessimistic() {
-            self.pessimistic_lock(mutations.iter().map(|m| Key::from(m.key().clone())), false)
+            self.pessimistic_lock(mutations.iter().map(|m| m.key().clone()), false)
                 .await?;
             for m in mutations {
                 self.buffer.mutate(m);
@@ -623,12 +620,11 @@ impl<PdC: PdClient> Transaction<PdC> {
         match self.options.kind {
             TransactionKind::Optimistic => {
                 for key in keys {
-                    self.buffer.lock(key.into());
+                    self.buffer.lock(key);
                 }
             }
             TransactionKind::Pessimistic(_) => {
-                self.pessimistic_lock(keys, false)
-                    .await?;
+                self.pessimistic_lock(keys, false).await?;
             }
         }
         Ok(())
@@ -771,6 +767,7 @@ impl<PdC: PdClient> Transaction<PdC> {
                 self.api_version,
             )
             .retry_multi_region(self.options.retry_options.region_backoff.clone())
+            .extract_error()
             .merge(CollectSingle)
             .post_process_default()
             .plan();
@@ -803,6 +800,7 @@ impl<PdC: PdClient> Transaction<PdC> {
                     let plan = PlanBuilder::new(rpc, api_version, request)
                         .resolve_lock(retry_options.lock_backoff, api_version)
                         .retry_multi_region(retry_options.region_backoff)
+                        .extract_error()
                         .merge(Collect)
                         .plan();
                     plan.execute()
@@ -863,6 +861,7 @@ impl<PdC: PdClient> Transaction<PdC> {
             )
             .preserve_shard()
             .retry_multi_region_preserve_results(self.options.retry_options.region_backoff.clone())
+            .extract_error()
             .merge(CollectWithShard)
             .plan();
         let pairs = plan.execute().await;
@@ -920,6 +919,7 @@ impl<PdC: PdClient> Transaction<PdC> {
             )
             .retry_multi_region(self.options.retry_options.region_backoff.clone())
             .extract_error()
+            .extract_error()
             .plan();
         plan.execute().await?;
 
@@ -989,6 +989,7 @@ impl<PdC: PdClient> Transaction<PdC> {
                 let plan = PlanBuilder::new(rpc.clone(), api_version, request)
                     .retry_multi_region(region_backoff.clone())
                     .merge(CollectSingle)
+                    .extract_error()
                     .plan();
                 plan.execute().await?;
             }
@@ -1042,7 +1043,6 @@ impl<PdC: PdClient> Drop for Transaction<PdC> {
         if self.get_status() == TransactionStatus::Active {
             match self.options.check_level {
                 CheckLevel::Panic => {
-                    dbg!(&self.timestamp);
                     panic!("Dropping an active transaction. Consider commit or rollback it.")
                 }
                 CheckLevel::Warn => {

diff --git a/tests/failpoint_tests.rs b/tests/failpoint_tests.rs
@@ -374,7 +374,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![].., 1024).await?;
+    let locks = client.scan_locks(&ts, .., 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(|l| l.key));
     Ok(locks_set.len())