feat(transaction): add Transaction::start_ts and scan_unbounded

Signed-off-by: Yijun Zhao <ariesdevil77@gmail.com>

Author: ariesdevil

src/transaction/snapshot.rs

@@ -54,6 +54,18 @@ impl Snapshot {
         self.transaction.scan(range, limit).await
     }
 
+    /// Scan a range, return ALL key-value pairs that lying in the range.
+    ///
+    /// Equivalent of client-go's `KVSnapshot.Iter(start, end)`; see
+    /// [`Transaction::scan_unbounded`].
+    pub async fn scan_unbounded(
+        &mut self,
+        range: impl Into<BoundRange>,
+    ) -> Result<impl Iterator<Item = KvPair>> {
+        debug!("invoking scan_unbounded request on snapshot");
+        self.transaction.scan_unbounded(range).await
+    }
+
     /// Scan a range, return at most `limit` keys that lying in the range.
     pub async fn scan_keys(
         &mut self,

src/transaction/transaction.rs

@@ -119,6 +119,19 @@ impl<PdC: PdClient> Transaction<PdC> {
         }
     }
 
+    /// Get the start timestamp of the transaction.
+    ///
+    /// This is the TSO timestamp at which the transaction reads. The
+    /// transaction's writes become visible at its *commit* timestamp
+    /// (returned by [`commit`](Transaction::commit)), not at `start_ts`.
+    ///
+    /// The `physical` component is a cluster-wide wall clock in
+    /// milliseconds, useful as a consistent clock reading across processes
+    /// (e.g. checking lease or deadline expiry).
+    pub fn start_ts(&self) -> Timestamp {
+        self.timestamp.clone()
+    }
+
     /// Create a new 'get' request
     ///
     /// Once resolved this request will result in the fetching of the value associated with the
@@ -448,6 +461,70 @@ impl<PdC: PdClient> Transaction<PdC> {
             .map(KvPair::into_key))
     }
 
+    /// Create a 'scan' request without a limit.
+    ///
+    /// Once resolved this request will result in a `Vec` of ALL key-value pairs
+    /// that lie in the specified range, ordered by key.
+    ///
+    /// Internally the range is fetched in batches of `SCAN_UNBOUNDED_BATCH_SIZE`
+    /// pairs, so no single RPC carries an unbounded limit and each response
+    /// message stays bounded in size. This is the equivalent of client-go's
+    /// `KVSnapshot.Iter(start, end)`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust,no_run
+    /// # use tikv_client::{Key, KvPair, 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 key1: Key = b"foo".to_vec().into();
+    /// let key2: Key = b"bar".to_vec().into();
+    /// let result: Vec<KvPair> = txn
+    ///     .scan_unbounded(key1..key2)
+    ///     .await
+    ///     .unwrap()
+    ///     .collect();
+    /// // Finish the transaction...
+    /// txn.commit().await.unwrap();
+    /// # });
+    /// ```
+    pub async fn scan_unbounded(
+        &mut self,
+        range: impl Into<BoundRange>,
+    ) -> Result<impl Iterator<Item = KvPair>> {
+        debug!("invoking transactional scan_unbounded request");
+        let (start, end) = range.into().into_keys();
+        let mut start = start;
+        let mut out: Vec<KvPair> = Vec::new();
+        loop {
+            let page: Vec<KvPair> = self
+                .scan((start.clone(), end.clone()), SCAN_UNBOUNDED_BATCH_SIZE)
+                .await?
+                .collect();
+            // Termination follows client-go's Scanner (txnsnapshot/scan.go):
+            // a short page means the range is exhausted — the plan layer fans
+            // a scan out to every region in the range and each region returns
+            // up to `limit` pairs, so a merged page smaller than the batch can
+            // only happen when no region has more data. A full page means more
+            // data may remain; advance past the last returned key with
+            // `next_key()` (the same key+'\x00' successor trick client-go
+            // uses) and fetch again. If the total is an exact multiple of the
+            // batch size, this costs one extra empty page fetch, same as
+            // client-go discovering EOF on its next getData.
+            let full_page = page.len() as u32 == SCAN_UNBOUNDED_BATCH_SIZE;
+            if let Some(last) = page.last() {
+                start = last.key().clone().next_key();
+            }
+            out.extend(page);
+            if !full_page {
+                break;
+            }
+        }
+        Ok(out.into_iter())
+    }
+
     /// Sets the value associated with the given key.
     ///
     /// # Examples
@@ -1128,6 +1205,9 @@ const DEFAULT_HEARTBEAT_INTERVAL: Duration = Duration::from_millis(MAX_TTL / 2);
 /// TiKV recommends each RPC packet should be less than around 1MB. We keep KV size of
 /// each request below 16KB.
 pub const TXN_COMMIT_BATCH_SIZE: u64 = 16 * 1024;
+
+/// Batch size used internally by `scan_unbounded` to paginate through a range.
+const SCAN_UNBOUNDED_BATCH_SIZE: u32 = 1024;
 const TTL_FACTOR: f64 = 6000.0;
 
 /// Optimistic or pessimistic transaction.
@@ -1701,6 +1781,7 @@ impl From<u8> for TransactionStatus {
 #[cfg(test)]
 mod tests {
     use std::any::Any;
+    use std::collections::BTreeMap;
     use std::io;
     use std::sync::atomic::AtomicUsize;
     use std::sync::atomic::Ordering;
@@ -1717,6 +1798,7 @@ mod tests {
     use crate::proto::pdpb::Timestamp;
     use crate::request::Keyspace;
     use crate::transaction::HeartbeatOption;
+    use crate::KvPair;
     use crate::TimestampExt;
     use crate::Transaction;
     use crate::TransactionOptions;
@@ -1870,4 +1952,85 @@ mod tests {
             "expected a lifecycle log carrying start_ts {start_ts}; captured: {logs:?}"
         );
     }
+
+    #[tokio::test]
+    async fn start_ts_returns_transaction_timestamp() {
+        let ts = Timestamp {
+            physical: 1_700_000_000_123,
+            logical: 42,
+            ..Default::default()
+        };
+        let txn = Transaction::new(
+            ts.clone(),
+            Arc::new(MockPdClient::default()),
+            TransactionOptions::new_optimistic().read_only(),
+            Keyspace::Disable,
+        );
+        assert_eq!(txn.start_ts(), ts);
+    }
+
+    #[tokio::test]
+    async fn scan_unbounded_paginates_through_range() {
+        // 2500 pairs force 3 pages at the internal batch size of 1024.
+        let data: BTreeMap<Vec<u8>, Vec<u8>> = (0..2500u32)
+            .map(|i| {
+                (
+                    format!("k{i:04}").into_bytes(),
+                    format!("v{i}").into_bytes(),
+                )
+            })
+            .collect();
+        let scan_calls = Arc::new(AtomicUsize::new(0));
+        let scan_calls_cloned = scan_calls.clone();
+        let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook(
+            move |req: &dyn Any| {
+                let scan = req
+                    .downcast_ref::<kvrpcpb::ScanRequest>()
+                    .expect("only scan requests are expected");
+                scan_calls_cloned.fetch_add(1, Ordering::SeqCst);
+                assert_eq!(scan.limit, super::SCAN_UNBOUNDED_BATCH_SIZE);
+                let mut pairs = Vec::new();
+                for (k, v) in data.range(scan.start_key.clone()..) {
+                    if !scan.end_key.is_empty() && k.as_slice() >= scan.end_key.as_slice() {
+                        break;
+                    }
+                    if pairs.len() >= scan.limit as usize {
+                        break;
+                    }
+                    pairs.push(kvrpcpb::KvPair {
+                        key: k.clone(),
+                        value: v.clone(),
+                        ..Default::default()
+                    });
+                }
+                Ok(Box::new(kvrpcpb::ScanResponse {
+                    pairs,
+                    ..Default::default()
+                }) as Box<dyn Any>)
+            },
+        )));
+
+        let mut txn = Transaction::new(
+            Timestamp::default(),
+            pd_client,
+            TransactionOptions::new_optimistic().read_only(),
+            Keyspace::Disable,
+        );
+        let result: Vec<KvPair> = txn
+            .scan_unbounded("k0000".to_owned()..="k9999".to_owned())
+            .await
+            .unwrap()
+            .collect();
+
+        assert_eq!(result.len(), 2500);
+        assert_eq!(scan_calls.load(Ordering::SeqCst), 3);
+        let keys: Vec<&[u8]> = result.iter().map(|p| p.key().into()).collect();
+        let mut sorted = keys.clone();
+        sorted.sort_unstable();
+        assert_eq!(keys, sorted, "pairs must come back in key order");
+        assert_eq!(keys.first().unwrap(), &b"k0000");
+        assert_eq!(keys.last().unwrap(), &b"k2499");
+        assert_eq!(result[0].value(), &b"v0".to_vec());
+        assert_eq!(result[2499].value(), &b"v2499".to_vec());
+    }
 }