Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/transaction/snapshot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
163 changes: 163 additions & 0 deletions src/transaction/transaction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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());
}
}