Optimize key_exists performance (#360)

* Optimize key_exists performance

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

* Fix clippy

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

* Add test case for key_exists

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

* Fix clippy for test

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

Author: yongman

src/raw/client.rs

@@ -744,11 +744,11 @@ impl<PdC: PdClient> Client<PdC> {
     }
 
     fn assert_non_atomic(&self) -> Result<()> {
-        (!self.atomic).then(|| ()).ok_or(Error::UnsupportedMode)
+        (!self.atomic).then_some(()).ok_or(Error::UnsupportedMode)
     }
 
     fn assert_atomic(&self) -> Result<()> {
-        self.atomic.then(|| ()).ok_or(Error::UnsupportedMode)
+        self.atomic.then_some(()).ok_or(Error::UnsupportedMode)
     }
 }
 

src/transaction/transaction.rs

@@ -208,8 +208,7 @@ impl<PdC: PdClient> Transaction<PdC> {
     /// ```
     pub async fn key_exists(&mut self, key: impl Into<Key>) -> Result<bool> {
         debug!(self.logger, "invoking transactional key_exists request");
-        let key = key.into();
-        Ok(self.scan_keys(key.clone()..=key, 1).await?.next().is_some())
+        Ok(self.get(key).await?.is_some())
     }
 
     /// Create a new 'batch get' request.

tests/integration_tests.rs

@@ -919,6 +919,29 @@ async fn txn_scan_reverse() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+#[serial]
+async fn txn_key_exists() -> Result<()> {
+    init().await?;
+    let client = TransactionClient::new_with_config(pd_addrs(), Default::default(), None).await?;
+    let key = "key".to_owned();
+    let value = "value".to_owned();
+    let mut t1 = client.begin_optimistic().await?;
+    t1.put(key.clone(), value.clone()).await?;
+    assert!(t1.key_exists(key.clone()).await?);
+    t1.commit().await?;
+
+    let mut t2 = client.begin_optimistic().await?;
+    assert!(t2.key_exists(key).await?);
+    t2.commit().await?;
+
+    let not_exists_key = "not_exists_key".to_owned();
+    let mut t3 = client.begin_optimistic().await?;
+    assert!(!t3.key_exists(not_exists_key).await?);
+    t3.commit().await?;
+    Ok(())
+}
+
 // helper function
 async fn get_u32(client: &RawClient, key: Vec<u8>) -> Result<u32> {
     let x = client.get(key).await?.unwrap();