De-async buffer

Signed-off-by: Nick Cameron <nrc@ncameron.org>

Author: nrc

README.md

@@ -8,7 +8,7 @@ This crate provides an easy-to-use client for [TiKV](https://github.com/tikv/tik
 
 This crate lets you connect to a TiKV cluster and use either a transactional or raw (simple get/put style without transactional consistency guarantees) API to access and update your data.
 
-This is an open source (Apache 2) project maintained by the TiKV Authors. We welcome community contributions, see below for more info.
+The TiKV Rust client is an open source (Apache 2) project maintained by the TiKV Authors. We welcome contributions, see below for more info.
 
 ## Getting started
 
@@ -19,7 +19,7 @@ The TiKV client is a Rust library (crate). To use this crate in your project, ad
 tikv-client = 0.1
 ```
 
-Note, that you will need `tikv-client = { git = "https://github.com/tikv/client-rust.git" }` until we publish the crate (should be any day now).
+Note, you will need `tikv-client = { git = "https://github.com/tikv/client-rust.git" }` until we publish the crate (should be any day now).
 
 The minimum supported version of Rust is 1.40. The minimum supported version of TiKV is 5.0.
 

src/transaction/buffer.rs

@@ -24,19 +24,19 @@ impl Buffer {
     }
 
     /// Get the primary key of the buffer.
-    pub async fn get_primary_key(&self) -> Option<Key> {
+    pub fn get_primary_key(&self) -> Option<Key> {
         self.primary_key.clone()
     }
 
     /// Set the primary key if it is not set
-    pub async fn primary_key_or(&mut self, key: &Key) {
+    pub fn primary_key_or(&mut self, key: &Key) {
         self.primary_key.get_or_insert_with(|| key.clone());
     }
 
     /// Get a value from the buffer.
     /// If the returned value is None, it means the key doesn't exist in buffer yet.
-    pub async fn get(&self, key: &Key) -> Option<Value> {
-        match self.get_from_mutations(key).await {
+    pub fn get(&self, key: &Key) -> Option<Value> {
+        match self.get_from_mutations(key) {
             MutationValue::Determined(value) => value,
             MutationValue::Undetermined => None,
         }
@@ -49,7 +49,7 @@ impl Buffer {
         F: FnOnce(Key) -> Fut,
         Fut: Future<Output = Result<Option<Value>>>,
     {
-        match self.get_from_mutations(&key).await {
+        match self.get_from_mutations(&key) {
             MutationValue::Determined(value) => Ok(value),
             MutationValue::Undetermined => {
                 let value = f(key.clone()).await?;
@@ -161,7 +161,7 @@ impl Buffer {
     }
 
     /// Lock the given key if necessary.
-    pub async fn lock(&mut self, key: Key) {
+    pub fn lock(&mut self, key: Key) {
         self.primary_key.get_or_insert_with(|| key.clone());
         let value = self
             .entry_map
@@ -175,7 +175,7 @@ impl Buffer {
     }
 
     /// Put a value into the buffer (does not write through).
-    pub async fn put(&mut self, key: Key, value: Value) {
+    pub fn put(&mut self, key: Key, value: Value) {
         let mut entry = self.entry_map.entry(key.clone());
         match entry {
             Entry::Occupied(ref mut o)
@@ -189,7 +189,7 @@ impl Buffer {
     }
 
     /// Mark a value as Insert mutation into the buffer (does not write through).
-    pub async fn insert(&mut self, key: Key, value: Value) {
+    pub fn insert(&mut self, key: Key, value: Value) {
         let mut entry = self.entry_map.entry(key.clone());
         match entry {
             Entry::Occupied(ref mut o) if matches!(o.get(), BufferEntry::Del) => {
@@ -200,7 +200,7 @@ impl Buffer {
     }
 
     /// Mark a value as deleted.
-    pub async fn delete(&mut self, key: Key) {
+    pub fn delete(&mut self, key: Key) {
         let is_pessimistic = self.is_pessimistic;
         let mut entry = self.entry_map.entry(key.clone());
 
@@ -217,14 +217,14 @@ impl Buffer {
     }
 
     /// Converts the buffered mutations to the proto buffer version
-    pub async fn to_proto_mutations(&self) -> Vec<kvrpcpb::Mutation> {
+    pub fn to_proto_mutations(&self) -> Vec<kvrpcpb::Mutation> {
         self.entry_map
             .iter()
             .filter_map(|(key, mutation)| mutation.to_proto_with_key(key))
             .collect()
     }
 
-    async fn get_from_mutations(&self, key: &Key) -> MutationValue {
+    fn get_from_mutations(&self, key: &Key) -> MutationValue {
         self.entry_map
             .get(&key)
             .map(BufferEntry::get_value)
@@ -365,15 +365,11 @@ mod tests {
     use futures::{executor::block_on, future::ready};
     use tikv_client_common::internal_err;
 
-    #[tokio::test]
-    async fn set_and_get_from_buffer() {
+    #[test]
+    fn set_and_get_from_buffer() {
         let mut buffer = Buffer::new(false);
-        buffer
-            .put(b"key1".to_vec().into(), b"value1".to_vec())
-            .await;
-        buffer
-            .put(b"key2".to_vec().into(), b"value2".to_vec())
-            .await;
+        buffer.put(b"key1".to_vec().into(), b"value1".to_vec());
+        buffer.put(b"key2".to_vec().into(), b"value2".to_vec());
         assert_eq!(
             block_on(
                 buffer.get_or_else(b"key1".to_vec().into(), move |_| ready(Err(internal_err!(
@@ -385,8 +381,8 @@ mod tests {
             b"value1".to_vec()
         );
 
-        buffer.delete(b"key2".to_vec().into()).await;
-        buffer.put(b"key1".to_vec().into(), b"value".to_vec()).await;
+        buffer.delete(b"key2".to_vec().into());
+        buffer.put(b"key1".to_vec().into(), b"value".to_vec());
         assert_eq!(
             block_on(buffer.batch_get_or_else(
                 vec![b"key2".to_vec().into(), b"key1".to_vec().into()].into_iter(),
@@ -401,15 +397,11 @@ mod tests {
         );
     }
 
-    #[tokio::test]
-    async fn insert_and_get_from_buffer() {
+    #[test]
+    fn insert_and_get_from_buffer() {
         let mut buffer = Buffer::new(false);
-        buffer
-            .insert(b"key1".to_vec().into(), b"value1".to_vec())
-            .await;
-        buffer
-            .insert(b"key2".to_vec().into(), b"value2".to_vec())
-            .await;
+        buffer.insert(b"key1".to_vec().into(), b"value1".to_vec());
+        buffer.insert(b"key2".to_vec().into(), b"value2".to_vec());
         assert_eq!(
             block_on(
                 buffer.get_or_else(b"key1".to_vec().into(), move |_| ready(Err(internal_err!(
@@ -421,10 +413,8 @@ mod tests {
             b"value1".to_vec()
         );
 
-        buffer.delete(b"key2".to_vec().into()).await;
-        buffer
-            .insert(b"key1".to_vec().into(), b"value".to_vec())
-            .await;
+        buffer.delete(b"key2".to_vec().into());
+        buffer.insert(b"key1".to_vec().into(), b"value".to_vec());
         assert_eq!(
             block_on(buffer.batch_get_or_else(
                 vec![b"key2".to_vec().into(), b"key1".to_vec().into()].into_iter(),
@@ -481,8 +471,8 @@ mod tests {
     }
 
     // Check that multiple writes to the same key combine in the correct way.
-    #[tokio::test]
-    async fn state_machine() {
+    #[test]
+    fn state_machine() {
         let mut buffer = Buffer::new(false);
 
         macro_rules! assert_entry {
@@ -493,28 +483,28 @@ mod tests {
 
         // Insert + Delete = CheckNotExists
         let key: Key = b"key1".to_vec().into();
-        buffer.insert(key.clone(), b"value1".to_vec()).await;
-        buffer.delete(key.clone()).await;
+        buffer.insert(key.clone(), b"value1".to_vec());
+        buffer.delete(key.clone());
         assert_entry!(key, BufferEntry::CheckNotExist);
 
         // CheckNotExists + Delete = CheckNotExists
-        buffer.delete(key.clone()).await;
+        buffer.delete(key.clone());
         assert_entry!(key, BufferEntry::CheckNotExist);
 
         // CheckNotExists + Put = Insert
-        buffer.put(key.clone(), b"value2".to_vec()).await;
+        buffer.put(key.clone(), b"value2".to_vec());
         assert_entry!(key, BufferEntry::Insert(_));
 
         // Insert + Put = Insert
         let key: Key = b"key2".to_vec().into();
-        buffer.insert(key.clone(), b"value1".to_vec()).await;
-        buffer.put(key.clone(), b"value2".to_vec()).await;
+        buffer.insert(key.clone(), b"value1".to_vec());
+        buffer.put(key.clone(), b"value2".to_vec());
         assert_entry!(key, BufferEntry::Insert(_));
 
         // Delete + Insert = Put
         let key: Key = b"key3".to_vec().into();
-        buffer.delete(key.clone()).await;
-        buffer.insert(key.clone(), b"value1".to_vec()).await;
+        buffer.delete(key.clone());
+        buffer.insert(key.clone(), b"value1".to_vec());
         assert_entry!(key, BufferEntry::Put(_));
     }
 }

src/transaction/transaction.rs

@@ -398,7 +398,7 @@ impl<PdC: PdClient> Transaction<PdC> {
             self.pessimistic_lock(iter::once(key.clone()), false)
                 .await?;
         }
-        self.buffer.put(key, value.into()).await;
+        self.buffer.put(key, value.into());
         Ok(())
     }
 
@@ -424,7 +424,7 @@ impl<PdC: PdClient> Transaction<PdC> {
     pub async fn insert(&mut self, key: impl Into<Key>, value: impl Into<Value>) -> Result<()> {
         self.check_allow_operation().await?;
         let key = key.into();
-        if self.buffer.get(&key).await.is_some() {
+        if self.buffer.get(&key).is_some() {
             return Err(Error::DuplicateKeyInsertion);
         }
         if self.is_pessimistic() {
@@ -434,7 +434,7 @@ impl<PdC: PdClient> Transaction<PdC> {
             )
             .await?;
         }
-        self.buffer.insert(key, value.into()).await;
+        self.buffer.insert(key, value.into());
         Ok(())
     }
 
@@ -462,7 +462,7 @@ impl<PdC: PdClient> Transaction<PdC> {
             self.pessimistic_lock(iter::once(key.clone()), false)
                 .await?;
         }
-        self.buffer.delete(key).await;
+        self.buffer.delete(key);
         Ok(())
     }
 
@@ -497,7 +497,7 @@ impl<PdC: PdClient> Transaction<PdC> {
         match self.options.kind {
             TransactionKind::Optimistic => {
                 for key in keys {
-                    self.buffer.lock(key.into()).await;
+                    self.buffer.lock(key.into());
                 }
             }
             TransactionKind::Pessimistic(_) => {
@@ -535,8 +535,8 @@ impl<PdC: PdClient> Transaction<PdC> {
             *status = TransactionStatus::StartedCommit;
         }
 
-        let primary_key = self.buffer.get_primary_key().await;
-        let mutations = self.buffer.to_proto_mutations().await;
+        let primary_key = self.buffer.get_primary_key();
+        let mutations = self.buffer.to_proto_mutations();
         if mutations.is_empty() {
             assert!(primary_key.is_none());
             return Ok(None);
@@ -595,8 +595,8 @@ impl<PdC: PdClient> Transaction<PdC> {
             *status = TransactionStatus::StartedRollback;
         }
 
-        let primary_key = self.buffer.get_primary_key().await;
-        let mutations = self.buffer.to_proto_mutations().await;
+        let primary_key = self.buffer.get_primary_key();
+        let mutations = self.buffer.to_proto_mutations();
         let res = Committer::new(
             primary_key,
             mutations,
@@ -619,7 +619,7 @@ impl<PdC: PdClient> Transaction<PdC> {
     /// Returns the TTL set on the transaction's locks by TiKV.
     pub async fn send_heart_beat(&mut self) -> Result<u64> {
         self.check_allow_operation().await?;
-        let primary_key = match self.buffer.get_primary_key().await {
+        let primary_key = match self.buffer.get_primary_key() {
             Some(k) => k,
             None => return Err(Error::NoPrimaryKey),
         };
@@ -695,7 +695,6 @@ impl<PdC: PdClient> Transaction<PdC> {
         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());
@@ -717,12 +716,12 @@ impl<PdC: PdClient> Transaction<PdC> {
         let pairs = plan.execute().await;
 
         // primary key will be set here if needed
-        self.buffer.primary_key_or(&first_key).await;
+        self.buffer.primary_key_or(&first_key);
 
         self.start_auto_heartbeat().await;
 
         for key in keys {
-            self.buffer.lock(key.key()).await;
+            self.buffer.lock(key.key());
         }
 
         pairs
@@ -755,7 +754,6 @@ impl<PdC: PdClient> Transaction<PdC> {
         let primary_key = self
             .buffer
             .get_primary_key()
-            .await
             .expect("Primary key should exist");
         let start_ts = self.timestamp.clone();
         let region_backoff = self.options.retry_options.region_backoff.clone();