From e0bf24b24d9776c1961ab7b949dcf665ed75e21e Mon Sep 17 00:00:00 2001
From: cfzjywxk <lsswxrxr@163.com>
Date: Tue, 16 Nov 2021 13:23:32 +0800
Subject: [PATCH] add newly inserted flag to optimize delete your write (#378)

Signed-off-by: cfzjywxk <lsswxrxr@163.com>
---
 integration_tests/2pc_test.go | 48 +++++++++++++++++++++++++++++++++++
 kv/keyflags.go                | 10 ++++++++
 txnkv/transaction/2pc.go      | 19 +++++++++++---
 3 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/integration_tests/2pc_test.go b/integration_tests/2pc_test.go
index ea4946de..bd24fc94 100644
--- a/integration_tests/2pc_test.go
+++ b/integration_tests/2pc_test.go
@@ -1550,3 +1550,51 @@ func (s *testCommitterSuite) TestCommitMultipleRegions() {
 	}
 	s.mustCommit(m)
 }
+
+func (s *testCommitterSuite) TestNewlyInsertedMemDBFlag() {
+	ctx := context.Background()
+	txn := s.begin()
+	memdb := txn.GetMemBuffer()
+	k0 := []byte("k0")
+	v0 := []byte("v0")
+	k1 := []byte("k1")
+	k2 := []byte("k2")
+	v1 := []byte("v1")
+	v2 := []byte("v2")
+
+	// Insert after delete, the newly inserted flag should not exist.
+	err := txn.Delete(k0)
+	s.Nil(err)
+	err = txn.Set(k0, v0)
+	s.Nil(err)
+	flags, err := memdb.GetFlags(k0)
+	s.Nil(err)
+	s.False(flags.HasNewlyInserted())
+
+	// Lock then insert, the newly inserted flag should exist.
+	lockCtx := &kv.LockCtx{ForUpdateTS: txn.StartTS(), WaitStartTime: time.Now()}
+	err = txn.LockKeys(context.Background(), lockCtx, k1)
+	s.Nil(err)
+	err = txn.GetMemBuffer().SetWithFlags(k1, v1, kv.SetNewlyInserted)
+	s.Nil(err)
+	flags, err = memdb.GetFlags(k1)
+	s.Nil(err)
+	s.True(flags.HasNewlyInserted())
+
+	// Lock then delete and insert, the newly inserted flag should not exist.
+	err = txn.LockKeys(ctx, lockCtx, k2)
+	s.Nil(err)
+	err = txn.Delete(k2)
+	s.Nil(err)
+	flags, err = memdb.GetFlags(k2)
+	s.Nil(err)
+	s.False(flags.HasNewlyInserted())
+	err = txn.Set(k2, v2)
+	s.Nil(err)
+	flags, err = memdb.GetFlags(k2)
+	s.Nil(err)
+	s.False(flags.HasNewlyInserted())
+
+	err = txn.Commit(ctx)
+	s.Nil(err)
+}
diff --git a/kv/keyflags.go b/kv/keyflags.go
index 133e5079..8bbba044 100644
--- a/kv/keyflags.go
+++ b/kv/keyflags.go
@@ -49,6 +49,7 @@ const (
 	flagPrewriteOnly
 	flagIgnoredIn2PC
 	flagReadable
+	flagNewlyInserted
 
 	persistentFlags = flagKeyLocked | flagKeyLockedValExist
 )
@@ -98,6 +99,11 @@ func (f KeyFlags) AndPersistent() KeyFlags {
 	return f & persistentFlags
 }
 
+// HasNewlyInserted returns whether the in-transaction key is generated by an "insert" operation.
+func (f KeyFlags) HasNewlyInserted() bool {
+	return f&flagNewlyInserted != 0
+}
+
 // ApplyFlagsOps applys flagspos to origin.
 func ApplyFlagsOps(origin KeyFlags, ops ...FlagsOp) KeyFlags {
 	for _, op := range ops {
@@ -126,6 +132,8 @@ func ApplyFlagsOps(origin KeyFlags, ops ...FlagsOp) KeyFlags {
 			origin |= flagIgnoredIn2PC
 		case SetReadable:
 			origin |= flagReadable
+		case SetNewlyInserted:
+			origin |= flagNewlyInserted
 		}
 	}
 	return origin
@@ -160,4 +168,6 @@ const (
 	SetIgnoredIn2PC
 	// SetReadable marks the key is readable by in-transaction read.
 	SetReadable
+	// SetNewlyInserted marks the key is newly inserted with value length greater than zero.
+	SetNewlyInserted
 )
diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go
index a4b55301..2a5c3671 100644
--- a/txnkv/transaction/2pc.go
+++ b/txnkv/transaction/2pc.go
@@ -439,10 +439,21 @@ func (c *twoPhaseCommitter) initKeysAndMutations() error {
 					checkCnt++
 					memBuf.UpdateFlags(key, kv.SetPrewriteOnly)
 				} else {
-					// normal delete keys in optimistic txn can be delete without not exists checking
-					// delete-your-writes keys in pessimistic txn can ensure must be no exists so can directly delete them
-					op = kvrpcpb.Op_Del
-					delCnt++
+					if flags.HasNewlyInserted() {
+						// The delete-your-write keys in pessimistic transactions, only lock needed keys and skip
+						// other deletes for example the secondary index delete.
+						// Here if `tidb_constraint_check_in_place` is enabled and the transaction is in optimistic mode,
+						// the logic is same as the pessimistic mode.
+						if flags.HasLocked() {
+							op = kvrpcpb.Op_Lock
+							lockCnt++
+						} else {
+							continue
+						}
+					} else {
+						op = kvrpcpb.Op_Del
+						delCnt++
+					}
 				}
 			}
 		}