From ffaaf7131a8df6ab4e858bf27e39cd7445cf7929 Mon Sep 17 00:00:00 2001
From: ekexium <eke@fastmail.com>
Date: Tue, 13 Sep 2022 13:15:14 +0800
Subject: [PATCH] Don't let pessimistic txn prewrite in the optimistic way
 (#584)

* don't let pessimistic txn prewrite in the optimistic way even if it does not lock any key

Signed-off-by: ekexium <eke@fastmail.com>

* add a todo

Signed-off-by: ekexium <eke@fastmail.com>

Signed-off-by: ekexium <eke@fastmail.com>
---
 txnkv/transaction/2pc.go      | 13 +++++++++++++
 txnkv/transaction/prewrite.go |  1 +
 2 files changed, 14 insertions(+)

diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go
index 888eb6c150..02a5173087 100644
--- a/txnkv/transaction/2pc.go
+++ b/txnkv/transaction/2pc.go
@@ -1398,6 +1398,18 @@ func (c *twoPhaseCommitter) execute(ctx context.Context) (err error) {
 		c.hasTriedOnePC = true
 	}
 
+	// if lazy uniqueness check is enabled in TiDB (@@constraint_check_in_place_pessimistic=0), for_update_ts might be
+	// zero for a pessimistic transaction. We set it to the start_ts to force the PrewritePessimistic path in TiKV.
+	// TODO: can we simply set for_update_ts = start_ts for all pessimistic transactions whose for_update_ts=0?
+	if c.forUpdateTS == 0 {
+		for i := 0; i < c.mutations.Len(); i++ {
+			if c.mutations.NeedConstraintCheckInPrewrite(i) {
+				c.forUpdateTS = c.startTS
+				break
+			}
+		}
+	}
+
 	// TODO(youjiali1995): It's better to use different maxSleep for different operations
 	// and distinguish permanent errors from temporary errors, for example:
 	//   - If all PDs are down, all requests to PD will fail due to network error.
@@ -1442,6 +1454,7 @@ func (c *twoPhaseCommitter) execute(ctx context.Context) (err error) {
 	}
 
 	start := time.Now()
+
 	err = c.prewriteMutations(bo, c.mutations)
 
 	if err != nil {
diff --git a/txnkv/transaction/prewrite.go b/txnkv/transaction/prewrite.go
index 7bf4e00dff..c7bc6daae6 100644
--- a/txnkv/transaction/prewrite.go
+++ b/txnkv/transaction/prewrite.go
@@ -75,6 +75,7 @@ func (c *twoPhaseCommitter) buildPrewriteRequest(batch batchMutations, txnSize u
 	m := batch.mutations
 	mutations := make([]*kvrpcpb.Mutation, m.Len())
 	pessimisticActions := make([]kvrpcpb.PrewriteRequest_PessimisticAction, m.Len())
+
 	for i := 0; i < m.Len(); i++ {
 		assertion := kvrpcpb.Assertion_None
 		if m.IsAssertExists(i) {