From e4e36581b3d781a78340cbe93994376fc15dbae3 Mon Sep 17 00:00:00 2001 From: Eduard Ralph Date: Tue, 20 Oct 2026 14:35:13 +0200 Subject: [PATCH] transaction: add a test-utils probe that drives 2PC one phase at a time Behind a new test-utils feature, Transaction::probe() returns a TransactionProbe that can set the primary key (one the transaction need not write), set an exact lock TTL, prewrite only chosen keys, and commit the primary alone. Nothing it does starts the heartbeat, commits secondaries or rolls back, so dropping the transaction afterwards leaves the locks a crashed client would. It mirrors client-go's CommitterProbe. Without the feature nothing changes. With it, Committer takes an optional lock TTL override; None keeps the computed TTL. Signed-off-by: Eduard Ralph --- Cargo.toml | 3 + Makefile | 6 +- src/transaction/mod.rs | 2 + src/transaction/transaction.rs | 483 ++++++++++++++++++++++++++++++++- tests/integration_tests.rs | 88 ++++++ 5 files changed, 575 insertions(+), 7 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ca1029f3..2d64d794 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,9 @@ prometheus = ["prometheus/push", "prometheus/process"] # Enable integration tests with a running TiKV and PD instance. # Use $PD_ADDRS, comma separated, to set the addresses the tests use. integration-tests = [] +# Test probes: drive a transaction's two-phase commit one phase at a time +# (`Transaction::probe`). For tests only. +test-utils = [] [lib] name = "tikv_client" diff --git a/Makefile b/Makefile index ef3fd397..0aa6bfa9 100644 --- a/Makefile +++ b/Makefile @@ -5,11 +5,11 @@ export RUSTFLAGS=-Dwarnings export PD_ADDRS ?= 127.0.0.1:2379 export MULTI_REGION ?= 1 -ALL_FEATURES := integration-tests +ALL_FEATURES := integration-tests test-utils NEXTEST_ARGS := --config-file $(shell pwd)/config/nextest.toml -INTEGRATION_TEST_ARGS := --features "integration-tests" --test-threads 1 +INTEGRATION_TEST_ARGS := --features "integration-tests test-utils" --test-threads 1 RUN_INTEGRATION_TEST := cargo nextest run ${NEXTEST_ARGS} --all ${INTEGRATION_TEST_ARGS} @@ -24,7 +24,7 @@ check: generate cargo clippy --all-targets --features "${ALL_FEATURES}" -- -D clippy::all unit-test: generate - cargo nextest run ${NEXTEST_ARGS} --all --no-default-features + cargo nextest run ${NEXTEST_ARGS} --all --no-default-features --features test-utils integration-test: integration-test-txn integration-test-raw diff --git a/src/transaction/mod.rs b/src/transaction/mod.rs index 86612a24..ba8efd9c 100644 --- a/src/transaction/mod.rs +++ b/src/transaction/mod.rs @@ -23,6 +23,8 @@ pub use transaction::Mutation; pub use transaction::Scanner; pub use transaction::Transaction; pub use transaction::TransactionOptions; +#[cfg(feature = "test-utils")] +pub use transaction::TransactionProbe; mod buffer; mod client; diff --git a/src/transaction/transaction.rs b/src/transaction/transaction.rs index 5bc36983..b5df61e5 100644 --- a/src/transaction/transaction.rs +++ b/src/transaction/transaction.rs @@ -95,6 +95,10 @@ pub struct Transaction { /// had started — which a rollback retry would otherwise need to know. prewritten: bool, start_instant: Instant, + /// What a [`TransactionProbe`] has set: kept on the transaction so it holds across + /// separate `probe()` calls. + #[cfg(feature = "test-utils")] + probe: ProbeState, } impl Transaction { @@ -119,6 +123,8 @@ impl Transaction { is_heartbeat_started: false, prewritten: false, start_instant: std::time::Instant::now(), + #[cfg(feature = "test-utils")] + probe: ProbeState::default(), } } @@ -1334,6 +1340,10 @@ struct Committer { undetermined: bool, write_size: u64, start_instant: Instant, + /// A lock TTL in milliseconds that replaces the computed one. Only a + /// [`TransactionProbe`] sets it. + #[new(default)] + lock_ttl: Option, } impl Committer { @@ -1413,20 +1423,25 @@ impl Committer { self.mutations.len() ); let primary_lock = self.primary_key.clone().unwrap(); - let elapsed = self.start_instant.elapsed().as_millis() as u64; - let lock_ttl = self.calc_txn_lock_ttl(); + let lock_ttl = match self.lock_ttl { + Some(ttl) => ttl, + None => { + let elapsed = self.start_instant.elapsed().as_millis() as u64; + self.calc_txn_lock_ttl() + elapsed + } + }; let mut request = match &self.options.kind { TransactionKind::Optimistic => new_prewrite_request( self.mutations.clone(), primary_lock, self.start_version.clone(), - lock_ttl + elapsed, + lock_ttl, ), TransactionKind::Pessimistic(for_update_ts) => new_pessimistic_prewrite_request( self.mutations.clone(), primary_lock, self.start_version.clone(), - lock_ttl + elapsed, + lock_ttl, for_update_ts.clone(), ), }; @@ -1719,6 +1734,212 @@ impl Committer { } } +/// Settings a [`TransactionProbe`] has made on its transaction. +#[cfg(feature = "test-utils")] +#[derive(Default)] +struct ProbeState { + primary_key: Option, + lock_ttl: Option, +} + +#[cfg(feature = "test-utils")] +impl Transaction { + /// Drive this transaction's two-phase commit one phase at a time, for tests. + /// + /// Available with the `test-utils` feature. See [`TransactionProbe`]. + pub fn probe(&mut self) -> TransactionProbe<'_, PdC> { + TransactionProbe { txn: self } + } +} + +/// Drives a [`Transaction`]'s two-phase commit one phase at a time, so a test can +/// build the states a crash or a slow client leaves behind: locks whose primary was +/// never written, or a committed primary with its secondaries still locked. +/// +/// Available with the `test-utils` feature, which is meant for tests only. A +/// transaction driven this way is not cleaned up: nothing here starts the heartbeat, +/// commits the secondaries or rolls back, and dropping the transaction after +/// [`prewrite`](Self::prewrite) leaves its locks in place, without the drop check +/// firing. Use an optimistic transaction (or [`HeartbeatOption::NoHeartbeat`]) when +/// the locks are meant to expire. +/// +/// The counterpart of client-go's `CommitterProbe` (`txnkv/transaction/test_probe.go`). +/// +/// # Examples +/// +/// Leave a lock on `"secondary"` whose primary, `"primary"`, is never written: +/// +/// ```rust,no_run +/// # use tikv_client::{TransactionClient, TransactionOptions, CheckLevel}; +/// # futures::executor::block_on(async { +/// # let client = TransactionClient::new(vec!["192.168.0.100"]).await.unwrap(); +/// let mut txn = client +/// .begin_with_options(TransactionOptions::new_optimistic().drop_check(CheckLevel::None)) +/// .await +/// .unwrap(); +/// txn.put("secondary".to_owned(), "value").await.unwrap(); +/// let mut probe = txn.probe(); +/// probe.set_primary_key("primary".to_owned()).set_lock_ttl(1000); +/// probe.prewrite(["secondary".to_owned()]).await.unwrap(); +/// drop(txn); // the crash: the lock stays until a reader resolves it +/// # }); +/// ``` +#[cfg(feature = "test-utils")] +pub struct TransactionProbe<'a, PdC: PdClient = PdRpcClient> { + txn: &'a mut Transaction, +} + +#[cfg(feature = "test-utils")] +impl TransactionProbe<'_, PdC> { + /// Use `key` as the primary key. It does not have to be a key this transaction + /// writes: a primary that is never prewritten leaves the secondaries' locks with no + /// record of the transaction's fate. Defaults to the transaction's own primary. + pub fn set_primary_key(&mut self, key: impl Into) -> &mut Self { + self.txn.probe.primary_key = + Some(key.into().encode_keyspace(self.txn.keyspace, KeyMode::Txn)); + self + } + + /// Prewrite with this lock TTL, in milliseconds, instead of the computed one. + pub fn set_lock_ttl(&mut self, ttl_ms: u64) -> &mut Self { + self.txn.probe.lock_ttl = Some(ttl_ms); + self + } + + /// Prewrite only `keys`, each of which this transaction must already have written + /// (`put`, `insert` or `delete`), and stop: the locks stay until something commits, + /// rolls back or resolves them. + /// + /// Returns the `min_commit_ts` TiKV reported, as [`Transaction::commit`]'s prewrite + /// does. After this call the transaction counts as committing: a later + /// [`Transaction::rollback`] cleans the locks up with `BatchRollback`. + /// + /// Refused for a transaction with 1PC or async commit: there the prewrite is the + /// commit point, so it could not stop between the phases. + pub async fn prewrite( + &mut self, + keys: impl IntoIterator>, + ) -> Result> { + let txn = &mut *self.txn; + txn.refuse_fast_commit_for_probe()?; + let keyspace = txn.keyspace; + let keys: Vec = keys + .into_iter() + .map(|k| k.into().encode_keyspace(keyspace, KeyMode::Txn)) + .collect(); + let mutations: Vec = txn + .buffer + .to_proto_mutations() + .into_iter() + .filter(|m| { + keys.iter() + .any(|k| AsRef::<[u8]>::as_ref(k) == m.key.as_slice()) + }) + .collect(); + if keys.is_empty() || mutations.len() != keys.len() { + return Err(Error::StringError(format!( + "probe: asked to prewrite {} key(s), but the transaction has written {} of \ + them; write each key before prewriting it", + keys.len(), + mutations.len() + ))); + } + let primary_key = txn.probe_primary_key()?; + txn.enter_commit_for_probe()?; + let mut committer = Committer::new( + Some(primary_key), + mutations, + txn.timestamp.clone(), + txn.rpc.clone(), + txn.options.clone(), + keyspace, + txn.buffer.get_write_size() as u64, + txn.start_instant, + ); + committer.lock_ttl = txn.probe.lock_ttl; + committer.prewrite().await + } + + /// Commit the primary key alone, at a fresh timestamp from PD, and return that + /// commit timestamp. The secondaries stay locked: resolving them is left to a + /// reader, which must roll them forward. + /// + /// A commit timestamp TiKV rejects as too small is retried with a newer one, as in + /// [`Transaction::commit`]. On success the transaction is committed, so + /// [`Transaction::rollback`] is refused from then on. When the outcome is unknown, + /// the error is an [`Error::UndeterminedError`], as from [`Transaction::commit`], and + /// rollback is refused too: the primary may be durable, so rolling back its + /// secondaries could tear the transaction. The caller has to find out what happened. + pub async fn commit_primary(&mut self) -> Result { + let txn = &mut *self.txn; + txn.refuse_fast_commit_for_probe()?; + let primary_key = txn.probe_primary_key()?; + txn.enter_commit_for_probe()?; + let mut committer = Committer::new( + Some(primary_key), + Vec::new(), + txn.timestamp.clone(), + txn.rpc.clone(), + txn.options.clone(), + txn.keyspace, + txn.buffer.get_write_size() as u64, + txn.start_instant, + ); + // Counted as committed from before the commit is sent, so a call cancelled in + // flight cannot leave rollback open while the primary may be durable. Only a + // clean failure, where the primary did not commit, reopens it. + txn.set_status(TransactionStatus::Committed); + match committer.commit_primary_with_retry().await { + Ok(commit_ts) => Ok(commit_ts), + Err(e) if committer.undetermined => Err(Error::UndeterminedError(Box::new(e))), + Err(e) => { + txn.set_status(TransactionStatus::StartedCommit); + Err(e) + } + } + } +} + +#[cfg(feature = "test-utils")] +impl Transaction { + fn probe_primary_key(&self) -> Result { + self.probe + .primary_key + .clone() + .or_else(|| self.buffer.get_primary_key()) + .ok_or(Error::NoPrimaryKey) + } + + /// The probe drives two-phase commit only: with 1PC or async commit the prewrite + /// is the commit point, so there is no phase boundary to stop at. + fn refuse_fast_commit_for_probe(&self) -> Result<()> { + if self.options.try_one_pc || self.options.async_commit { + return Err(Error::StringError( + "probe: drives two-phase commit only; this transaction uses 1PC or async commit" + .to_owned(), + )); + } + Ok(()) + } + + /// Move to `StartedCommit`, as [`Transaction::commit`] does before it prewrites. + fn enter_commit_for_probe(&mut self) -> Result<()> { + if !self.transit_status( + |status| { + matches!( + status, + TransactionStatus::StartedCommit | TransactionStatus::Active + ) + }, + TransactionStatus::StartedCommit, + ) { + return Err(Error::OperationAfterCommitError); + } + self.prewritten = true; + Ok(()) + } +} + /// Is the outcome of a commit-point request unknown? /// /// Two ways to lose certainty: we failed to receive a response (a gRPC error — the @@ -1966,4 +2187,258 @@ mod tests { }) ))); } + + /// Every request a mock TiKV received, by type, for the probe tests. + #[cfg(feature = "test-utils")] + fn recording_pd_client( + seen: Arc>>, + prewrites: Arc>>, + commits: Arc>>, + ) -> Arc { + Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + move |req: &dyn Any| { + if let Some(r) = req.downcast_ref::() { + seen.lock().unwrap().push("prewrite".to_owned()); + prewrites.lock().unwrap().push(r.clone()); + Ok(Box::::default() as Box) + } else if let Some(r) = req.downcast_ref::() { + seen.lock().unwrap().push("commit".to_owned()); + commits.lock().unwrap().push(r.clone()); + Ok(Box::::default() as Box) + } else if req.downcast_ref::().is_some() { + seen.lock().unwrap().push("heartbeat".to_owned()); + Ok(Box::::default() as Box) + } else { + seen.lock().unwrap().push("other".to_owned()); + Ok(Box::::default() as Box) + } + }, + ))) + } + + #[cfg(feature = "test-utils")] + #[tokio::test] + async fn probe_prewrites_only_the_named_keys_under_the_named_primary() { + let seen = Arc::new(Mutex::new(Vec::new())); + let prewrites = Arc::new(Mutex::new(Vec::new())); + let commits = Arc::new(Mutex::new(Vec::new())); + let pd_client = recording_pd_client(seen.clone(), prewrites.clone(), commits.clone()); + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + TransactionOptions::new_optimistic(), + Keyspace::Disable, + ); + txn.put("secondary".to_owned(), "v").await.unwrap(); + txn.put("other".to_owned(), "w").await.unwrap(); + + let mut probe = txn.probe(); + probe + .set_primary_key("primary".to_owned()) + .set_lock_ttl(1234); + probe.prewrite(["secondary".to_owned()]).await.unwrap(); + + // One prewrite, and nothing else: no commit, no heartbeat, no rollback. + assert_eq!(*seen.lock().unwrap(), vec!["prewrite".to_owned()]); + let req = prewrites.lock().unwrap().pop().unwrap(); + assert_eq!(req.primary_lock, b"primary".to_vec()); + assert_eq!(req.lock_ttl, 1234, "the TTL is exactly the one set"); + let keys: Vec<_> = req.mutations.iter().map(|m| m.key.clone()).collect(); + assert_eq!(keys, vec![b"secondary".to_vec()]); + assert_eq!(req.secondaries, vec![b"secondary".to_vec()]); + + // Dropping it now is the crash: no cleanup, and no drop-check panic either. + drop(txn); + assert_eq!(*seen.lock().unwrap(), vec!["prewrite".to_owned()]); + } + + #[cfg(feature = "test-utils")] + #[tokio::test] + async fn probe_refuses_a_key_the_transaction_did_not_write() { + let seen = Arc::new(Mutex::new(Vec::new())); + let pd_client = recording_pd_client( + seen.clone(), + Arc::new(Mutex::new(Vec::new())), + Arc::new(Mutex::new(Vec::new())), + ); + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + TransactionOptions::new_optimistic().drop_check(crate::CheckLevel::None), + Keyspace::Disable, + ); + txn.put("written".to_owned(), "v").await.unwrap(); + let err = txn + .probe() + .prewrite(["written".to_owned(), "never-written".to_owned()]) + .await + .unwrap_err(); + assert!(err.to_string().contains("has written 1 of them"), "{err}"); + assert!(seen.lock().unwrap().is_empty(), "nothing was sent"); + // A refused prewrite leaves the transaction usable. + txn.rollback().await.unwrap(); + } + + #[cfg(feature = "test-utils")] + #[tokio::test] + async fn probe_commits_the_primary_alone() { + let seen = Arc::new(Mutex::new(Vec::new())); + let commits = Arc::new(Mutex::new(Vec::new())); + let pd_client = recording_pd_client( + seen.clone(), + Arc::new(Mutex::new(Vec::new())), + commits.clone(), + ); + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + TransactionOptions::new_optimistic(), + Keyspace::Disable, + ); + txn.put("primary".to_owned(), "p").await.unwrap(); + txn.put("secondary".to_owned(), "s").await.unwrap(); + let mut probe = txn.probe(); + probe.set_primary_key("primary".to_owned()); + probe + .prewrite(["primary".to_owned(), "secondary".to_owned()]) + .await + .unwrap(); + probe.commit_primary().await.unwrap(); + + assert_eq!( + *seen.lock().unwrap(), + vec!["prewrite".to_owned(), "commit".to_owned()] + ); + let req = commits.lock().unwrap().pop().unwrap(); + assert_eq!( + req.keys, + vec![b"primary".to_vec()], + "the primary, and only it" + ); + + // The primary is durable, so rolling back the secondaries is refused. + assert!(matches!( + txn.rollback().await, + Err(crate::Error::OperationAfterCommitError) + )); + assert_eq!(seen.lock().unwrap().len(), 2, "no rollback was sent"); + } + + #[cfg(feature = "test-utils")] + #[tokio::test] + async fn probe_reports_a_lost_primary_commit_as_undetermined() { + let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + |req: &dyn Any| { + if req.downcast_ref::().is_some() { + Ok(Box::::default() as Box) + } else { + // The commit's answer never arrives. + Err(crate::Error::GrpcAPI(tonic::Status::unavailable("lost"))) + } + }, + ))); + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + TransactionOptions::new_optimistic().drop_check(crate::CheckLevel::None), + Keyspace::Disable, + ); + txn.put("primary".to_owned(), "p").await.unwrap(); + let mut probe = txn.probe(); + probe.prewrite(["primary".to_owned()]).await.unwrap(); + let err = probe.commit_primary().await.unwrap_err(); + assert!(matches!(err, crate::Error::UndeterminedError(_)), "{err:?}"); + // The primary may be durable: rolling back its secondaries is refused. + assert!(matches!( + txn.rollback().await, + Err(crate::Error::OperationAfterCommitError) + )); + } + + #[cfg(feature = "test-utils")] + #[tokio::test] + async fn probe_commit_primary_rejected_cleanly_leaves_rollback_open() { + let pd_client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + |req: &dyn Any| { + if req.downcast_ref::().is_some() { + Ok(Box::new(kvrpcpb::CommitResponse { + error: Some(kvrpcpb::KeyError { + abort: "rejected".to_owned(), + ..Default::default() + }), + ..Default::default() + }) as Box) + } else if req.downcast_ref::().is_some() { + Ok(Box::::default() as Box) + } else { + Ok(Box::::default() as Box) + } + }, + ))); + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + TransactionOptions::new_optimistic().drop_check(crate::CheckLevel::None), + Keyspace::Disable, + ); + txn.put("primary".to_owned(), "p").await.unwrap(); + let mut probe = txn.probe(); + probe.prewrite(["primary".to_owned()]).await.unwrap(); + let err = probe.commit_primary().await.unwrap_err(); + assert!( + !matches!(err, crate::Error::UndeterminedError(_)), + "{err:?}" + ); + // The primary did not commit, so the locks can still be rolled back. + txn.rollback().await.unwrap(); + } + + #[cfg(feature = "test-utils")] + #[tokio::test] + async fn probe_refuses_one_pc_and_async_commit() { + for options in [ + TransactionOptions::new_optimistic().try_one_pc(), + TransactionOptions::new_optimistic().use_async_commit(), + ] { + let seen = Arc::new(Mutex::new(Vec::new())); + let pd_client = recording_pd_client( + seen.clone(), + Arc::new(Mutex::new(Vec::new())), + Arc::new(Mutex::new(Vec::new())), + ); + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + options.drop_check(crate::CheckLevel::None), + Keyspace::Disable, + ); + txn.put("k".to_owned(), "v").await.unwrap(); + let err = txn.probe().prewrite(["k".to_owned()]).await.unwrap_err(); + assert!(err.to_string().contains("two-phase commit only"), "{err}"); + let err = txn.probe().commit_primary().await.unwrap_err(); + assert!(err.to_string().contains("two-phase commit only"), "{err}"); + assert!(seen.lock().unwrap().is_empty(), "nothing was sent"); + } + } + + #[cfg(feature = "test-utils")] + #[tokio::test] + async fn without_the_probe_the_lock_ttl_is_still_computed() { + let prewrites = Arc::new(Mutex::new(Vec::new())); + let pd_client = recording_pd_client( + Arc::new(Mutex::new(Vec::new())), + prewrites.clone(), + Arc::new(Mutex::new(Vec::new())), + ); + let mut txn = Transaction::new( + Timestamp::default(), + pd_client, + TransactionOptions::new_optimistic().heartbeat_option(HeartbeatOption::NoHeartbeat), + Keyspace::Disable, + ); + txn.put("k".to_owned(), "v").await.unwrap(); + txn.commit().await.unwrap(); + let req = prewrites.lock().unwrap().pop().unwrap(); + assert!(req.lock_ttl >= super::DEFAULT_LOCK_TTL, "{}", req.lock_ttl); + } } diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index de1d111a..cd2e2bb8 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -1623,3 +1623,91 @@ async fn txn_unsafe_destroy_range() -> Result<()> { Ok(()) } + +/// The probe leaves the lock a crashed client would: a secondary whose primary was never +/// written. Once its TTL has passed, a reader resolves it as rolled back. +#[cfg(feature = "test-utils")] +#[tokio::test] +#[serial] +async fn txn_probe_leaves_an_orphaned_secondary_lock() -> Result<()> { + init().await?; + let client = + TransactionClient::new_with_config(pd_addrs(), Config::default().with_default_keyspace()) + .await?; + let primary = "probe-orphan/a-primary".to_owned(); + let secondary = "probe-orphan/z-secondary".to_owned(); + + let mut txn = client + .begin_with_options( + TransactionOptions::new_optimistic().drop_check(tikv_client::CheckLevel::None), + ) + .await?; + txn.put(secondary.clone(), "orphan").await?; + let mut probe = txn.probe(); + probe.set_primary_key(primary.clone()).set_lock_ttl(1000); + probe.prewrite([secondary.clone()]).await?; + drop(txn); + + let locks = client + .scan_locks( + &client.current_timestamp().await?, + "probe-orphan/".to_owned().., + 16, + ) + .await?; + assert_eq!(locks.len(), 1, "{locks:?}"); + assert_eq!(locks[0].key, secondary.clone().into_bytes()); + assert_eq!(locks[0].primary_lock, primary.into_bytes()); + + tokio::time::sleep(std::time::Duration::from_millis(1500)).await; + let mut reader = client.begin_optimistic().await?; + assert_eq!(reader.get(secondary).await?, None); + reader.rollback().await?; + Ok(()) +} + +/// With the primary committed and the secondary still locked, a reader must roll the +/// secondary forward. +#[cfg(feature = "test-utils")] +#[tokio::test] +#[serial] +async fn txn_probe_commit_primary_leaves_the_secondary_to_roll_forward() -> Result<()> { + init().await?; + let client = + TransactionClient::new_with_config(pd_addrs(), Config::default().with_default_keyspace()) + .await?; + let primary = "probe-forward/a-primary".to_owned(); + let secondary = "probe-forward/z-secondary".to_owned(); + + let mut txn = client + .begin_with_options( + TransactionOptions::new_optimistic().drop_check(tikv_client::CheckLevel::None), + ) + .await?; + txn.put(primary.clone(), "p").await?; + txn.put(secondary.clone(), "s").await?; + let mut probe = txn.probe(); + probe.set_primary_key(primary.clone()); + probe.prewrite([primary.clone(), secondary.clone()]).await?; + probe.commit_primary().await?; + drop(txn); + + let locks = client + .scan_locks( + &client.current_timestamp().await?, + "probe-forward/".to_owned().., + 16, + ) + .await?; + assert_eq!( + locks.len(), + 1, + "only the secondary is still locked: {locks:?}" + ); + assert_eq!(locks[0].key, secondary.clone().into_bytes()); + + let mut reader = client.begin_optimistic().await?; + assert_eq!(reader.get(secondary).await?, Some(b"s".to_vec())); + reader.rollback().await?; + Ok(()) +}