diff --git a/src/storage/mvcc/mod.rs b/src/storage/mvcc/mod.rs index 13c6d6b3389..8c723b986ee 100644 --- a/src/storage/mvcc/mod.rs +++ b/src/storage/mvcc/mod.rs @@ -271,7 +271,7 @@ pub mod tests { Ok(()) } - fn must_prewrite_put_impl( + pub fn must_prewrite_put_impl( engine: &E, key: &[u8], value: &[u8], @@ -336,7 +336,7 @@ pub mod tests { ts: u64, for_update_ts: u64, is_pessimistic_lock: bool, - ) { + ) -> Error { let ctx = Context::default(); let snapshot = engine.snapshot(&ctx).unwrap(); let mut txn = MvccTxn::new(snapshot, ts, true).unwrap(); @@ -344,10 +344,10 @@ pub mod tests { options.for_update_ts = for_update_ts; let mutation = Mutation::Put((Key::from_raw(key), value.to_vec())); if for_update_ts == 0 { - txn.prewrite(mutation, pk, &options).unwrap_err(); + txn.prewrite(mutation, pk, &options).unwrap_err() } else { txn.pessimistic_prewrite(mutation, pk, is_pessimistic_lock, &options) - .unwrap_err(); + .unwrap_err() } } @@ -357,8 +357,8 @@ pub mod tests { value: &[u8], pk: &[u8], ts: u64, - ) { - must_prewrite_put_err_impl(engine, key, value, pk, ts, 0, false); + ) -> Error { + must_prewrite_put_err_impl(engine, key, value, pk, ts, 0, false) } pub fn must_pessimistic_prewrite_put_err( @@ -369,7 +369,7 @@ pub mod tests { ts: u64, for_update_ts: u64, is_pessimistic_lock: bool, - ) { + ) -> Error { must_prewrite_put_err_impl( engine, key, @@ -378,7 +378,7 @@ pub mod tests { ts, for_update_ts, is_pessimistic_lock, - ); + ) } fn must_prewrite_delete_impl( @@ -516,14 +516,14 @@ pub mod tests { pk: &[u8], start_ts: u64, for_update_ts: u64, - ) { + ) -> Error { let ctx = Context::default(); let snapshot = engine.snapshot(&ctx).unwrap(); let mut txn = MvccTxn::new(snapshot, start_ts, true).unwrap(); let mut options = Options::default(); options.for_update_ts = for_update_ts; txn.acquire_pessimistic_lock(Key::from_raw(key), pk, false, &options) - .unwrap_err(); + .unwrap_err() } pub fn must_pessimistic_rollback( @@ -587,11 +587,16 @@ pub mod tests { write(engine, &ctx, txn.into_modifies()); } - pub fn must_cleanup_err(engine: &E, key: &[u8], start_ts: u64, current_ts: u64) { + pub fn must_cleanup_err( + engine: &E, + key: &[u8], + start_ts: u64, + current_ts: u64, + ) -> Error { let ctx = Context::default(); let snapshot = engine.snapshot(&ctx).unwrap(); let mut txn = MvccTxn::new(snapshot, start_ts, true).unwrap(); - assert!(txn.cleanup(Key::from_raw(key), current_ts).is_err()); + txn.cleanup(Key::from_raw(key), current_ts).unwrap_err() } pub fn must_txn_heart_beat( diff --git a/src/storage/mvcc/txn.rs b/src/storage/mvcc/txn.rs index 9db11a99b18..32ca7112601 100644 --- a/src/storage/mvcc/txn.rs +++ b/src/storage/mvcc/txn.rs @@ -256,7 +256,7 @@ impl MvccTxn { info.set_lock_version(lock.ts); info.set_key(key.into_raw()?); info.set_lock_ttl(lock.ttl); - info.set_txn_size(options.txn_size); + info.set_txn_size(lock.txn_size); return Err(Error::KeyIsLocked(info)); } if lock.lock_type != LockType::Pessimistic { @@ -797,7 +797,7 @@ mod tests { use crate::storage::kv::Engine; use crate::storage::mvcc::tests::*; use crate::storage::mvcc::WriteType; - use crate::storage::mvcc::{MvccReader, MvccTxn}; + use crate::storage::mvcc::{Error, MvccReader, MvccTxn}; use crate::storage::{ Key, Mutation, Options, ScanMode, TestEngineBuilder, SHORT_VALUE_MAX_LEN, }; @@ -2052,4 +2052,48 @@ mod tests { must_get(&engine, k, 19, v); assert!(try_prewrite_insert(&engine, k, v, k, 20).is_err()); } + + #[test] + fn test_lock_info_validation() { + let engine = TestEngineBuilder::new().build().unwrap(); + + let k = b"k"; + let v = b"v"; + + let mut options = Options::default(); + options.lock_ttl = 3; + options.txn_size = 10; + + let mut expected_lock_info = kvproto::kvrpcpb::LockInfo::default(); + expected_lock_info.set_primary_lock(k.to_vec()); + expected_lock_info.set_lock_version(10); + expected_lock_info.set_key(k.to_vec()); + expected_lock_info.set_lock_ttl(options.lock_ttl); + expected_lock_info.set_txn_size(options.txn_size); + + let assert_lock_info_eq = |e, expected_lock_info: &kvproto::kvrpcpb::LockInfo| match e { + Error::KeyIsLocked(info) => assert_eq!(info, *expected_lock_info), + _ => panic!("unexpected error"), + }; + + must_prewrite_put_impl(&engine, k, v, k, 10, false, options); + + assert_lock_info_eq( + must_prewrite_put_err(&engine, k, v, k, 20), + &expected_lock_info, + ); + + assert_lock_info_eq( + must_acquire_pessimistic_lock_err(&engine, k, k, 30, 30), + &expected_lock_info, + ); + + assert_lock_info_eq(must_cleanup_err(&engine, k, 10, 1), &expected_lock_info); + + expected_lock_info.set_lock_ttl(0); + assert_lock_info_eq( + must_pessimistic_prewrite_put_err(&engine, k, v, k, 40, 40, false), + &expected_lock_info, + ); + } }