txn: Do constraint check when handling repeated acqurie_pessimsitic_lock request

src/storage/txn/actions/acquire_pessimistic_lock.rs

@@ -142,10 +142,22 @@ pub fn acquire_pessimistic_lock<S: Snapshot>(
                 None
             };
 
-        if need_load_value {
-            val = reader.get(&key, for_update_ts)?;
-        } else if need_check_existence {
-            val = reader.get_write(&key, for_update_ts)?.map(|_| vec![]);
+        if need_load_value || need_check_existence || should_not_exist {
+            let write = reader.get_write_with_commit_ts(&key, for_update_ts)?;
+            if let Some((write, commit_ts)) = write {
+                // Here `get_write_with_commit_ts` returns only the latest PUT if it exists and
+                // is not deleted. It's still ok to pass it into `check_data_constraint`.
+                // In case we are going to lock it with write conflict, we do not check it since
+                // the statement will then retry.
+                if locked_with_conflict_ts.is_none() {
+                    check_data_constraint(reader, should_not_exist, &write, commit_ts, &key)?;
+                }
+                if need_load_value {
+                    val = Some(reader.load_data(&key, write)?);
+                } else if need_check_existence {
+                    val = Some(vec![]);
+                }
+            }
         }
         // Pervious write is not loaded.
         let (prev_write_loaded, prev_write) = (false, None);
@@ -1832,4 +1844,134 @@ pub mod tests {
         must_pessimistic_rollback(&mut engine, b"k1", 10, 50);
         must_unlocked(&mut engine, b"k1");
     }
+
+    #[test]
+    fn test_repeated_request_check_should_not_exist() {
+        let mut engine = TestEngineBuilder::new().build().unwrap();
+
+        for &(return_values, check_existence) in
+            &[(false, false), (false, true), (true, false), (true, true)]
+        {
+            let key = &[b'k', (return_values as u8 * 2) + check_existence as u8] as &[u8];
+
+            // An empty key.
+            must_succeed(&mut engine, key, key, 10, 10);
+            let res = must_succeed_impl(
+                &mut engine,
+                key,
+                key,
+                10,
+                true,
+                1000,
+                10,
+                return_values,
+                check_existence,
+                15,
+                false,
+            );
+            assert!(res.is_none());
+            must_pessimistic_prewrite_lock(&mut engine, key, key, 10, 10, DoPessimisticCheck);
+            must_commit(&mut engine, key, 10, 19);
+
+            // The key has one record: Lock(10, 19)
+            must_succeed(&mut engine, key, key, 20, 20);
+            let res = must_succeed_impl(
+                &mut engine,
+                key,
+                key,
+                20,
+                true,
+                1000,
+                20,
+                return_values,
+                check_existence,
+                25,
+                false,
+            );
+            assert!(res.is_none());
+            must_pessimistic_prewrite_put(&mut engine, key, b"v1", key, 20, 20, DoPessimisticCheck);
+            must_commit(&mut engine, key, 20, 29);
+
+            // The key has records:
+            // Lock(10, 19), Put(20, 29)
+            must_succeed(&mut engine, key, key, 30, 30);
+            let error = must_err_impl(
+                &mut engine,
+                key,
+                key,
+                30,
+                true,
+                30,
+                return_values,
+                check_existence,
+                35,
+                false,
+            );
+            assert!(matches!(
+                error,
+                MvccError(box ErrorInner::AlreadyExist { .. })
+            ));
+            must_pessimistic_prewrite_lock(&mut engine, key, key, 30, 30, DoPessimisticCheck);
+            must_commit(&mut engine, key, 30, 39);
+
+            // Lock(10, 19), Put(20, 29), Lock(30, 39)
+            must_succeed(&mut engine, key, key, 40, 40);
+            let error = must_err_impl(
+                &mut engine,
+                key,
+                key,
+                40,
+                true,
+                40,
+                return_values,
+                check_existence,
+                45,
+                false,
+            );
+            assert!(matches!(
+                error,
+                MvccError(box ErrorInner::AlreadyExist { .. })
+            ));
+            must_pessimistic_prewrite_delete(&mut engine, key, key, 40, 40, DoPessimisticCheck);
+            must_commit(&mut engine, key, 40, 49);
+
+            // Lock(10, 19), Put(20, 29), Lock(30, 39), Delete(40, 49)
+            must_succeed(&mut engine, key, key, 50, 50);
+            let res = must_succeed_impl(
+                &mut engine,
+                key,
+                key,
+                50,
+                true,
+                1000,
+                50,
+                return_values,
+                check_existence,
+                55,
+                false,
+            );
+            assert!(res.is_none());
+            must_pessimistic_prewrite_lock(&mut engine, key, key, 50, 50, DoPessimisticCheck);
+            must_commit(&mut engine, key, 50, 59);
+
+            // Lock(10, 19), Put(20, 29), Lock(30, 39), Delete(40, 49), Lock(50, 59)
+            must_succeed(&mut engine, key, key, 60, 60);
+            let res = must_succeed_impl(
+                &mut engine,
+                key,
+                key,
+                60,
+                true,
+                1000,
+                60,
+                return_values,
+                check_existence,
+                65,
+                false,
+            );
+            assert!(res.is_none());
+            must_pessimistic_prewrite_lock(&mut engine, key, key, 60, 60, DoPessimisticCheck);
+            must_commit(&mut engine, key, 60, 69);
+        }
+    }
 }