Merge pull request #233 from ziyi-yan/support-mutation-insert-checkne

Support mutation Insert and CheckNotExists

Author: nrc

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,9 +201,30 @@ 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);
+        let mut mutations = self.mutations.lock().await;
+        let value = mutations
+            .entry_map
+            .entry(key.clone())
+            .or_insert(BufferEntry::Del);
+
+        let new_value: BufferEntry;
+        if let BufferEntry::Insert(_) = value {
+            new_value = BufferEntry::CheckNotExist
+        } else {
+            new_value = BufferEntry::Del
+        }
+
+        mutations.insert(key, new_value);
     }
 
     /// Converts the buffered mutations to the proto buffer version
@@ -237,6 +267,12 @@ impl Buffer {
             Some(BufferEntry::Del) => {
                 assert!(value.is_none());
             }
+            Some(BufferEntry::Insert(v)) => {
+                assert!(value.as_ref() == Some(v))
+            }
+            Some(BufferEntry::CheckNotExist) => {
+                assert!(value.is_none());
+            }
         }
     }
 }
@@ -247,6 +283,8 @@ impl Buffer {
 // Mutations:
 //   - `Put`
 //   - `Del`
+//   - `Insert`
+//   - `CheckNotExist`, a constraint to ensure the key doesn't exist. See https://github.com/pingcap/tidb/pull/14968.
 // 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 +312,10 @@ enum BufferEntry {
     Put(Value),
     // Value has been deleted.
     Del,
+    // Key should be check not exists before.
+    Insert(Value),
+    // Key should be check not exists before.
+    CheckNotExist,
 }
 
 impl BufferEntry {
@@ -287,6 +329,11 @@ 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());
+            }
+            BufferEntry::CheckNotExist => pb.set_op(kvrpcpb::Op::CheckNotExists),
         };
         pb.set_key(key.clone().into());
         Some(pb)
@@ -299,6 +346,8 @@ 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())),
+            BufferEntry::CheckNotExist => MutationValue::Determined(None),
         }
     }
 }
@@ -359,6 +408,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.

tests/integration_tests.rs

@@ -132,6 +132,31 @@ async fn crud() -> Result<()> {
     Ok(())
 }
 
+// Tests transactional insert and delete-your-writes cases
+#[tokio::test]
+#[serial]
+async fn insert_duplicate_keys() -> Result<()> {
+    clear_tikv().await?;
+
+    let client = TransactionClient::new(pd_addrs()).await?;
+    // Initialize TiKV store with {foo => bar}
+    let mut txn = client.begin_optimistic().await?;
+    txn.put("foo".to_owned(), "bar".to_owned()).await?;
+    txn.commit().await?;
+    // Try insert foo again
+    let mut txn = client.begin_optimistic().await?;
+    txn.insert("foo".to_owned(), "foo".to_owned()).await?;
+    assert!(txn.commit().await.is_err());
+
+    // Delete-your-writes
+    let mut txn = client.begin_optimistic().await?;
+    txn.insert("foo".to_owned(), "foo".to_owned()).await?;
+    txn.delete("foo".to_owned()).await?;
+    assert!(txn.commit().await.is_err());
+
+    Ok(())
+}
+
 #[tokio::test]
 #[serial]
 async fn pessimistic() -> Result<()> {

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,