tikv
/
client-rust
mirror of https://github.com/tikv/client-rust.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/atomic' into atomic

This commit is contained in:
ekexium 2021-04-20 17:17:14 +08:00
parent 1c6815abaf 5fb1dc071a
commit 873fde5e00
4 changed files with 64 additions and 93 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

123
src/transaction/buffer.rs
View File

@ -1,58 +1,36 @@
// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
use crate::{BoundRange, Key, KvPair, Result, Value}; use crate::{BoundRange, Key, KvPair, Result, Value};
use derive_new::new;
use std::{ use std::{
collections::{btree_map::Entry, BTreeMap, HashMap}, collections::{btree_map::Entry, BTreeMap, HashMap},
future::Future, future::Future,
}; };
use tikv_client_proto::kvrpcpb; use tikv_client_proto::kvrpcpb;
use tokio::sync::{Mutex, MutexGuard};
#[derive(new)]
struct InnerBuffer {
#[new(default)]
primary_key: Option<Key>,
#[new(default)]
entry_map: BTreeMap<Key, BufferEntry>,
is_pessimistic: bool,
}
impl InnerBuffer {
fn insert(&mut self, key: impl Into<Key>, entry: BufferEntry) {
let key = key.into();
if !matches!(entry, BufferEntry::Cached(_) | BufferEntry::CheckNotExist) {
self.primary_key.get_or_insert_with(|| key.clone());
}
self.entry_map.insert(key, entry);
}
/// 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());
}
}
/// A caching layer which buffers reads and writes in a transaction. /// A caching layer which buffers reads and writes in a transaction.
pub struct Buffer { pub struct Buffer {
inner: Mutex<InnerBuffer>, primary_key: Option<Key>,
entry_map: BTreeMap<Key, BufferEntry>,
is_pessimistic: bool,
} }
impl Buffer { impl Buffer {
pub fn new(is_pessimistic: bool) -> Buffer { pub fn new(is_pessimistic: bool) -> Buffer {
Buffer { Buffer {
inner: Mutex::new(InnerBuffer::new(is_pessimistic)), primary_key: None,
entry_map: BTreeMap::new(),
is_pessimistic,
} }
} }
/// Get the primary key of the buffer. /// Get the primary key of the buffer.
pub async fn get_primary_key(&self) -> Option<Key> { pub async fn get_primary_key(&self) -> Option<Key> {
self.inner.lock().await.primary_key.clone() self.primary_key.clone()
} }
/// Set the primary key if it is not set /// Set the primary key if it is not set
pub async fn primary_key_or(&self, key: &Key) { pub async fn primary_key_or(&mut self, key: &Key) {
self.inner.lock().await.primary_key_or(key); self.primary_key.get_or_insert_with(|| key.clone());
} }
/// Get a value from the buffer. /// Get a value from the buffer.
@ -66,7 +44,7 @@ impl Buffer {
/// Get a value from the buffer. If the value is not present, run `f` to get /// Get a value from the buffer. If the value is not present, run `f` to get
/// the value. /// the value.
pub async fn get_or_else<F, Fut>(&self, key: Key, f: F) -> Result<Option<Value>> pub async fn get_or_else<F, Fut>(&mut self, key: Key, f: F) -> Result<Option<Value>>
where where
F: FnOnce(Key) -> Fut, F: FnOnce(Key) -> Fut,
Fut: Future<Output = Result<Option<Value>>>, Fut: Future<Output = Result<Option<Value>>>,
@ -75,8 +53,7 @@ impl Buffer {
MutationValue::Determined(value) => Ok(value), MutationValue::Determined(value) => Ok(value),
MutationValue::Undetermined => { MutationValue::Undetermined => {
let value = f(key.clone()).await?; let value = f(key.clone()).await?;
let mut mutations = self.inner.lock().await; self.update_cache(key, value.clone());
Self::update_cache(&mut mutations, key, value.clone());
Ok(value) Ok(value)
} }
} }
@ -87,7 +64,7 @@ impl Buffer {
/// ///
/// only used for snapshot read (i.e. not for `batch_get_for_update`) /// only used for snapshot read (i.e. not for `batch_get_for_update`)
pub async fn batch_get_or_else<F, Fut>( pub async fn batch_get_or_else<F, Fut>(
&self, &mut self,
keys: impl Iterator<Item = Key>, keys: impl Iterator<Item = Key>,
f: F, f: F,
) -> Result<impl Iterator<Item = KvPair>> ) -> Result<impl Iterator<Item = KvPair>>
@ -96,7 +73,6 @@ impl Buffer {
Fut: Future<Output = Result<Vec<KvPair>>>, Fut: Future<Output = Result<Vec<KvPair>>>,
{ {
let (cached_results, undetermined_keys) = { let (cached_results, undetermined_keys) = {
let mutations = self.inner.lock().await;
// Partition the keys into those we have buffered and those we have to // Partition the keys into those we have buffered and those we have to
// get from the store. // get from the store.
let (undetermined_keys, cached_results): ( let (undetermined_keys, cached_results): (
@ -104,7 +80,7 @@ impl Buffer {
Vec<(Key, MutationValue)>, Vec<(Key, MutationValue)>,
) = keys ) = keys
.map(|key| { .map(|key| {
let value = mutations let value = self
.entry_map .entry_map
.get(&key) .get(&key)
.map(BufferEntry::get_value) .map(BufferEntry::get_value)
@ -122,11 +98,10 @@ impl Buffer {
}; };
let fetched_results = f(Box::new(undetermined_keys)).await?; let fetched_results = f(Box::new(undetermined_keys)).await?;
let mut mutations = self.inner.lock().await;
for kvpair in &fetched_results { for kvpair in &fetched_results {
let key = kvpair.0.clone(); let key = kvpair.0.clone();
let value = Some(kvpair.1.clone()); let value = Some(kvpair.1.clone());
Self::update_cache(&mut mutations, key, value); self.update_cache(key, value);
} }
let results = cached_results.chain(fetched_results.into_iter()); let results = cached_results.chain(fetched_results.into_iter());
@ -135,7 +110,7 @@ impl Buffer {
/// Run `f` to fetch entries in `range` from TiKV. Combine them with mutations in local buffer. Returns the results. /// Run `f` to fetch entries in `range` from TiKV. Combine them with mutations in local buffer. Returns the results.
pub async fn scan_and_fetch<F, Fut>( pub async fn scan_and_fetch<F, Fut>(
&self, &mut self,
range: BoundRange, range: BoundRange,
limit: u32, limit: u32,
f: F, f: F,
@ -145,8 +120,7 @@ impl Buffer {
Fut: Future<Output = Result<Vec<KvPair>>>, Fut: Future<Output = Result<Vec<KvPair>>>,
{ {
// read from local buffer // read from local buffer
let mut mutations = self.inner.lock().await; let mutation_range = self.entry_map.range(range.clone());
let mutation_range = mutations.entry_map.range(range.clone());
// fetch from TiKV // fetch from TiKV
// fetch more entries because some of them may be deleted. // fetch more entries because some of them may be deleted.
@ -177,7 +151,7 @@ impl Buffer {
// update local buffer // update local buffer
for (k, v) in &results { for (k, v) in &results {
Self::update_cache(&mut mutations, k.clone(), Some(v.clone())); self.update_cache(k.clone(), Some(v.clone()));
} }
let mut res = results let mut res = results
@ -190,10 +164,9 @@ impl Buffer {
} }
/// Lock the given key if necessary. /// Lock the given key if necessary.
pub async fn lock(&self, key: Key) { pub async fn lock(&mut self, key: Key) {
let mutations = &mut self.inner.lock().await; self.primary_key.get_or_insert_with(|| key.clone());
mutations.primary_key.get_or_insert_with(|| key.clone()); let value = self
let value = mutations
.entry_map .entry_map
.entry(key) .entry(key)
// Mutated keys don't need a lock. // Mutated keys don't need a lock.
@ -205,27 +178,25 @@ impl Buffer {
} }
/// Insert a value into the buffer (does not write through). /// Insert a value into the buffer (does not write through).
pub async fn put(&self, key: Key, value: Value) { pub async fn put(&mut self, key: Key, value: Value) {
self.inner.lock().await.insert(key, BufferEntry::Put(value)); self.insert_entry(key, BufferEntry::Put(value));
} }
/// Mark a value as Insert mutation into the buffer (does not write through). /// Mark a value as Insert mutation into the buffer (does not write through).
pub async fn insert(&self, key: Key, value: Value) { pub async fn insert(&mut self, key: Key, value: Value) {
let mut mutations = self.inner.lock().await; let mut entry = self.entry_map.entry(key.clone());
let mut entry = mutations.entry_map.entry(key.clone());
match entry { match entry {
Entry::Occupied(ref mut o) if matches!(o.get(), BufferEntry::Del) => { Entry::Occupied(ref mut o) if matches!(o.get(), BufferEntry::Del) => {
o.insert(BufferEntry::Put(value)); o.insert(BufferEntry::Put(value));
} }
_ => mutations.insert(key, BufferEntry::Insert(value)), _ => self.insert_entry(key, BufferEntry::Insert(value)),
} }
} }
/// Mark a value as deleted. /// Mark a value as deleted.
pub async fn delete(&self, key: Key) { pub async fn delete(&mut self, key: Key) {
let mut mutations = self.inner.lock().await; let is_pessimistic = self.is_pessimistic;
let is_pessimistic = mutations.is_pessimistic; let mut entry = self.entry_map.entry(key.clone());
let mut entry = mutations.entry_map.entry(key.clone());
match entry { match entry {
Entry::Occupied(ref mut o) Entry::Occupied(ref mut o)
@ -233,40 +204,32 @@ impl Buffer {
{ {
o.insert(BufferEntry::CheckNotExist); o.insert(BufferEntry::CheckNotExist);
} }
_ => mutations.insert(key, BufferEntry::Del), _ => self.insert_entry(key, BufferEntry::Del),
} }
} }
/// Converts the buffered mutations to the proto buffer version /// Converts the buffered mutations to the proto buffer version
pub async fn to_proto_mutations(&self) -> Vec<kvrpcpb::Mutation> { pub async fn to_proto_mutations(&self) -> Vec<kvrpcpb::Mutation> {
self.inner self.entry_map
.lock()
.await
.entry_map
.iter() .iter()
.filter_map(|(key, mutation)| mutation.to_proto_with_key(key)) .filter_map(|(key, mutation)| mutation.to_proto_with_key(key))
.collect() .collect()
} }
async fn get_from_mutations(&self, key: &Key) -> MutationValue { async fn get_from_mutations(&self, key: &Key) -> MutationValue {
self.inner self.entry_map
.lock()
.await
.entry_map
.get(&key) .get(&key)
.map(BufferEntry::get_value) .map(BufferEntry::get_value)
.unwrap_or(MutationValue::Undetermined) .unwrap_or(MutationValue::Undetermined)
} }
fn update_cache(buffer: &mut MutexGuard<InnerBuffer>, key: Key, value: Option<Value>) { fn update_cache(&mut self, key: Key, value: Option<Value>) {
match buffer.entry_map.get(&key) { match self.entry_map.get(&key) {
Some(BufferEntry::Locked(None)) => { Some(BufferEntry::Locked(None)) => {
buffer self.entry_map.insert(key, BufferEntry::Locked(Some(value)));
.entry_map
.insert(key, BufferEntry::Locked(Some(value)));
} }
None => { None => {
buffer.entry_map.insert(key, BufferEntry::Cached(value)); self.entry_map.insert(key, BufferEntry::Cached(value));
} }
Some(BufferEntry::Cached(v)) | Some(BufferEntry::Locked(Some(v))) => { Some(BufferEntry::Cached(v)) | Some(BufferEntry::Locked(Some(v))) => {
assert!(&value == v); assert!(&value == v);
@ -285,6 +248,14 @@ impl Buffer {
} }
} }
} }
fn insert_entry(&mut self, key: impl Into<Key>, entry: BufferEntry) {
let key = key.into();
if !matches!(entry, BufferEntry::Cached(_) | BufferEntry::CheckNotExist) {
self.primary_key.get_or_insert_with(|| key.clone());
}
self.entry_map.insert(key, entry);
}
} }
// The state of a key-value pair in the buffer. // The state of a key-value pair in the buffer.
@ -388,7 +359,7 @@ mod tests {
#[tokio::test] #[tokio::test]
#[allow(unreachable_code)] #[allow(unreachable_code)]
async fn set_and_get_from_buffer() { async fn set_and_get_from_buffer() {
let buffer = Buffer::new(false); let mut buffer = Buffer::new(false);
buffer buffer
.put(b"key1".to_vec().into(), b"value1".to_vec()) .put(b"key1".to_vec().into(), b"value1".to_vec())
.await; .await;
@ -421,7 +392,7 @@ mod tests {
#[tokio::test] #[tokio::test]
#[allow(unreachable_code)] #[allow(unreachable_code)]
async fn insert_and_get_from_buffer() { async fn insert_and_get_from_buffer() {
let buffer = Buffer::new(false); let mut buffer = Buffer::new(false);
buffer buffer
.insert(b"key1".to_vec().into(), b"value1".to_vec()) .insert(b"key1".to_vec().into(), b"value1".to_vec())
.await; .await;
@ -463,13 +434,13 @@ mod tests {
let v2: Value = b"value2".to_vec(); let v2: Value = b"value2".to_vec();
let v2_ = v2.clone(); let v2_ = v2.clone();
let buffer = Buffer::new(false); let mut buffer = Buffer::new(false);
let r1 = block_on(buffer.get_or_else(k1.clone(), move |_| ready(Ok(Some(v1_))))); let r1 = block_on(buffer.get_or_else(k1.clone(), move |_| ready(Ok(Some(v1_)))));
let r2 = block_on(buffer.get_or_else(k1.clone(), move |_| ready(panic!()))); let r2 = block_on(buffer.get_or_else(k1.clone(), move |_| ready(panic!())));
assert_eq!(r1.unwrap().unwrap(), v1); assert_eq!(r1.unwrap().unwrap(), v1);
assert_eq!(r2.unwrap().unwrap(), v1); assert_eq!(r2.unwrap().unwrap(), v1);
let buffer = Buffer::new(false); let mut buffer = Buffer::new(false);
let r1 = block_on( let r1 = block_on(
buffer.batch_get_or_else(vec![k1.clone(), k2.clone()].into_iter(), move |_| { buffer.batch_get_or_else(vec![k1.clone(), k2.clone()].into_iter(), move |_| {
ready(Ok(vec![(k1_, v1__).into(), (k2_, v2_).into()])) ready(Ok(vec![(k1_, v1__).into(), (k2_, v2_).into()]))

12
src/transaction/snapshot.rs
View File

@ -19,18 +19,18 @@ pub struct Snapshot {
impl Snapshot { impl Snapshot {
/// Get the value associated with the given key. /// Get the value associated with the given key.
pub async fn get(&self, key: impl Into<Key>) -> Result<Option<Value>> { pub async fn get(&mut self, key: impl Into<Key>) -> Result<Option<Value>> {
self.transaction.get(key).await self.transaction.get(key).await
} }
/// Check whether the key exists. /// Check whether the key exists.
pub async fn key_exists(&self, key: impl Into<Key>) -> Result<bool> { pub async fn key_exists(&mut self, key: impl Into<Key>) -> Result<bool> {
self.transaction.key_exists(key).await self.transaction.key_exists(key).await
} }
/// Get the values associated with the given keys. /// Get the values associated with the given keys.
pub async fn batch_get( pub async fn batch_get(
&self, &mut self,
keys: impl IntoIterator<Item = impl Into<Key>>, keys: impl IntoIterator<Item = impl Into<Key>>,
) -> Result<impl Iterator<Item = KvPair>> { ) -> Result<impl Iterator<Item = KvPair>> {
self.transaction.batch_get(keys).await self.transaction.batch_get(keys).await
@ -38,7 +38,7 @@ impl Snapshot {
/// Scan a range, return at most `limit` key-value pairs that lying in the range. /// Scan a range, return at most `limit` key-value pairs that lying in the range.
pub async fn scan( pub async fn scan(
&self, &mut self,
range: impl Into<BoundRange>, range: impl Into<BoundRange>,
limit: u32, limit: u32,
) -> Result<impl Iterator<Item = KvPair>> { ) -> Result<impl Iterator<Item = KvPair>> {
@ -47,7 +47,7 @@ impl Snapshot {
/// Scan a range, return at most `limit` keys that lying in the range. /// Scan a range, return at most `limit` keys that lying in the range.
pub async fn scan_keys( pub async fn scan_keys(
&self, &mut self,
range: impl Into<BoundRange>, range: impl Into<BoundRange>,
limit: u32, limit: u32,
) -> Result<impl Iterator<Item = Key>> { ) -> Result<impl Iterator<Item = Key>> {
@ -56,7 +56,7 @@ impl Snapshot {
/// Unimplemented. Similar to scan, but in the reverse direction. /// Unimplemented. Similar to scan, but in the reverse direction.
#[allow(dead_code)] #[allow(dead_code)]
fn scan_reverse(&self, range: impl RangeBounds<Key>) -> BoxStream<Result<KvPair>> { fn scan_reverse(&mut self, range: impl RangeBounds<Key>) -> BoxStream<Result<KvPair>> {
self.transaction.scan_reverse(range) self.transaction.scan_reverse(range)
} }
} }

12
src/transaction/transaction.rs
View File

@ -93,7 +93,7 @@ impl<PdC: PdClient> Transaction<PdC> {
/// txn.commit().await.unwrap(); /// txn.commit().await.unwrap();
/// # }); /// # });
/// ``` /// ```
pub async fn get(&self, key: impl Into<Key>) -> Result<Option<Value>> { pub async fn get(&mut self, key: impl Into<Key>) -> Result<Option<Value>> {
self.check_allow_operation().await?; self.check_allow_operation().await?;
let timestamp = self.timestamp.clone(); let timestamp = self.timestamp.clone();
let rpc = self.rpc.clone(); let rpc = self.rpc.clone();
@ -184,7 +184,7 @@ impl<PdC: PdClient> Transaction<PdC> {
/// txn.commit().await.unwrap(); /// txn.commit().await.unwrap();
/// # }); /// # });
/// ``` /// ```
pub async fn key_exists(&self, key: impl Into<Key>) -> Result<bool> { pub async fn key_exists(&mut self, key: impl Into<Key>) -> Result<bool> {
let key = key.into(); let key = key.into();
Ok(self.scan_keys(key.clone()..=key, 1).await?.next().is_some()) Ok(self.scan_keys(key.clone()..=key, 1).await?.next().is_some())
} }
@ -216,7 +216,7 @@ impl<PdC: PdClient> Transaction<PdC> {
/// # }); /// # });
/// ``` /// ```
pub async fn batch_get( pub async fn batch_get(
&self, &mut self,
keys: impl IntoIterator<Item = impl Into<Key>>, keys: impl IntoIterator<Item = impl Into<Key>>,
) -> Result<impl Iterator<Item = KvPair>> { ) -> Result<impl Iterator<Item = KvPair>> {
self.check_allow_operation().await?; self.check_allow_operation().await?;
@ -307,7 +307,7 @@ impl<PdC: PdClient> Transaction<PdC> {
/// # }); /// # });
/// ``` /// ```
pub async fn scan( pub async fn scan(
&self, &mut self,
range: impl Into<BoundRange>, range: impl Into<BoundRange>,
limit: u32, limit: u32,
) -> Result<impl Iterator<Item = KvPair>> { ) -> Result<impl Iterator<Item = KvPair>> {
@ -341,7 +341,7 @@ impl<PdC: PdClient> Transaction<PdC> {
/// # }); /// # });
/// ``` /// ```
pub async fn scan_keys( pub async fn scan_keys(
&self, &mut self,
range: impl Into<BoundRange>, range: impl Into<BoundRange>,
limit: u32, limit: u32,
) -> Result<impl Iterator<Item = Key>> { ) -> Result<impl Iterator<Item = Key>> {
@ -602,7 +602,7 @@ impl<PdC: PdClient> Transaction<PdC> {
} }
async fn scan_inner( async fn scan_inner(
&self, &mut self,
range: impl Into<BoundRange>, range: impl Into<BoundRange>,
limit: u32, limit: u32,
key_only: bool, key_only: bool,

10
tests/integration_tests.rs
View File

@ -103,7 +103,7 @@ async fn crud() -> Result<()> {
txn.commit().await?; txn.commit().await?;
// Read again from TiKV // Read again from TiKV
let snapshot = client.snapshot( let mut snapshot = client.snapshot(
client.current_timestamp().await?, client.current_timestamp().await?,
// TODO needed because pessimistic does not check locks (#235) // TODO needed because pessimistic does not check locks (#235)
TransactionOptions::new_optimistic(), TransactionOptions::new_optimistic(),
@ -339,9 +339,9 @@ async fn txn_bank_transfer() -> Result<()> {
.await?; .await?;
let chosen_people = people.iter().choose_multiple(&mut rng, 2); let chosen_people = people.iter().choose_multiple(&mut rng, 2);
let alice = chosen_people[0]; let alice = chosen_people[0];
let mut alice_balance = get_txn_u32(&txn, alice.clone()).await?; let mut alice_balance = get_txn_u32(&mut txn, alice.clone()).await?;
let bob = chosen_people[1]; let bob = chosen_people[1];
let mut bob_balance = get_txn_u32(&txn, bob.clone()).await?; let mut bob_balance = get_txn_u32(&mut txn, bob.clone()).await?;
if alice_balance == 0 { if alice_balance == 0 {
txn.rollback().await?; txn.rollback().await?;
continue; continue;
@ -360,7 +360,7 @@ async fn txn_bank_transfer() -> Result<()> {
let mut new_sum = 0; let mut new_sum = 0;
let mut txn = client.begin_optimistic().await?; let mut txn = client.begin_optimistic().await?;
for person in people.iter() { for person in people.iter() {
new_sum += get_txn_u32(&txn, person.clone()).await?; new_sum += get_txn_u32(&mut txn, person.clone()).await?;
} }
assert_eq!(sum, new_sum); assert_eq!(sum, new_sum);
txn.commit().await?; txn.commit().await?;
@ -827,7 +827,7 @@ async fn get_u32(client: &RawClient, key: Vec<u8>) -> Result<u32> {
} }
// helper function // helper function
async fn get_txn_u32(txn: &Transaction, key: Vec<u8>) -> Result<u32> { async fn get_txn_u32(txn: &mut Transaction, key: Vec<u8>) -> Result<u32> {
let x = txn.get(key).await?.unwrap(); let x = txn.get(key).await?.unwrap();
let boxed_slice = x.into_boxed_slice(); let boxed_slice = x.into_boxed_slice();
let array: Box<[u8; 4]> = boxed_slice let array: Box<[u8; 4]> = boxed_slice