diff --git a/src/transaction/snapshot.rs b/src/transaction/snapshot.rs index f9cb6b9d..a44ceca2 100644 --- a/src/transaction/snapshot.rs +++ b/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, + ) -> Result> { + 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, diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index 61b87423..e6bee423 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -119,6 +119,19 @@ impl Transaction { } } + /// 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 Transaction { .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 = 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, + ) -> Result> { + debug!("invoking transactional scan_unbounded request"); + let (start, end) = range.into().into_keys(); + let mut start = start; + let mut out: Vec = Vec::new(); + loop { + let page: Vec = 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 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> = (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::() + .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) + }, + ))); + + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + TransactionOptions::new_optimistic().read_only(), + Keyspace::Disable, + ); + let result: Vec = 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()); + } }