refactor: introduce an atomic mode for raw client

Signed-off-by: ekexium <[email protected]>

src/raw/client.rs

@@ -24,6 +24,8 @@ const MAX_RAW_KV_SCAN_LIMIT: u32 = 10240;
 pub struct Client {
  rpc: Arc<PdRpcClient>,
  cf: Option<ColumnFamily>,
+ /// Whether to use the [`atomic mode`](Client::with_atomic).
+ atomic: bool,
 }
 
 impl Client {
@@ -63,7 +65,11 @@ impl Client {
  ) -> Result<Client> {
  let pd_endpoints: Vec<String> = pd_endpoints.into_iter().map(Into::into).collect();
  let rpc = Arc::new(PdRpcClient::connect(&pd_endpoints, &config, false).await?);
- Ok(Client { rpc, cf: None })
+ Ok(Client {
+ rpc,
+ cf: None,
+ atomic: false,
+ })
  }
 
  /// Set the column family of requests.
@@ -89,6 +95,22 @@ impl Client {
  Client {
  rpc: self.rpc.clone(),
  cf: Some(cf),
+ atomic: self.atomic,
+ }
+ }
+
+ /// Set to use the atomic mode.
+ ///
+ /// The only reason of using atomic mode is the
+ /// [`compare_and_swap`](Client::compare_and_swap) operation. To guarantee
+ /// the atomicity of CAS, write operations like [`put`](Client::put) or
+ /// [`delete`](Client::delete) in atomic mode are more expensive. Some
+ /// operations are not supported in the mode.
+ pub fn with_atomic(&self) -> Client {
+ Client {
+ rpc: self.rpc.clone(),
+ cf: self.cf.clone(),
+ atomic: true,
  }
  }
 
@@ -171,15 +193,15 @@ impl Client {
  /// # });
  /// ```
  pub async fn put(&self, key: impl Into<Key>, value: impl Into<Value>) -> Result<()> {
- self.put_inner(key, value, false).await
- }
-
- /// Create a new *atomic* 'put' request.
- /// Atomic operations can block each other on the same key.
- ///
- /// Once resolved this request will result in the setting of the value associated with the given key.
- pub async fn atomic_put(&self, key: impl Into<Key>, value: impl Into<Value>) -> Result<()> {
- self.put_inner(key, value, true).await
+ let request = new_raw_put_request(key.into(), value.into(), self.cf.clone(), self.atomic);
+  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
+ .single_region()
+  .await?
+  .retry_region(DEFAULT_REGION_BACKOFF)
+  .extract_error()
+  .plan();
+  plan.execute().await?;
+ Ok(())
  }
 
  /// Create a new 'batch put' request.
@@ -203,19 +225,18 @@ impl Client {
  &self,
  pairs: impl IntoIterator<Item = impl Into<KvPair>>,
  ) -> Result<()> {
- self.batch_put_inner(pairs, false).await
- }
-
- /// Create a new *atomic* 'batch put' request.
- /// Atomic operations can block each other on the same key.
- ///
- /// Once resolved this request will result in the setting of the values
- /// associated with the given keys.
- pub async fn atomic_batch_put(
- &self,
- pairs: impl IntoIterator<Item = impl Into<KvPair>>,
- ) -> Result<()> {
- self.batch_put_inner(pairs, true).await
+ let request = new_raw_batch_put_request(
+ pairs.into_iter().map(Into::into),
+ self.cf.clone(),
+ self.atomic,
+ );
+ let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
+ .multi_region()
+ .retry_region(DEFAULT_REGION_BACKOFF)
+ .extract_error()
+ .plan();
+ plan.execute().await?;
+ Ok(())
  }
 
  /// Create a new 'delete' request.
@@ -236,29 +257,15 @@ impl Client {
  /// # });
  /// ```
  pub async fn delete(&self, key: impl Into<Key>) -> Result<()> {
- self.delete_inner(key, false).await
- }
-
- /// Create a new *atomic* 'delete' request.
- /// Atomic operations can block each other on the same key.
- ///
- /// Once resolved this request will result in the deletion of the given key.
- ///
- /// It does not return an error if the key does not exist in TiKV.
- ///
- /// # Examples
- /// ```rust,no_run
- /// # use tikv_client::{Key, Config, RawClient};
- /// # use futures::prelude::*;
- /// # futures::executor::block_on(async {
- /// # let client = RawClient::new(vec!["192.168.0.100"]).await.unwrap();
- /// let key = "TiKV".to_owned();
- /// let req = client.delete(key);
- /// let result: () = req.await.unwrap();
- /// # });
- /// ```
- pub async fn atomic_delete(&self, key: impl Into<Key>) -> Result<()> {
- self.delete_inner(key, true).await
+ let request = new_raw_delete_request(key.into(), self.cf.clone(), self.atomic);
+ let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
+ .single_region()
+ .await?
+ .retry_region(DEFAULT_REGION_BACKOFF)
+ .extract_error()
+ .plan();
+ plan.execute().await?;
+ Ok(())
  }
 
  /// Create a new 'batch delete' request.
@@ -279,6 +286,7 @@ impl Client {
  /// # });
  /// ```
  pub async fn batch_delete(&self, keys: impl IntoIterator<Item = impl Into<Key>>) -> Result<()> {
+ self.assert_non_atomic()?;
  let request =
  new_raw_batch_delete_request(keys.into_iter().map(Into::into), self.cf.clone());
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
@@ -306,6 +314,7 @@ impl Client {
  /// # });
  /// ```
  pub async fn delete_range(&self, range: impl Into<BoundRange>) -> Result<()> {
+ self.assert_non_atomic()?;
  let request = new_raw_delete_range_request(range.into(), self.cf.clone());
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
  .multi_region()
@@ -436,18 +445,23 @@ impl Client {
 
  /// Create a new *atomic* 'compare and set' request.
  ///
- /// Once resolved this request will result in an atomic `compare and set' operation for the given key.
+ /// Once resolved this request will result in an atomic `compare and set'
+ /// operation for the given key.
  ///
- /// If the value retrived is equal to `current_value`, `new_value` is written.
+ /// If the value retrived is equal to `current_value`, `new_value` is
+ /// written.
  ///
  /// # Return Value
- /// A tuple is returned if successful: the previous value and whether the value is swapped
- pub async fn atomic_compare_and_swap(
+ ///
+ /// A tuple is returned if successful: the previous value and whether the
+ /// value is swapped
+ pub async fn compare_and_swap(
  &self,
  key: impl Into<Key>,
  previous_value: impl Into<Option<Value>>,
  new_value: impl Into<Value>,
  ) -> Result<(Option<Value>, bool)> {
+ self.assert_atomic()?;
  let req = new_cas_request(
  key.into(),
  new_value.into(),
@@ -516,48 +530,15 @@ impl Client {
  plan.execute().await
  }
 
- async fn put_inner(
- &self,
- key: impl Into<Key>,
- value: impl Into<Value>,
- atomic: bool,
- ) -> Result<()> {
- let request = new_raw_put_request(key.into(), value.into(), self.cf.clone(), atomic);
- let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
- .single_region()
- .await?
- .retry_region(DEFAULT_REGION_BACKOFF)
- .extract_error()
- .plan();
- plan.execute().await?;
- Ok(())
- }
-
- async fn batch_put_inner(
- &self,
- pairs: impl IntoIterator<Item = impl Into<KvPair>>,
- atomic: bool,
- ) -> Result<()> {
- let request =
- new_raw_batch_put_request(pairs.into_iter().map(Into::into), self.cf.clone(), atomic);
- let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
- .multi_region()
- .retry_region(DEFAULT_REGION_BACKOFF)
- .extract_error()
- .plan();
- plan.execute().await?;
- Ok(())
+ fn assert_non_atomic(&self) -> Result<()> {
+ (!self.atomic)
+ .then(|| ())
+ .ok_or(Error::UnsupportedInAtomicMode)
  }
 
- async fn delete_inner(&self, key: impl Into<Key>, atomic: bool) -> Result<()> {
- let request = new_raw_delete_request(key.into(), self.cf.clone(), atomic);
- let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
- .single_region()
- .await?
- .retry_region(DEFAULT_REGION_BACKOFF)
- .extract_error()
- .plan();
- plan.execute().await?;
- Ok(())
+ fn assert_atomic(&self) -> Result<()> {
+ self.atomic
+ .then(|| ())
+ .ok_or(Error::UnsupportedInNonAtomicMode)
  }
 }

tests/integration_tests.rs

@@ -11,7 +11,8 @@ use std::{
  iter,
 };
 use tikv_client::{
- Key, KvPair, RawClient, Result, Transaction, TransactionClient, TransactionOptions, Value,
+ Error, Key, KvPair, RawClient, Result, Transaction, TransactionClient, TransactionOptions,
+ Value,
 };
 
 // Parameters used in test
@@ -759,19 +760,19 @@ async fn pessimistic_heartbeat() -> Result<()> {
 #[serial]
 async fn raw_cas() -> Result<()> {
  clear_tikv().await;
- let client = RawClient::new(pd_addrs()).await?;
+ let client = RawClient::new(pd_addrs()).await?.with_atomic();
  let key = "key".to_owned();
  let value = "value".to_owned();
  let new_value = "new value".to_owned();
 
- client.atomic_put(key.clone(), value.clone()).await?;
+ client.put(key.clone(), value.clone()).await?;
  assert_eq!(
  client.get(key.clone()).await?.unwrap(),
  value.clone().as_bytes()
  );
 
  client
- .atomic_compare_and_swap(
+ .compare_and_swap(
  key.clone(),
  Some("another_value".to_owned()).map(|v| v.into()),
  new_value.clone(),
@@ -783,7 +784,7 @@ async fn raw_cas() -> Result<()> {
  );
 
  client
- .atomic_compare_and_swap(
+ .compare_and_swap(
  key.clone(),
  Some(value.to_owned()).map(|v| v.into()),
  new_value.clone(),
@@ -794,8 +795,24 @@ async fn raw_cas() -> Result<()> {
  new_value.clone().as_bytes()
  );
 
- client.atomic_delete(key.clone()).await?;
+ client.delete(key.clone()).await?;
  assert!(client.get(key.clone()).await?.is_none());
+
+ // check unsupported operations
+ assert!(matches!(
+ client.batch_delete(vec![key.clone()]).await.err().unwrap(),
+ Error::UnsupportedInAtomicMode
+ ));
+ let client = RawClient::new(pd_addrs()).await?;
+ assert!(matches!(
+ client
+ .compare_and_swap(key.clone(), None, vec![])
+ .await
+ .err()
+ .unwrap(),
+ Error::UnsupportedInNonAtomicMode
+ ));
+
  Ok(())
 }
 

tikv-client-common/src/errors.rs

@@ -28,6 +28,12 @@ pub enum Error {
  /// An operation requires a primary key, but the transaction was empty.
  #[error("transaction has no primary key")]
  NoPrimaryKey,
+ #[error(
+ "The operation does is not supported in raw-atomic mode, please consider using an atomic client"
+ )]
+ UnsupportedInAtomicMode,
+ #[error("The operation is only supported in raw-atomic mode, please consider using a non-atomic raw client")]
+ UnsupportedInNonAtomicMode,
  /// Wraps a `std::io::Error`.
  #[error("IO error: {0}")]
  Io(#[from] std::io::Error),