txn: fix wrong txn_size in acquire_pessimistic_lock (tikv#5681)

* txn: fix wrong txn_size in acquire_pessimistic_lock

Signed-off-by: youjiali1995 <[email protected]>

src/storage/mvcc/mod.rs

@@ -271,7 +271,7 @@ pub mod tests {
         Ok(())
     }
 
-    fn must_prewrite_put_impl<E: Engine>(
+    pub fn must_prewrite_put_impl<E: Engine>(
         engine: &E,
         key: &[u8],
         value: &[u8],
@@ -336,18 +336,18 @@ 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();
         let mut options = Options::default();
         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<E: Engine>(
@@ -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<E: Engine>(
@@ -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<E: Engine>(
@@ -587,11 +587,16 @@ pub mod tests {
         write(engine, &ctx, txn.into_modifies());
     }
 
-    pub fn must_cleanup_err<E: Engine>(engine: &E, key: &[u8], start_ts: u64, current_ts: u64) {
+    pub fn must_cleanup_err<E: Engine>(
+        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<E: Engine>(

src/storage/mvcc/txn.rs

@@ -256,7 +256,7 @@ impl<S: Snapshot> MvccTxn<S> {
                 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,
+        );
+    }
 }