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

src/lib.rs

@@ -164,6 +164,8 @@ pub use crate::transaction::Client as TransactionClient;
 #[doc(inline)]
 pub use crate::transaction::ProtoLockInfo;
 #[doc(inline)]
+pub use crate::transaction::Scanner;
+#[doc(inline)]
 pub use crate::transaction::Snapshot;
 #[doc(inline)]
 pub use crate::transaction::SyncSnapshot;

src/transaction/mod.rs

@@ -20,6 +20,7 @@ pub use transaction::CheckLevel;
 #[doc(hidden)]
 pub use transaction::HeartbeatOption;
 pub use transaction::Mutation;
+pub use transaction::Scanner;
 pub use transaction::Transaction;
 pub use transaction::TransactionOptions;
 

src/transaction/snapshot.rs

@@ -7,6 +7,7 @@ use crate::BoundRange;
 use crate::Key;
 use crate::KvPair;
 use crate::Result;
+use crate::Scanner;
 use crate::Transaction;
 use crate::Value;
 
@@ -54,6 +55,15 @@ impl Snapshot {
         self.transaction.scan(range, limit).await
     }
 
+    /// Scan a range lazily, buffering at most one batch of pairs in memory.
+    ///
+    /// Equivalent of client-go's `KVSnapshot.Iter(start, end)`; see
+    /// [`Transaction::scan_unbounded`].
+    pub fn scan_unbounded(&mut self, range: impl Into<BoundRange>) -> Scanner<'_> {
+        debug!("creating scan_unbounded scanner on snapshot");
+        self.transaction.scan_unbounded(range)
+    }
+
     /// Scan a range, return at most `limit` keys that lying in the range.
     pub async fn scan_keys(
         &mut self,

src/transaction/transaction.rs

@@ -448,6 +448,41 @@ impl<PdC: PdClient> Transaction<PdC> {
             .map(KvPair::into_key))
     }
 
+    /// Create a lazy scanner over a range, the equivalent of client-go's
+    /// `KVSnapshot.Iter(start, end)`.
+    ///
+    /// The returned [`Scanner`] holds at most one batch of
+    /// `SCAN_UNBOUNDED_BATCH_SIZE` pairs in memory and fetches the next batch
+    /// only as it is advanced, so no RPC carries an unbounded limit and client
+    /// memory stays bounded regardless of the range size.
+    ///
+    /// Construction performs no I/O; the first fetch happens on the first
+    /// [`Scanner::next`] call.
+    ///
+    /// # 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 mut scanner = txn.scan_unbounded(key1..key2);
+    /// while let Some(_pair) = scanner.next().await.unwrap() {
+    ///     // Process the pair...
+    /// }
+    /// // Finish the transaction...
+    /// txn.commit().await.unwrap();
+    /// # });
+    /// ```
+    pub fn scan_unbounded(&mut self, range: impl Into<BoundRange>) -> Scanner<'_, PdC> {
+        debug!("creating transactional scan_unbounded scanner");
+        let (start, end) = range.into().into_keys();
+        Scanner::new(self, start, end)
+    }
+
     /// Sets the value associated with the given key.
     ///
     /// # Examples
@@ -1128,6 +1163,76 @@ 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;
+
+/// A lazy range scanner: the equivalent of client-go's `Scanner`
+/// (txnkv/txnsnapshot/scan.go). It buffers at most one batch of
+/// `SCAN_UNBOUNDED_BATCH_SIZE` pairs and fetches the next batch only when the
+/// buffered one is exhausted, so scanning a large range cannot balloon client
+/// memory the way a collect-all API can.
+///
+/// Created by [`Transaction::scan_unbounded`]. It borrows the transaction
+/// mutably until dropped; client-go's `Scanner.Close` maps to Rust's drop.
+pub struct Scanner<'a, PdC: PdClient = PdRpcClient> {
+    txn: &'a mut Transaction<PdC>,
+    end: Option<Key>,
+    next_start: Key,
+    cache: std::vec::IntoIter<KvPair>,
+    eof: bool,
+}
+
+impl<'a, PdC: PdClient> Scanner<'a, PdC> {
+    fn new(txn: &'a mut Transaction<PdC>, start: Key, end: Option<Key>) -> Scanner<'a, PdC> {
+        Scanner {
+            txn,
+            end,
+            next_start: start,
+            cache: Vec::new().into_iter(),
+            eof: false,
+        }
+    }
+
+    /// Advance to the next pair, fetching the next batch from TiKV only when
+    /// the buffered batch is exhausted. Returns `Ok(None)` once the range is
+    /// exhausted.
+    pub async fn next(&mut self) -> Result<Option<KvPair>> {
+        loop {
+            if let Some(pair) = self.cache.next() {
+                return Ok(Some(pair));
+            }
+            if self.eof {
+                return Ok(None);
+            }
+            let page: Vec<KvPair> = self
+                .txn
+                .scan(
+                    (self.next_start.clone(), self.end.clone()),
+                    SCAN_UNBOUNDED_BATCH_SIZE,
+                )
+                .await?
+                .collect();
+            // Termination follows client-go's Scanner (txnsnapshot/scan.go):
+            // a short page means no region in the range has more data — the
+            // plan layer fans a scan out to every region and each returns up
+            // to `limit` pairs, so a merged page smaller than the batch can
+            // only happen at the end of the range. 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 on the next call. If the total is an exact
+            // multiple of the batch size, the final call costs one extra empty
+            // page fetch, same as client-go discovering EOF on its next
+            // getData.
+            self.eof = (page.len() as u32) < SCAN_UNBOUNDED_BATCH_SIZE;
+            if let Some(last) = page.last() {
+                self.next_start = last.key().clone().next_key();
+            }
+            self.cache = page.into_iter();
+        }
+    }
+}
+
+/// Batch size used internally by `scan_unbounded` to paginate through a range.
+/// Matches client-go's `DefaultScanBatchSize` (txnkv/txnsnapshot/snapshot.go).
+const SCAN_UNBOUNDED_BATCH_SIZE: u32 = 256;
 const TTL_FACTOR: f64 = 6000.0;
 
 /// Optimistic or pessimistic transaction.
@@ -1701,6 +1806,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 +1823,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 +1977,92 @@ mod tests {
             "expected a lifecycle log carrying start_ts {start_ts}; captured: {logs:?}"
         );
     }
+
+    #[tokio::test]
+    async fn start_timestamp_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_timestamp(), ts);
+    }
+
+    #[tokio::test]
+    async fn scan_unbounded_paginates_through_range() {
+        // 2500 pairs force 10 pages (9 full + 1 short) at the internal batch size of 256.
+        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 mut scanner = txn.scan_unbounded("k0000".to_owned()..="k9999".to_owned());
+        // Construction is lazy: no RPC before the first next().
+        assert_eq!(scan_calls.load(Ordering::SeqCst), 0);
+        let mut result: Vec<KvPair> = Vec::new();
+        while let Some(pair) = scanner.next().await.unwrap() {
+            result.push(pair);
+            // The scanner buffers at most one batch: after consuming exactly
+            // one batch, only one fetch has happened.
+            if result.len() as u32 == super::SCAN_UNBOUNDED_BATCH_SIZE {
+                assert_eq!(scan_calls.load(Ordering::SeqCst), 1);
+            }
+        }
+
+        assert_eq!(result.len(), 2500);
+        assert_eq!(scan_calls.load(Ordering::SeqCst), 10);
+        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());
+    }
 }