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:
Trigger condition is:
1. 1PC committed, memory lock is not released.
2. Concurrent reads encounter memory lock, use `max_ts` as the read timestamp.
3. `async-prewrite-apply` is enabled and the read-after-write could not be ensured from a same session.

Labling it `major` severity as the `async-prewrite-apply` is not enabled by default.

Fix patch:
diff --git a/components/txn_types/src/lock.rs b/components/txn_types/src/lock.rs
index 103353318e0..8a90132e09b 100644
--- a/components/txn_types/src/lock.rs
+++ b/components/txn_types/src/lock.rs
@@ -78,6 +78,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>>,
@@ -154,6 +158,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_ts: TimeStamp::zero(),
@@ -438,7 +443,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.
@@ -1122,6 +1131,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_ts: 8.into(),
diff --git a/src/storage/txn/actions/prewrite.rs b/src/storage/txn/actions/prewrite.rs
index de2cf481c94..6dcf4c22ca3 100644
--- a/src/storage/txn/actions/prewrite.rs
+++ b/src/storage/txn/actions/prewrite.rs
@@ -427,9 +427,9 @@ impl<'a> PrewriteMutation<'a> {
                             continue;
                         }
 
-                        warn!("conflicting write was found, pessimistic lock must be lost for the corresponding row key"; 
-                            "key" => %self.key, 
-                            "start_ts" => self.txn_props.start_ts, 
+                        warn!("conflicting write was found, pessimistic lock must be lost for the corresponding row key";
+                            "key" => %self.key,
+                            "start_ts" => self.txn_props.start_ts,
                             "for_update_ts" => for_update_ts,
                             "conflicting start_ts" => write.start_ts,
                             "conflicting commit_ts" => commit_ts);
@@ -489,6 +489,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 {
@@ -503,6 +507,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();
             }
@@ -783,6 +788,8 @@ fn amend_pessimistic_lock<S: Snapshot>(
 }
 
 pub mod tests {
+    #[cfg(test)]
+    use std::borrow::Cow;
     #[cfg(test)]
     use std::sync::Arc;
 
@@ -793,7 +800,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)]
@@ -1127,6 +1134,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>(
@@ -2510,4 +2519,57 @@ pub mod tests {
         assert_eq!(lock.versions_to_last_change, 0);
         must_rollback(&mut engine, key, 40, false);
     }
+
+    #[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,
+        )
+        .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,
+        )
+        .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 7f6dfd258c5..cade9b05e38 100644
--- a/src/storage/txn/commands/prewrite.rs
+++ b/src/storage/txn/commands/prewrite.rs
@@ -870,7 +870,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, _) in std::mem::take(&mut txn.locks_for_1pc) {
+    for (key, mut lock, _) in std::mem::take(&mut txn.locks_for_1pc) {
+        lock.use_one_pc = false;
         txn.put_lock(key, &lock);
     }
 }


Test patch: