From 5ed11f62d7b09f795a6b75abfdc4324a0d8bd585 Mon Sep 17 00:00:00 2001
From: Andrei Matei <andrei@cockroachlabs.com>
Date: Tue, 8 Oct 2019 14:13:35 -0400
Subject: [PATCH] kv: fix cluster_logical_timestamp()

Before this patch, the txn.CommitTimestamp() function (which powers
the cluster_logical_timestamp(), among others) was failing to take into
consideration possible refreshes that might have happened since the
current epoch began (i.e. txn.RefreshedTimestamp).

Fixes #41424

Release note (bug fix): Fix a bug causing the
cluster_logical_timestamp() function to sometimes return incorrect
results.
---
 pkg/kv/txn_coord_sender.go |  5 +++-
 pkg/kv/txn_test.go         | 54 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/pkg/kv/txn_coord_sender.go b/pkg/kv/txn_coord_sender.go
index c691d905f965..09f01c374c5d 100644
--- a/pkg/kv/txn_coord_sender.go
+++ b/pkg/kv/txn_coord_sender.go
@@ -1012,8 +1012,11 @@ func (tc *TxnCoordSender) OrigTimestamp() hlc.Timestamp {
 func (tc *TxnCoordSender) CommitTimestamp() hlc.Timestamp {
 	tc.mu.Lock()
 	defer tc.mu.Unlock()
+	txn := &tc.mu.txn
 	tc.mu.txn.OrigTimestampWasObserved = true
-	return tc.mu.txn.OrigTimestamp
+	commitTS := txn.OrigTimestamp
+	commitTS.Forward(txn.RefreshedTimestamp)
+	return commitTS
 }
 
 // CommitTimestampFixed is part of the client.TxnSender interface.
diff --git a/pkg/kv/txn_test.go b/pkg/kv/txn_test.go
index 9ff4bbab5145..78955e5cb76a 100644
--- a/pkg/kv/txn_test.go
+++ b/pkg/kv/txn_test.go
@@ -21,8 +21,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/pkg/errors"
-
 	"github.com/cockroachdb/cockroach/pkg/internal/client"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -31,7 +29,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/tscache"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/localtestcluster"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/require"
 )
 
 // TestTxnDBBasics verifies that a simple transaction can be run and
@@ -702,3 +703,52 @@ func TestTxnResolveIntentsFromMultipleEpochs(t *testing.T) {
 		}
 	}
 }
+
+// Test that txn.CommitTimestamp() reflects refreshes.
+func TestTxnCommitTimestampAdvancedByRefresh(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	ctx := context.Background()
+
+	// We're going to inject an uncertainty error, expect the refresh to succeed,
+	// and then check that the txn.CommitTimestamp() value reflects the refresh.
+	injected := false
+	var refreshTS hlc.Timestamp
+	errKey := roachpb.Key("inject_err")
+	s := createTestDBWithContextAndKnobs(t, client.DefaultDBContext(), &storage.StoreTestingKnobs{
+		TestingRequestFilter: func(ba roachpb.BatchRequest) *roachpb.Error {
+			if g, ok := ba.GetArg(roachpb.Get); ok && g.(*roachpb.GetRequest).Key.Equal(errKey) {
+				if injected {
+					return nil
+				}
+				injected = true
+				txn := ba.Txn.Clone()
+				refreshTS = txn.Timestamp.Add(0, 1)
+				pErr := roachpb.NewReadWithinUncertaintyIntervalError(
+					txn.OrigTimestamp,
+					refreshTS,
+					&txn)
+				return roachpb.NewErrorWithTxn(pErr, &txn)
+			}
+			return nil
+		},
+	})
+	defer s.Stop()
+
+	err := s.DB.Txn(ctx, func(ctx context.Context, txn *client.Txn) error {
+		_, err := txn.Get(ctx, errKey)
+		if err != nil {
+			return err
+		}
+		if !injected {
+			return errors.Errorf("didn't inject err")
+		}
+		commitTS := txn.CommitTimestamp()
+		// We expect to have refreshed just after the timestamp injected by the error.
+		expTS := refreshTS.Add(0, 1)
+		if !commitTS.Equal(expTS) {
+			return errors.Errorf("expected refreshTS: %s, got: %s", refreshTS, commitTS)
+		}
+		return nil
+	})
+	require.NoError(t, err)
+}