Add scan_reverse and scan_keys_reverse support for txnkv (#354)

Co-authored-by: iosmanthus <MyOsmanthusTree@gmail.com>

src/transaction/buffer.rs

@@ -111,6 +111,7 @@ impl Buffer {
         range: BoundRange,
         limit: u32,
         update_cache: bool,
+        reverse: bool,
         f: F,
     ) -> Result<impl Iterator<Item = KvPair>>
     where
@@ -158,7 +159,13 @@ impl Buffer {
             .into_iter()
             .map(|(k, v)| KvPair::new(k, v))
             .collect::<Vec<_>>();
-        res.sort_by_cached_key(|x| x.key().clone());
+
+        // TODO: use `BTreeMap` instead of `HashMap` to avoid sorting.
+        if reverse {
+            res.sort_unstable_by(|a, b| b.key().cmp(a.key()));
+        } else {
+            res.sort_unstable_by(|a, b| a.key().cmp(b.key()));
+        }
 
         Ok(res.into_iter().take(limit as usize))
     }

src/transaction/lowering.rs

@@ -23,6 +23,7 @@ pub fn new_scan_request(
     timestamp: Timestamp,
     limit: u32,
     key_only: bool,
+    reverse: bool,
 ) -> kvrpcpb::ScanRequest {
     let (start_key, end_key) = range.into_keys();
     requests::new_scan_request(
@@ -31,6 +32,7 @@ pub fn new_scan_request(
         timestamp.version(),
         limit,
         key_only,
+        reverse,
     )
 }
 

src/transaction/requests.rs

@@ -125,13 +125,20 @@ pub fn new_scan_request(
     timestamp: u64,
     limit: u32,
     key_only: bool,
+    reverse: bool,
 ) -> kvrpcpb::ScanRequest {
     let mut req = kvrpcpb::ScanRequest::default();
-    req.set_start_key(start_key);
-    req.set_end_key(end_key);
+    if !reverse {
+        req.set_start_key(start_key);
+        req.set_end_key(end_key);
+    } else {
+        req.set_start_key(end_key);
+        req.set_end_key(start_key);
+    }
     req.set_limit(limit);
     req.set_key_only(key_only);
     req.set_version(timestamp);
+    req.set_reverse(reverse);
     req
 }
 

src/transaction/snapshot.rs

@@ -2,9 +2,7 @@
 
 use crate::{BoundRange, Key, KvPair, Result, Transaction, Value};
 use derive_new::new;
-use futures::stream::BoxStream;
 use slog::Logger;
-use std::ops::RangeBounds;
 
 /// A read-only transaction which reads at the given timestamp.
 ///
@@ -61,10 +59,26 @@ impl Snapshot {
         self.transaction.scan_keys(range, limit).await
     }
 
-    /// Unimplemented. Similar to scan, but in the reverse direction.
-    #[allow(dead_code)]
-    fn scan_reverse(&mut self, range: impl RangeBounds<Key>) -> BoxStream<Result<KvPair>> {
+    /// Similar to scan, but in the reverse direction.
+    pub async fn scan_reverse(
+        &mut self,
+        range: impl Into<BoundRange>,
+        limit: u32,
+    ) -> Result<impl Iterator<Item = KvPair>> {
         debug!(self.logger, "invoking scan_reverse request on snapshot");
-        self.transaction.scan_reverse(range)
+        self.transaction.scan_reverse(range, limit).await
+    }
+
+    /// Similar to scan_keys, but in the reverse direction.
+    pub async fn scan_keys_reverse(
+        &mut self,
+        range: impl Into<BoundRange>,
+        limit: u32,
+    ) -> Result<impl Iterator<Item = Key>> {
+        debug!(
+            self.logger,
+            "invoking scan_keys_reverse request on snapshot"
+        );
+        self.transaction.scan_keys_reverse(range, limit).await
     }
 }

src/transaction/transaction.rs

@@ -12,9 +12,9 @@ use crate::{
 };
 use derive_new::new;
 use fail::fail_point;
-use futures::{prelude::*, stream::BoxStream};
+use futures::prelude::*;
 use slog::Logger;
-use std::{iter, ops::RangeBounds, sync::Arc, time::Instant};
+use std::{iter, sync::Arc, time::Instant};
 use tikv_client_proto::{kvrpcpb, pdpb::Timestamp};
 use tokio::{sync::RwLock, time::Duration};
 
@@ -344,7 +344,7 @@ impl<PdC: PdClient> Transaction<PdC> {
         limit: u32,
     ) -> Result<impl Iterator<Item = KvPair>> {
         debug!(self.logger, "invoking transactional scan request");
-        self.scan_inner(range, limit, false).await
+        self.scan_inner(range, limit, false, false).await
     }
 
     /// Create a new 'scan' request that only returns the keys.
@@ -381,7 +381,7 @@ impl<PdC: PdClient> Transaction<PdC> {
     ) -> Result<impl Iterator<Item = Key>> {
         debug!(self.logger, "invoking transactional scan_keys request");
         Ok(self
-            .scan_inner(range, limit, true)
+            .scan_inner(range, limit, true, false)
             .await?
             .map(KvPair::into_key))
     }
@@ -389,9 +389,31 @@ impl<PdC: PdClient> Transaction<PdC> {
     /// Create a 'scan_reverse' request.
     ///
     /// Similar to [`scan`](Transaction::scan), but scans in the reverse direction.
-    pub(crate) fn scan_reverse(&self, _range: impl RangeBounds<Key>) -> BoxStream<Result<KvPair>> {
+    pub async fn scan_reverse(
+        &mut self,
+        range: impl Into<BoundRange>,
+        limit: u32,
+    ) -> Result<impl Iterator<Item = KvPair>> {
         debug!(self.logger, "invoking transactional scan_reverse request");
-        unimplemented!()
+        self.scan_inner(range, limit, false, true).await
+    }
+
+    /// Create a 'scan_keys_reverse' request.
+    ///
+    /// Similar to [`scan`](Transaction::scan_keys), but scans in the reverse direction.
+    pub async fn scan_keys_reverse(
+        &mut self,
+        range: impl Into<BoundRange>,
+        limit: u32,
+    ) -> Result<impl Iterator<Item = Key>> {
+        debug!(
+            self.logger,
+            "invoking transactional scan_keys_reverse request"
+        );
+        Ok(self
+            .scan_inner(range, limit, true, true)
+            .await?
+            .map(KvPair::into_key))
     }
 
     /// Sets the value associated with the given key.
@@ -680,6 +702,7 @@ impl<PdC: PdClient> Transaction<PdC> {
         range: impl Into<BoundRange>,
         limit: u32,
         key_only: bool,
+        reverse: bool,
     ) -> Result<impl Iterator<Item = KvPair>> {
         self.check_allow_operation().await?;
         let timestamp = self.timestamp.clone();
@@ -691,8 +714,10 @@ impl<PdC: PdClient> Transaction<PdC> {
                 range.into(),
                 limit,
                 !key_only,
+                reverse,
                 move |new_range, new_limit| async move {
-                    let request = new_scan_request(new_range, timestamp, new_limit, key_only);
+                    let request =
+                        new_scan_request(new_range, timestamp, new_limit, key_only, reverse);
                     let plan = PlanBuilder::new(rpc, request)
                         .resolve_lock(retry_options.lock_backoff)
                         .retry_multi_region(retry_options.region_backoff)

tests/integration_tests.rs

@@ -22,7 +22,7 @@ use std::{
     iter,
 };
 use tikv_client::{
-    transaction::HeartbeatOption, Error, Key, KvPair, RawClient, Result, Transaction,
+    transaction::HeartbeatOption, BoundRange, Error, Key, KvPair, RawClient, Result, Transaction,
     TransactionClient, TransactionOptions, Value,
 };
 
@@ -883,6 +883,42 @@ async fn txn_scan() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+#[serial]
+async fn txn_scan_reverse() -> Result<()> {
+    init().await?;
+    let client = TransactionClient::new_with_config(pd_addrs(), Default::default(), None).await?;
+
+    let k1 = b"a1".to_vec();
+    let k2 = b"a2".to_vec();
+    let v1 = b"b1".to_vec();
+    let v2 = b"b2".to_vec();
+
+    let reverse_resp = vec![
+        (Key::from(k2.clone()), v2.clone()),
+        (Key::from(k1.clone()), v1.clone()),
+    ];
+
+    // pessimistic
+    let option = TransactionOptions::new_pessimistic().drop_check(tikv_client::CheckLevel::Warn);
+    let mut t = client.begin_with_options(option.clone()).await?;
+    t.put(k1.clone(), v1).await?;
+    t.put(k2.clone(), v2).await?;
+    t.commit().await?;
+
+    let mut t2 = client.begin_with_options(option).await?;
+    let bound_range: BoundRange = (k1..=k2).into();
+    let resp = t2
+        .scan_reverse(bound_range, 2)
+        .await?
+        .map(|kv| (kv.0, kv.1))
+        .collect::<Vec<(Key, Vec<u8>)>>();
+    assert_eq!(resp, reverse_resp);
+    t2.commit().await?;
+
+    Ok(())
+}
+
 // helper function
 async fn get_u32(client: &RawClient, key: Vec<u8>) -> Result<u32> {
     let x = client.get(key).await?.unwrap();