txn: add insert() to Transaction

Signed-off-by: Ziyi Yan <ziyi.yan@foxmail.com>

Author: ziyi-yan

src/transaction/buffer.rs

@@ -44,6 +44,15 @@ impl Buffer {
         self.mutations.lock().await.get_primary_key_or(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 {
+            MutationValue::Determined(value) => value,
+            MutationValue::Undetermined => None,
+        }
+    }
+
     /// Get a value from the buffer. If the value is not present, run `f` to get
     /// the value.
     pub async fn get_or_else<F, Fut>(&self, key: Key, f: F) -> Result<Option<Value>>
@@ -192,6 +201,14 @@ impl Buffer {
             .insert(key, BufferEntry::Put(value));
     }
 
+    /// Mark a value as Insert mutation into the buffer (does not write through).
+    pub async fn insert(&self, key: Key, value: Value) {
+        self.mutations
+            .lock()
+            .await
+            .insert(key, BufferEntry::Insert(value));
+    }
+
     /// Mark a value as deleted.
     pub async fn delete(&self, key: Key) {
         self.mutations.lock().await.insert(key, BufferEntry::Del);
@@ -237,6 +254,9 @@ impl Buffer {
             Some(BufferEntry::Del) => {
                 assert!(value.is_none());
             }
+            Some(BufferEntry::Insert(v)) => {
+                assert!(value.as_ref() == Some(v))
+            }
         }
     }
 }
@@ -247,6 +267,7 @@ impl Buffer {
 // Mutations:
 //   - `Put`
 //   - `Del`
+//   - `Insert`
 // Cache of read requests:
 //   - `Cached`, generated by normal read requests
 //   - `ReadLockCached`, generated by lock commands (`lock_keys`, `get_for_update`) and optionally read requests
@@ -274,6 +295,8 @@ enum BufferEntry {
     Put(Value),
     // Value has been deleted.
     Del,
+    // Key should be check not exists before.
+    Insert(Value),
 }
 
 impl BufferEntry {
@@ -287,6 +310,10 @@ impl BufferEntry {
             }
             BufferEntry::Del => pb.set_op(kvrpcpb::Op::Del),
             BufferEntry::Locked(_) => pb.set_op(kvrpcpb::Op::Lock),
+            BufferEntry::Insert(v) => {
+                pb.set_op(kvrpcpb::Op::Insert);
+                pb.set_value(v.clone());
+            }
         };
         pb.set_key(key.clone().into());
         Some(pb)
@@ -299,6 +326,7 @@ impl BufferEntry {
             BufferEntry::Del => MutationValue::Determined(None),
             BufferEntry::Locked(None) => MutationValue::Undetermined,
             BufferEntry::Locked(Some(value)) => MutationValue::Determined(value.clone()),
+            BufferEntry::Insert(value) => MutationValue::Determined(Some(value.clone())),
         }
     }
 }
@@ -359,6 +387,38 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    #[allow(unreachable_code)]
+    async fn insert_and_get_from_buffer() {
+        let buffer = Buffer::default();
+        buffer
+            .insert(b"key1".to_vec().into(), b"value1".to_vec())
+            .await;
+        buffer
+            .insert(b"key2".to_vec().into(), b"value2".to_vec())
+            .await;
+        assert_eq!(
+            block_on(buffer.get_or_else(b"key1".to_vec().into(), move |_| ready(panic!())))
+                .unwrap()
+                .unwrap(),
+            b"value1".to_vec()
+        );
+
+        buffer.delete(b"key2".to_vec().into()).await;
+        buffer
+            .insert(b"key1".to_vec().into(), b"value".to_vec())
+            .await;
+        assert_eq!(
+            block_on(buffer.batch_get_or_else(
+                vec![b"key2".to_vec().into(), b"key1".to_vec().into()].into_iter(),
+                move |_| ready(Ok(vec![])),
+            ))
+            .unwrap()
+            .collect::<Vec<_>>(),
+            vec![KvPair(Key::from(b"key1".to_vec()), b"value".to_vec()),]
+        );
+    }
+
     #[test]
     #[allow(unreachable_code)]
     fn repeat_reads_are_cached() {

src/transaction/transaction.rs

@@ -362,6 +362,37 @@ impl Transaction {
         Ok(())
     }
 
+    /// Inserts the value associated with the given key.
+    /// It has a constraint that key should not exist before.
+    ///
+    /// # Examples
+    /// ```rust,no_run
+    /// # use tikv_client::{Key, Value, Config, TransactionClient};
+    /// # use futures::prelude::*;
+    /// # 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_optimistic().await.unwrap();
+    /// let key = "TiKV".to_owned();
+    /// let val = "TiKV".to_owned();
+    /// txn.insert(key, val);
+    /// // Finish the transaction...
+    /// txn.commit().await.unwrap();
+    /// # });
+    /// ```
+    pub async fn insert(&mut self, key: impl Into<Key>, value: impl Into<Value>) -> Result<()> {
+        self.check_allow_operation()?;
+        let key = key.into();
+        if self.buffer.get(&key).await.is_some() {
+            return Err(Error::DuplicateKeyInsertion);
+        }
+        if self.is_pessimistic() {
+            self.pessimistic_lock(iter::once(key.clone()), false)
+                .await?;
+        }
+        self.buffer.insert(key, value.into()).await;
+        Ok(())
+    }
+
     /// Deletes the given key.
     ///
     /// Deleting a non-existent key will not result in an error.

tikv-client-common/src/errors.rs

@@ -10,6 +10,9 @@ pub enum Error {
     /// Feature is not implemented.
     #[error("Unimplemented feature")]
     Unimplemented,
+    /// Duplicate key insertion happens.
+    #[error("Duplicate key insertion")]
+    DuplicateKeyInsertion,
     /// Failed to resolve a lock
     #[error("Failed to resolve lock")]
     ResolveLockError,