From 545eeb85278b8d87b3f54524e59e724e242b9b7f Mon Sep 17 00:00:00 2001 From: Yijun Zhao Date: Sun, 30 Aug 2026 23:56:04 +0800 Subject: [PATCH] fix(transaction): recover expired async commit locks Recover expired async-commit locks by checking all secondary locks before resolving, instead of retrying CheckTxnStatus forever (#528). Also bundled: - cleanup_locks: take max with the primary lock's min_commit_ts when computing the commit version from secondaries - CheckSecondaryLocks merge: return an error instead of panicking on conflicting commit TS across regions Signed-off-by: Yijun Zhao --- src/common/errors.rs | 4 + src/transaction/lock.rs | 733 +++++++++++++++++++++++++++++++++--- src/transaction/requests.rs | 377 +++++++++++++++++-- tests/failpoint_tests.rs | 96 ++++- 4 files changed, 1115 insertions(+), 95 deletions(-) diff --git a/src/common/errors.rs b/src/common/errors.rs index f9a99d2f..3137fbd7 100644 --- a/src/common/errors.rs +++ b/src/common/errors.rs @@ -117,6 +117,10 @@ pub enum Error { KvError { message: String }, #[error("{}", message)] InternalError { message: String }, + /// The server returned a response that violates a protocol invariant. Acting on such + /// a response could corrupt transaction state, so the operation is aborted instead. + #[error("protocol violation: {}", message)] + ProtocolViolation { message: String }, #[error("{0}")] StringError(String), #[error("PessimisticLock error: {:?}", inner)] diff --git a/src/transaction/lock.rs b/src/transaction/lock.rs index 98bb3683..bd254434 100644 --- a/src/transaction/lock.rs +++ b/src/transaction/lock.rs @@ -16,6 +16,7 @@ use crate::backoff::DEFAULT_REGION_BACKOFF; use crate::backoff::OPTIMISTIC_BACKOFF; use crate::kv::HexRepr; use crate::pd::PdClient; +use crate::Key; use crate::proto::kvrpcpb; use crate::proto::kvrpcpb::TxnInfo; @@ -78,6 +79,10 @@ pub(crate) fn reject_shared_locks(locks: &[kvrpcpb::LockInfo]) -> Result<()> { /// which means the key is finally either committed or rolled back, before we read the value of /// the key. We first use `CheckTxnStatus` to get the transaction's final status (committed or /// rolled back), then use `ResolveLock` to resolve the remaining locks in the transaction. +/// +/// An expired async-commit lock needs an extra step in between: TiKV refuses to roll it back +/// via `CheckTxnStatus`, so the transaction's final status is first recovered from all of its +/// secondary locks (`CheckSecondaryLocks`), and then every lock of the transaction is resolved. pub async fn resolve_locks( locks: Vec, timestamp: Timestamp, @@ -107,24 +112,11 @@ pub async fn resolve_locks( // This matches the client-go `LockResolver.ResolveLocksWithOpts` flow: query txn status for // each encountered lock, then resolve immediately when the status is final. for lock in locks { - let region_ver_id = pd_client - .region_for_key(&lock.key.clone().into()) - .await? - .ver_id(); - // skip if the region is cleaned - if clean_regions - .get(&lock.lock_version) - .map(|regions| regions.contains(®ion_ver_id)) - .unwrap_or(false) - { - continue; - } - + let mut keys = Vec::new(); let commit_version = match commit_versions.get(&lock.lock_version) { - Some(&commit_version) => Some(commit_version), + Some(&commit_version) => commit_version, None => { - // TODO: handle primary mismatch error. - let status = lock_resolver + let status = match lock_resolver .get_txn_status_from_lock( OPTIMISTIC_BACKOFF, &lock, @@ -134,28 +126,101 @@ pub async fn resolve_locks( pd_client.clone(), keyspace, ) - .await?; - match &status.kind { - TransactionStatusKind::Committed(ts) => { - let commit_version = ts.version(); - commit_versions.insert(lock.lock_version, commit_version); - Some(commit_version) + .await + { + Ok(status) => status, + Err(Error::KeyError(key_err)) + if key_err.primary_mismatch.is_some() + && lock.lock_type == kvrpcpb::Op::PessimisticLock as i32 => + { + // The encountered pessimistic lock points at a stale primary: the + // transaction changed its primary after writing this lock + // (pingcap/tidb#42937). Roll back only this stale lock — the + // transaction itself may still be alive, and a region-wide + // ResolveLock could roll back its other, legitimate locks. + let for_update_ts = if lock.lock_for_update_ts == 0 { + u64::MAX + } else { + lock.lock_for_update_ts + }; + let req = requests::new_pessimistic_rollback_request( + vec![lock.key.clone()], + lock.lock_version, + for_update_ts, + ); + let plan = + crate::request::PlanBuilder::new(pd_client.clone(), keyspace, req) + .retry_multi_region(DEFAULT_REGION_BACKOFF) + .extract_error() + .plan(); + plan.execute().await?; + continue; } - TransactionStatusKind::RolledBack => { - commit_versions.insert(lock.lock_version, 0); - Some(0) + Err(err) => return Err(err), + }; + match &status.kind { + TransactionStatusKind::Committed(ts) => ts.version(), + TransactionStatusKind::RolledBack => 0, + TransactionStatusKind::Locked(_, primary_lock) + if status.is_expired && primary_lock.use_async_commit => + { + // TiKV will not roll back an async-commit primary. Inspect every + // secondary to recover its decision, then resolve all its regions. + let secondary_status = lock_resolver + .check_all_secondaries( + pd_client.clone(), + keyspace, + primary_lock.secondaries.clone(), + lock.lock_version, + ) + .await?; + let commit_version = if secondary_status.fallback_2pc { + let fallback_status = lock_resolver + .get_txn_status_from_lock( + OPTIMISTIC_BACKOFF, + &lock, + caller_start_ts, + current_ts, + true, + pd_client.clone(), + keyspace, + ) + .await?; + match &fallback_status.kind { + TransactionStatusKind::Committed(ts) => ts.version(), + TransactionStatusKind::RolledBack => 0, + TransactionStatusKind::Locked(_, lock_info) => { + live_locks.push(lock_info.clone()); + continue; + } + } + } else { + secondary_status.resolved_commit_version(primary_lock.min_commit_ts)? + }; + keys.extend(primary_lock.secondaries.iter().cloned()); + keys.push(primary_lock.key.clone()); + commit_version } TransactionStatusKind::Locked(_, lock_info) => { live_locks.push(lock_info.clone()); - None + continue; } } } }; + commit_versions.insert(lock.lock_version, commit_version); - if let Some(commit_version) = commit_version { + // ResolveLock sweeps the transaction's locks throughout a region, so each + // region needs only one successful request even when several keys point to it. + for key in keys.iter().chain([&lock.key]) { + let key = key.into(); + let region_ver_id = pd_client.region_for_key(key).await?.ver_id(); + let regions = clean_regions.entry(lock.lock_version).or_default(); + if regions.contains(®ion_ver_id) { + continue; + } let cleaned_region = resolve_lock_with_retry( - &lock.key, + key, lock.lock_version, commit_version, lock.is_txn_file, @@ -164,17 +229,14 @@ pub async fn resolve_locks( OPTIMISTIC_BACKOFF, ) .await?; - clean_regions - .entry(lock.lock_version) - .or_default() - .insert(cleaned_region); + regions.insert(cleaned_region); } } Ok(live_locks) } async fn resolve_lock_with_retry( - #[allow(clippy::ptr_arg)] key: &Vec, + key: &Key, start_version: u64, commit_version: u64, is_txn_file: bool, @@ -187,7 +249,7 @@ async fn resolve_lock_with_retry( loop { attempt += 1; debug!("resolving locks: attempt {}", attempt); - let store = pd_client.clone().store_for_key(key.into()).await?; + let store = pd_client.clone().store_for_key(key).await?; let ver_id = store.region_with_leader.ver_id(); let request = requests::new_resolve_lock_request(start_version, commit_version, is_txn_file); @@ -379,10 +441,7 @@ impl LockResolver { debug!( "secondary status, txn_id:{}, commit_ts:{:?}, min_commit_version:{}, fallback_2pc:{}", txn_id, - secondary_status - .commit_ts - .as_ref() - .map_or(0, |ts| ts.version()), + secondary_status.commit_ts, secondary_status.min_commit_ts, secondary_status.fallback_2pc, ); @@ -404,11 +463,8 @@ impl LockResolver { ) .await?; } else { - let commit_ts = if let Some(commit_ts) = &secondary_status.commit_ts { - commit_ts.version() - } else { - secondary_status.min_commit_ts - }; + let commit_ts = + secondary_status.resolved_commit_version(lock_info.min_commit_ts)?; txn_infos.insert(txn_id, (commit_ts, l.is_txn_file)); continue; } @@ -500,17 +556,20 @@ impl LockResolver { .plan(); let mut status: TransactionStatus = match plan.execute().await { Ok(status) => status, - Err(Error::ExtractedErrors(mut errors)) => match errors.pop() { - Some(Error::KeyError(key_err)) => { - if let Some(txn_not_found) = key_err.txn_not_found { - return Err(Error::TxnNotFound(txn_not_found)); + Err(Error::ExtractedErrors(mut errors)) | Err(Error::MultipleKeyErrors(mut errors)) => { + match errors.pop() { + Some(Error::KeyError(key_err)) => { + if let Some(txn_not_found) = key_err.txn_not_found { + return Err(Error::TxnNotFound(txn_not_found)); + } + // A PrimaryMismatch error propagates to `resolve_locks`, which rolls + // back the stale pessimistic lock it was reported for. + return Err(Error::KeyError(key_err)); } - // TODO: handle primary mismatch error. - return Err(Error::KeyError(key_err)); + Some(err) => return Err(err), + None => unreachable!(), } - Some(err) => return Err(err), - None => unreachable!(), - }, + } Err(err) => return Err(err), }; @@ -532,6 +591,7 @@ impl LockResolver { ) -> Result { let req = new_check_secondary_locks_request(keys, txn_id); let plan = crate::request::PlanBuilder::new(pd_client.clone(), keyspace, req) + .preserve_shard() .retry_multi_region(DEFAULT_REGION_BACKOFF) .extract_error() .merge(Collect) @@ -641,6 +701,8 @@ pub fn lock_until_expired_ms(lock_version: u64, ttl: u64, current: Timestamp) -> #[cfg(test)] mod tests { use std::any::Any; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicU64; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -651,6 +713,80 @@ mod tests { use crate::mock::MockKvClient; use crate::mock::MockPdClient; use crate::proto::errorpb; + use crate::proto::metapb; + use crate::region::RegionWithLeader; + use crate::request::EncodeKeyspace; + use crate::request::KeyMode; + use crate::Key; + + /// A transaction ID whose lock is always expired under the mock PD clock: + /// `Timestamp::from_version` casts the version to `i64` before shifting, so + /// `u64::MAX` yields a physical time of -1 (production timestamp arithmetic, not + /// mock behavior), while `MockPdClient::get_timestamp` stands at physical 0 — + /// any positive TTL has therefore already lapsed. + const EXPIRED_TXN_VERSION: u64 = u64::MAX; + + /// The async-commit recovery tests place these on both sides of the mock region + /// boundary (`MockPdClient::region1` ends at `[10]`), so a full recovery must + /// resolve two regions. + const PRIMARY_KEY: &[u8] = &[1]; // mock region 1 + const SECONDARY_KEY: &[u8] = &[11]; // mock region 2 + + fn encoded(key: &[u8], keyspace: Keyspace) -> Vec { + Key::from(key.to_vec()) + .encode_keyspace(keyspace, KeyMode::Txn) + .into() + } + + /// The primary lock of an expired async-commit transaction — the shape returned + /// by `CheckTxnStatus` and encountered by readers. + fn async_commit_primary_lock( + primary_key: &[u8], + min_commit_ts: u64, + secondaries: Vec>, + ) -> kvrpcpb::LockInfo { + kvrpcpb::LockInfo { + key: primary_key.to_vec(), + primary_lock: primary_key.to_vec(), + lock_version: EXPIRED_TXN_VERSION, + lock_ttl: 1, + min_commit_ts, + use_async_commit: true, + secondaries, + ..Default::default() + } + } + + /// A `CheckTxnStatusResponse` reporting the transaction still locked by `primary`. + fn still_locked_response(primary: kvrpcpb::LockInfo) -> kvrpcpb::CheckTxnStatusResponse { + kvrpcpb::CheckTxnStatusResponse { + lock_ttl: 1, + lock_info: Some(primary), + action: kvrpcpb::Action::NoAction as i32, + ..Default::default() + } + } + + /// A still-live async-commit secondary lock. + fn async_commit_secondary_lock(key: &[u8], min_commit_ts: u64) -> kvrpcpb::LockInfo { + kvrpcpb::LockInfo { + key: key.to_vec(), + lock_version: EXPIRED_TXN_VERSION, + min_commit_ts, + use_async_commit: true, + ..Default::default() + } + } + + /// A secondary lock that fell back from async commit to plain 2PC. + fn fallback_2pc_secondary_lock(key: &[u8]) -> kvrpcpb::LockInfo { + kvrpcpb::LockInfo { + key: key.to_vec(), + lock_version: EXPIRED_TXN_VERSION, + use_async_commit: false, + ..Default::default() + } + } #[test] fn shared_locks_are_refused_never_misresolved() { @@ -714,10 +850,17 @@ mod tests { let key = vec![1]; let region1 = MockPdClient::region1(); - let resolved_region = - resolve_lock_with_retry(&key, 1, 2, false, client.clone(), keyspace, backoff.clone()) - .await - .unwrap(); + let resolved_region = resolve_lock_with_retry( + (&key).into(), + 1, + 2, + false, + client.clone(), + keyspace, + backoff.clone(), + ) + .await + .unwrap(); assert_eq!(region1.ver_id(), resolved_region); // Test resolve lock over retry limit @@ -727,7 +870,7 @@ mod tests { ) .unwrap(); let key = vec![100]; - resolve_lock_with_retry(&key, 3, 4, false, client, keyspace, backoff) + resolve_lock_with_retry((&key).into(), 3, 4, false, client, keyspace, backoff) .await .expect_err("should return error"); } @@ -774,6 +917,482 @@ mod tests { assert_eq!(resolve_lock_count.load(Ordering::SeqCst), 1); } + #[rstest::rstest] + // With `Keyspace::Enable` every key gains the keyspace prefix, which places them + // all in mock region 2 — the recovery then resolves one region instead of two. + #[case(Keyspace::Disable, 2, &[PRIMARY_KEY])] + #[case(Keyspace::Enable { keyspace_id: 0 }, 1, &[PRIMARY_KEY])] + #[case(Keyspace::Disable, 2, &[SECONDARY_KEY, PRIMARY_KEY])] + #[case(Keyspace::Enable { keyspace_id: 0 }, 1, &[SECONDARY_KEY, PRIMARY_KEY])] + #[tokio::test] + #[serial] + async fn test_resolve_locks_recovers_expired_async_commit( + #[case] keyspace: Keyspace, + #[case] expected_resolved_regions: usize, + #[case] encountered_keys: &[&[u8]], + ) { + let primary_key = encoded(PRIMARY_KEY, keyspace); + let secondary_key = encoded(SECONDARY_KEY, keyspace); + + let check_secondary_count = Arc::new(AtomicUsize::new(0)); + let resolve_lock_count = Arc::new(AtomicUsize::new(0)); + let resolved_commit_version = Arc::new(AtomicU64::new(0)); + + let check_secondary_count_captured = check_secondary_count.clone(); + let resolve_lock_count_captured = resolve_lock_count.clone(); + let resolved_commit_version_captured = resolved_commit_version.clone(); + let primary_key_captured = primary_key.clone(); + let secondary_key_captured = secondary_key.clone(); + let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + move |req: &dyn Any| { + if let Some(req) = req.downcast_ref::() { + assert_eq!(req.primary_key, primary_key_captured); + return Ok(Box::new(still_locked_response(async_commit_primary_lock( + &primary_key_captured, + 44, + vec![secondary_key_captured.clone()], + ))) as Box); + } + if let Some(req) = req.downcast_ref::() { + check_secondary_count_captured.fetch_add(1, Ordering::SeqCst); + // The secondaries listed in the primary lock are already encoded and + // must be sent verbatim — re-encoding them would corrupt the keys. + assert_eq!(req.keys, vec![secondary_key_captured.clone()]); + let resp = kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![async_commit_secondary_lock(&secondary_key_captured, 43)], + ..Default::default() + }; + return Ok(Box::new(resp) as Box); + } + if let Some(req) = req.downcast_ref::() { + resolve_lock_count_captured.fetch_add(1, Ordering::SeqCst); + resolved_commit_version_captured.store(req.commit_version, Ordering::SeqCst); + return Ok(Box::::default() as Box); + } + panic!("unexpected request type: {:?}", req.type_id()); + }, + ))); + + let locks = encountered_keys + .iter() + .map(|key| kvrpcpb::LockInfo { + key: encoded(key, keyspace), + primary_lock: primary_key.clone(), + lock_version: EXPIRED_TXN_VERSION, + lock_ttl: 1, + use_async_commit: true, + ..Default::default() + }) + .collect(); + + let live_locks = resolve_locks(locks, Timestamp::default(), client, keyspace) + .await + .unwrap(); + + assert!(live_locks.is_empty()); + assert_eq!(check_secondary_count.load(Ordering::SeqCst), 1); + // Every region the transaction wrote to must be resolved, exactly once each. + assert_eq!( + resolve_lock_count.load(Ordering::SeqCst), + expected_resolved_regions + ); + // The secondary's min_commit_ts (43) is below the primary's (44): the recovered + // commit version must be the maximum across ALL locks, i.e. the primary's. + assert_eq!(resolved_commit_version.load(Ordering::SeqCst), 44); + } + + #[tokio::test] + #[serial] + async fn test_resolve_locks_recovers_missing_async_secondary_as_rollback() { + let check_secondary_count = Arc::new(AtomicUsize::new(0)); + let resolve_lock_count = Arc::new(AtomicUsize::new(0)); + let resolved_commit_version = Arc::new(AtomicU64::new(u64::MAX)); + + let check_secondary_count_captured = check_secondary_count.clone(); + let resolve_lock_count_captured = resolve_lock_count.clone(); + let resolved_commit_version_captured = resolved_commit_version.clone(); + let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + move |req: &dyn Any| { + if req.is::() { + return Ok(Box::new(still_locked_response(async_commit_primary_lock( + PRIMARY_KEY, + 42, + vec![SECONDARY_KEY.to_vec()], + ))) as Box); + } + if req.is::() { + check_secondary_count_captured.fetch_add(1, Ordering::SeqCst); + // The requested secondary is absent and commit_ts is zero: TiKV has + // established a rollback tombstone, so the transaction must roll back. + return Ok( + Box::::default() as Box + ); + } + if let Some(req) = req.downcast_ref::() { + resolve_lock_count_captured.fetch_add(1, Ordering::SeqCst); + resolved_commit_version_captured.store(req.commit_version, Ordering::SeqCst); + return Ok(Box::::default() as Box); + } + panic!("unexpected request type: {:?}", req.type_id()); + }, + ))); + + let lock = async_commit_primary_lock(PRIMARY_KEY, 42, vec![SECONDARY_KEY.to_vec()]); + + let live_locks = resolve_locks(vec![lock], Timestamp::default(), client, Keyspace::Disable) + .await + .unwrap(); + + assert!(live_locks.is_empty()); + assert_eq!(check_secondary_count.load(Ordering::SeqCst), 1); + // The rollback must reach every region the transaction wrote to: the + // secondary's region and the primary's. + assert_eq!(resolve_lock_count.load(Ordering::SeqCst), 2); + assert_eq!(resolved_commit_version.load(Ordering::SeqCst), 0); + } + + #[rstest::rstest] + #[case(55, 0)] + #[case(0, 0)] + #[case(0, 1000)] + #[tokio::test] + #[serial] + async fn test_resolve_locks_falls_back_to_2pc_with_real_current_ts( + #[case] commit_version: u64, + #[case] lock_ttl: u64, + ) { + let check_txn_status_count = Arc::new(AtomicUsize::new(0)); + let force_sync_count = Arc::new(AtomicUsize::new(0)); + let resolve_lock_count = Arc::new(AtomicUsize::new(0)); + let resolved_commit_version = Arc::new(AtomicU64::new(0)); + + let check_txn_status_count_captured = check_txn_status_count.clone(); + let force_sync_count_captured = force_sync_count.clone(); + let resolve_lock_count_captured = resolve_lock_count.clone(); + let resolved_commit_version_captured = resolved_commit_version.clone(); + let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + move |req: &dyn Any| { + if let Some(req) = req.downcast_ref::() { + check_txn_status_count_captured.fetch_add(1, Ordering::SeqCst); + if req.force_sync_commit { + force_sync_count_captured.fetch_add(1, Ordering::SeqCst); + assert_ne!(req.current_ts, u64::MAX); + return Ok(Box::new(kvrpcpb::CheckTxnStatusResponse { + commit_version, + lock_ttl, + lock_info: (lock_ttl != 0).then(|| kvrpcpb::LockInfo { + key: PRIMARY_KEY.to_vec(), + primary_lock: PRIMARY_KEY.to_vec(), + lock_version: EXPIRED_TXN_VERSION, + lock_ttl, + ..Default::default() + }), + action: kvrpcpb::Action::NoAction as i32, + ..Default::default() + }) as Box); + } + return Ok(Box::new(still_locked_response(async_commit_primary_lock( + PRIMARY_KEY, + 42, + vec![SECONDARY_KEY.to_vec()], + ))) as Box); + } + if req.is::() { + return Ok(Box::new(kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![fallback_2pc_secondary_lock(SECONDARY_KEY)], + ..Default::default() + }) as Box); + } + if let Some(req) = req.downcast_ref::() { + resolve_lock_count_captured.fetch_add(1, Ordering::SeqCst); + resolved_commit_version_captured.store(req.commit_version, Ordering::SeqCst); + return Ok(Box::::default() as Box); + } + panic!("unexpected request type: {:?}", req.type_id()); + }, + ))); + + let lock = async_commit_primary_lock(PRIMARY_KEY, 42, vec![SECONDARY_KEY.to_vec()]); + + let live_locks = resolve_locks(vec![lock], Timestamp::default(), client, Keyspace::Disable) + .await + .unwrap(); + + assert_eq!(check_txn_status_count.load(Ordering::SeqCst), 2); + assert_eq!(force_sync_count.load(Ordering::SeqCst), 1); + if lock_ttl == 0 { + assert!(live_locks.is_empty()); + assert_eq!(resolve_lock_count.load(Ordering::SeqCst), 2); + assert_eq!( + resolved_commit_version.load(Ordering::SeqCst), + commit_version + ); + } else { + assert_eq!(resolve_lock_count.load(Ordering::SeqCst), 0); + assert_eq!(live_locks.len(), 1); + assert_eq!(live_locks[0].lock_ttl, lock_ttl); + assert!(!live_locks[0].use_async_commit); + } + } + + /// A region error on `CheckSecondaryLocks` makes the plan re-shard against fresh + /// region boundaries. Every retried sub-request must be paired with its OWN keys + /// (`preserve_shard` + `Collect::merge` contract): pairing a stale shard would make + /// the shorter per-region lock lists below look like missing locks and roll back a + /// committable transaction. + #[tokio::test] + #[serial] + async fn test_check_secondary_locks_reshards_on_region_error() { + // The transaction spans primary [1] and secondaries [2] and [3]. Every key + // starts out in one region; after the simulated split, [2] and [3] live in two. + fn mock_region(id: u64, start_key: Vec, end_key: Vec) -> RegionWithLeader { + let mut region = RegionWithLeader::default(); + region.region.id = id; + region.region.start_key = start_key; + region.region.end_key = end_key; + region.region.region_epoch = Some(metapb::RegionEpoch { + conf_ver: 0, + version: 1, + }); + region.leader = Some(metapb::Peer { + store_id: 41, + ..Default::default() + }); + region + } + + let split = Arc::new(AtomicBool::new(false)); + let check_secondary_count = Arc::new(AtomicUsize::new(0)); + let resolve_lock_count = Arc::new(AtomicUsize::new(0)); + let resolved_commit_version = Arc::new(AtomicU64::new(u64::MAX)); + + let split_in_dispatch = split.clone(); + let check_secondary_count_captured = check_secondary_count.clone(); + let resolve_lock_count_captured = resolve_lock_count.clone(); + let resolved_commit_version_captured = resolved_commit_version.clone(); + let split_in_region_hook = split.clone(); + let client = Arc::new( + MockPdClient::new(MockKvClient::with_dispatch_hook(move |req: &dyn Any| { + if req.is::() { + return Ok(Box::new(still_locked_response(async_commit_primary_lock( + &[1], + 40, + vec![vec![2], vec![3]], + ))) as Box); + } + if let Some(req) = req.downcast_ref::() { + check_secondary_count_captured.fetch_add(1, Ordering::SeqCst); + if req.keys == vec![vec![2], vec![3]] { + // First attempt, before the split: fail the whole shard. + split_in_dispatch.store(true, Ordering::SeqCst); + let resp = kvrpcpb::CheckSecondaryLocksResponse { + region_error: Some(errorpb::Error::default()), + ..Default::default() + }; + return Ok(Box::new(resp) as Box); + } + let min_commit_ts = match req.keys.as_slice() { + [key] if key.as_slice() == [2] => 41, + [key] if key.as_slice() == [3] => 42, + keys => panic!("unexpected CheckSecondaryLocks shard: {:?}", keys), + }; + let resp = kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![async_commit_secondary_lock(&req.keys[0], min_commit_ts)], + ..Default::default() + }; + return Ok(Box::new(resp) as Box); + } + if let Some(req) = req.downcast_ref::() { + resolve_lock_count_captured.fetch_add(1, Ordering::SeqCst); + resolved_commit_version_captured.store(req.commit_version, Ordering::SeqCst); + return Ok(Box::::default() as Box); + } + panic!("unexpected request type: {:?}", req.type_id()); + })) + .with_region_for_key_hook(move |key: &Key| { + let key: &[u8] = key.into(); + if !split_in_region_hook.load(Ordering::SeqCst) { + Ok(mock_region(1, vec![], vec![10])) + } else if key < &[3][..] { + Ok(mock_region(1, vec![], vec![3])) + } else { + Ok(mock_region(4, vec![3], vec![10])) + } + }), + ); + + let lock = async_commit_primary_lock(&[1], 40, vec![vec![2], vec![3]]); + + let live_locks = resolve_locks(vec![lock], Timestamp::default(), client, Keyspace::Disable) + .await + .unwrap(); + + assert!(live_locks.is_empty()); + // One failed pre-split request plus one per post-split region. + assert_eq!(check_secondary_count.load(Ordering::SeqCst), 3); + // Both locks survived, so the transaction must COMMIT at the maximum + // min_commit_ts (42) — a stale shard pairing would have inferred a missing + // lock instead and rolled the transaction back (commit version 0). + assert_eq!(resolved_commit_version.load(Ordering::SeqCst), 42); + assert_eq!(resolve_lock_count.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + #[serial] + async fn test_resolve_locks_retries_expired_missing_primary_with_rollback() { + let check_txn_status_count = Arc::new(AtomicUsize::new(0)); + let resolve_lock_count = Arc::new(AtomicUsize::new(0)); + + let check_txn_status_count_captured = check_txn_status_count.clone(); + let resolve_lock_count_captured = resolve_lock_count.clone(); + let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + move |req: &dyn Any| { + if let Some(req) = req.downcast_ref::() { + check_txn_status_count_captured.fetch_add(1, Ordering::SeqCst); + if !req.rollback_if_not_exist { + return Ok(Box::new(kvrpcpb::CheckTxnStatusResponse { + error: Some(kvrpcpb::KeyError { + txn_not_found: Some(kvrpcpb::TxnNotFound { + start_ts: req.lock_ts, + primary_key: req.primary_key.clone(), + }), + ..Default::default() + }), + ..Default::default() + }) as Box); + } + return Ok(Box::::default() as Box); + } + if req.is::() { + resolve_lock_count_captured.fetch_add(1, Ordering::SeqCst); + return Ok(Box::::default() as Box); + } + panic!("unexpected request type: {:?}", req.type_id()); + }, + ))); + + let lock = kvrpcpb::LockInfo { + key: vec![1], + primary_lock: vec![2], + lock_version: EXPIRED_TXN_VERSION, + lock_ttl: 1, + ..Default::default() + }; + + let live_locks = resolve_locks(vec![lock], Timestamp::default(), client, Keyspace::Disable) + .await + .unwrap(); + + assert!(live_locks.is_empty()); + assert_eq!(check_txn_status_count.load(Ordering::SeqCst), 2); + assert_eq!(resolve_lock_count.load(Ordering::SeqCst), 1); + } + + /// A pessimistic lock whose recorded primary is stale (the transaction changed its + /// primary, pingcap/tidb#42937) makes TiKV answer CheckTxnStatus with PrimaryMismatch. + /// Only that stale lock may be rolled back — the transaction itself may still be alive, + /// so a region-wide ResolveLock is out of the question: each stale lock gets its own + /// single-key PessimisticRollback, and a second stale lock of the same transaction in + /// the same region must NOT be skipped or swept along. + #[tokio::test] + #[serial] + async fn test_resolve_locks_rolls_back_stale_pessimistic_lock_on_primary_mismatch() { + let rolled_back_keys = Arc::new(std::sync::Mutex::new(Vec::new())); + + let rolled_back_keys_captured = rolled_back_keys.clone(); + let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + move |req: &dyn Any| { + if req.is::() { + return Ok(Box::new(kvrpcpb::CheckTxnStatusResponse { + error: Some(kvrpcpb::KeyError { + primary_mismatch: Some(kvrpcpb::PrimaryMismatch { + lock_info: Some(kvrpcpb::LockInfo::default()), + }), + ..Default::default() + }), + ..Default::default() + }) as Box); + } + if let Some(req) = req.downcast_ref::() { + assert_eq!(req.start_version, EXPIRED_TXN_VERSION); + assert_eq!(req.for_update_ts, u64::MAX); + assert_eq!(req.keys.len(), 1, "rollback must target a single key"); + rolled_back_keys_captured + .lock() + .unwrap() + .push(req.keys[0].clone()); + return Ok( + Box::::default() as Box + ); + } + if req.is::() { + panic!("must not sweep a region for a stale pessimistic lock"); + } + panic!("unexpected request type: {:?}", req.type_id()); + }, + ))); + + let make_lock = |key: u8| kvrpcpb::LockInfo { + key: vec![key], + primary_lock: vec![9], + lock_version: EXPIRED_TXN_VERSION, + lock_ttl: 1, + lock_type: kvrpcpb::Op::PessimisticLock as i32, + ..Default::default() + }; + // Two stale pessimistic locks of the same transaction in the same mock region. + let locks = vec![make_lock(1), make_lock(2)]; + + let live_locks = resolve_locks(locks, Timestamp::default(), client, Keyspace::Disable) + .await + .unwrap(); + + assert!(live_locks.is_empty()); + assert_eq!( + *rolled_back_keys.lock().unwrap(), + vec![vec![1], vec![2]], + "each stale lock must be rolled back individually" + ); + } + + /// A PrimaryMismatch for a non-pessimistic lock is unexpected (client-go treats it as an + /// error too) and must propagate instead of rolling anything back. + #[tokio::test] + #[serial] + async fn test_resolve_locks_propagates_primary_mismatch_for_non_pessimistic_lock() { + let client = Arc::new(MockPdClient::new(MockKvClient::with_dispatch_hook( + move |req: &dyn Any| { + if req.is::() { + return Ok(Box::new(kvrpcpb::CheckTxnStatusResponse { + error: Some(kvrpcpb::KeyError { + primary_mismatch: Some(kvrpcpb::PrimaryMismatch { + lock_info: Some(kvrpcpb::LockInfo::default()), + }), + ..Default::default() + }), + ..Default::default() + }) as Box); + } + if req.is::() { + panic!("must not resolve a lock on unexpected primary mismatch"); + } + panic!("unexpected request type: {:?}", req.type_id()); + }, + ))); + + let lock = kvrpcpb::LockInfo { + key: PRIMARY_KEY.to_vec(), + primary_lock: vec![2], + lock_version: EXPIRED_TXN_VERSION, + lock_ttl: 1, + ..Default::default() + }; + + let result = + resolve_locks(vec![lock], Timestamp::default(), client, Keyspace::Disable).await; + assert!(matches!(result, Err(Error::KeyError(_)))); + } + #[test] fn format_key_for_log_hex_encodes_the_prefix() { assert_eq!(format_key_for_log(b"hello"), "len=5, prefix=68656C6C6F"); diff --git a/src/transaction/requests.rs b/src/transaction/requests.rs index db524f16..765fb0a0 100644 --- a/src/transaction/requests.rs +++ b/src/transaction/requests.rs @@ -1,6 +1,7 @@ // Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0. use std::cmp; +use std::collections::HashSet; use std::iter; use std::sync::Arc; @@ -43,9 +44,11 @@ use crate::store::Request; use crate::store::Store; use crate::store::{region_stream_for_keys, region_stream_for_range}; use crate::timestamp::TimestampExt; +use crate::transaction::lock::format_key_for_log; use crate::transaction::requests::kvrpcpb::prewrite_request::PessimisticAction; use crate::transaction::HasLocks; use crate::util::iter::FlatMapOkIterExt; +use crate::Error; use crate::KvPair; use crate::Result; use crate::Value; @@ -725,26 +728,14 @@ impl TransactionStatus { } } - // is_cacheable checks whether the transaction status is certain. - // If transaction is already committed, the result could be cached. - // Otherwise: - // If l.LockType is pessimistic lock type: - // - if its primary lock is pessimistic too, the check txn status result should not be cached. - // - if its primary lock is prewrite lock type, the check txn status could be cached. - // If l.lockType is prewrite lock type: - // - always cache the check txn status result. - // For prewrite locks, their primary keys should ALWAYS be the correct one and will NOT change. + // Only final states are cacheable. A Locked result is not final even when its TTL expired: + // async-commit recovery still has to inspect every secondary, and force-sync fallback must be + // able to issue a fresh CheckTxnStatus request. pub fn is_cacheable(&self) -> bool { - match &self.kind { - TransactionStatusKind::RolledBack | TransactionStatusKind::Committed(..) => true, - TransactionStatusKind::Locked(..) if self.is_expired => matches!( - self.action, - kvrpcpb::Action::NoAction - | kvrpcpb::Action::LockNotExistRollback - | kvrpcpb::Action::TtlExpireRollback - ), - _ => false, - } + matches!( + self.kind, + TransactionStatusKind::RolledBack | TransactionStatusKind::Committed(..) + ) } } @@ -775,47 +766,124 @@ impl KvRequest for kvrpcpb::CheckSecondaryLocksRequest { shardable_keys!(kvrpcpb::CheckSecondaryLocksRequest); -impl Merge for Collect { +/// Merge the per-region responses of a sharded `CheckSecondaryLocks` request. Each shard +/// (`Vec>`) is the list of secondary keys sent to one region; pairing every response +/// with its own key list is what makes missing-lock detection possible — TiKV only returns +/// the locks it found, never the keys that no longer hold one. +impl Merge>>> for Collect { type Out = SecondaryLocksStatus; - fn merge(&self, input: Vec>) -> Result { + fn merge( + &self, + input: Vec>>>>, + ) -> Result { let mut out = SecondaryLocksStatus { commit_ts: None, min_commit_ts: 0, fallback_2pc: false, }; + for resp in input { - let resp = resp?; - for lock in resp.locks.into_iter() { + let ResponseWithShard(resp, requested_keys) = resp?; + if resp.locks.len() > requested_keys.len() { + return Err(Error::ProtocolViolation { + message: format!( + "CheckSecondaryLocks returned {} locks for {} requested keys", + resp.locks.len(), + requested_keys.len() + ), + }); + } + + // TiKV checks the requested keys one by one and stops at the first key that no + // longer holds a lock of this transaction, making the transaction's fate durable + // on the way: unless that key is already committed, a protected rollback is + // written for it. The decision is reported through `commit_ts` — the commit TS, + // or zero for a rollback — and the returned locks then no longer cover every + // requested key. A short lock list therefore means the transaction is decided. + let response_missing_lock = resp.locks.len() < requested_keys.len(); + if !response_missing_lock && resp.commit_ts != 0 { + return Err(Error::ProtocolViolation { + message: format!( + "CheckSecondaryLocks returned commit TS {} although every requested lock is still present", + resp.commit_ts + ), + }); + } + + let mut remaining: HashSet<&Vec> = requested_keys.iter().collect(); + for lock in &resp.locks { + if !remaining.remove(&lock.key) { + return Err(Error::ProtocolViolation { + message: format!( + "CheckSecondaryLocks returned an unrequested or duplicate lock ({})", + format_key_for_log(&lock.key) + ), + }); + } if !lock.use_async_commit { out.fallback_2pc = true; - return Ok(out); } out.min_commit_ts = cmp::max(out.min_commit_ts, lock.min_commit_ts); } - out.commit_ts = match ( - out.commit_ts.take(), - Timestamp::try_from_version(resp.commit_ts), - ) { - (Some(a), Some(b)) => { - assert_eq!(a, b); - Some(a) + + if response_missing_lock { + if out.commit_ts.is_some_and(|ts| ts != resp.commit_ts) { + return Err(Error::ProtocolViolation { + message: format!( + "CheckSecondaryLocks reported conflicting commit TS ({:?} and {}) for one transaction", + out.commit_ts, + resp.commit_ts + ), + }); } - (Some(a), None) => Some(a), - (None, Some(b)) => Some(b), - (None, None) => None, - }; + out.commit_ts = Some(resp.commit_ts); + } } + Ok(out) } } +/// The aggregated outcome of `CheckSecondaryLocks` over the secondary keys of an +/// async-commit transaction. pub struct SecondaryLocksStatus { - pub commit_ts: Option, + /// A missing lock's durable decision: `Some(0)` for rollback, `Some(ts)` for commit. + /// `None` means every requested lock is still present. + pub commit_ts: Option, + /// The maximum `min_commit_ts` across the locks that are still alive. pub min_commit_ts: u64, + /// True when a surviving lock fell back from async commit to 2PC: the transaction's + /// fate then belongs to its primary lock, not to the secondaries. pub fallback_2pc: bool, } +impl SecondaryLocksStatus { + /// The version this transaction must be resolved with: a positive commit version to + /// commit, or zero to roll back — the same encoding `TxnInfo.status` uses on the wire. + /// + /// While every lock is still alive the transaction is committable, and the commit + /// version is the maximum `min_commit_ts` across the primary and all secondary locks — + /// exactly the value the transaction's own committer would compute. Once a lock is + /// missing, TiKV has already made the decision durable and `commit_ts` carries it. + /// + /// Returns an error when TiKV reports a commit TS below a surviving lock's + /// `min_commit_ts`: every lock promised its readers no commit below that point. + pub fn resolved_commit_version(&self, primary_min_commit_ts: u64) -> Result { + let min_commit_ts = cmp::max(primary_min_commit_ts, self.min_commit_ts); + let commit_version = self.commit_ts.unwrap_or(min_commit_ts); + if commit_version != 0 && commit_version < min_commit_ts { + return Err(Error::ProtocolViolation { + message: format!( + "CheckSecondaryLocks reported commit TS {} below a surviving lock's min_commit_ts {}", + commit_version, min_commit_ts + ), + }); + } + Ok(commit_version) + } +} + pair_locks!(kvrpcpb::BatchGetResponse); pair_locks!(kvrpcpb::ScanResponse); error_locks!(kvrpcpb::GetResponse); @@ -891,14 +959,251 @@ impl Merge for Collect { #[cfg(test)] mod tests { + use crate::common::Error; use crate::common::Error::PessimisticLockError; use crate::common::Error::ResolveLockError; use crate::proto::kvrpcpb; + use crate::proto::pdpb::Timestamp; use crate::request::plan::Merge; + use crate::request::Collect; use crate::request::CollectWithShard; use crate::request::ResponseWithShard; + use crate::timestamp::TimestampExt; use crate::KvPair; + use super::TransactionStatus; + use super::TransactionStatusKind; + + /// A still-live async-commit lock, as returned inside `CheckSecondaryLocksResponse`. + fn async_commit_lock(key: &[u8], min_commit_ts: u64) -> kvrpcpb::LockInfo { + kvrpcpb::LockInfo { + key: key.to_vec(), + use_async_commit: true, + min_commit_ts, + ..Default::default() + } + } + + #[rstest::rstest] + #[case(7, 8)] + #[case(0, 8)] + #[case(7, 0)] + fn check_secondary_conflicting_commit_ts_is_a_protocol_violation( + #[case] first: u64, + #[case] second: u64, + ) { + let result = Collect.merge(vec![ + Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + commit_ts: first, + ..Default::default() + }, + vec![b"a".to_vec()], + )), + Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + commit_ts: second, + ..Default::default() + }, + vec![b"b".to_vec()], + )), + ]); + + assert!(matches!(result, Err(Error::ProtocolViolation { .. }))); + } + + #[test] + fn check_secondary_rejects_more_locks_than_requested_keys() { + let result = Collect.merge(vec![Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![async_commit_lock(b"a", 1), async_commit_lock(b"b", 2)], + ..Default::default() + }, + vec![b"a".to_vec()], + ))]); + + assert!(matches!(result, Err(Error::ProtocolViolation { .. }))); + } + + #[test] + fn check_secondary_rejects_a_lock_that_was_not_requested() { + let result = Collect.merge(vec![Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![async_commit_lock(b"b", 1)], + ..Default::default() + }, + vec![b"a".to_vec()], + ))]); + + assert!(matches!(result, Err(Error::ProtocolViolation { .. }))); + } + + #[test] + fn check_secondary_rejects_a_duplicate_lock_key() { + let result = Collect.merge(vec![Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![async_commit_lock(b"a", 1), async_commit_lock(b"a", 2)], + ..Default::default() + }, + vec![b"a".to_vec(), b"b".to_vec()], + ))]); + + assert!(matches!(result, Err(Error::ProtocolViolation { .. }))); + } + + #[test] + fn check_secondary_all_locks_present_uses_max_min_commit_ts() { + let result = Collect + .merge(vec![Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![async_commit_lock(b"secondary", 70)], + ..Default::default() + }, + vec![b"secondary".to_vec()], + ))]) + .unwrap(); + + assert_eq!(result.commit_ts, None); + assert_eq!(result.resolved_commit_version(80).unwrap(), 80); + } + + #[test] + fn check_secondary_missing_lock_preserves_exact_commit_ts() { + let result = Collect + .merge(vec![Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + commit_ts: 77, + ..Default::default() + }, + vec![b"missing".to_vec()], + ))]) + .unwrap(); + + assert_eq!(result.commit_ts, Some(77)); + assert_eq!(result.resolved_commit_version(70).unwrap(), 77); + } + + #[test] + fn check_secondary_missing_lock_with_zero_commit_ts_resolves_as_rollback() { + let result = Collect + .merge(vec![Ok(ResponseWithShard( + // No lock and no commit TS: TiKV wrote a protected rollback for the key. + kvrpcpb::CheckSecondaryLocksResponse::default(), + vec![b"missing".to_vec()], + ))]) + .unwrap(); + + assert_eq!(result.commit_ts, Some(0)); + assert_eq!(result.resolved_commit_version(80).unwrap(), 0); + } + + #[rstest::rstest] + #[case(0)] + #[case(77)] + fn check_secondary_missing_lock_keeps_decision_across_live_shards(#[case] commit_ts: u64) { + let missing = ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + commit_ts, + ..Default::default() + }, + vec![b"missing".to_vec()], + ); + let live = ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![async_commit_lock(b"live", 70)], + ..Default::default() + }, + vec![b"live".to_vec()], + ); + for responses in [ + vec![Ok(missing.clone()), Ok(live.clone())], + vec![Ok(live), Ok(missing)], + ] { + let result = Collect.merge(responses).unwrap(); + assert_eq!(result.resolved_commit_version(75).unwrap(), commit_ts); + } + } + + #[test] + fn check_secondary_rejects_commit_ts_below_primary_min_commit_ts() { + let result = Collect + .merge(vec![Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + commit_ts: 77, + ..Default::default() + }, + vec![b"missing".to_vec()], + ))]) + .unwrap(); + + assert!(matches!( + result.resolved_commit_version(80), + Err(Error::ProtocolViolation { .. }) + )); + } + + #[test] + fn check_secondary_rejects_commit_ts_below_locked_min_commit_ts() { + let result = Collect + .merge(vec![ + Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + locks: vec![async_commit_lock(b"locked", 80)], + ..Default::default() + }, + vec![b"locked".to_vec()], + )), + Ok(ResponseWithShard( + kvrpcpb::CheckSecondaryLocksResponse { + commit_ts: 77, + ..Default::default() + }, + vec![b"missing".to_vec()], + )), + ]) + .unwrap(); + + // The merge only aggregates; the min_commit_ts gate lives in + // `resolved_commit_version`, which every commit-path caller goes through. + assert!(matches!( + result.resolved_commit_version(0), + Err(Error::ProtocolViolation { .. }) + )); + } + + #[test] + fn only_final_transaction_statuses_are_cacheable() { + let committed = TransactionStatus { + kind: TransactionStatusKind::Committed(Timestamp::from_version(5)), + action: kvrpcpb::Action::NoAction, + is_expired: false, + }; + assert!(committed.is_cacheable()); + + let rolled_back = TransactionStatus { + kind: TransactionStatusKind::RolledBack, + action: kvrpcpb::Action::NoAction, + is_expired: false, + }; + assert!(rolled_back.is_cacheable()); + + // A `Locked` status is a snapshot, never a fact — not even once the TTL has + // expired: async-commit recovery must inspect the secondaries afresh, and the + // force-sync fallback must be able to issue a new CheckTxnStatus request. + let expired_async_commit_lock = TransactionStatus { + kind: TransactionStatusKind::Locked( + 1, + kvrpcpb::LockInfo { + use_async_commit: true, + ..Default::default() + }, + ), + action: kvrpcpb::Action::NoAction, + is_expired: true, + }; + assert!(!expired_async_commit_lock.is_cacheable()); + } + #[tokio::test] async fn test_merge_pessimistic_lock_response() { let (key1, key2, key3, key4) = (b"key1", b"key2", b"key3", b"key4"); diff --git a/tests/failpoint_tests.rs b/tests/failpoint_tests.rs index 550b3d8a..4c876866 100644 --- a/tests/failpoint_tests.rs +++ b/tests/failpoint_tests.rs @@ -20,6 +20,8 @@ use tikv_client::CheckLevel; use tikv_client::Config; use tikv_client::Result; use tikv_client::RetryOptions; +use tikv_client::Timestamp; +use tikv_client::TimestampExt; use tikv_client::TransactionClient; use tikv_client::TransactionOptions; @@ -133,8 +135,12 @@ async fn txn_cleanup_locks_batch_size() -> Result<()> { let client = TransactionClient::new_with_config(pd_addrs(), Config::default().with_default_keyspace()) .await?; + let write_start_ts = client.current_timestamp().await?; let keys = write_data(&client, true, true).await?; - assert_eq!(count_locks(&client).await?, keys.len()); + assert_eq!( + count_locks_since(&client, &write_start_ts).await?, + keys.len() + ); let safepoint = client.current_timestamp().await?; let options = ResolveLocksOptions { @@ -148,6 +154,14 @@ async fn txn_cleanup_locks_batch_size() -> Result<()> { assert_eq!(res.resolved_locks, keys.len()); assert_eq!(count_locks(&client).await?, keys.len()); + // Clean up the locks this case deliberately left behind, so later cases start + // from a clean cluster instead of tripping over them (#509). + fail::cfg("before-cleanup-locks", "off").unwrap(); + client + .cleanup_locks(full_range, &safepoint, Default::default()) + .await?; + assert_eq!(count_locks(&client).await?, 0); + scenario.teardown(); Ok(()) } @@ -272,8 +286,12 @@ async fn txn_cleanup_range_async_commit_locks() -> Result<()> { let client = TransactionClient::new_with_config(pd_addrs(), Config::default().with_default_keyspace()) .await?; + let write_start_ts = client.current_timestamp().await?; let keys = write_data(&client, true, true).await?; - assert_eq!(count_locks(&client).await?, keys.len()); + assert_eq!( + count_locks_since(&client, &write_start_ts).await?, + keys.len() + ); info!("total keys' count {}", keys.len()); let mut sorted_keys: Vec> = Vec::from_iter(keys.clone()); @@ -350,6 +368,62 @@ async fn txn_resolve_locks() -> Result<()> { Ok(()) } +// Regression test for #528: an async-commit transaction can finish prewrite and +// disappear before its secondary commit task starts. Once the lock TTL expires, +// an ordinary read must recover the transaction from all secondary locks instead +// of returning the same ResolveLockError forever. +#[tokio::test] +#[serial] +async fn txn_read_recovers_expired_async_commit_locks() -> Result<()> { + init().await?; + let scenario = FailScenario::setup(); + + fail::cfg("after-prewrite", "return").unwrap(); + defer! {{ + fail::cfg("after-prewrite", "off").unwrap(); + }} + + let client = + TransactionClient::new_with_config(pd_addrs(), Config::default().with_default_keyspace()) + .await?; + // `init` deliberately splits the u32 keyspace when MULTI_REGION is set, so + // these endpoints exercise the sharded CheckSecondaryLocks collector in CI. + let primary = 1u32.to_be_bytes().to_vec(); + let secondary = (u32::MAX - 1).to_be_bytes().to_vec(); + let primary_value = b"recovered-primary-value".to_vec(); + let secondary_value = b"recovered-secondary-value".to_vec(); + + let mut txn = client + .begin_with_options( + TransactionOptions::new_optimistic() + .use_async_commit() + .heartbeat_option(HeartbeatOption::NoHeartbeat) + .drop_check(CheckLevel::Warn), + ) + .await?; + txn.put(primary.clone(), primary_value.clone()).await?; + txn.put(secondary.clone(), secondary_value.clone()).await?; + assert!(txn.commit().await.is_err()); + fail::cfg("after-prewrite", "off").unwrap(); + + let ts = client.current_timestamp().await?; + assert!(!client.scan_locks(&ts, .., 1024).await?.is_empty()); + + // The default prewrite lock TTL is three seconds. + tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; + + let mut reader = client.begin_optimistic().await?; + assert_eq!(reader.get(primary).await?, Some(primary_value)); + assert_eq!(reader.get(secondary).await?, Some(secondary_value)); + reader.rollback().await?; + + let ts = client.current_timestamp().await?; + assert_eq!(client.scan_locks(&ts, .., 1024).await?.len(), 0); + + scenario.teardown(); + Ok(()) +} + // Regression test for #545: a pessimistic transaction whose commit fails after // prewrite has placed its 2PC lock must have that lock cleared by `rollback()`. // Previously the terminal pessimistic rollback sent `PessimisticRollback`, which @@ -557,6 +631,24 @@ async fn count_locks(client: &TransactionClient) -> Result { count_locks_in_range(client, b"", b"").await } +/// Count locks created at or after `since`, de-duplicated. +/// +/// Unlike `count_locks`, this ignores locks left behind by earlier test cases — on a +/// churning cluster (hundreds of tiny regions splitting and merging) such locks can +/// disappear from one scan and re-appear in a later one, which made the bare +/// `count_locks == keys.len()` assertion flaky (#509). +async fn count_locks_since(client: &TransactionClient, since: &Timestamp) -> Result { + let ts = client.current_timestamp().await.unwrap(); + let locks = client.scan_locks(&ts, .., 65536).await?; + let locks_set: HashSet> = HashSet::from_iter( + locks + .into_iter() + .filter(|l| l.lock_version >= since.version()) + .map(|l| l.key), + ); + Ok(locks_set.len()) +} + async fn count_locks_in_range( client: &TransactionClient, start_key: &[u8],