Merge branch 'master' into sync-api-txclient

Author: odecode

Cargo.toml

@@ -14,7 +14,6 @@ prometheus = ["prometheus/push", "prometheus/process"]
 # Enable integration tests with a running TiKV and PD instance.
 # Use $PD_ADDRS, comma separated, to set the addresses the tests use.
 integration-tests = []
-apiv2-no-prefix = []
 
 [lib]
 name = "tikv_client"

Makefile

@@ -5,11 +5,11 @@ export RUSTFLAGS=-Dwarnings
 export PD_ADDRS     ?= 127.0.0.1:2379
 export MULTI_REGION ?= 1
 
-ALL_FEATURES := integration-tests apiv2-no-prefix
+ALL_FEATURES := integration-tests
 
 NEXTEST_ARGS := --config-file $(shell pwd)/config/nextest.toml
 
-INTEGRATION_TEST_ARGS := --features "integration-tests apiv2-no-prefix" --test-threads 1
+INTEGRATION_TEST_ARGS := --features "integration-tests" --test-threads 1
 
 RUN_INTEGRATION_TEST := cargo nextest run ${NEXTEST_ARGS} --all ${INTEGRATION_TEST_ARGS}
 

README.md

@@ -17,7 +17,7 @@ The TiKV client is a Rust library (crate). To use this crate in your project, ad
 
 ```toml
 [dependencies]
-tikv-client = "0.3"
+tikv-client = "0.4"
 ```
 
 ### Prerequisites

src/common/errors.rs

@@ -8,6 +8,20 @@ use crate::proto::kvrpcpb;
 use crate::region::RegionVerId;
 use crate::BoundRange;
 
+/// Protobuf-generated region-level error returned by TiKV.
+///
+/// This type is generated from TiKV's protobuf definitions and may change in a
+/// future release even if the wire format is compatible.
+#[doc(inline)]
+pub use crate::proto::errorpb::Error as ProtoRegionError;
+
+/// Protobuf-generated per-key error returned by TiKV.
+///
+/// This type is generated from TiKV's protobuf definitions and may change in a
+/// future release even if the wire format is compatible.
+#[doc(inline)]
+pub use crate::proto::kvrpcpb::KeyError as ProtoKeyError;
+
 /// An error originating from the TiKV client or dependencies.
 #[derive(Debug, Error)]
 #[allow(clippy::large_enum_variant)]
@@ -63,13 +77,13 @@ pub enum Error {
     Canceled(#[from] futures::channel::oneshot::Canceled),
     /// Errors caused by changes of region information
     #[error("Region error: {0:?}")]
-    RegionError(Box<crate::proto::errorpb::Error>),
+    RegionError(Box<ProtoRegionError>),
     /// Whether the transaction is committed or not is undetermined
     #[error("Whether the transaction is committed or not is undetermined")]
     UndeterminedError(Box<Error>),
-    /// Wraps `crate::proto::kvrpcpb::KeyError`
+    /// Wraps a per-key error returned by TiKV.
     #[error("{0:?}")]
-    KeyError(Box<crate::proto::kvrpcpb::KeyError>),
+    KeyError(Box<ProtoKeyError>),
     /// Multiple errors generated from the ExtractError plan.
     #[error("Multiple errors: {0:?}")]
     ExtractedErrors(Vec<Error>),
@@ -122,14 +136,14 @@ Use TransactionClient instead or move SyncTransactionClient usage outside the as
     NestedRuntimeError(String),
 }
 
-impl From<crate::proto::errorpb::Error> for Error {
-    fn from(e: crate::proto::errorpb::Error) -> Error {
+impl From<ProtoRegionError> for Error {
+    fn from(e: ProtoRegionError) -> Error {
         Error::RegionError(Box::new(e))
     }
 }
 
-impl From<crate::proto::kvrpcpb::KeyError> for Error {
-    fn from(e: crate::proto::kvrpcpb::KeyError) -> Error {
+impl From<ProtoKeyError> for Error {
+    fn from(e: ProtoKeyError) -> Error {
         Error::KeyError(Box::new(e))
     }
 }

src/common/mod.rs

@@ -4,4 +4,6 @@ mod errors;
 pub mod security;
 
 pub use self::errors::Error;
+pub use self::errors::ProtoKeyError;
+pub use self::errors::ProtoRegionError;
 pub use self::errors::Result;

src/kv/mod.rs

@@ -14,7 +14,7 @@ pub use key::KvPairTTL;
 pub use kvpair::KvPair;
 pub use value::Value;
 
-struct HexRepr<'a>(pub &'a [u8]);
+pub struct HexRepr<'a>(pub &'a [u8]);
 
 impl fmt::Display for HexRepr<'_> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

src/lib.rs

@@ -122,6 +122,10 @@ pub use common::security::SecurityManager;
 #[doc(inline)]
 pub use common::Error;
 #[doc(inline)]
+pub use common::ProtoKeyError;
+#[doc(inline)]
+pub use common::ProtoRegionError;
+#[doc(inline)]
 pub use common::Result;
 #[doc(inline)]
 pub use config::Config;
@@ -157,6 +161,8 @@ pub use crate::transaction::CheckLevel;
 #[doc(inline)]
 pub use crate::transaction::Client as TransactionClient;
 #[doc(inline)]
+pub use crate::transaction::ProtoLockInfo;
+#[doc(inline)]
 pub use crate::transaction::Snapshot;
 #[doc(inline)]
 pub use crate::transaction::SyncSnapshot;

src/request/keyspace.rs

@@ -13,6 +13,7 @@ pub const TXN_KEY_PREFIX: u8 = b'x';
 pub const KEYSPACE_PREFIX_LEN: usize = 4;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[non_exhaustive]
 pub enum Keyspace {
     Disable,
     Enable {
@@ -23,7 +24,6 @@ pub enum Keyspace {
     /// This mode is intended for **server-side embedding** use cases (e.g. embedding this client in
     /// `tikv-server`) where keys are already in API V2 "logical key bytes" form and must be passed
     /// through unchanged.
-    #[cfg(feature = "apiv2-no-prefix")]
     ApiV2NoPrefix,
 }
 
@@ -38,7 +38,6 @@ impl Keyspace {
         match self {
             Keyspace::Disable => kvrpcpb::ApiVersion::V1,
             Keyspace::Enable { .. } => kvrpcpb::ApiVersion::V2,
-            #[cfg(feature = "apiv2-no-prefix")]
             Keyspace::ApiV2NoPrefix => kvrpcpb::ApiVersion::V2,
         }
     }
@@ -183,6 +182,32 @@ impl TruncateKeyspace for Vec<crate::proto::kvrpcpb::LockInfo> {
     }
 }
 
+impl EncodeKeyspace for Vec<crate::proto::kvrpcpb::LockInfo> {
+    fn encode_keyspace(mut self, keyspace: Keyspace, key_mode: KeyMode) -> Self {
+        if !matches!(keyspace, Keyspace::Enable { .. }) {
+            return self;
+        }
+        for lock in &mut self {
+            take_mut::take(&mut lock.key, |key| {
+                Key::from(key).encode_keyspace(keyspace, key_mode).into()
+            });
+            take_mut::take(&mut lock.primary_lock, |primary| {
+                Key::from(primary)
+                    .encode_keyspace(keyspace, key_mode)
+                    .into()
+            });
+            for secondary in lock.secondaries.iter_mut() {
+                take_mut::take(secondary, |secondary| {
+                    Key::from(secondary)
+                        .encode_keyspace(keyspace, key_mode)
+                        .into()
+                });
+            }
+        }
+        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 {
@@ -274,26 +299,105 @@ mod tests {
             mutation.encode_keyspace(keyspace, key_mode),
             expected_mutation
         );
+
+        let key_mode = KeyMode::Txn;
+        let lock = crate::proto::kvrpcpb::LockInfo {
+            key: vec![b'k', b'1'],
+            primary_lock: vec![b'p', b'1'],
+            secondaries: vec![vec![b's', b'1'], vec![b's', b'2']],
+            ..Default::default()
+        };
+        let locks = vec![lock].encode_keyspace(keyspace, key_mode);
+        assert_eq!(locks.len(), 1);
+        assert_eq!(locks[0].key, vec![b'x', 0, 0xDE, 0xAD, b'k', b'1']);
+        assert_eq!(locks[0].primary_lock, vec![b'x', 0, 0xDE, 0xAD, b'p', b'1']);
+        assert_eq!(
+            locks[0].secondaries,
+            vec![
+                vec![b'x', 0, 0xDE, 0xAD, b's', b'1'],
+                vec![b'x', 0, 0xDE, 0xAD, b's', b'2']
+            ]
+        );
     }
 
     #[test]
     fn test_truncate_version() {
-        let key = Key::from(vec![b'r', 0, 0xDE, 0xAD, 0xBE, 0xEF]);
         let keyspace = Keyspace::Enable {
             keyspace_id: 0xDEAD,
         };
+
+        let key = Key::from(vec![b'r', 0, 0xDE, 0xAD, 0xBE, 0xEF]);
         let expected_key = Key::from(vec![0xBE, 0xEF]);
         assert_eq!(key.truncate_keyspace(keyspace), expected_key);
 
         let key = Key::from(vec![b'x', 0, 0xDE, 0xAD, 0xBE, 0xEF]);
-        let keyspace = Keyspace::Enable {
-            keyspace_id: 0xDEAD,
-        };
         let expected_key = Key::from(vec![0xBE, 0xEF]);
         assert_eq!(key.truncate_keyspace(keyspace), expected_key);
+
+        let pair = KvPair(Key::from(vec![b'x', 0, 0xDE, 0xAD, b'k']), vec![b'v']);
+        let expected_pair = KvPair(Key::from(vec![b'k']), vec![b'v']);
+        assert_eq!(pair.truncate_keyspace(keyspace), expected_pair);
+
+        let range = Range {
+            start: Key::from(vec![b'x', 0, 0xDE, 0xAD, b'a']),
+            end: Key::from(vec![b'x', 0, 0xDE, 0xAD, b'b']),
+        };
+        let expected_range = Range {
+            start: Key::from(vec![b'a']),
+            end: Key::from(vec![b'b']),
+        };
+        assert_eq!(range.truncate_keyspace(keyspace), expected_range);
+
+        let ranges = vec![
+            Range {
+                start: Key::from(vec![b'x', 0, 0xDE, 0xAD, b'a']),
+                end: Key::from(vec![b'x', 0, 0xDE, 0xAD, b'b']),
+            },
+            Range {
+                start: Key::from(vec![b'x', 0, 0xDE, 0xAD, b'c']),
+                end: Key::from(vec![b'x', 0, 0xDE, 0xAD, b'd']),
+            },
+        ];
+        let expected_ranges = vec![
+            Range {
+                start: Key::from(vec![b'a']),
+                end: Key::from(vec![b'b']),
+            },
+            Range {
+                start: Key::from(vec![b'c']),
+                end: Key::from(vec![b'd']),
+            },
+        ];
+        assert_eq!(ranges.truncate_keyspace(keyspace), expected_ranges);
+
+        let pairs = vec![
+            KvPair(Key::from(vec![b'x', 0, 0xDE, 0xAD, b'k']), vec![b'v']),
+            KvPair(
+                Key::from(vec![b'x', 0, 0xDE, 0xAD, b'k', b'2']),
+                vec![b'v', b'2'],
+            ),
+        ];
+        let expected_pairs = vec![
+            KvPair(Key::from(vec![b'k']), vec![b'v']),
+            KvPair(Key::from(vec![b'k', b'2']), vec![b'v', b'2']),
+        ];
+        assert_eq!(pairs.truncate_keyspace(keyspace), expected_pairs);
+
+        let lock = crate::proto::kvrpcpb::LockInfo {
+            key: vec![b'x', 0, 0xDE, 0xAD, b'k'],
+            primary_lock: vec![b'x', 0, 0xDE, 0xAD, b'p'],
+            secondaries: vec![vec![b'x', 0, 0xDE, 0xAD, b's']],
+            ..Default::default()
+        };
+        let expected_lock = crate::proto::kvrpcpb::LockInfo {
+            key: vec![b'k'],
+            primary_lock: vec![b'p'],
+            secondaries: vec![vec![b's']],
+            ..Default::default()
+        };
+        assert_eq!(vec![lock].truncate_keyspace(keyspace), vec![expected_lock]);
     }
 
-    #[cfg(feature = "apiv2-no-prefix")]
     #[test]
     fn test_apiv2_no_prefix_api_version() {
         assert_eq!(
@@ -302,7 +406,6 @@ mod tests {
         );
     }
 
-    #[cfg(feature = "apiv2-no-prefix")]
     #[test]
     fn test_apiv2_no_prefix_encode_is_noop() {
         let keyspace = Keyspace::ApiV2NoPrefix;
@@ -323,9 +426,29 @@ mod tests {
             mutation.clone().encode_keyspace(keyspace, key_mode),
             mutation
         );
+
+        let lock = crate::proto::kvrpcpb::LockInfo {
+            key: vec![b'x', 0, 0, 0, b'k'],
+            primary_lock: vec![b'x', 0, 0, 0, b'p'],
+            secondaries: vec![vec![b'x', 0, 0, 0, b's']],
+            ..Default::default()
+        };
+        let locks = vec![lock];
+        assert_eq!(locks.clone().encode_keyspace(keyspace, key_mode), locks);
+
+        let lock = crate::proto::kvrpcpb::LockInfo {
+            key: vec![b'k', b'1'],
+            primary_lock: vec![b'p', b'1'],
+            secondaries: vec![vec![b's', b'1']],
+            ..Default::default()
+        };
+        let locks = vec![lock.clone()];
+        assert_eq!(
+            locks.clone().encode_keyspace(Keyspace::Disable, key_mode),
+            locks
+        );
     }
 
-    #[cfg(feature = "apiv2-no-prefix")]
     #[test]
     fn test_apiv2_no_prefix_truncate_is_noop() {
         let keyspace = Keyspace::ApiV2NoPrefix;

src/request/plan.rs

@@ -88,7 +88,7 @@ impl<Req: KvRequest + StoreRequest> StoreRequest for Dispatch<Req> {
 const MULTI_REGION_CONCURRENCY: usize = 16;
 const MULTI_STORES_CONCURRENCY: usize = 16;
 
-fn is_grpc_error(e: &Error) -> bool {
+pub(crate) fn is_grpc_error(e: &Error) -> bool {
     matches!(e, Error::GrpcAPI(_) | Error::Grpc(_))
 }
 

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,11 +160,15 @@ 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(e) => Some(vec![Error::StringError(e.to_string())]),
+            Err(Error::MultipleKeyErrors(errs)) => Some(std::mem::take(errs)),
+            Err(e) => Some(vec![std::mem::replace(
+                e,
+                Error::StringError("".to_string()), // placeholder, no use.
+            )]),
         }
     }
 }