Repository: tikv/tikv (language: rust)

Issue:
Reads with max-ts may not see previous writes when async-apply-prewrite is enabled
## Bug Report

<!-- Thanks for your bug report! Don't worry if you can't fill out all the sections. -->

### What version of TiKV are you using?
<!-- You can run `tikv-server --version` -->
master

### What operating system and CPU are you using?
<!-- If you're using Linux, you can run `cat /proc/cpuinfo` -->

### Steps to reproduce
<!-- If possible, provide a recipe for reproducing the error. A complete runnable program is good. -->
ref https://github.com/pingcap/tidb/issues/58677

### What did you expect?

No anomaly detected by jepsen.

### What did happened?

The jepsen test failed.



Discussion / hints:

Fix patch:
diff --git a/components/txn_types/src/lock.rs b/components/txn_types/src/lock.rs
index 26dc493262e..98650f1a8e4 100644
--- a/components/txn_types/src/lock.rs
+++ b/components/txn_types/src/lock.rs
@@ -84,6 +84,10 @@ pub struct Lock {
     pub txn_size: u64,
     pub min_commit_ts: TimeStamp,
     pub use_async_commit: bool,
+    // This field is only valid for in-memory locks and does not need to be persisted because:
+    //   1. the lock should be converted to a write directly when 1pc succeeds.
+    //   2. the field should be reverted to false (default value) when 1pc fails.
+    pub use_one_pc: bool,
     // Only valid when `use_async_commit` is true, and the lock is primary. Do not set
     // `secondaries` for secondaries.
     pub secondaries: Vec<Vec<u8>>,
@@ -171,6 +175,7 @@ impl Lock {
             txn_size,
             min_commit_ts,
             use_async_commit: false,
+            use_one_pc: false,
             secondaries: Vec::default(),
             rollback_ts: Vec::default(),
             last_change: LastChange::default(),
@@ -489,7 +494,11 @@ impl Lock {
             )));
         }
 
-        if ts == TimeStamp::max() && raw_key == lock.primary && !lock.use_async_commit {
+        if ts == TimeStamp::max()
+            && raw_key == lock.primary
+            && !lock.use_async_commit
+            && !lock.use_one_pc
+        {
             // When `ts == TimeStamp::max()` (which means to get latest committed version
             // for primary key), and current key is the primary key, we ignore
             // this lock.
@@ -1326,6 +1335,7 @@ mod tests {
             txn_size: 0,
             min_commit_ts: 20.into(),
             use_async_commit: false,
+            use_one_pc: false,
             secondaries: vec![],
             rollback_ts: vec![],
             last_change: LastChange::make_exist(8.into(), 2),
diff --git a/src/storage/txn/actions/prewrite.rs b/src/storage/txn/actions/prewrite.rs
index 9e95a3ef63f..b618f63892b 100644
--- a/src/storage/txn/actions/prewrite.rs
+++ b/src/storage/txn/actions/prewrite.rs
@@ -618,6 +618,10 @@ impl<'a> PrewriteMutation<'a> {
         if let Some(secondary_keys) = self.secondary_keys {
             lock.use_async_commit = true;
             lock.secondaries = secondary_keys.to_owned();
+        } else if try_one_pc {
+            // Set `use_one_pc` to true to prevent the in-memory lock from being skipped
+            // when reading with max-ts.
+            lock.use_one_pc = true;
         }
 
         let final_min_commit_ts = if lock.use_async_commit || try_one_pc {
@@ -632,6 +636,7 @@ impl<'a> PrewriteMutation<'a> {
             fail_point!("after_calculate_min_commit_ts");
             if let Err(Error(box ErrorInner::CommitTsTooLarge { .. })) = &res {
                 try_one_pc = false;
+                lock.use_one_pc = false;
                 lock.use_async_commit = false;
                 lock.secondaries = Vec::new();
             }
@@ -916,6 +921,8 @@ fn amend_pessimistic_lock<S: Snapshot>(
 }
 
 pub mod tests {
+    #[cfg(test)]
+    use std::borrow::Cow;
     #[cfg(test)]
     use std::sync::Arc;
 
@@ -926,7 +933,7 @@ pub mod tests {
     #[cfg(test)]
     use tikv_kv::RocksEngine;
     #[cfg(test)]
-    use txn_types::OldValue;
+    use txn_types::{OldValue, TsSet};
 
     use super::*;
     #[cfg(test)]
@@ -1275,6 +1282,8 @@ pub mod tests {
         // success 1pc prewrite needs to be transformed to locks
         assert!(!must_locked(&mut engine, b"k1", 10).use_async_commit);
         assert!(!must_locked(&mut engine, b"k2", 10).use_async_commit);
+        assert!(!must_locked(&mut engine, b"k1", 10).use_one_pc);
+        assert!(!must_locked(&mut engine, b"k2", 10).use_one_pc);
     }
 
     pub fn try_pessimistic_prewrite_check_not_exists<E: Engine>(
@@ -2757,4 +2766,59 @@ pub mod tests {
         prewrite_err(&mut engine, key, value, key, 120, 130, Some(130));
         must_unlocked(&mut engine, key);
     }
+
+    #[test]
+    fn test_1pc_set_lock_use_one_pc() {
+        let mut engine = crate::storage::TestEngineBuilder::new().build().unwrap();
+        let cm = ConcurrencyManager::new(42.into());
+
+        let snapshot = engine.snapshot(Default::default()).unwrap();
+
+        let mut txn = MvccTxn::new(10.into(), cm.clone());
+        let mut reader = SnapshotReader::new(10.into(), snapshot, false);
+
+        let k1 = b"k1";
+        let k2 = b"k2";
+
+        prewrite(
+            &mut txn,
+            &mut reader,
+            &optimistic_async_props(k1, 10.into(), 50.into(), 2, true),
+            Mutation::make_put(Key::from_raw(k1), b"v1".to_vec()),
+            &None,
+            SkipPessimisticCheck,
+            None,
+        )
+        .unwrap();
+        prewrite(
+            &mut txn,
+            &mut reader,
+            &optimistic_async_props(k1, 10.into(), 50.into(), 1, true),
+            Mutation::make_put(Key::from_raw(k2), b"v2".to_vec()),
+            &None,
+            SkipPessimisticCheck,
+            None,
+        )
+        .unwrap();
+
+        // lock.use_one_pc should be set to true when using 1pc.
+        assert_eq!(txn.guards.len(), 2);
+        txn.guards[0].with_lock(|l| assert!(l.as_ref().unwrap().use_one_pc));
+        txn.guards[1].with_lock(|l| assert!(l.as_ref().unwrap().use_one_pc));
+
+        // read with max_ts should be blocked by the lock.
+        for &key in &[k1, k2] {
+            let k = Key::from_raw(key);
+            let res = cm.read_key_check(&k, |l| {
+                Lock::check_ts_conflict(
+                    Cow::Borrowed(l),
+                    &k,
+                    TimeStamp::max(),
+                    &TsSet::Empty,
+                    crate::storage::IsolationLevel::Si,
+                )
+            });
+            assert!(res.is_err());
+        }
+    }
 }
diff --git a/src/storage/txn/commands/prewrite.rs b/src/storage/txn/commands/prewrite.rs
index 861083fb117..d9c7a65921b 100644
--- a/src/storage/txn/commands/prewrite.rs
+++ b/src/storage/txn/commands/prewrite.rs
@@ -988,7 +988,8 @@ fn handle_1pc_locks(txn: &mut MvccTxn, commit_ts: TimeStamp) -> ReleasedLocks {
 
 /// Change all 1pc locks in txn to 2pc locks.
 pub(in crate::storage::txn) fn fallback_1pc_locks(txn: &mut MvccTxn) {
-    for (key, lock, remove_pessimistic_lock) in std::mem::take(&mut txn.locks_for_1pc) {
+    for (key, mut lock, remove_pessimistic_lock) in std::mem::take(&mut txn.locks_for_1pc) {
+        lock.use_one_pc = false;
         let is_new_lock = !remove_pessimistic_lock;
         txn.put_lock(key, &lock, is_new_lock);
     }


Test patch: