txnkv: fix SetCommitWaitUntilTSOTimeout does not work in causal consistency (tikv#1847)

fix tikv#1846

Signed-off-by: Chao Wang <cclcwangchao@hotmail.com>

Author: lcwangchao

integration_tests/option_test.go

@@ -119,10 +119,6 @@ func (s *testOptionSuite) GetHistogramMetricSampleCount(m any) uint64 {
 }
 
 func (s *testOptionSuite) TestSetCommitWaitUntilTSO() {
-	if *withTiKV {
-		s.T().Skip("TestSetCommitWaitUntilTSO only runs on local mock storage")
-	}
-
 	s.Nil(failpoint.Enable("tikvclient/InjectStartTSForCommit", "return(true)"))
 	defer func() {
 		s.Nil(failpoint.Disable("tikvclient/InjectStartTSForCommit"))
@@ -151,6 +147,7 @@ func (s *testOptionSuite) TestSetCommitWaitUntilTSO() {
 		//	}
 		mockCommitTSO  []uint64
 		setWaitTimeout time.Duration
+		extraPrepare   func(transaction.TxnProbe)
 		err            bool
 	}{
 		{
@@ -181,11 +178,38 @@ func (s *testOptionSuite) TestSetCommitWaitUntilTSO() {
 			setWaitTimeout: time.Millisecond,
 			err:            true,
 		},
+		{
+			name:          "should also check for 1pc",
+			commitWaitTSO: oracle.ComposeTS(10000, 0),
+			mockCommitTSO: []uint64{100},
+			extraPrepare: func(txn transaction.TxnProbe) {
+				txn.SetEnableAsyncCommit(true)
+				txn.SetEnable1PC(true)
+			},
+			err: true,
+		},
+		{
+			name:          "should also check for causal consistency",
+			commitWaitTSO: oracle.ComposeTS(10000, 0),
+			mockCommitTSO: []uint64{100},
+			extraPrepare: func(txn transaction.TxnProbe) {
+				txn.SetEnableAsyncCommit(true)
+				txn.SetEnable1PC(true)
+				txn.SetCausalConsistency(true)
+			},
+			err: true,
+		},
 	} {
 		s.Run(tc.name, func() {
 			txn, err := s.store.Begin()
 			s.NoError(err)
 			s.NoError(txn.Set([]byte("somekey:"+uuid.NewString()), []byte("somevalue")))
+			// by default, 1pc and async commit are disabled, but it can be overridden in extraPrepare
+			txn.SetEnableAsyncCommit(false)
+			txn.SetEnable1PC(false)
+			if tc.extraPrepare != nil {
+				tc.extraPrepare(txn)
+			}
 
 			txn.SetCommitWaitUntilTSO(tc.commitWaitTSO + txn.StartTS())
 			if tc.setWaitTimeout > 0 {

internal/unionstore/memdb_art.go

@@ -166,6 +166,6 @@ func (db *artDBWithContext) SnapshotIterReverse(upper, lower []byte) Iterator {
 }
 
 // SnapshotGetter returns a Getter for a snapshot of MemBuffer.
-func (db *artDBWithContext) SnapshotGetter() Getter {
+func (db *artDBWithContext) SnapshotGetter() kv.Getter {
 	return db.ART.SnapshotGetter()
 }

internal/unionstore/memdb_rbt.go

@@ -176,6 +176,6 @@ func (db *rbtDBWithContext) SnapshotIterReverse(upper, lower []byte) Iterator {
 }
 
 // SnapshotGetter returns a Getter for a snapshot of MemBuffer.
-func (db *rbtDBWithContext) SnapshotGetter() Getter {
+func (db *rbtDBWithContext) SnapshotGetter() kv.Getter {
 	return db.RBT.SnapshotGetter()
 }

txnkv/transaction/2pc.go

@@ -1747,7 +1747,8 @@ func (c *twoPhaseCommitter) execute(ctx context.Context) (err error) {
 	// all nodes, we have to make sure the commit TS of this transaction is greater
 	// than the snapshot TS of all existent readers. So we get a new timestamp
 	// from PD and plus one as our MinCommitTS.
-	if commitTSMayBeCalculated && c.needLinearizability() {
+	// If `commitWaitUntilTSO > 0`, we also need to get a new timestamp from TSO to ensure the constraint satisfied.
+	if commitTSMayBeCalculated && (c.needLinearizability() || c.txn.commitWaitUntilTSO > 0) {
 		util.EvalFailpoint("getMinCommitTSFromTSO")
 		start := time.Now()
 		latestTS, err := c.txn.GetTimestampForCommit(bo, c.txn.GetScope())