Merge remote-tracking branch 'upstream/master' into buffer-mutex

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

src/raw/client.rs

@@ -3,7 +3,7 @@
 use tikv_client_common::Error;
 
 use crate::{
- backoff::{DEFAULT_REGION_BACKOFF, OPTIMISTIC_BACKOFF},
+ backoff::DEFAULT_REGION_BACKOFF,
  config::Config,
  pd::PdRpcClient,
  raw::lowering::*,
@@ -115,7 +115,6 @@ impl Client {
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
  .single_region()
  .await?
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .retry_region(DEFAULT_REGION_BACKOFF)
  .post_process_default()
  .plan();
@@ -146,7 +145,6 @@ impl Client {
  ) -> Result<Vec<KvPair>> {
  let request = new_raw_batch_get_request(keys.into_iter().map(Into::into), self.cf.clone());
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .multi_region()
  .retry_region(DEFAULT_REGION_BACKOFF)
  .merge(Collect)
@@ -177,7 +175,6 @@ impl Client {
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
  .single_region()
  .await?
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .retry_region(DEFAULT_REGION_BACKOFF)
  .extract_error()
  .plan();
@@ -208,7 +205,6 @@ impl Client {
  ) -> Result<()> {
  let request = new_raw_batch_put_request(pairs.into_iter().map(Into::into), self.cf.clone());
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .multi_region()
  .retry_region(DEFAULT_REGION_BACKOFF)
  .extract_error()
@@ -239,7 +235,6 @@ impl Client {
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
  .single_region()
  .await?
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .retry_region(DEFAULT_REGION_BACKOFF)
  .extract_error()
  .plan();
@@ -268,7 +263,6 @@ impl Client {
  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)
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .multi_region()
  .retry_region(DEFAULT_REGION_BACKOFF)
  .extract_error()
@@ -295,7 +289,6 @@ impl Client {
  pub async fn delete_range(&self, range: impl Into<BoundRange>) -> Result<()> {
  let request = new_raw_delete_range_request(range.into(), self.cf.clone());
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .multi_region()
  .retry_region(DEFAULT_REGION_BACKOFF)
  .extract_error()
@@ -437,7 +430,6 @@ impl Client {
 
  let request = new_raw_scan_request(range.into(), limit, key_only, self.cf.clone());
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .multi_region()
  .retry_region(DEFAULT_REGION_BACKOFF)
  .merge(Collect)
@@ -469,7 +461,6 @@ impl Client {
  self.cf.clone(),
  );
  let plan = crate::request::PlanBuilder::new(self.rpc.clone(), request)
- .resolve_lock(OPTIMISTIC_BACKOFF)
  .multi_region()
  .retry_region(DEFAULT_REGION_BACKOFF)
  .merge(Collect)

src/transaction/buffer.rs

@@ -20,14 +20,15 @@ struct InnerBuffer {
 impl InnerBuffer {
  fn insert(&mut self, key: impl Into<Key>, entry: BufferEntry) {
  let key = key.into();
- if !matches!(entry, BufferEntry::Cached(_)) {
+ if !matches!(entry, BufferEntry::Cached(_) | BufferEntry::CheckNotExist) {
  self.primary_key.get_or_insert_with(|| key.clone());
  }
  self.entry_map.insert(key, entry);
  }
 
- pub fn get_primary_key_or(&mut self, key: &Key) -> &Key {
- self.primary_key.get_or_insert(key.clone())
+ /// Set the primary key if it is not set
+ pub fn primary_key_or(&mut self, key: &Key) {
+ self.primary_key.get_or_insert(key.clone());
  }
 }
 
@@ -48,9 +49,9 @@ impl Buffer {
  self.inner.primary_key.clone()
  }
 
- /// Get the primary key of the buffer, if not exists, use `key` as the primary key.
- pub async fn get_primary_key_or(&mut self, key: &Key) -> Key {
- self.inner.get_primary_key_or(key).clone()
+ /// Set the primary key if it is not set
+ pub async fn primary_key_or(&mut self, key: &Key) {
+ self.inner.primary_key_or(key);
  }
 
  /// Get a value from the buffer.
@@ -186,7 +187,7 @@ impl Buffer {
 
  /// Lock the given key if necessary.
  pub async fn lock(&mut self, key: Key) {
- self.inner.primary_key.get_or_insert(key.clone());
+ self.inner.primary_key.get_or_insert_with(|| key.clone());
  let value = self
  .inner
  .entry_map

src/transaction/transaction.rs

@@ -116,18 +116,30 @@ impl<PdC: PdClient> Transaction<PdC> {
  }
 
  /// Create a `get for udpate` request.
- /// Once resolved this request will pessimistically lock and fetch the latest
- /// value associated with the given key at **current timestamp**.
  ///
- /// The "current timestamp" (also called `for_update_ts` of the request) is fetched immediately from PD.
+ /// The request reads and "locks" a key. It is similar to `SELECT ... FOR
+ /// UPDATE` in TiDB, and has different behavior in optimistic and
+ /// pessimistic transactions.
  ///
- /// Note: The behavior of this command does not follow snapshot isolation. It is similar to `select for update` in TiDB,
- /// which is similar to that in MySQL. It reads the latest value (using current timestamp),
- /// and the value is not cached in the local buffer.
- /// So normal `get`-like commands after `get_for_update` will not be influenced, they still read values at `start_ts`.
+ /// # Optimistic transaction
  ///
+ /// It reads at the "start timestamp" and caches the value, just like normal
+ /// get requests. The lock is written in prewrite and commit, so it cannot
+ /// prevent concurrent transactions from writing the same key, but can only
+ /// prevent itself from committing.
  ///
- /// It can only be used in pessimistic mode.
+ /// # Pessimistic transaction
+ ///
+ /// It reads at the "current timestamp" and thus does not cache the value.
+ /// So following read requests won't be affected by the `get_for_udpate`.
+ /// A lock will be acquired immediately with this request, which prevents
+ /// concurrent transactions from mutating the keys.
+ ///
+ /// The "current timestamp" (also called `for_update_ts` of the request) is
+ /// fetched immediately from the timestamp oracle.
+ ///
+ /// Note: The behavior of the request under pessimistic transaction does not
+ /// follow snapshot isolation.
  ///
  /// # Examples
  /// ```rust,no_run
@@ -146,7 +158,9 @@ impl<PdC: PdClient> Transaction<PdC> {
  pub async fn get_for_update(&mut self, key: impl Into<Key>) -> Result<Option<Value>> {
  self.check_allow_operation().await?;
  if !self.is_pessimistic() {
- Err(Error::InvalidTransactionType)
+ let key = key.into();
+ self.lock_keys(iter::once(key.clone())).await?;
+ self.get(key).await
  } else {
  let mut pairs = self.pessimistic_lock(iter::once(key.into()), true).await?;
  debug_assert!(pairs.len() <= 1);
@@ -228,33 +242,25 @@ impl<PdC: PdClient> Transaction<PdC> {
 
  /// Create a new 'batch get for update' request.
  ///
- /// Once resolved this request will pessimistically lock the keys and
- /// fetch the values associated with the given keys.
- ///
- /// Note: The behavior of this command does not follow snapshot isolation. It is similar to `select for update` in TiDB,
- /// which is similar to that in MySQL. It reads the latest value (using current timestamp),
- /// and the value is not cached in the local buffer.
- /// So normal `get`-like commands after `batch_get_for_update` will not be influenced, they still read values at `start_ts`.
- ///
- /// Non-existent entries will not appear in the result. The order of the keys is not retained in the result.
+ /// Similar [`get_for_update`](Transaction::get_for_update), but it works
+ /// for a batch of keys.
  ///
- /// It can only be used in pessimistic mode.
+ /// Non-existent entries will not appear in the result. The order of the
+ /// keys is not retained in the result.
  ///
  /// # Examples
  /// ```rust,no_run
- /// # use tikv_client::{Key, Value, Config, TransactionClient};
+ /// # use tikv_client::{Key, Value, Config, TransactionClient, KvPair};
  /// # use futures::prelude::*;
  /// # use std::collections::HashMap;
  /// # 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_pessimistic().await.unwrap();
  /// let keys = vec!["TiKV".to_owned(), "TiDB".to_owned()];
- /// let result: HashMap<Key, Value> = txn
+ /// let result: Vec<KvPair> = txn
  /// .batch_get_for_update(keys)
  /// .await
- /// .unwrap()
- /// .map(|pair| (pair.0, pair.1))
- /// .collect();
+ /// .unwrap();
  /// // now "TiKV" and "TiDB" are both locked
  /// // Finish the transaction...
  /// txn.commit().await.unwrap();
@@ -263,13 +269,15 @@ impl<PdC: PdClient> Transaction<PdC> {
  pub async fn batch_get_for_update(
  &mut self,
  keys: impl IntoIterator<Item = impl Into<Key>>,
- ) -> Result<impl Iterator<Item = KvPair>> {
+ ) -> Result<Vec<KvPair>> {
  self.check_allow_operation().await?;
+ let keys: Vec<Key> = keys.into_iter().map(|k| k.into()).collect();
  if !self.is_pessimistic() {
- return Err(Error::InvalidTransactionType);
+ self.lock_keys(keys.clone()).await?;
+ Ok(self.batch_get(keys).await?.collect())
+ } else {
+ self.pessimistic_lock(keys, true).await
  }
- let keys: Vec<Key> = keys.into_iter().map(|it| it.into()).collect();
- Ok(self.pessimistic_lock(keys, true).await?.into_iter())
  }
 
  /// Create a new 'scan' request.
@@ -473,8 +481,8 @@ impl<PdC: PdClient> Transaction<PdC> {
  }
  }
  TransactionKind::Pessimistic(_) => {
- self.pessimistic_lock(keys.into_iter().map(|k| k.into()), false)
-  .await?;
+ let keys: Vec<Key> = keys.into_iter().map(|k| k.into()).collect();
+ self.pessimistic_lock(keys.into_iter(), false).await?;
  }
  }
  Ok(())
@@ -649,7 +657,13 @@ impl<PdC: PdClient> Transaction<PdC> {
  }
 
  let first_key = keys[0].clone().key();
- let primary_lock = self.buffer.get_primary_key_or(&first_key).await;
+ // we do not set the primary key here, because pessimistic lock request
+ // can fail, in which case the keys may not be part of the transaction.
+ let primary_lock = self
+ .buffer
+ .get_primary_key()
+ .await
+ .unwrap_or_else(|| first_key.clone());
  let for_update_ts = self.rpc.clone().get_timestamp().await?;
  self.options.push_for_update_ts(for_update_ts.clone());
  let request = new_pessimistic_lock_request(
@@ -669,6 +683,9 @@ impl<PdC: PdClient> Transaction<PdC> {
  .plan();
  let pairs = plan.execute().await;
 
+ // primary key will be set here if needed
+ self.buffer.primary_key_or(&first_key).await;
+
  self.start_auto_heartbeat().await;
 
  for key in keys {

tests/integration_tests.rs

@@ -305,10 +305,7 @@ async fn txn_write_million() -> Result<()> {
  let res = txn.batch_get(keys.clone()).await?.collect::<Vec<_>>();
  assert_eq!(res.len(), keys.len());
 
- let res = txn
- .batch_get_for_update(keys.clone())
- .await?
- .collect::<Vec<_>>();
+ let res = txn.batch_get_for_update(keys.clone()).await?;
  assert_eq!(res.len(), keys.len());
 
  txn.commit().await?;
@@ -686,7 +683,8 @@ async fn get_for_update() -> Result<()> {
 
  let mut t1 = client.begin_pessimistic().await?;
  let mut t2 = client.begin_pessimistic().await?;
-
+ let mut t3 = client.begin_optimistic().await?;
+ let mut t4 = client.begin_optimistic().await?;
  let mut t0 = client.begin_pessimistic().await?;
  t0.put(key1.clone(), value1).await?;
  t0.put(key2.clone(), value2).await?;
@@ -700,12 +698,19 @@ async fn get_for_update() -> Result<()> {
  let res: HashMap<_, _> = t2
  .batch_get_for_update(keys.clone())
  .await?
+ .into_iter()
  .map(From::from)
  .collect();
  t2.commit().await?;
- assert!(res.get(&key1.into()).unwrap() == &value1);
+ assert!(res.get(&key1.clone().into()).unwrap() == &value1);
  assert!(res.get(&key2.into()).unwrap() == &value2);
 
+ assert!(t3.get_for_update(key1).await?.is_none());
+ assert!(t3.commit().await.is_err());
+
+ assert!(t4.batch_get_for_update(keys).await?.len() == 0);
+ assert!(t4.commit().await.is_err());
+
  Ok(())
 }