txn: Support check for_update_ts when prewriting

components/resolved_ts/src/cmd.rs

@@ -424,6 +424,7 @@ mod tests {
             Mutation::make_put(k1.clone(), b"v4".to_vec()),
             &None,
             SkipPessimisticCheck,
+            None,
         )
         .unwrap();
         one_pc_commit(true, &mut txn, 10.into());

src/storage/mod.rs

@@ -8254,6 +8254,7 @@ mod tests {
                     None,
                     false,
                     AssertionLevel::Off,
+                    vec![],
                     Context::default(),
                 ),
                 expect_ok_callback(tx.clone(), 0),
@@ -9620,6 +9621,7 @@ mod tests {
                     Some(vec![b"e".to_vec()]),
                     false,
                     AssertionLevel::Off,
+                    vec![],
                     Context::default(),
                 ),
                 Box::new(move |res| {
@@ -9718,6 +9720,7 @@ mod tests {
                     None,
                     false,
                     AssertionLevel::Off,
+                    vec![],
                     Default::default(),
                 ),
                 expect_ok_callback(tx.clone(), 0),
@@ -9768,6 +9771,7 @@ mod tests {
                     Some(vec![k2.to_vec()]),
                     false,
                     AssertionLevel::Off,
+                    vec![],
                     Default::default(),
                 ),
                 expect_ok_callback(tx.clone(), 0),
@@ -10604,6 +10608,7 @@ mod tests {
                     None,
                     false,
                     AssertionLevel::Off,
+                    vec![],
                     Context::default(),
                 ),
                 Box::new(move |res| {
@@ -10662,6 +10667,7 @@ mod tests {
                     None,
                     false,
                     AssertionLevel::Off,
+                    vec![],
                     Context::default(),
                 ),
                 Box::new(move |res| {

src/storage/mvcc/mod.rs

@@ -132,10 +132,14 @@ pub enum ErrorInner {
     KeyVersion,
 
     #[error(
-        "pessimistic lock not found, start_ts:{}, key:{}",
-        .start_ts, log_wrappers::Value::key(.key)
+        "pessimistic lock not found, start_ts:{}, key:{}, reason: {:?}",
+        .start_ts, log_wrappers::Value::key(.key), .reason
     )]
-    PessimisticLockNotFound { start_ts: TimeStamp, key: Vec<u8> },
+    PessimisticLockNotFound {
+        start_ts: TimeStamp,
+        key: Vec<u8>,
+        reason: PessimisticLockNotFoundReason,
+    },
 
     #[error(
         "min_commit_ts {} is larger than max_commit_ts {}, start_ts: {}",
@@ -257,12 +261,15 @@ impl ErrorInner {
                     key: key.to_owned(),
                 })
             }
-            ErrorInner::PessimisticLockNotFound { start_ts, key } => {
-                Some(ErrorInner::PessimisticLockNotFound {
-                    start_ts: *start_ts,
-                    key: key.to_owned(),
-                })
-            }
+            ErrorInner::PessimisticLockNotFound {
+                start_ts,
+                key,
+                reason,
+            } => Some(ErrorInner::PessimisticLockNotFound {
+                start_ts: *start_ts,
+                key: key.to_owned(),
+                reason: *reason,
+            }),
             ErrorInner::CommitTsTooLarge {
                 start_ts,
                 min_commit_ts,
@@ -421,6 +428,15 @@ pub fn default_not_found_error(key: Vec<u8>, hint: &str) -> Error {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+pub enum PessimisticLockNotFoundReason {
+    LockTsMismatch,
+    LockMissingAmendFail,
+    LockForUpdateTsMismatch,
+    NonLockKeyConflict,
+    FailpointInjected,
+}
+
 pub mod tests {
     use std::borrow::Cow;
 

src/storage/mvcc/reader/reader.rs

@@ -886,6 +886,7 @@ pub mod tests {
                 m,
                 &None,
                 SkipPessimisticCheck,
+                None,
             )
             .unwrap();
             self.write(txn.into_modifies());
@@ -910,6 +911,7 @@ pub mod tests {
                 m,
                 &None,
                 DoPessimisticCheck,
+                None,
             )
             .unwrap();
             self.write(txn.into_modifies());

src/storage/mvcc/txn.rs

@@ -9,7 +9,7 @@ use kvproto::kvrpcpb::LockInfo;
 use txn_types::{Key, Lock, PessimisticLock, TimeStamp, Value};
 
 use super::metrics::{GC_DELETE_VERSIONS_HISTOGRAM, MVCC_VERSIONS_HISTOGRAM};
-use crate::storage::kv::Modify;
+use crate::storage::{kv::Modify, mvcc::PessimisticLockNotFoundReason};
 
 pub const MAX_TXN_WRITE_SIZE: usize = 32 * 1024;
 
@@ -306,6 +306,7 @@ pub(crate) fn make_txn_error(
             "pessimisticlocknotfound" => ErrorInner::PessimisticLockNotFound {
                 start_ts,
                 key: key.to_raw().unwrap(),
+                reason: PessimisticLockNotFoundReason::FailpointInjected,
             },
             _ => ErrorInner::Other(box_err!("unexpected error string")),
         }
@@ -815,6 +816,7 @@ pub(crate) mod tests {
             Mutation::make_put(key.clone(), v.to_vec()),
             &None,
             SkipPessimisticCheck,
+            None,
         )
         .unwrap();
         assert!(txn.write_size() > 0);
@@ -859,6 +861,7 @@ pub(crate) mod tests {
             Mutation::make_put(Key::from_raw(key), value.to_vec()),
             &None,
             SkipPessimisticCheck,
+            None,
         )
         .unwrap_err();
 
@@ -872,6 +875,7 @@ pub(crate) mod tests {
             Mutation::make_put(Key::from_raw(key), value.to_vec()),
             &None,
             SkipPessimisticCheck,
+            None,
         )
         .unwrap();
     }
@@ -1312,6 +1316,7 @@ pub(crate) mod tests {
                 mutation,
                 &Some(vec![b"key1".to_vec(), b"key2".to_vec(), b"key3".to_vec()]),
                 SkipPessimisticCheck,
+                None,
             )
             .unwrap();
             let modifies = txn.into_modifies();
@@ -1370,6 +1375,7 @@ pub(crate) mod tests {
                 mutation,
                 &Some(vec![b"key1".to_vec(), b"key2".to_vec(), b"key3".to_vec()]),
                 DoPessimisticCheck,
+                None,
             )
             .unwrap();
             let modifies = txn.into_modifies();
@@ -1439,6 +1445,7 @@ pub(crate) mod tests {
             mutation,
             &Some(vec![b"key1".to_vec(), b"key2".to_vec(), b"key3".to_vec()]),
             DoPessimisticCheck,
+            None,
         )
         .unwrap();
         assert_eq!(min_commit_ts.into_inner(), 100);

src/storage/txn/actions/acquire_pessimistic_lock.rs

@@ -452,7 +452,6 @@ pub mod tests {
         TestEngineBuilder,
     };
 
-    #[cfg(test)]
     pub fn acquire_pessimistic_lock_allow_lock_with_conflict<E: Engine>(
         engine: &mut E,
         key: &[u8],
@@ -496,7 +495,6 @@ pub mod tests {
         res.map(|r| r.0)
     }
 
-    #[cfg(test)]
     pub fn must_succeed_allow_lock_with_conflict<E: Engine>(
         engine: &mut E,
         key: &[u8],