txn_coord_sender_test.go

// Copyright 2014 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package kvcoord_test

import (
	"bytes"
	"context"
	"fmt"
	"reflect"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/pkg/base"
	"github.com/cockroachdb/cockroach/pkg/kv"
	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/isolation"
	"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities/tenantcapabilitiesauthorizer"
	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
	"github.com/cockroachdb/cockroach/pkg/storage"
	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
	"github.com/cockroachdb/cockroach/pkg/testutils"
	"github.com/cockroachdb/cockroach/pkg/testutils/kvclientutils"
	"github.com/cockroachdb/cockroach/pkg/testutils/localtestcluster"
	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
	"github.com/cockroachdb/cockroach/pkg/util/hlc"
	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
	"github.com/cockroachdb/cockroach/pkg/util/log"
	"github.com/cockroachdb/cockroach/pkg/util/metric"
	"github.com/cockroachdb/cockroach/pkg/util/stop"
	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
	"github.com/cockroachdb/errors"
	"github.com/stretchr/testify/require"
	"golang.org/x/sync/errgroup"
)

// createTestDB creates a local test server and starts it. The caller
// is responsible for stopping the test server.
func createTestDB(t testing.TB) *localtestcluster.LocalTestCluster {
	return createTestDBWithKnobs(t, nil)
}

func createTestDBWithKnobs(
	t testing.TB, knobs *kvserver.StoreTestingKnobs,
) *localtestcluster.LocalTestCluster {
	if knobs == nil {
		knobs = &kvserver.StoreTestingKnobs{}
	}
	if knobs.TenantRateKnobs.Authorizer == nil {
		knobs.TenantRateKnobs.Authorizer = tenantcapabilitiesauthorizer.NewAllowEverythingAuthorizer()
	}
	s := &localtestcluster.LocalTestCluster{
		StoreTestingKnobs: knobs,
	}
	s.Start(t, kvcoord.InitFactoryForLocalTestCluster)
	return s
}

// makeTS creates a new timestamp.
func makeTS(walltime int64, logical int32) hlc.Timestamp {
	return hlc.Timestamp{
		WallTime: walltime,
		Logical:  logical,
	}
}

// TestTxnCoordSenderAssignsTxnKey verifies that a transaction is assigned a
// transaction record key corresponding to the first key it writes.
func TestTxnCoordSenderAssignsTxnKey(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDB(t)
	defer s.Stop()
	ctx := context.Background()

	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	txn.TestingSetPriority(10)
	txn.SetDebugName("test txn")
	require.NoError(t, txn.SetIsoLevel(isolation.ReadCommitted))
	require.NoError(t, txn.Put(ctx, "key", []byte("value")))

	proto := txn.TestingCloneTxn()
	require.Equal(t, []byte("key"), proto.Key)
	require.Equal(t, isolation.ReadCommitted, proto.IsoLevel)
	require.Equal(t, enginepb.TxnPriority(10), proto.Priority)
	require.Equal(t, "test txn", proto.Name)
}

// TestTxnCoordSenderKeyRanges verifies that multiple requests to same or
// overlapping key ranges causes the coordinator to keep track only of
// the minimum number of ranges.
func TestTxnCoordSenderKeyRanges(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	ctx := context.Background()
	ranges := []struct {
		start, end roachpb.Key
	}{
		{roachpb.Key("a"), roachpb.Key(nil)},
		{roachpb.Key("a"), roachpb.Key(nil)},
		{roachpb.Key("aa"), roachpb.Key(nil)},
		{roachpb.Key("b"), roachpb.Key(nil)},
		{roachpb.Key("aa"), roachpb.Key("c")},
		{roachpb.Key("b"), roachpb.Key("c")},
	}

	s := createTestDB(t)
	defer s.Stop()

	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	// Disable txn pipelining so that all write spans are immediately
	// added to the transaction's lock footprint.
	if err := txn.DisablePipelining(); err != nil {
		t.Fatal(err)
	}
	tc := txn.Sender().(*kvcoord.TxnCoordSender)

	for _, rng := range ranges {
		if rng.end != nil {
			if _, err := txn.DelRange(ctx, rng.start, rng.end, false /* returnKeys */); err != nil {
				t.Fatal(err)
			}
		} else {
			if err := txn.Put(ctx, rng.start, []byte("value")); err != nil {
				t.Fatal(err)
			}
		}
	}

	// Verify that the transaction coordinator is only tracking two lock
	// spans. "a" and range "aa"-"c".
	lockSpans := tc.TestingGetLockFootprint(true)
	if len(lockSpans) != 2 {
		t.Errorf("expected 2 entries in keys range group; got %v", lockSpans)
	}
}

// Test that the heartbeat loop detects aborted transactions and stops.
func TestTxnCoordSenderHeartbeat(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDBWithKnobs(t, &kvserver.StoreTestingKnobs{
		DisableScanner:    true,
		DisableSplitQueue: true,
		DisableMergeQueue: true,
	})
	defer s.Stop()
	ctx := context.Background()

	keyA := roachpb.Key("a")
	keyC := roachpb.Key("c")
	splitKey := roachpb.Key("b")
	if err := s.DB.AdminSplit(
		ctx,
		splitKey,
		hlc.MaxTimestamp, /* expirationTime */
	); err != nil {
		t.Fatal(err)
	}

	ambient := s.AmbientCtx
	// Make a db with a short heartbeat interval.
	tsf := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			// Short heartbeat interval.
			HeartbeatInterval: time.Millisecond,
			Settings:          s.Cfg.Settings,
			Clock:             s.Clock,
			Stopper:           s.Stopper(),
		},
		kvcoord.NewDistSenderForLocalTestCluster(
			ctx,
			s.Cfg.Settings, &roachpb.NodeDescriptor{NodeID: 1},
			ambient.Tracer, s.Clock, s.Latency, s.Stores, s.Stopper(), s.Gossip,
		),
	)
	quickHeartbeatDB := kv.NewDB(ambient, tsf, s.Clock, s.Stopper())

	// We're going to test twice. In both cases the heartbeat is supposed to
	// notice that its transaction is aborted, but:
	// - once the abort span is populated on the txn's range.
	// - once the abort span is not populated.
	// The two conditions are created by either clearing an intent from the txn's
	// range or not (i.e. clearing an intent from another range).
	// The difference is supposed to be immaterial for the heartbeat loop (that's
	// what we're testing). As of June 2018, HeartbeatTxnRequests don't check the
	// abort span.
	for _, pusherKey := range []roachpb.Key{keyA, keyC} {
		t.Run(fmt.Sprintf("pusher:%s", pusherKey), func(t *testing.T) {
			// Make a db with a short heartbeat interval.
			initialTxn := kv.NewTxn(ctx, quickHeartbeatDB, 0 /* gatewayNodeID */)
			tc := initialTxn.Sender().(*kvcoord.TxnCoordSender)

			if err := initialTxn.Put(ctx, keyA, []byte("value")); err != nil {
				t.Fatal(err)
			}
			if err := initialTxn.Put(ctx, keyC, []byte("value")); err != nil {
				t.Fatal(err)
			}

			// Verify 3 heartbeats.
			var heartbeatTS hlc.Timestamp
			for i := 0; i < 3; i++ {
				testutils.SucceedsSoon(t, func() error {
					txn, pErr := getTxn(ctx, initialTxn)
					if pErr != nil {
						t.Fatal(pErr)
					}
					// Advance clock by 1ns.
					s.Manual.Advance(1)
					if lastActive := txn.LastActive(); heartbeatTS.Less(lastActive) {
						heartbeatTS = lastActive
						return nil
					}
					return errors.Errorf("expected heartbeat")
				})
			}

			// Push our txn with another high-priority txn.
			{
				if err := s.DB.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
					if err := txn.SetUserPriority(roachpb.MaxUserPriority); err != nil {
						return err
					}
					return txn.Put(ctx, pusherKey, []byte("pusher val"))
				}); err != nil {
					t.Fatal(err)
				}
			}

			// Verify that the abort is discovered and the heartbeat discontinued.
			// This relies on the heartbeat loop stopping once it figures out that the txn
			// has been aborted.
			testutils.SucceedsSoon(t, func() error {
				if tc.IsTracking() {
					return fmt.Errorf("transaction is not aborted")
				}
				return nil
			})

			// Trying to do something else should give us a TransactionAbortedError.
			_, err := initialTxn.Get(ctx, "a")
			assertTransactionAbortedError(t, err)
		})
	}
}

// Verify the txn sees a retryable error without using the handle: Normally the
// caller uses the txn handle directly, and if there is a retryable error then
// Send() fails and the handle gets "poisoned", meaning, the coordinator will be
// in state txnRetryableError instead of txnPending. But sometimes the handle
// can be idle and aborted by a heartbeat failure. This test verifies that in
// those cases the state of the handle ends up as txnRetryableError.
// This is important to verify because if the handle stays in txnPending then
// GetRetryableErr() returns nil, and PrepareForRetry() will not reset the
// handle.
func TestDB_PrepareForRetryAfterHeartbeatFailure(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	// Create a DB with a short heartbeat interval.
	s := createTestDB(t)
	defer s.Stop()
	ctx := context.Background()
	ambient := s.AmbientCtx
	tsf := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			// Short heartbeat interval.
			HeartbeatInterval: time.Millisecond,
			Settings:          s.Cfg.Settings,
			Clock:             s.Clock,
			Stopper:           s.Stopper(),
		},
		kvcoord.NewDistSenderForLocalTestCluster(
			ctx,
			s.Cfg.Settings, &roachpb.NodeDescriptor{NodeID: 1},
			ambient.Tracer, s.Clock, s.Latency, s.Stores, s.Stopper(), s.Gossip,
		),
	)
	db := kv.NewDB(ambient, tsf, s.Clock, s.Stopper())

	// Create a txn which will be aborted by a high priority txn.
	txn := kv.NewTxn(ctx, db, 0)

	// We first write to one range, then a high priority txn will abort our txn,
	// then we will try to read from another range until we see that the
	// transaction is poisoned because of a heartbeat failure.
	// Note that if we read from the same range then we will check the AbortSpan
	// and fail immediately (Send() will fail), even before the heartbeat failure,
	// which is not the case we want to test here.
	keyA := roachpb.Key("a")
	keyC := roachpb.Key("c")
	splitKey := roachpb.Key("b")
	require.NoError(t, s.DB.AdminSplit(
		ctx,
		splitKey,
		hlc.MaxTimestamp, /* expirationTime */
	))
	require.NoError(t, txn.Put(ctx, keyA, "1"))

	{
		// High priority txn - will abort the other txn.
		hpTxn := kv.NewTxn(ctx, db, 0)
		require.NoError(t, hpTxn.SetUserPriority(roachpb.MaxUserPriority))
		require.NoError(t, hpTxn.Put(ctx, keyA, "hp txn"))
		require.NoError(t, hpTxn.Commit(ctx))
	}

	tc := txn.Sender().(*kvcoord.TxnCoordSender)

	// Wait until we know that the handle was poisoned due to a heartbeat failure.
	testutils.SucceedsSoon(t, func() error {
		_, err := txn.Get(ctx, keyC)
		if err == nil {
			return errors.New("the handle is not poisoned yet")
		}
		require.IsType(t, &kvpb.TransactionRetryWithProtoRefreshError{}, err)
		return nil
	})

	// At this point the handle should be in state txnRetryableError - verify we
	// can read the error.
	pErr := tc.GetRetryableErr(ctx)
	require.NotNil(t, pErr)
	require.Equal(t, txn.ID(), pErr.PrevTxnID)
	// The transaction was aborted, therefore we should have a new transaction ID.
	require.NotEqual(t, pErr.PrevTxnID, pErr.NextTransaction.ID)
}

// getTxn fetches the requested key and returns the transaction info.
func getTxn(ctx context.Context, txn *kv.Txn) (*roachpb.Transaction, *kvpb.Error) {
	txnMeta := txn.TestingCloneTxn().TxnMeta
	qt := &kvpb.QueryTxnRequest{
		RequestHeader: kvpb.RequestHeader{
			Key: txnMeta.Key,
		},
		Txn: txnMeta,
	}

	ba := &kvpb.BatchRequest{}
	ba.Timestamp = txnMeta.WriteTimestamp
	ba.Add(qt)

	db := txn.DB()
	sender := db.NonTransactionalSender()

	br, pErr := sender.Send(ctx, ba)
	if pErr != nil {
		return nil, pErr
	}
	return &br.Responses[0].GetInner().(*kvpb.QueryTxnResponse).QueriedTxn, nil
}

func verifyCleanup(
	key roachpb.Key, eng storage.Engine, t *testing.T, coords ...*kvcoord.TxnCoordSender,
) {
	ctx := context.Background()
	testutils.SucceedsSoon(t, func() error {
		for _, coord := range coords {
			if coord.IsTracking() {
				return fmt.Errorf("expected no heartbeat")
			}
		}
		intentRes, err := storage.MVCCGet(ctx, eng, key, hlc.MaxTimestamp, storage.MVCCGetOptions{
			Inconsistent: true,
		})
		require.NoError(t, err)
		if intentRes.Intent != nil {
			return fmt.Errorf("found unexpected write intent: %s", intentRes.Intent)
		}
		return nil
	})
}

// TestTxnCoordSenderEndTxn verifies that ending a transaction
// sends resolve write intent requests.
func TestTxnCoordSenderEndTxn(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDB(t)
	defer s.Stop()
	ctx := context.Background()

	// 4 cases: no deadline, past deadline, equal deadline, future deadline.
	for i := 0; i < 4; i++ {
		key := roachpb.Key("key: " + strconv.Itoa(i))
		txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
		// Initialize the transaction.
		if pErr := txn.Put(ctx, key, []byte("value")); pErr != nil {
			t.Fatal(pErr)
		}
		// Conflicting transaction that pushes the above transaction.
		conflictTxn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
		conflictTxn.TestingSetPriority(enginepb.MaxTxnPriority)
		if _, pErr := conflictTxn.Get(ctx, key); pErr != nil {
			t.Fatal(pErr)
		}

		// The transaction was pushed at least to conflictTxn's timestamp (but
		// it could have been pushed more - the push takes a timestamp off the
		// HLC).
		pusheeTxn, pErr := getTxn(ctx, txn)
		if pErr != nil {
			t.Fatal(pErr)
		}
		pushedTimestamp := pusheeTxn.WriteTimestamp

		{
			var err error
			switch i {
			case 0:
				// No deadline.

			case 1:
				// Past deadline.
				err := txn.UpdateDeadline(ctx, pushedTimestamp.Prev())
				require.NoError(t, err, "Deadline update to past failed")
			case 2:
				// Equal deadline.
				err := txn.UpdateDeadline(ctx, pushedTimestamp)
				require.NoError(t, err, "Deadline update to equal failed")

			case 3:
				// Future deadline.
				err := txn.UpdateDeadline(ctx, pushedTimestamp.Next())
				require.NoError(t, err, "Deadline update to future failed")
			}
			if err = txn.Commit(ctx); err != nil {
				require.NoError(t, txn.Rollback(ctx))
			}

			switch i {
			case 0:
				// No deadline.
				if err != nil {
					t.Fatal(err)
				}

			case 1:
				// Past deadline.
				fallthrough
			case 2:
				// Equal deadline.
				assertTransactionRetryError(t, err)
				if !testutils.IsError(err, "RETRY_COMMIT_DEADLINE_EXCEEDED") {
					t.Fatalf("expected deadline exceeded, got: %s", err)
				}
			case 3:
				// Future deadline.
				if err != nil {
					t.Fatal(err)
				}
			}
		}
		verifyCleanup(key, s.Eng, t, txn.Sender().(*kvcoord.TxnCoordSender))
	}
}

// TestTxnCoordSenderCommitCanceled is a regression test for
// https://github.com/cockroachdb/cockroach/issues/68643. It makes sure that an
// EndTxn(commit=false) sent by the caller in response to a client context
// cancellation isn't passed through TxnCoordSender concurrently with an
// asynchronous EndTxn(commit=true) request sent by txnCommitter to make an
// implicitly committed transaction explicit.
func TestTxnCoordSenderCommitCanceled(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	ctx := context.Background()

	// blockCommits is used to block commit responses for a given txn. The key is
	// a txn ID, and the value is a ready channel (chan struct) that will be
	// closed when the commit has been received and blocked.
	var blockCommits sync.Map
	responseFilter := func(_ context.Context, ba *kvpb.BatchRequest, _ *kvpb.BatchResponse) *kvpb.Error {
		if arg, ok := ba.GetArg(kvpb.EndTxn); ok && ba.Txn != nil {
			et := arg.(*kvpb.EndTxnRequest)
			readyC, ok := blockCommits.Load(ba.Txn.ID)
			if ok && et.Commit && len(et.InFlightWrites) == 0 {
				close(readyC.(chan struct{})) // notify test that commit is received and blocked
				<-ctx.Done()                  // wait for test to complete (NB: not the passed context)
			}
		}
		return nil
	}

	s := createTestDBWithKnobs(t, &kvserver.StoreTestingKnobs{
		TestingResponseFilter: responseFilter,
	})
	defer s.Stop()
	ctx, _ = s.Stopper().WithCancelOnQuiesce(ctx)

	// Set up a new txn, and write a couple of values.
	txn := kv.NewTxn(ctx, s.DB, 0)
	require.NoError(t, txn.Put(ctx, "a", "1"))
	require.NoError(t, txn.Put(ctx, "b", "2"))

	// Read back a. This is crucial to reproduce the original bug. We need
	// txnPipeliner to record the lock in its lock footprint, but it doesn't do
	// that if the intents are proven together with the commit EndTxn request
	// (because it incorrectly assumes no further requests will be sent). If the
	// lock footprint isn't updated, the TxnCoordSender will incorrectly believe
	// the txn hasn't taken out any locks, and will elide the final
	// EndTxn(commit=false) rollback request. For details, see:
	// https://github.com/cockroachdb/cockroach/issues/68643
	_, err := txn.Get(ctx, "a")
	require.NoError(t, err)

	// Commit the transaction, but ask the response filter to block the final
	// async commit sent by txnCommitter to make the implicit commit explicit.
	readyC := make(chan struct{})
	blockCommits.Store(txn.ID(), readyC)
	require.NoError(t, txn.Commit(ctx))
	<-readyC

	// From the TxnCoordSender's point of view, the txn is implicitly committed,
	// and the commit response is on its way back up the stack. However, if the
	// client were to disconnect before receiving the response (canceling the
	// context), and something rolls back the transaction because of that, then
	// txn.Rollback() would send an asynchronous rollback request using a separate
	// context.
	//
	// However, this is hard to test since txn.Rollback() in this case sends the
	// EndTxn(commit=false) async. We instead replicate what Txn.Rollback() would
	// do here (i.e. send a EndTxn(commit=false)) and assert that we receive the
	// expected error.
	ba := &kvpb.BatchRequest{}
	ba.Add(&kvpb.EndTxnRequest{Commit: false})
	_, pErr := txn.Send(ctx, ba)
	require.NotNil(t, pErr)
	require.IsType(t, &kvpb.TransactionStatusError{}, pErr.GetDetail())
	txnErr := pErr.GetDetail().(*kvpb.TransactionStatusError)
	require.Equal(t, kvpb.TransactionStatusError_REASON_TXN_COMMITTED, txnErr.Reason)
}

// TestTxnCoordSenderAddLockOnError verifies that locks are tracked if the
// transaction is, even on error.
func TestTxnCoordSenderAddLockOnError(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDB(t)
	defer s.Stop()

	ctx := context.Background()

	// Create a transaction with intent at "x".
	key := roachpb.Key("x")
	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	tc := txn.Sender().(*kvcoord.TxnCoordSender)

	// Write so that the coordinator begins tracking this txn.
	if err := txn.Put(ctx, "x", "y"); err != nil {
		t.Fatal(err)
	}
	{
		err := txn.CPut(ctx, key, []byte("x"), kvclientutils.StrToCPutExistingValue("born to fail"))
		if !errors.HasType(err, (*kvpb.ConditionFailedError)(nil)) {
			t.Fatal(err)
		}
	}
	lockSpans := tc.TestingGetLockFootprint(true)
	expSpans := []roachpb.Span{{Key: key, EndKey: []byte("")}}
	equal := !reflect.DeepEqual(lockSpans, expSpans)
	if err := txn.Rollback(ctx); err != nil {
		t.Fatal(err)
	}
	if !equal {
		t.Fatalf("expected stored locks %v, got %v", expSpans, lockSpans)
	}
}

func assertTransactionRetryError(t *testing.T, e error) {
	t.Helper()
	if retErr := (*kvpb.TransactionRetryWithProtoRefreshError)(nil); errors.As(e, &retErr) {
		if !testutils.IsError(retErr, "TransactionRetryError") {
			t.Fatalf("expected the cause to be TransactionRetryError, but got %s",
				retErr)
		}
	} else {
		t.Fatalf("expected a retryable error, but got %s (%T)", e, e)
	}
}

func assertTransactionAbortedError(t *testing.T, e error) {
	if retErr := (*kvpb.TransactionRetryWithProtoRefreshError)(nil); errors.As(e, &retErr) {
		if !testutils.IsError(retErr, "TransactionAbortedError") {
			t.Fatalf("expected the cause to be TransactionAbortedError, but got %s",
				retErr)
		}
	} else {
		t.Fatalf("expected a retryable error, but got %s (%T)", e, e)
	}
}

// TestTxnCoordSenderCleanupOnAborted verifies that if a txn receives a
// TransactionAbortedError, the coordinator cleans up the transaction.
func TestTxnCoordSenderCleanupOnAborted(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDB(t)
	defer s.Stop()
	ctx := context.Background()

	// Create a transaction with intent at "a".
	key := roachpb.Key("a")
	txn1 := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	if err := txn1.Put(ctx, key, []byte("value")); err != nil {
		t.Fatal(err)
	}

	// Push the transaction (by writing key "a" with higher priority) to abort it.
	txn2 := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	if err := txn2.SetUserPriority(roachpb.MaxUserPriority); err != nil {
		t.Fatal(err)
	}
	if err := txn2.Put(ctx, key, []byte("value2")); err != nil {
		t.Fatal(err)
	}

	// Now end the transaction and verify we've cleanup up, even though
	// end transaction failed.
	err := txn1.Commit(ctx)
	assertTransactionAbortedError(t, err)
	require.NoError(t, txn2.Commit(ctx))
	verifyCleanup(key, s.Eng, t, txn1.Sender().(*kvcoord.TxnCoordSender), txn2.Sender().(*kvcoord.TxnCoordSender))
}

// TestTxnCoordSenderCleanupOnCommitAfterRestart verifies that if a txn restarts
// at a higher epoch and then commits before it has acquired any locks in the new
// epoch, the coordinator still cleans up the transaction. In #40466, we saw that
// this case could be detected as a 1PC transaction and the cleanup during the
// commit could be omitted.
func TestTxnCoordSenderCleanupOnCommitAfterRestart(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDB(t)
	defer s.Stop()
	ctx := context.Background()

	// Create a transaction with intent at "a".
	key := roachpb.Key("a")
	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	if err := txn.Put(ctx, key, []byte("value")); err != nil {
		t.Fatal(err)
	}

	// Restart the transaction with a new epoch.
	require.Error(t, txn.Sender().GenerateForcedRetryableErr(ctx, s.Clock.Now(), true /* mustRestart */, "force retry"))
	require.NoError(t, txn.Sender().ClearRetryableErr(ctx))

	// Now immediately commit.
	require.NoError(t, txn.Commit(ctx))
	verifyCleanup(key, s.Eng, t, txn.Sender().(*kvcoord.TxnCoordSender))
}

// TestTxnCoordSenderGCWithAmbiguousResultErr verifies that the coordinator
// cleans up extant transactions and locks after an ambiguous result error is
// observed, even if the error is on the first request.
func TestTxnCoordSenderGCWithAmbiguousResultErr(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	testutils.RunTrueAndFalse(t, "errOnFirst", func(t *testing.T, errOnFirst bool) {
		key := roachpb.Key("a")
		are := kvpb.NewAmbiguousResultErrorf("very ambiguous")
		knobs := &kvserver.StoreTestingKnobs{
			TestingResponseFilter: func(ctx context.Context, ba *kvpb.BatchRequest, br *kvpb.BatchResponse) *kvpb.Error {
				for _, req := range ba.Requests {
					if putReq, ok := req.GetInner().(*kvpb.PutRequest); ok && putReq.Key.Equal(key) {
						return kvpb.NewError(are)
					}
				}
				return nil
			},
		}

		s := createTestDBWithKnobs(t, knobs)
		defer s.Stop()

		ctx := context.Background()
		txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
		tc := txn.Sender().(*kvcoord.TxnCoordSender)
		if !errOnFirst {
			otherKey := roachpb.Key("other")
			if err := txn.Put(ctx, otherKey, []byte("value")); err != nil {
				t.Fatal(err)
			}
		}
		if err := txn.Put(ctx, key, []byte("value")); !testutils.IsError(err, "result is ambiguous") {
			t.Fatalf("expected error %v, found %v", are, err)
		}

		if err := txn.Rollback(ctx); err != nil {
			t.Fatal(err)
		}

		testutils.SucceedsSoon(t, func() error {
			// Locking the TxnCoordSender to prevent a data race.
			if tc.IsTracking() {
				return errors.Errorf("expected garbage collection")
			}
			return nil
		})

		verifyCleanup(key, s.Eng, t, tc)
	})
}

// TestTxnCoordSenderTxnUpdatedOnError verifies that errors adjust the
// response transaction's timestamp and priority as appropriate.
func TestTxnCoordSenderTxnUpdatedOnError(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	isolation.RunEachLevel(t, testTxnCoordSenderTxnUpdatedOnError)
}

func testTxnCoordSenderTxnUpdatedOnError(t *testing.T, isoLevel isolation.Level) {
	ctx := context.Background()
	origTS := makeTS(123, 0)
	plus10 := origTS.Add(10, 10)
	plus20 := origTS.Add(20, 0)
	testCases := []struct {
		// The test's name.
		name                  string
		pErrGen               func(txn *roachpb.Transaction) *kvpb.Error
		callPrepareForRetry   bool
		expEpoch              enginepb.TxnEpoch
		expPri                enginepb.TxnPriority
		expWriteTS, expReadTS hlc.Timestamp
		// Is set, we're expecting that the Transaction proto is re-initialized (as
		// opposed to just having the epoch incremented).
		expNewTransaction bool
	}{
		{
			// No error, so nothing interesting either.
			name:       "nil",
			pErrGen:    func(_ *roachpb.Transaction) *kvpb.Error { return nil },
			expEpoch:   0,
			expPri:     1,
			expWriteTS: origTS,
			expReadTS:  origTS,
		},
		{
			// On uncertainty error, the read timestamp moves ahead of the existing
			// write and LocalUncertaintyLimit, if one exists. Also, a new epoch
			// begins for isolation levels that cannot move their read timestamp
			// between operations.
			name: "ReadWithinUncertaintyIntervalError without LocalUncertaintyLimit",
			pErrGen: func(txn *roachpb.Transaction) *kvpb.Error {
				pErr := kvpb.NewErrorWithTxn(
					kvpb.NewReadWithinUncertaintyIntervalError(
						origTS,               // readTS
						hlc.ClockTimestamp{}, // localUncertaintyLimit
						txn,                  // txn
						plus10,               // valueTS
						hlc.ClockTimestamp{}, // localTS
					),
					txn)
				return pErr
			},
			expEpoch: func() enginepb.TxnEpoch {
				if isoLevel == isolation.ReadCommitted {
					return 0 // PerStatementReadSnapshot
				}
				return 1
			}(),
			expPri:     1,
			expWriteTS: plus10.Next(),
			expReadTS:  plus10.Next(),
		},
		{
			// On uncertainty error, the read timestamp moves ahead of the existing
			// write and LocalUncertaintyLimit, if one exists. Also, a new epoch
			// begins for isolation levels that cannot move their read timestamp
			// between operations.
			name: "ReadWithinUncertaintyIntervalError with LocalUncertaintyLimit",
			pErrGen: func(txn *roachpb.Transaction) *kvpb.Error {
				pErr := kvpb.NewErrorWithTxn(
					kvpb.NewReadWithinUncertaintyIntervalError(
						origTS,                     // readTS
						hlc.ClockTimestamp(plus20), // localUncertaintyLimit
						txn,                        // txn
						plus10,                     // valueTS
						hlc.ClockTimestamp{},       // localTS
					),
					txn)
				return pErr
			},
			expEpoch: func() enginepb.TxnEpoch {
				if isoLevel == isolation.ReadCommitted {
					return 0 // PerStatementReadSnapshot
				}
				return 1
			}(),
			expPri:     1,
			expWriteTS: plus20,
			expReadTS:  plus20,
		},
		{
			// On abort, nothing changes - we are left with a poisoned txn (unless we
			// call PrepareForRetry as in the next test case).
			name: "TransactionAbortedError",
			pErrGen: func(txn *roachpb.Transaction) *kvpb.Error {
				txn.WriteTimestamp = plus20
				txn.Priority = 10
				return kvpb.NewErrorWithTxn(&kvpb.TransactionAbortedError{}, txn)
			},
			expPri:     1,
			expWriteTS: origTS,
			expReadTS:  origTS,
		},
		{
			// On abort, reset the txn by calling PrepareForRetry, and then we get a
			// new priority to use for the next attempt.
			name: "TransactionAbortedError with PrepareForRetry",
			pErrGen: func(txn *roachpb.Transaction) *kvpb.Error {
				txn.WriteTimestamp = plus20
				txn.Priority = 10
				return kvpb.NewErrorWithTxn(&kvpb.TransactionAbortedError{}, txn)
			},
			callPrepareForRetry: true,
			expNewTransaction:   true,
			expPri:              10,
			expWriteTS:          plus20,
			expReadTS:           plus20,
		},
		{
			// On failed push, read timestamp advances just past the pushed timestamp.
			// Additionally, priority ratchets up to just below the pusher's. Finally,
			// a new epoch begins for isolation levels that cannot move their read
			// timestamp between operations.
			name: "TransactionPushError",
			pErrGen: func(txn *roachpb.Transaction) *kvpb.Error {
				return kvpb.NewErrorWithTxn(&kvpb.TransactionPushError{
					PusheeTxn: roachpb.Transaction{
						TxnMeta: enginepb.TxnMeta{WriteTimestamp: plus10, Priority: 10},
					},
				}, txn)
			},
			expEpoch: func() enginepb.TxnEpoch {
				if isoLevel == isolation.ReadCommitted {
					return 0
				}
				return 1
			}(),
			expPri:     9,
			expWriteTS: plus10,
			expReadTS:  plus10,
		},
		{
			// On retry, advance timestamp and priority. Also, restart with new epoch
			// for isolation levels that cannot move their read timestamp between
			// operations.
			name: "TransactionRetryError",
			pErrGen: func(txn *roachpb.Transaction) *kvpb.Error {
				txn.WriteTimestamp = plus10
				txn.Priority = 10
				return kvpb.NewErrorWithTxn(&kvpb.TransactionRetryError{}, txn)
			},
			expEpoch: func() enginepb.TxnEpoch {
				if isoLevel == isolation.ReadCommitted {
					return 0
				}
				return 1
			}(),
			expPri:     10,
			expWriteTS: plus10,
			expReadTS:  plus10,
		},
		{
			// On retry, advance timestamp and priority. Also, restart with new epoch,
			// even for isolation levels that can move their read timestamp between
			// operations. This is because PrepareForRetry is called instead of
			// PrepareForPartialRetry, indicating that the client wants to promote
			// what could be a partial retry to a full retry.
			name: "TransactionRetryError with PrepareForRetry",
			pErrGen: func(txn *roachpb.Transaction) *kvpb.Error {
				txn.WriteTimestamp = plus10
				txn.Priority = 10
				return kvpb.NewErrorWithTxn(&kvpb.TransactionRetryError{}, txn)
			},
			expEpoch:            1,
			callPrepareForRetry: true,
			expPri:              10,
			expWriteTS:          plus10,
			expReadTS:           plus10,
		},
	}

	for _, test := range testCases {
		t.Run(test.name, func(t *testing.T) {
			stopper := stop.NewStopper()

			manual := timeutil.NewManualTime(origTS.GoTime())
			clock := hlc.NewClockForTesting(manual)

			var senderFn kv.SenderFunc = func(
				_ context.Context, ba *kvpb.BatchRequest,
			) (*kvpb.BatchResponse, *kvpb.Error) {
				var reply *kvpb.BatchResponse
				pErr := test.pErrGen(ba.Txn)
				if pErr == nil {
					reply = ba.CreateReply()
					reply.Txn = ba.Txn
				} else if txn := pErr.GetTxn(); txn != nil {
					// Update the manual clock to simulate an
					// error updating a local hlc clock.
					manual.AdvanceTo(txn.WriteTimestamp.GoTime())
				}
				return reply, pErr
			}
			ambient := log.MakeTestingAmbientCtxWithNewTracer()
			tsf := kvcoord.NewTxnCoordSenderFactory(
				kvcoord.TxnCoordSenderFactoryConfig{
					AmbientCtx: ambient,
					Clock:      clock,
					Stopper:    stopper,
				},
				senderFn,
			)
			db := kv.NewDB(ambient, tsf, clock, stopper)
			key := roachpb.Key("test-key")
			now := clock.NowAsClockTimestamp()
			origTxnProto := roachpb.MakeTransaction("test txn", key, isoLevel, roachpb.UserPriority(0), now.ToTimestamp(), clock.MaxOffset().Nanoseconds(), 0, 0, false /* omitInRangefeeds */)
			// TODO(andrei): I've monkeyed with the priorities on this initial
			// Transaction to keep the test happy from a previous version in which the
			// Transaction was not initialized before use (which became insufficient
			// when we started testing that TransactionAbortedError's properly
			// re-initializes the proto), but this deserves cleanup. I think this test
			// is strict in what updated priorities it expects and also our mechanism
			// for assigning exact priorities doesn't work properly when faced with
			// updates.
			origTxnProto.Priority = 1
			txn := kv.NewTxnFromProto(ctx, db, 0 /* gatewayNodeID */, now, kv.RootTxn, &origTxnProto)
			txn.TestingSetPriority(1)
			// Enable stepping so that Read Committed transactions don't advance their
			// read timestamp on each operation.
			txn.ConfigureStepping(ctx, kv.SteppingEnabled)

			err := txn.Put(ctx, key, []byte("value"))
			stopper.Stop(ctx)

			if test.callPrepareForRetry {
				if err := txn.PrepareForRetry(ctx); err != nil {
					t.Fatal(err)
				}
			}
			if test.name != "nil" && err == nil {
				t.Fatalf("expected an error")
			}
			proto := txn.TestingCloneTxn()
			txnReset := origTxnProto.ID != proto.ID
			if txnReset != test.expNewTransaction {
				t.Fatalf("expected txn reset: %t and got: %t", test.expNewTransaction, txnReset)
			}
			if proto.Epoch != test.expEpoch {
				t.Errorf("expected epoch = %d; got %d",
					test.expEpoch, proto.Epoch)
			}
			if proto.Priority != test.expPri {
				t.Errorf("expected priority = %d; got %d",
					test.expPri, proto.Priority)
			}
			if proto.WriteTimestamp != test.expWriteTS {
				t.Errorf("expected timestamp to be %s; got %s",
					test.expWriteTS, proto.WriteTimestamp)
			}
			if proto.ReadTimestamp != test.expReadTS {
				t.Errorf("expected orig timestamp to be %s; got %s",
					test.expReadTS, proto.ReadTimestamp)
			}
		})
	}
}

// TestWTOBitTerminatedOnErrorResponses is a regression test for #85711. It
// ensures that when batch request errors have the WTO bit set, subsequent
// request don't carry that bit (something that's asserted on in client-side
// interceptors).
func TestWTOBitTerminatedOnErrorResponses(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()

	keyA := roachpb.Key("a")
	keyB := roachpb.Key("b")

	s, _, db := serverutils.StartServer(t, base.TestServerArgs{})
	defer s.Stopper().Stop(ctx)

	// Split to ensure batch requests get split through the distsender.
	require.NoError(t, db.AdminSplit(
		ctx,
		keyB,             /* splitKey */
		hlc.MaxTimestamp, /* expirationTime */
	))

	// Write a key that the txn-al CPut below doesn't expect, failing the batch
	// request.
	require.NoError(t, db.Put(ctx, keyB, []byte("b-unexpected")))

	txn := db.NewTxn(ctx, "root")

	// Read a key to pin the read timestamp. We'll use it to trigger to WTO
	// error below.
	b0 := txn.NewBatch()
	b0.Get(keyB)
	require.NoError(t, txn.Run(ctx, b0))

	// Write to keyA out-of-band to induce a WTO condition in the txn-al Put
	// below.
	require.NoError(t, db.Put(ctx, keyA, "a-unexpected"))

	// Send a batch (as part of a leaf txn) that will be split into two
	// sub-batches: [Put(a), CPut(b, nil)]. Put(a) should observe the WTO bit
	// set in the batch response, whereas the CPut induces an error response.
	// Since these are two separate requests, the error response is combined
	// with the batch request such that the error itself has the WTO bit set. In
	// #85711 we observed that the bit was not terminated on the client side,
	// and subsequent requests were issued with the WTO bit set, which tripped
	// up assertions.
	b1 := txn.NewBatch()
	b1.Put(keyA, "a")
	b1.CPut(keyB, "b", nil /* expValue */)
	require.True(t, testutils.IsError(txn.Run(ctx, b1), "unexpected value"))
	require.False(t, txn.TestingCloneTxn().WriteTooOld) // WTO bit is terminated

	b2 := txn.NewBatch()
	b2.Put(keyB, "b")
	require.NoError(t, txn.Run(ctx, b2))
	require.False(t, txn.TestingCloneTxn().WriteTooOld) // WTO bit is terminated
	require.NoError(t, txn.Commit(ctx))
}

// TestTxnMultipleCoord checks that multiple txn coordinators can be
// used for reads by a single transaction, and their state can be combined.
func TestTxnMultipleCoord(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDB(t)
	defer s.Stop()

	ctx := context.Background()
	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)

	// Start the transaction.
	key := roachpb.Key("a")
	if _, err := txn.Get(ctx, key); err != nil {
		t.Fatal(err)
	}

	// New create a second, leaf coordinator.
	leafInputState, err := txn.GetLeafTxnInputState(ctx)
	require.NoError(t, err)
	txn2 := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState)

	// Start the second transaction.
	key2 := roachpb.Key("b")
	if _, err := txn2.Get(ctx, key2); err != nil {
		t.Fatal(err)
	}

	// Augment txn with txn2's meta & commit.
	tfs, err := txn2.GetLeafTxnFinalState(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if err := txn.UpdateRootWithLeafFinalState(ctx, tfs); err != nil {
		t.Fatal(err)
	}

	// Verify presence of both locks.
	tcs := txn.Sender().(*kvcoord.TxnCoordSender)
	refreshSpans := tcs.TestingGetRefreshFootprint()
	require.Equal(t, []roachpb.Span{{Key: key}, {Key: key2}}, refreshSpans)

	ba := txn.NewBatch()
	ba.AddRawRequest(&kvpb.EndTxnRequest{Commit: true})
	if err := txn.Run(ctx, ba); err != nil {
		t.Fatal(err)
	}
}

// TestTxnCoordSenderNoDuplicateLockSpans verifies that TxnCoordSender does not
// generate duplicate lock spans and that it merges lock spans that have
// overlapping ranges.
func TestTxnCoordSenderNoDuplicateLockSpans(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	skip.IgnoreLint(t, "WIP")
	ctx := context.Background()
	stopper := stop.NewStopper()
	clock := hlc.NewClockForTesting(timeutil.NewManualTime(timeutil.Unix(0, 123)))

	var expectedLockSpans []roachpb.Span

	var senderFn kv.SenderFunc = func(_ context.Context, ba *kvpb.BatchRequest) (
		*kvpb.BatchResponse, *kvpb.Error) {
		br := ba.CreateReply()
		br.Txn = ba.Txn.Clone()
		if rArgs, ok := ba.GetArg(kvpb.EndTxn); ok {
			et := rArgs.(*kvpb.EndTxnRequest)
			if !reflect.DeepEqual(et.LockSpans, expectedLockSpans) {
				t.Errorf("Invalid lock spans: %+v; expected %+v", et.LockSpans, expectedLockSpans)
			}
			br.Txn.Status = roachpb.COMMITTED
		}
		return br, nil
	}
	ambient := log.MakeTestingAmbientCtxWithNewTracer()

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		senderFn,
	)
	defer stopper.Stop(ctx)

	db := kv.NewDB(ambient, factory, clock, stopper)
	txn := kv.NewTxn(ctx, db, 0 /* gatewayNodeID */)

	// Acquire locks on a-b, c, m, u-w before the final batch.
	_, pErr := txn.ScanForShare(
		ctx, roachpb.Key("a"), roachpb.Key("b"), 0, kvpb.GuaranteedDurability,
	)
	if pErr != nil {
		t.Fatal(pErr)
	}
	pErr = txn.Put(ctx, roachpb.Key("c"), []byte("value"))
	if pErr != nil {
		t.Fatal(pErr)
	}
	_, pErr = txn.GetForUpdate(ctx, roachpb.Key("m"), kvpb.GuaranteedDurability)
	if pErr != nil {
		t.Fatal(pErr)
	}
	_, pErr = txn.DelRange(ctx, roachpb.Key("u"), roachpb.Key("w"), false /* returnKeys */)
	if pErr != nil {
		t.Fatal(pErr)
	}

	// The final batch overwrites key c, reads key n, and overlaps part of the a-b and u-w ranges.
	b := txn.NewBatch()
	b.Put(roachpb.Key("b"), []byte("value"))
	b.Put(roachpb.Key("c"), []byte("value"))
	b.Put(roachpb.Key("d"), []byte("value"))
	b.GetForUpdate(roachpb.Key("n"), kvpb.GuaranteedDurability)
	b.ReverseScanForShare(roachpb.Key("v"), roachpb.Key("z"), kvpb.GuaranteedDurability)

	// The expected locks are a-b, c, m, n, and v-z.
	//
	// A note about the v-z span -- because the DeleteRange request did not
	// actually delete any keys, we'll not track anything for it in the lock
	// footprint for this transaction. The v-z range comes from the
	// ReverseScanForShare request in the final batch.
	expectedLockSpans = []roachpb.Span{
		{Key: roachpb.Key("a"), EndKey: roachpb.Key("b").Next()},
		{Key: roachpb.Key("c"), EndKey: nil},
		{Key: roachpb.Key("d"), EndKey: nil},
		{Key: roachpb.Key("m"), EndKey: nil},
		{Key: roachpb.Key("n"), EndKey: nil},
		{Key: roachpb.Key("v"), EndKey: roachpb.Key("z")},
	}

	pErr = txn.CommitInBatch(ctx, b)
	if pErr != nil {
		t.Fatal(pErr)
	}
}

// checkTxnMetrics verifies that the provided Sender's transaction metrics match the expected
// values. This is done through a series of retries with increasing backoffs, to work around
// the TxnCoordSender's asynchronous updating of metrics after a transaction ends.
func checkTxnMetrics(
	t *testing.T,
	metrics kvcoord.TxnMetrics,
	name string,
	commits, commits1PC, commitsReadOnly, aborts, restarts int64,
) {
	testutils.SucceedsSoon(t, func() error {
		return checkTxnMetricsOnce(metrics, name, commits, commits1PC, commitsReadOnly, aborts, restarts)
	})
}

func checkTxnMetricsOnce(
	metrics kvcoord.TxnMetrics,
	name string,
	commits, commits1PC, commitReadOnly, aborts, restarts int64,
) error {
	durationCounts, _ := metrics.Durations.CumulativeSnapshot().Total()
	restartsCounts, _ := metrics.Restarts.CumulativeSnapshot().Total()
	testcases := []struct {
		name string
		a, e int64
	}{
		{"commits", metrics.Commits.Count(), commits},
		{"commits1PC", metrics.Commits1PC.Count(), commits1PC},
		{"commitsReadOnly", metrics.CommitsReadOnly.Count(), commitReadOnly},
		{"aborts", metrics.Aborts.Count(), aborts},
		{"durations", durationCounts, commits + aborts},
		{"restarts", restartsCounts, restarts},
	}

	for _, tc := range testcases {
		if tc.a != tc.e {
			return errors.Errorf("%s: actual %s %d != expected %d", name, tc.name, tc.a, tc.e)
		}
	}

	return nil
}

// setupMetricsTest sets the txn coord sender factory's metrics to
// have a faster sample interval and returns a cleanup function to be
// executed by callers.
func setupMetricsTest(
	t *testing.T,
) (*localtestcluster.LocalTestCluster, kvcoord.TxnMetrics, func()) {
	s := &localtestcluster.LocalTestCluster{
		// Liveness heartbeat txns mess up the metrics.
		DisableLivenessHeartbeat: true,
		DontCreateSystemRanges:   true,
	}
	s.Start(t, kvcoord.InitFactoryForLocalTestCluster)

	metrics := kvcoord.MakeTxnMetrics(metric.TestSampleInterval)
	s.DB.GetFactory().(*kvcoord.TxnCoordSenderFactory).TestingSetMetrics(metrics)
	return s, metrics, s.Stop
}

// Test a normal transaction. This and the other metrics tests below use real KV operations,
// because it took far too much mucking with TxnCoordSender internals to mock out the sender
// function as other tests do.
func TestTxnCommit(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s, metrics, cleanupFn := setupMetricsTest(t)
	defer cleanupFn()
	value := []byte("value")

	// Test a write txn commit.
	if err := s.DB.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
		key := []byte("key-commit")
		return txn.Put(ctx, key, value)
	}); err != nil {
		t.Fatal(err)
	}
	checkTxnMetrics(t, metrics, "commit txn", 1 /* commits */, 0 /* commits1PC */, 0 /* commitsReadOnly */, 0, 0)

	// Test a read-only get txn.
	if err := s.DB.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
		key := []byte("key-commit")
		_, err := txn.Get(ctx, key)
		return err
	}); err != nil {
		t.Fatal(err)
	}

	checkTxnMetrics(t, metrics, "commit txn", 2 /* commits */, 0 /* commits1PC */, 1 /* commitsReadOnly */, 0, 0)

	// Test a read-only scan txn.
	if err := s.DB.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
		key := []byte("key-commit")
		endKey := []byte("key-commit2")
		_, err := txn.Scan(ctx, key, endKey, 5)
		return err
	}); err != nil {
		t.Fatal(err)
	}

	checkTxnMetrics(t, metrics, "commit txn", 3 /* commits */, 0 /* commits1PC */, 2 /* commitsReadOnly */, 0, 0)
}

// TestTxnOnePhaseCommit verifies that 1PC metric tracking works.
func TestTxnOnePhaseCommit(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s, metrics, cleanupFn := setupMetricsTest(t)
	defer cleanupFn()

	value := []byte("value")

	ctx := context.Background()
	if err := s.DB.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
		key := []byte("key-commit")
		b := txn.NewBatch()
		b.Put(key, value)
		return txn.CommitInBatch(ctx, b)
	}); err != nil {
		t.Fatal(err)
	}
	kv, err := s.DB.Get(ctx, []byte("key-commit"))
	if err != nil {
		t.Fatal(err)
	}
	if kv.Value == nil {
		t.Fatal("expected value not found")
	}
	val, err := kv.Value.GetBytes()
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(val, value) {
		t.Fatalf("expected: %s, got: %s", value, val)
	}
	checkTxnMetrics(t, metrics, "commit 1PC txn", 1 /* commits */, 1 /* 1PC */, 0, 0, 0)
}

func TestTxnAbortCount(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	s, metrics, cleanupFn := setupMetricsTest(t)
	defer cleanupFn()
	const intentionalErrText = "intentional error to cause abort"

	// Test aborted read-write transaction.
	if err := s.DB.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
		err := txn.Put(ctx, "key", "value")
		require.NoError(t, err)

		return errors.New(intentionalErrText)
	}); !testutils.IsError(err, intentionalErrText) {
		t.Fatalf("unexpected error: %v", err)
	}
	checkTxnMetrics(t, metrics, "abort txn", 0, 0, 0, 1 /* aborts */, 0)

	// Test aborted read-only transaction.
	if err := s.DB.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
		_, err := txn.Get(ctx, "key")
		require.NoError(t, err)

		return errors.New(intentionalErrText)
	}); !testutils.IsError(err, intentionalErrText) {
		t.Fatalf("unexpected error: %v", err)
	}
	checkTxnMetrics(t, metrics, "abort txn", 0, 0, 0, 2 /* aborts */, 0)
}

func TestTxnRestartCount(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	readKey := []byte("read")
	writeKey := []byte("write")
	value := []byte("value")
	ctx := context.Background()

	s, metrics, cleanupFn := setupMetricsTest(t)
	defer cleanupFn()

	// Start a transaction and read a key that we're going to modify outside the
	// txn. This ensures that refreshing the txn will not succeed, so a restart
	// will be necessary.
	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	if _, err := txn.Get(ctx, readKey); err != nil {
		t.Fatal(err)
	}

	// Write the read key outside of the transaction, at a higher timestamp, which
	// will necessitate a txn restart when the original read key span is updated.
	if err := s.DB.Put(ctx, readKey, value); err != nil {
		t.Fatal(err)
	}

	// Outside of the transaction, read the same key as will be
	// written within the transaction. This means that future
	// attempts to write will forward the txn timestamp.
	if _, err := s.DB.Get(ctx, writeKey); err != nil {
		t.Fatal(err)
	}

	// This put will lay down an intent, txn write timestamp will increase beyond
	// the read timestamp.
	if err := txn.Put(ctx, writeKey, value); err != nil {
		t.Fatal(err)
	}
	proto := txn.TestingCloneTxn()
	if proto.WriteTimestamp.LessEq(proto.ReadTimestamp) {
		t.Errorf("expected timestamp to increase: %s", proto)
	}

	// Wait for heartbeat to start.
	tc := txn.Sender().(*kvcoord.TxnCoordSender)
	testutils.SucceedsSoon(t, func() error {
		if !tc.IsTracking() {
			return errors.New("expected heartbeat to start")
		}
		return nil
	})

	// Commit (should cause restart metric to increase).
	err := txn.Commit(ctx)
	if err != nil {
		require.NoError(t, txn.Rollback(ctx))
	}
	assertTransactionRetryError(t, err)
	checkTxnMetrics(t, metrics, "restart txn", 0, 0, 0, 1 /* aborts */, 1 /* restarts */)
}

func TestTxnDurations(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s, metrics, cleanupFn := setupMetricsTest(t)
	manual := s.Manual
	defer cleanupFn()
	const puts = 10

	const incr time.Duration = 1000
	for i := 0; i < puts; i++ {
		key := roachpb.Key(fmt.Sprintf("key-txn-durations-%d", i))
		if err := s.DB.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
			if err := txn.Put(ctx, key, []byte("val")); err != nil {
				return err
			}
			manual.Advance(incr)
			return nil
		}); err != nil {
			t.Fatal(err)
		}
	}

	checkTxnMetrics(t, metrics, "txn durations", puts, 0, 0, 0, 0)

	hist := metrics.Durations
	// The clock is a bit odd in these tests, so I can't test the mean without
	// introducing spurious errors or being overly lax.
	//
	// TODO(cdo): look into cause of variance.
	count, _ := hist.CumulativeSnapshot().Total()
	if a, e := count, int64(puts); a != e {
		t.Fatalf("durations %d != expected %d", a, e)
	}

	for _, b := range hist.ToPrometheusMetric().GetHistogram().GetBucket() {
		thresh := incr.Nanoseconds()
		if *b.UpperBound < float64(thresh) && *b.CumulativeCount != 0 {
			t.Fatalf("expected no values in bucket: %f", *b.UpperBound)
		}
	}

}

// TestTxnCommitWait tests the commit-wait sleep phase of transactions under
// various conditions. It verifies that transactions sleep for the correct
// amount of time and verifies that metrics are correctly updated to reflect
// these commit-wait sleeps. For more, see TxnCoordSender.maybeCommitWait.
func TestTxnCommitWait(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()

	// Testing parameters:
	//
	// - linearizable: is linearizable mode enabled?
	// - commit:       does the transaction commit or rollback?
	// - readOnly:     does the transaction perform any writes?
	// - futureTime:   does the transaction commit in the future?
	// - deferred:     does the caller assume responsibility for commit waiting?
	//
	testFn := func(t *testing.T, linearizable, commit, readOnly, futureTime, deferred bool) {
		s, metrics, cleanupFn := setupMetricsTest(t)
		defer cleanupFn()
		s.DB.GetFactory().(*kvcoord.TxnCoordSenderFactory).TestingSetLinearizable(linearizable)
		commitWaitC := make(chan struct{})
		s.DB.GetFactory().(*kvcoord.TxnCoordSenderFactory).TestingSetCommitWaitFilter(func() {
			close(commitWaitC)
		})

		// maxClockOffset defines the maximum clock offset between nodes in the
		// cluster. When in linearizable mode, all writing transactions must
		// wait this additional duration after committing to ensure that their
		// commit timestamp is below the current HLC clock time of any other
		// node in the system. In doing so, all causally dependent transactions
		// are guaranteed to start with higher timestamps, regardless of the
		// gateway they use. This ensures that all causally dependent
		// transactions commit with higher timestamps, even if their read and
		// writes sets do not conflict with the original transaction's. This, in
		// turn, prevents the "causal reverse" anomaly which can be observed by
		// a third, concurrent transaction. See the following blog post for
		// more: https://www.cockroachlabs.com/blog/consistency-model/.
		maxClockOffset := s.Clock.MaxOffset()
		// futureOffset defines how far in the future the test transaction will
		// commit, if futureTime is true. We set it to less than the maximum
		// clock offset as a convenience, so that a value this far in the future
		// falls within the uncertainty interval of a read, which allows the
		// test to cause a read-only transaction to commit with a future
		// timestamp, if readOnly is true.
		futureOffset := maxClockOffset / 2

		txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
		readyC := make(chan struct{})
		deferredWaitC := make(chan struct{})
		errC := make(chan error, 1)
		go func() {
			var commitWaitFn func(context.Context) error
			if deferred {
				// If the test wants the caller to assume responsibility for commit
				// waiting, we expect the transaction to return immediately after a
				// commit. However, the commit-wait function returned here will then
				// block for the duration of the commit-wait.
				commitWaitFn = txn.DeferCommitWait(ctx)
			}

			// Run the transaction.
			err := func() error {
				key := roachpb.Key("a")

				// If we want the transaction to commit in the future, we
				// perform a write in its near future. If the transaction is
				// read-only, it will restart due to uncertainty. If the
				// transaction is read-write, it will need to bump its write
				// timestamp above the other value.
				if futureTime {
					ts := txn.TestingCloneTxn().WriteTimestamp.Add(futureOffset.Nanoseconds(), 0)
					h := kvpb.Header{Timestamp: ts}
					put := kvpb.NewPut(key, roachpb.Value{})
					if _, pErr := kv.SendWrappedWith(ctx, s.DB.NonTransactionalSender(), h, put); pErr != nil {
						return pErr.GoError()
					}
				}

				if readOnly {
					if _, err := txn.Get(ctx, key); err != nil {
						return err
					}
				} else {
					if err := txn.Put(ctx, key, "val"); err != nil {
						return err
					}
				}

				close(readyC)
				if !commit {
					return txn.Rollback(ctx)
				}
				return txn.Commit(ctx)
			}()

			if commitWaitFn != nil {
				close(deferredWaitC)
				_ = commitWaitFn(ctx) // NOTE: blocks
			}

			errC <- err
		}()

		// Wait until the transaction is about to commit / rollback.
		<-readyC

		// If the test deferred the commit wait, the transaction commit should
		// return immediately, without waiting. However, the caller will then
		// wait when it calls the commit-wait function.
		if deferred {
			<-deferredWaitC
		}

		if !commit {
			// If the transaction rolled back, it should immediately return,
			// without waiting.
			require.NoError(t, <-errC)
			require.Equal(t, int64(0), metrics.CommitWaits.Count())
			return
		}

		// If the transaction committed, it will need to wait at least until its
		// commit timestamp is below the HLC clock. If linearizable mode is
		// enabled and this transaction wrote then it will need to wait an
		// additional max_offset.
		expWait := time.Duration(0)
		if linearizable && !readOnly {
			expWait += maxClockOffset
		}
		if futureTime {
			expWait += futureOffset
		}
		expMetric := int64(0)
		if expWait > 0 {
			expMetric = 1
		}

		// Advance the manual clock slowly. If the commit-wait sleep completes
		// too early, we'll catch it with the require.Empty. If it completes too
		// late, we'll stall when pulling from the channel. Before doing so, wait
		// until the transaction has begun its commit-wait.
		for expWait > 0 {
			<-commitWaitC
			require.Empty(t, errC)

			adv := futureOffset / 5
			expWait -= adv
			s.Manual.Advance(adv)
		}
		require.NoError(t, <-errC)
		require.Equal(t, expMetric, metrics.CommitWaits.Count())

		// Validate the transaction's commit timestamp in relation to the local
		// HLC clock.
		minClockTS := txn.TestingCloneTxn().WriteTimestamp
		if linearizable && !readOnly {
			minClockTS = minClockTS.Add(maxClockOffset.Nanoseconds(), 0)
		}
		require.True(t, minClockTS.Less(s.Clock.Now()))
	}
	testutils.RunTrueAndFalse(t, "linearizable", func(t *testing.T, linearizable bool) {
		testutils.RunTrueAndFalse(t, "commit", func(t *testing.T, commit bool) {
			testutils.RunTrueAndFalse(t, "readOnly", func(t *testing.T, readOnly bool) {
				testutils.RunTrueAndFalse(t, "futureTime", func(t *testing.T, futureTime bool) {
					testutils.RunTrueAndFalse(t, "deferred", func(t *testing.T, deferred bool) {
						testFn(t, linearizable, commit, readOnly, futureTime, deferred)
					})
				})
			})
		})
	})
}

// TestAbortTransactionOnCommitErrors verifies that transactions are
// aborted on the correct errors.
func TestAbortTransactionOnCommitErrors(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)

	testCases := []struct {
		err        error
		errFn      func(roachpb.Transaction) *kvpb.Error
		asyncAbort bool
	}{
		{
			errFn: func(txn roachpb.Transaction) *kvpb.Error {
				const nodeID = 0
				// ReadWithinUncertaintyIntervalErrors need a clock to have been
				// recorded on the origin.
				txn.UpdateObservedTimestamp(nodeID, makeTS(123, 0).UnsafeToClockTimestamp())
				return kvpb.NewErrorWithTxn(
					kvpb.NewReadWithinUncertaintyIntervalError(
						hlc.Timestamp{}, hlc.ClockTimestamp{}, nil, hlc.Timestamp{}, hlc.ClockTimestamp{}),
					&txn)
			},
			asyncAbort: false},
		{err: &kvpb.TransactionAbortedError{}, asyncAbort: true},
		{err: &kvpb.TransactionPushError{}, asyncAbort: false},
		{err: &kvpb.TransactionRetryError{}, asyncAbort: false},
		{err: &kvpb.RangeNotFoundError{}, asyncAbort: false},
		{err: &kvpb.RangeKeyMismatchError{}, asyncAbort: false},
		{err: &kvpb.TransactionStatusError{}, asyncAbort: false},
	}

	for _, test := range testCases {
		t.Run(fmt.Sprintf("%T", test.err), func(t *testing.T) {
			var commit, abort atomic.Value
			commit.Store(false)
			abort.Store(false)

			stopper := stop.NewStopper()
			defer stopper.Stop(ctx)
			var senderFn kv.SenderFunc = func(
				_ context.Context, ba *kvpb.BatchRequest,
			) (*kvpb.BatchResponse, *kvpb.Error) {
				br := ba.CreateReply()
				br.Txn = ba.Txn.Clone()

				if et, hasET := ba.GetArg(kvpb.EndTxn); hasET {
					if et.(*kvpb.EndTxnRequest).Commit {
						commit.Store(true)
						if test.errFn != nil {
							return nil, test.errFn(*ba.Txn)
						}
						return nil, kvpb.NewErrorWithTxn(test.err, ba.Txn)
					}
					abort.Store(true)
				}
				return br, nil
			}
			ambient := log.MakeTestingAmbientCtxWithNewTracer()
			factory := kvcoord.NewTxnCoordSenderFactory(
				kvcoord.TxnCoordSenderFactoryConfig{
					AmbientCtx: ambient,
					Clock:      clock,
					Stopper:    stopper,
					Settings:   cluster.MakeTestingClusterSettings(),
				},
				senderFn,
			)

			db := kv.NewDB(ambient, factory, clock, stopper)
			txn := kv.NewTxn(ctx, db, 0 /* gatewayNodeID */)
			if pErr := txn.Put(ctx, "a", "b"); pErr != nil {
				t.Fatalf("put failed: %s", pErr)
			}
			require.Error(t, txn.Commit(ctx), "unexpected commit success")
			require.NoError(t, txn.Rollback(ctx))

			if !commit.Load().(bool) {
				t.Errorf("%T: failed to find initial commit request", test.err)
			}
			if !test.asyncAbort && !abort.Load().(bool) {
				t.Errorf("%T: failed to find expected synchronous abort", test.err)
			} else {
				testutils.SucceedsSoon(t, func() error {
					if !abort.Load().(bool) {
						return errors.Errorf("%T: failed to find expected asynchronous abort", test.err)
					}
					return nil
				})
			}
		})
	}
}

// mockSender is a client.Sender implementation that passes requests to a list
// of provided matchers, in sequence. The first matcher that returns either a
// response or an error is used to provide the result for the request.
type mockSender struct {
	matchers []matcher
}

var _ kv.Sender = &mockSender{}

type matcher func(*kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error)

// match adds a matcher to the list of matchers.
func (s *mockSender) match(m matcher) {
	s.matchers = append(s.matchers, m)
}

// Send implements the client.Sender interface.
func (s *mockSender) Send(
	_ context.Context, ba *kvpb.BatchRequest,
) (*kvpb.BatchResponse, *kvpb.Error) {
	for _, m := range s.matchers {
		br, pErr := m(ba)
		if br != nil || pErr != nil {
			return br, pErr
		}
	}
	// If none of the matchers triggered, just create an empty reply.
	br := ba.CreateReply()
	br.Txn = ba.Txn.Clone()
	return br, nil
}

// Test that a rollback sent to the TxnCoordSender stops the heartbeat loop even
// if it encounters an error. As of June 2018, there's a separate code path for
// handling errors on rollback in this regard.
func TestRollbackErrorStopsHeartbeat(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: log.MakeTestingAmbientCtxWithNewTracer(),
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	db := kv.NewDB(ambient, factory, clock, stopper)

	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		if _, ok := ba.GetArg(kvpb.EndTxn); !ok {
			resp := ba.CreateReply()
			resp.Txn = ba.Txn
			return resp, nil
		}
		return nil, kvpb.NewErrorf("injected err")
	})

	txn := kv.NewTxn(ctx, db, roachpb.NodeID(1))
	txnHeader := kvpb.Header{
		Txn: txn.TestingCloneTxn(),
	}
	if _, pErr := kv.SendWrappedWith(
		ctx, txn, txnHeader, &kvpb.PutRequest{
			RequestHeader: kvpb.RequestHeader{
				Key: roachpb.Key("a"),
			},
		},
	); pErr != nil {
		t.Fatal(pErr)
	}
	if !txn.Sender().(*kvcoord.TxnCoordSender).IsTracking() {
		t.Fatalf("expected TxnCoordSender to be tracking after the write")
	}

	if _, pErr := kv.SendWrappedWith(
		ctx, txn, txnHeader,
		&kvpb.EndTxnRequest{Commit: false},
	); !testutils.IsPError(pErr, "injected err") {
		t.Fatal(pErr)
	}

	testutils.SucceedsSoon(t, func() error {
		if txn.Sender().(*kvcoord.TxnCoordSender).IsTracking() {
			return fmt.Errorf("still tracking")
		}
		return nil
	})
}

// Test that lock tracking behaves correctly for transactions that attempt to
// run a batch containing an EndTxn. Since in case of an error it's not easy to
// determine whether any locks have been laid down (i.e. in case the batch was
// split by the DistSender and then there was mixed success for the sub-batches,
// or in case a retriable error is returned), the test verifies that all
// possible locks are properly tracked and attached to a subsequent EndTxn.
func TestOnePCErrorTracking(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: log.MakeTestingAmbientCtxWithNewTracer(),
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	db := kv.NewDB(ambient, factory, clock, stopper)
	keyA, keyB, keyC := roachpb.Key("a"), roachpb.Key("b"), roachpb.Key("c")

	// Register a matcher catching the commit attempt.
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		if et, ok := ba.GetArg(kvpb.EndTxn); !ok {
			return nil, nil
		} else if !et.(*kvpb.EndTxnRequest).Commit {
			return nil, nil
		}
		return nil, kvpb.NewErrorf("injected err")
	})
	// Register a matcher catching the rollback attempt.
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		et, ok := ba.GetArg(kvpb.EndTxn)
		if !ok {
			return nil, nil
		}
		etReq := et.(*kvpb.EndTxnRequest)
		if etReq.Commit {
			return nil, nil
		}
		expLocks := []roachpb.Span{{Key: keyA}, {Key: keyB, EndKey: keyC}}
		locks := etReq.LockSpans
		if !reflect.DeepEqual(locks, expLocks) {
			return nil, kvpb.NewErrorf("expected locks %s, got: %s", expLocks, locks)
		}
		resp := ba.CreateReply()
		// Set the response's txn to the Aborted status (as the server would). This
		// will make the TxnCoordSender stop the heartbeat loop.
		resp.Txn = ba.Txn.Clone()
		resp.Txn.Status = roachpb.ABORTED
		return resp, nil
	})

	txn := kv.NewTxn(ctx, db, roachpb.NodeID(1))
	txnHeader := kvpb.Header{
		Txn: txn.TestingCloneTxn(),
	}
	b := txn.NewBatch()
	b.Put(keyA, "test value")
	b.ScanForUpdate(keyB, keyC, kvpb.BestEffort)
	if err := txn.CommitInBatch(ctx, b); !testutils.IsError(err, "injected err") {
		t.Fatal(err)
	}

	// Now send a rollback and verify that the TxnCoordSender attaches the locks
	// to it.
	if _, pErr := kv.SendWrappedWith(
		ctx, txn, txnHeader,
		&kvpb.EndTxnRequest{Commit: false},
	); pErr != nil {
		t.Fatal(pErr)
	}

	// As always, check that the rollback we just sent stops the heartbeat loop.
	testutils.SucceedsSoon(t, func() error {
		if txn.Sender().(*kvcoord.TxnCoordSender).IsTracking() {
			return fmt.Errorf("still tracking")
		}
		return nil
	})
}

// TestCommitReadOnlyTransaction verifies that a read-only does not send an
// EndTxnRequest.
func TestCommitReadOnlyTransaction(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	var calls []kvpb.Method
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		calls = append(calls, ba.Methods()...)
		return nil, nil
	})

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	testutils.RunTrueAndFalse(t, "explicit txn", func(t *testing.T, explicitTxn bool) {
		testutils.RunTrueAndFalse(t, "with get", func(t *testing.T, withGet bool) {
			calls = nil
			db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)
			if err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
				b := txn.NewBatch()
				if withGet {
					b.Get("foo")
				}
				if explicitTxn {
					return txn.CommitInBatch(ctx, b)
				}
				return txn.Run(ctx, b)
			}); err != nil {
				t.Fatal(err)
			}

			expectedCalls := []kvpb.Method(nil)
			if withGet {
				expectedCalls = append(expectedCalls, kvpb.Get)
			}
			if !reflect.DeepEqual(expectedCalls, calls) {
				t.Fatalf("expected %s, got %s", expectedCalls, calls)
			}
		})
	})
}

// TestCommitMutatingTransaction verifies that a transaction is committed
// upon successful invocation of the retryable func.
func TestCommitMutatingTransaction(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	var calls []kvpb.Method
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		br := ba.CreateReply()
		br.Txn = ba.Txn.Clone()

		calls = append(calls, ba.Methods()...)
		if !bytes.Equal(ba.Txn.Key, roachpb.Key("a")) {
			t.Errorf("expected transaction key to be \"a\"; got %s", ba.Txn.Key)
		}

		if _, ok := ba.GetArg(kvpb.DeleteRange); ok {
			// Simulate deleting a single key for DeleteRange requests. Unlike other
			// point writes, pipelined DeleteRange writes are tracked by looking at
			// the batch response.
			resp := br.Responses[0].GetInner()
			resp.(*kvpb.DeleteRangeResponse).Keys = []roachpb.Key{roachpb.Key("a")}
		}

		if et, ok := ba.GetArg(kvpb.EndTxn); ok {
			if !et.(*kvpb.EndTxnRequest).Commit {
				t.Errorf("expected commit to be true")
			}
			br.Txn.Status = roachpb.COMMITTED
		}
		return br, nil
	})

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)

	// Test all transactional write methods.
	testArgs := []struct {
		f         func(ctx context.Context, txn *kv.Txn) error
		expMethod kvpb.Method
		// All retryable functions below involve writing to exactly one key. The
		// write will be pipelined if it's not in the same batch as the
		// EndTxnRequest, in which case we expect a single QueryIntent request to be
		// added when trying to commit the transaction.
		expQueryIntent bool
	}{
		{
			f:              func(ctx context.Context, txn *kv.Txn) error { return txn.Put(ctx, "a", "b") },
			expMethod:      kvpb.Put,
			expQueryIntent: true,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				b := txn.NewBatch()
				b.Put("a", "b")
				return txn.CommitInBatch(ctx, b)
			},
			expMethod: kvpb.Put,
		},
		{
			f:              func(ctx context.Context, txn *kv.Txn) error { return txn.CPut(ctx, "a", "b", nil) },
			expMethod:      kvpb.ConditionalPut,
			expQueryIntent: true,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				b := txn.NewBatch()
				b.CPut("a", "b", nil)
				return txn.CommitInBatch(ctx, b)
			},
			expMethod: kvpb.ConditionalPut,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				_, err := txn.Inc(ctx, "a", 1)
				return err
			},
			expMethod:      kvpb.Increment,
			expQueryIntent: true,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				b := txn.NewBatch()
				b.Inc("a", 1)
				return txn.CommitInBatch(ctx, b)
			},
			expMethod: kvpb.Increment,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				_, err := txn.Del(ctx, "a")
				return err
			},
			expMethod:      kvpb.Delete,
			expQueryIntent: true,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				b := txn.NewBatch()
				b.Del("a")
				return txn.CommitInBatch(ctx, b)
			},
			expMethod: kvpb.Delete,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				_, err := txn.DelRange(ctx, "a", "b", false /* returnKeys */)
				return err
			},
			expMethod:      kvpb.DeleteRange,
			expQueryIntent: true,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				b := txn.NewBatch()
				b.DelRange("a", "b", false /* returnKeys */)
				return txn.CommitInBatch(ctx, b)
			},
			expMethod:      kvpb.DeleteRange,
			expQueryIntent: false,
		},
	}
	for _, test := range testArgs {
		t.Run(test.expMethod.String(), func(t *testing.T) {
			calls = nil
			db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)
			require.NoError(t, db.Txn(ctx, test.f))
			expectedCalls := []kvpb.Method{test.expMethod}
			if test.expQueryIntent {
				expectedCalls = append(expectedCalls, kvpb.QueryIntent)
			}
			expectedCalls = append(expectedCalls, kvpb.EndTxn)
			require.Equal(t, expectedCalls, calls)
		})
	}
}

// TestCommitNoopDeleteRangeTransaction ensures that committing a no-op
// DeleteRange transaction works correctly. In particular, the EndTxn request is
// elided, if possible.
func TestCommitNoopDeleteRangeTransaction(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	var calls []kvpb.Method
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		br := ba.CreateReply()
		br.Txn = ba.Txn.Clone()

		calls = append(calls, ba.Methods()...)
		if !bytes.Equal(ba.Txn.Key, roachpb.Key("a")) {
			t.Errorf("expected transaction key to be \"a\"; got %s", ba.Txn.Key)
		}
		if et, ok := ba.GetArg(kvpb.EndTxn); ok {
			if !et.(*kvpb.EndTxnRequest).Commit {
				t.Errorf("expected commit to be true")
			}
			br.Txn.Status = roachpb.COMMITTED
		}

		// Don't return any keys in the DeleteRange response, which simulates what a
		// no-op DeleteRange looks like from the client's perspective.
		return br, nil
	})

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)

	testCases := []struct {
		f           func(ctx context.Context, txn *kv.Txn) error
		elideEndTxn bool
	}{
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				_, err := txn.DelRange(ctx, "a", "d", true /* returnKeys */)
				return err
			},
			elideEndTxn: true,
		},
		{
			f: func(ctx context.Context, txn *kv.Txn) error {
				b := txn.NewBatch()
				b.DelRange("a", "b", true /* returnKeys */)
				return txn.CommitInBatch(ctx, b)
			},
			elideEndTxn: false,
		},
	}

	for i, test := range testCases {
		t.Run(fmt.Sprintf("#%d", i), func(t *testing.T) {
			calls = nil
			db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)
			require.NoError(t, db.Txn(ctx, test.f))
			expectedCalls := []kvpb.Method{kvpb.DeleteRange}
			if !test.elideEndTxn {
				expectedCalls = append(expectedCalls, kvpb.EndTxn)
			}
			require.Equal(t, expectedCalls, calls)
		})
	}
}

// TestAbortReadOnlyTransaction verifies that aborting a read-only
// transaction does not prompt an EndTxn call.
func TestAbortReadOnlyTransaction(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	var calls []kvpb.Method
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		calls = append(calls, ba.Methods()...)
		return nil, nil
	})

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)
	if err := db.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
		return errors.New("foo")
	}); err == nil {
		t.Fatal("expected error on abort")
	}

	if calls != nil {
		t.Fatalf("expected no calls, got %s", calls)
	}
}

// TestEndWriteRestartReadOnlyTransaction verifies that if
// a transaction writes, then restarts and turns read-only,
// an explicit EndTxn call is still sent if retry- able
// didn't, regardless of whether there is an error or not.
func TestEndWriteRestartReadOnlyTransaction(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	var calls []kvpb.Method
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		br := ba.CreateReply()
		br.Txn = ba.Txn.Clone()

		calls = append(calls, ba.Methods()...)
		switch ba.Requests[0].GetInner().Method() {
		case kvpb.Put, kvpb.Scan:
			return nil, kvpb.NewErrorWithTxn(
				kvpb.NewTransactionRetryError(kvpb.RETRY_SERIALIZABLE, "test err"),
				ba.Txn)
		case kvpb.EndTxn:
			br.Txn.Status = roachpb.COMMITTED
		}
		return br, nil
	})

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
			TestingKnobs: kvcoord.ClientTestingKnobs{
				// Disable span refresh, otherwise it kicks and retries batches by
				// itself.
				MaxTxnRefreshAttempts: -1,
			},
		},
		sender,
	)
	db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)

	testutils.RunTrueAndFalse(t, "write", func(t *testing.T, write bool) {
		testutils.RunTrueAndFalse(t, "success", func(t *testing.T, success bool) {
			calls = nil
			firstIter := true
			if err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
				var err error
				if firstIter {
					firstIter = false
					if write {
						err = txn.Put(ctx, "consider", "phlebas")
					} else /* locking read */ {
						_, err = txn.ScanForUpdate(ctx, "a", "b", 0, kvpb.BestEffort)
					}
					if err == nil {
						t.Fatal("missing injected retriable error")
					}
				}
				if !success {
					return errors.New("aborting on purpose")
				}
				return err
			}); err == nil != success {
				t.Fatalf("expected error: %t, got error: %v", !success, err)
			}

			var expCalls []kvpb.Method
			if write {
				expCalls = []kvpb.Method{kvpb.Put, kvpb.EndTxn}
			} else {
				expCalls = []kvpb.Method{kvpb.Scan, kvpb.EndTxn}
			}
			if !reflect.DeepEqual(expCalls, calls) {
				t.Fatalf("expected %v, got %v", expCalls, calls)
			}
		})
	})
}

// TestTransactionKeyNotChangedInRestart verifies that if the transaction
// already has a key (we're in a restart), the key in the transaction request is
// not changed.
func TestTransactionKeyNotChangedInRestart(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	keys := []string{"first", "second"}
	attempt := 0
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		br := ba.CreateReply()
		br.Txn = ba.Txn.Clone()

		// Ignore the final EndTxnRequest.
		if _, ok := ba.GetArg(kvpb.EndTxn); ok {
			br.Txn.Status = roachpb.COMMITTED
			return br, nil
		}

		// Both attempts should have a PutRequest.
		if _, ok := ba.GetArg(kvpb.Put); !ok {
			t.Fatalf("failed to find a put request: %v", ba)
		}

		// In the first attempt, the transaction key is the key of the first write command.
		// This key is retained between restarts, so we see the same key in the second attempt.
		if expectedKey := []byte(keys[0]); !bytes.Equal(expectedKey, ba.Txn.Key) {
			t.Fatalf("expected transaction key %v, got %v", expectedKey, ba.Txn.Key)
		}

		if attempt == 0 {
			return nil, kvpb.NewErrorWithTxn(
				kvpb.NewTransactionRetryError(kvpb.RETRY_SERIALIZABLE, "test err"),
				ba.Txn)
		}
		return br, nil
	})
	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)

	if err := db.Txn(context.Background(), func(ctx context.Context, txn *kv.Txn) error {
		defer func() { attempt++ }()
		b := txn.NewBatch()
		b.Put(keys[attempt], "b")
		return txn.Run(ctx, b)
	}); err != nil {
		t.Errorf("unexpected error on commit: %s", err)
	}
	minimumAttempts := 2
	if attempt < minimumAttempts {
		t.Errorf("expected attempt count >= %d, got %d", minimumAttempts, attempt)
	}
}

// TestSequenceNumbers verifies Requests are given sequence numbers and that
// they are incremented on successive commands.
func TestSequenceNumbers(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	var expSequence enginepb.TxnSeq
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		for _, ru := range ba.Requests {
			args := ru.GetInner()
			if args.Method() == kvpb.QueryIntent {
				// QueryIntent requests don't have sequence numbers.
				continue
			}
			expSequence++
			if seq := args.Header().Sequence; expSequence != seq {
				t.Errorf("expected Request sequence %d; got %d. request: %T",
					expSequence, seq, args)
			}
		}
		br := ba.CreateReply()
		br.Txn = ba.Txn
		return br, nil
	})

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)
	txn := kv.NewTxn(ctx, db, 0 /* gatewayNodeID */)

	for i := 0; i < 5; i++ {
		ba := &kvpb.BatchRequest{}
		for j := 0; j < i; j++ {
			ba.Add(kvpb.NewPut(roachpb.Key("a"), roachpb.MakeValueFromString("foo")).(*kvpb.PutRequest))
		}
		if _, pErr := txn.Send(ctx, ba); pErr != nil {
			t.Fatal(pErr)
		}
	}
}

// TestConcurrentTxnRequests verifies that multiple requests cannot be executed
// on a transaction at the same time from multiple goroutines.
func TestConcurrentTxnRequestsProhibited(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	clock := hlc.NewClockForTesting(nil)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	putSync := make(chan struct{})
	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		if _, ok := ba.GetArg(kvpb.Put); ok {
			// Block the Put until the Get runs.
			putSync <- struct{}{}
			<-putSync
		}
		br := ba.CreateReply()
		br.Txn = ba.Txn.Clone()
		if _, ok := ba.GetArg(kvpb.EndTxn); ok {
			br.Txn.Status = roachpb.COMMITTED
		}
		return br, nil
	})

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	db := kv.NewDB(ambient, factory, clock, stopper)

	err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
		g, gCtx := errgroup.WithContext(ctx)
		g.Go(func() error {
			return txn.Put(gCtx, "test_put", "val")
		})
		g.Go(func() error {
			// Wait for the Put to be blocked.
			<-putSync
			_, err := txn.Get(gCtx, "test_get")
			// Unblock the Put.
			putSync <- struct{}{}
			return err
		})
		return g.Wait()
	})
	require.Regexp(t, "concurrent txn use detected", err)
}

// TestTxnRequestTxnTimestamp verifies response txn timestamp is
// always upgraded on successive requests.
func TestTxnRequestTxnTimestamp(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	manual := timeutil.NewManualTime(timeutil.Unix(0, 1))
	clock := hlc.NewClockForTesting(manual)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)

	curReq := 0
	requests := []struct {
		expRequestTS, responseTS hlc.Timestamp
	}{
		{hlc.Timestamp{WallTime: 5, Logical: 0}, hlc.Timestamp{WallTime: 10, Logical: 0}},
		{hlc.Timestamp{WallTime: 10, Logical: 0}, hlc.Timestamp{WallTime: 10, Logical: 1}},
		{hlc.Timestamp{WallTime: 10, Logical: 1}, hlc.Timestamp{WallTime: 10, Logical: 0}},
		{hlc.Timestamp{WallTime: 10, Logical: 1}, hlc.Timestamp{WallTime: 20, Logical: 1}},
		{hlc.Timestamp{WallTime: 20, Logical: 1}, hlc.Timestamp{WallTime: 20, Logical: 1}},
		{hlc.Timestamp{WallTime: 20, Logical: 1}, hlc.Timestamp{WallTime: 19, Logical: 0}},
		{hlc.Timestamp{WallTime: 20, Logical: 1}, hlc.Timestamp{WallTime: 20, Logical: 1}},
	}

	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		req := requests[curReq]
		if req.expRequestTS != ba.Txn.WriteTimestamp {
			return nil, kvpb.NewErrorf("%d: expected ts %s got %s",
				curReq, req.expRequestTS, ba.Txn.WriteTimestamp)
		}

		br := ba.CreateReply()
		br.Txn = ba.Txn.Clone()
		br.Txn.WriteTimestamp.Forward(requests[curReq].responseTS)
		return br, nil
	})

	manual.AdvanceTo(requests[0].expRequestTS.GoTime())

	if err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
		for curReq = range requests {
			if _, err := txn.Get(ctx, "k"); err != nil {
				return err
			}
		}
		manual.AdvanceTo(txn.ProvisionalCommitTimestamp().GoTime())
		return nil
	}); err != nil {
		t.Fatal(err)
	}
}

// TestReadOnlyTxnObeysDeadline tests that read-only transactions obey the
// deadline. Read-only transactions have their EndTxn elided, so the enforcement
// of the deadline is done in the client.
func TestReadOnlyTxnObeysDeadline(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	manual := timeutil.NewManualTime(timeutil.Unix(0, 123))
	clock := hlc.NewClockForTesting(manual)
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	sender := &mockSender{}
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	sender.match(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
		if _, ok := ba.GetArg(kvpb.Get); ok {
			manual.Advance(100)
			br := ba.CreateReply()
			br.Txn = ba.Txn.Clone()
			br.Txn.WriteTimestamp.Forward(clock.Now())
			return br, nil
		}
		return nil, nil
	})

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		sender,
	)
	db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)

	// We're going to run two tests: one where the EndTxn is by itself in a
	// batch, one where it is not. As of June 2018, the EndTxn is elided in
	// different ways in the two cases.

	t.Run("standalone commit", func(t *testing.T) {
		txn := kv.NewTxn(ctx, db, 0 /* gatewayNodeID */)
		// Set a deadline. We'll generate a retriable error with a higher timestamp.
		err := txn.UpdateDeadline(ctx, clock.Now())
		require.NoError(t, err, "Deadline update to now failed")
		if _, err := txn.Get(ctx, "k"); err != nil {
			t.Fatal(err)
		}
		err = txn.Commit(ctx)
		assertTransactionRetryError(t, err)
		if !testutils.IsError(err, "RETRY_COMMIT_DEADLINE_EXCEEDED") {
			t.Fatalf("expected deadline exceeded, got: %s", err)
		}
	})

	t.Run("commit in batch", func(t *testing.T) {
		txn := kv.NewTxn(ctx, db, 0 /* gatewayNodeID */)
		// Set a deadline. We'll generate a retriable error with a higher timestamp.
		err := txn.UpdateDeadline(ctx, clock.Now())
		require.NoError(t, err, "Deadline update to now failed")
		b := txn.NewBatch()
		b.Get("k")
		err = txn.CommitInBatch(ctx, b)
		assertTransactionRetryError(t, err)
		if !testutils.IsError(err, "RETRY_COMMIT_DEADLINE_EXCEEDED") {
			t.Fatalf("expected deadline exceeded, got: %s", err)
		}
	})
}

// TestTxnCoordSenderPipelining verifies that transactional pipelining of writes
// is enabled by default in a transaction and is disabled after
// DisablePipelining is called. It also verifies that DisablePipelining returns
// an error if the transaction has already performed an operation.
func TestTxnCoordSenderPipelining(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	ctx := context.Background()
	s := createTestDB(t)
	defer s.Stop()
	distSender := s.DB.GetFactory().(*kvcoord.TxnCoordSenderFactory).NonTransactionalSender()

	var calls []kvpb.Method
	var senderFn kv.SenderFunc = func(
		ctx context.Context, ba *kvpb.BatchRequest,
	) (*kvpb.BatchResponse, *kvpb.Error) {
		calls = append(calls, ba.Methods()...)
		if et, ok := ba.GetArg(kvpb.EndTxn); ok {
			// Ensure that no transactions enter a STAGING state.
			et.(*kvpb.EndTxnRequest).InFlightWrites = nil
		}
		return distSender.Send(ctx, ba)
	}

	tsf := kvcoord.NewTxnCoordSenderFactory(kvcoord.TxnCoordSenderFactoryConfig{
		AmbientCtx: s.AmbientCtx,
		Settings:   s.Cfg.Settings,
		Clock:      s.Clock,
		Stopper:    s.Stopper(),
		// Disable transaction heartbeats so that they don't disrupt our attempt to
		// track the requests issued by the transactions.
		HeartbeatInterval: -1,
	}, senderFn)
	db := kv.NewDB(s.AmbientCtx, tsf, s.Clock, s.Stopper())

	err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
		return txn.Put(ctx, "key", "val")
	})
	if err != nil {
		t.Fatal(err)
	}

	err = db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
		if err := txn.DisablePipelining(); err != nil {
			return err
		}
		return txn.Put(ctx, "key", "val")
	})
	if err != nil {
		t.Fatal(err)
	}

	require.Equal(t, []kvpb.Method{
		kvpb.Put, kvpb.QueryIntent, kvpb.EndTxn,
		kvpb.Put, kvpb.EndTxn,
	}, calls)

	for _, action := range []func(ctx context.Context, txn *kv.Txn) error{
		func(ctx context.Context, txn *kv.Txn) error { return txn.Put(ctx, "key", "val") },
		func(ctx context.Context, txn *kv.Txn) error { _, err := txn.Get(ctx, "key"); return err },
	} {
		err = db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
			if err := action(ctx, txn); err != nil {
				t.Fatal(err)
			}
			return txn.DisablePipelining()
		})
		if exp := "cannot disable pipelining on a running transaction"; !testutils.IsError(err, exp) {
			t.Fatalf("expected %q error, but got %v", exp, err)
		}
	}
}

// Test that a txn's anchor is set to the first write key in batches mixing
// reads with writes.
func TestAnchorKey(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	ctx := context.Background()
	clock := hlc.NewClockForTesting(timeutil.NewManualTime(timeutil.Unix(0, 123)))
	ambient := log.MakeTestingAmbientCtxWithNewTracer()
	stopper := stop.NewStopper()
	defer stopper.Stop(ctx)

	key1 := roachpb.Key("a")
	key2 := roachpb.Key("b")

	var senderFn kv.SenderFunc = func(
		ctx context.Context, ba *kvpb.BatchRequest,
	) (*kvpb.BatchResponse, *kvpb.Error) {
		if !roachpb.Key(ba.Txn.Key).Equal(key2) {
			t.Fatalf("expected anchor %q, got %q", key2, ba.Txn.Key)
		}
		br := ba.CreateReply()
		br.Txn = ba.Txn.Clone()
		if _, ok := ba.GetArg(kvpb.EndTxn); ok {
			br.Txn.Status = roachpb.COMMITTED
		}
		return br, nil
	}

	factory := kvcoord.NewTxnCoordSenderFactory(
		kvcoord.TxnCoordSenderFactoryConfig{
			AmbientCtx: ambient,
			Clock:      clock,
			Stopper:    stopper,
			Settings:   cluster.MakeTestingClusterSettings(),
		},
		senderFn,
	)
	db := kv.NewDB(log.MakeTestingAmbientCtxWithNewTracer(), factory, clock, stopper)

	if err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
		ba := txn.NewBatch()
		ba.Get(key1)
		ba.Put(key2, "val")
		return txn.Run(ctx, ba)
	}); err != nil {
		t.Fatal(err)
	}
}

// Test that a leaf txn returns a raw error when "rejecting a client" (a client
// sending something after the txn is known to be aborted), not a
// TransactionRetryWithProtoRefreshError. This is important as leaves are not supposed to create
// "handled" errors; instead the DistSQL infra knows to recognize raw retryable
// errors and feed them to the root txn.
func TestLeafTxnClientRejectError(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	// We're going to inject an error so that a leaf txn is "poisoned". This can
	// happen, for example, if the leaf is used concurrently by multiple requests,
	// where the first one gets a TransactionAbortedError.
	errKey := roachpb.Key("a")
	knobs := &kvserver.StoreTestingKnobs{
		TestingRequestFilter: func(_ context.Context, ba *kvpb.BatchRequest) *kvpb.Error {
			if g, ok := ba.GetArg(kvpb.Get); ok && g.(*kvpb.GetRequest).Key.Equal(errKey) {
				txn := ba.Txn.Clone()
				txn.Status = roachpb.ABORTED
				return kvpb.NewErrorWithTxn(
					kvpb.NewTransactionAbortedError(kvpb.ABORT_REASON_UNKNOWN), txn,
				)
			}
			return nil
		},
	}

	s := createTestDBWithKnobs(t, knobs)
	defer s.Stop()

	ctx := context.Background()
	rootTxn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
	require.NoError(t, err)

	// New create a second, leaf coordinator.
	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState)

	if _, err := leafTxn.Get(ctx, errKey); !testutils.IsError(err, "TransactionAbortedError") {
		t.Fatalf("expected injected err, got: %v", err)
	}

	// Now use the leaf and check the error. At the TxnCoordSender level, the
	// pErr will be TransactionAbortedError. When pErr.GoError() is called, that's
	// transformed into an UnhandledRetryableError. For our purposes, what this
	// test is interested in demonstrating is that it's not a
	// TransactionRetryWithProtoRefreshError.
	_, err = leafTxn.Get(ctx, roachpb.Key("a"))
	if !errors.HasType(err, (*kvpb.UnhandledRetryableError)(nil)) {
		t.Fatalf("expected UnhandledRetryableError(TransactionAbortedError), got: (%T) %v", err, err)
	}
}

// Check that ingesting an Aborted txn record is a no-op. The TxnCoordSender is
// supposed to reject such updates because they risk putting it into an
// inconsistent state. See comments in TxnCoordSender.UpdateRootWithLeafFinalState().
func TestUpdateRootWithLeafFinalStateInAbortedTxn(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDBWithKnobs(t, nil /* knobs */)
	defer s.Stop()
	ctx := context.Background()

	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	leafInputState, err := txn.GetLeafTxnInputState(ctx)
	require.NoError(t, err)
	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0, leafInputState)

	finalState, err := leafTxn.GetLeafTxnFinalState(ctx)
	if err != nil {
		t.Fatal(err)
	}
	finalState.Txn.Status = roachpb.ABORTED
	if err := txn.UpdateRootWithLeafFinalState(ctx, finalState); err != nil {
		t.Fatal(err)
	}

	// Check that the transaction was not updated.
	leafInputState2, err := txn.GetLeafTxnInputState(ctx)
	require.NoError(t, err)
	if leafInputState2.Txn.Status != roachpb.PENDING {
		t.Fatalf("expected PENDING txn, got: %s", leafInputState2.Txn.Status)
	}
}

// Test that evaluating a request within a txn with the STAGING status works
// fine. It's unusual for the server to receive a request in the STAGING status
// (other than a single EndTxn which transitions from STAGING->COMMITTED),
// because the committer interceptor reverts the status from STAGING->PENDING if
// the batch moving to STAGING has been split and some part of it failed.
// However, it can still happen that the server receives a request with the
// STAGING record when the DistSender splits a batch and doesn't send all the
// sub-batches in parallel; in that case it's possible that a sub-batch with the
// EndTxn succeeds, and then the DistSender will send the remaining sub-batches
// with a STAGING status.
func TestPutsInStagingTxn(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()

	keyA := roachpb.Key("a")
	keyB := roachpb.Key("b")

	var putInStagingSeen bool
	var storeKnobs kvserver.StoreTestingKnobs
	storeKnobs.TestingRequestFilter = func(ctx context.Context, ba *kvpb.BatchRequest) *kvpb.Error {
		put, ok := ba.GetArg(kvpb.Put)
		if !ok || !put.(*kvpb.PutRequest).Key.Equal(keyB) {
			return nil
		}
		txn := ba.Txn
		if txn == nil {
			return nil
		}
		if txn.Status == roachpb.STAGING {
			putInStagingSeen = true
		}
		return nil
	}

	// Disable the DistSender concurrency so that sub-batches split by the
	// DistSender are send serially and the transaction is updated from one to
	// another. See below.
	settings := cluster.MakeTestingClusterSettings()
	kvcoord.TestingSenderConcurrencyLimit.Override(ctx, &settings.SV, 0)

	s, _, db := serverutils.StartServer(t,
		base.TestServerArgs{
			Settings: settings,
			Knobs:    base.TestingKnobs{Store: &storeKnobs},
		})
	defer s.Stopper().Stop(ctx)

	require.NoError(t, db.AdminSplit(
		ctx,
		keyB,             /* splitKey */
		hlc.MaxTimestamp, /* expirationTime */
	))

	txn := db.NewTxn(ctx, "test")

	// Cause a write too old condition for the upcoming txn writes, to spicy up
	// the test.
	require.NoError(t, db.Put(ctx, keyB, "b"))

	// Send a batch that will be split into two sub-batches: [Put(a)+EndTxn,
	// Put(b)] (the EndTxn is grouped with the first write). These sub-batches are
	// sent serially since we've inhibited the DistSender's concurrency. The first
	// one will transition the txn to STAGING, and the DistSender will use that
	// updated txn when sending the 2nd sub-batch.
	b := txn.NewBatch()
	b.Put(keyA, "a")
	b.Put(keyB, "b")
	require.NoError(t, txn.CommitInBatch(ctx, b))
	// Verify that the test isn't fooling itself by checking that we've indeed
	// seen a batch with the STAGING status.
	require.True(t, putInStagingSeen)
}

// TestTxnCoordSenderSetFixedTimestamp tests that SetFixedTimestamp cannot be
// called after a transaction has already been used in the current epoch to read
// or write.
func TestTxnCoordSenderSetFixedTimestamp(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()

	for _, test := range []struct {
		name   string
		before func(*testing.T, *kv.Txn)
		expErr string
	}{
		{
			name:   "nothing before",
			before: func(t *testing.T, txn *kv.Txn) {},
		},
		{
			name: "read before",
			before: func(t *testing.T, txn *kv.Txn) {
				_, err := txn.Get(ctx, "k")
				require.NoError(t, err)
			},
			expErr: "cannot set fixed timestamp, .* already performed reads",
		},
		{
			name: "write before",
			before: func(t *testing.T, txn *kv.Txn) {
				require.NoError(t, txn.Put(ctx, "k", "v"))
			},
			expErr: "cannot set fixed timestamp, .* already performed writes",
		},
		{
			name: "read and write before",
			before: func(t *testing.T, txn *kv.Txn) {
				_, err := txn.Get(ctx, "k")
				require.NoError(t, err)
				require.NoError(t, txn.Put(ctx, "k", "v"))
			},
			expErr: "cannot set fixed timestamp, .* already performed reads",
		},
		{
			name: "read before, in prior epoch",
			before: func(t *testing.T, txn *kv.Txn) {
				_, err := txn.Get(ctx, "k")
				require.NoError(t, err)
				require.Error(t, txn.Sender().GenerateForcedRetryableErr(ctx, txn.ReadTimestamp().Next(), true /* mustRestart */, "force retry"))
				require.NoError(t, txn.Sender().ClearRetryableErr(ctx))
			},
		},
		{
			name: "write before, in prior epoch",
			before: func(t *testing.T, txn *kv.Txn) {
				require.NoError(t, txn.Put(ctx, "k", "v"))
				require.Error(t, txn.Sender().GenerateForcedRetryableErr(ctx, txn.ReadTimestamp().Next(), true /* mustRestart */, "force retry"))
				require.NoError(t, txn.Sender().ClearRetryableErr(ctx))
			},
		},
		{
			name: "read and write before, in prior epoch",
			before: func(t *testing.T, txn *kv.Txn) {
				_, err := txn.Get(ctx, "k")
				require.NoError(t, err)
				require.NoError(t, txn.Put(ctx, "k", "v"))
				require.Error(t, txn.Sender().GenerateForcedRetryableErr(ctx, txn.ReadTimestamp().Next(), true /* mustRestart */, "force retry"))
				require.NoError(t, txn.Sender().ClearRetryableErr(ctx))
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			s := createTestDB(t)
			defer s.Stop()

			txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
			test.before(t, txn)

			ts := s.Clock.Now()
			err := txn.SetFixedTimestamp(ctx, ts)
			if test.expErr != "" {
				require.Error(t, err)
				require.Regexp(t, test.expErr, err)
				require.False(t, txn.ReadTimestampFixed())
			} else {
				require.NoError(t, err)
				require.True(t, txn.ReadTimestampFixed())
				require.Equal(t, ts, txn.ReadTimestamp())
			}
		})
	}
}

// TestTxnTypeCompatibleWithBatchRequest tests if a transaction type and a batch
// request are compatible.
func TestTxnTypeCompatibleWithBatchRequest(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	ctx := context.Background()
	s := createTestDB(t)
	defer s.Stop()

	rootTxn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
	require.NoError(t, err)
	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState)

	// A LeafTxn is not compatible with locking requests.
	// 1. Locking Get requests.
	_, err = leafTxn.GetForUpdate(ctx, roachpb.Key("a"), kvpb.GuaranteedDurability)
	require.Error(t, err)
	require.Regexp(t, "LeafTxn .* incompatible with locking request .*", err)
	_, err = leafTxn.GetForShare(ctx, roachpb.Key("a"), kvpb.GuaranteedDurability)
	require.Error(t, err)
	require.Regexp(t, "LeafTxn .* incompatible with locking request .*", err)
	// 2. Locking Scan requests.
	_, err = leafTxn.ScanForUpdate(
		ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */, kvpb.GuaranteedDurability,
	)
	require.Error(t, err)
	require.Regexp(t, "LeafTxn .* incompatible with locking request .*", err)
	_, err = leafTxn.ScanForShare(
		ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */, kvpb.GuaranteedDurability,
	)
	require.Error(t, err)
	require.Regexp(t, "LeafTxn .* incompatible with locking request .*", err)
	// 3. Locking ReverseScan requests.
	_, err = leafTxn.ReverseScanForUpdate(
		ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */, kvpb.GuaranteedDurability,
	)
	require.Error(t, err)
	require.Regexp(t, "LeafTxn .* incompatible with locking request .*", err)
	_, err = leafTxn.ReverseScanForShare(
		ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */, kvpb.GuaranteedDurability,
	)
	require.Error(t, err)
	require.Regexp(t, "LeafTxn .* incompatible with locking request .*", err)
	// 4. Writes.
	err = leafTxn.Put(ctx, roachpb.Key("a"), []byte("b"))
	require.Error(t, err)
	require.Regexp(t, "LeafTxn .* incompatible with locking request .*", err)
	// A LeafTxn is compatible with non-locking requests.
	_, err = leafTxn.Get(ctx, roachpb.Key("a"))
	require.NoError(t, err)
	_, err = leafTxn.Scan(ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */)
	require.NoError(t, err)
	_, err = leafTxn.ReverseScan(ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */)
	require.NoError(t, err)

	// A RootTxn is compatible with all requests.
	// 1. All types of Get requests.
	_, err = rootTxn.GetForUpdate(ctx, roachpb.Key("a"), kvpb.GuaranteedDurability)
	require.NoError(t, err)
	_, err = rootTxn.GetForShare(ctx, roachpb.Key("a"), kvpb.GuaranteedDurability)
	require.NoError(t, err)
	_, err = rootTxn.Get(ctx, roachpb.Key("a"))
	require.NoError(t, err)
	// 2. All types of Scan requests.
	_, err = rootTxn.ScanForUpdate(
		ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */, kvpb.GuaranteedDurability,
	)
	require.NoError(t, err)
	_, err = rootTxn.ScanForShare(
		ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */, kvpb.GuaranteedDurability,
	)
	require.NoError(t, err)
	_, err = rootTxn.Scan(ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */)
	require.NoError(t, err)
	// 3. All types of ReverseScan requests.
	_, err = rootTxn.ReverseScanForUpdate(
		ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */, kvpb.GuaranteedDurability,
	)
	require.NoError(t, err)
	_, err = rootTxn.ReverseScanForShare(
		ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */, kvpb.GuaranteedDurability,
	)
	require.NoError(t, err)
	_, err = rootTxn.ReverseScan(ctx, roachpb.Key("a"), roachpb.Key("d"), 0 /* maxRows */)
	require.NoError(t, err)
	// 4. And writes.
	require.NoError(t, err)
	err = rootTxn.Put(ctx, roachpb.Key("a"), []byte("b"))
	require.NoError(t, err)
}

// TestTxnSetIsoLevel tests setting and getting the isolation level of a
// transaction. It also tests that isolation levels can only be set before
// the transaction is active.
func TestTxnSetIsoLevel(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)
	s := createTestDB(t)
	defer s.Stop()
	ctx := context.Background()
	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)

	defaultLevel := isolation.Serializable
	levels := isolation.Levels()

	// The default isolation level is Serializable.
	require.Equal(t, defaultLevel, txn.IsoLevel())

	// Can set and get all isolation levels.
	for _, isoLevel := range levels {
		require.NoError(t, txn.SetIsoLevel(isoLevel))
		require.Equal(t, isoLevel, txn.IsoLevel())
	}

	// Can't set isolation level after the transaction is active.
	require.NoError(t, txn.Put(ctx, "k", "v"))
	for _, isoLevel := range levels {
		prev := txn.IsoLevel()
		err := txn.SetIsoLevel(isoLevel)
		if isoLevel == prev {
			// If the isolation level does not need to change, no error is returned.
			require.NoError(t, err)
		} else {
			// If the isolation level needs to change, an error is returned.
			require.Error(t, err)
			require.Regexp(t, "cannot change the isolation level of a running transaction", err)
		}
		// The isolation level should not have changed.
		require.Equal(t, prev, txn.IsoLevel())
	}
}

// TestRefreshWithSavepoint is an integration test that ensures the correct
// behavior of refreshes under savepoint rollback. The test sets up a write-skew
// example where txn1 reads keyA and writes to keyB, while concurrently txn2
// reads keyB and writes to keyA. The two txns can't be serialized so one is
// expected to get a serialization error upon commit.
//
// However, with the old behavior of discarding refresh spans upon savepoint
// rollback, the read corresponding to the discarded refresh span is not
// refreshed, so the conflict goes unnoticed and both txns commit successfully.
// See #111228 for more details.
func TestRefreshWithSavepoint(t *testing.T) {
	defer leaktest.AfterTest(t)()
	defer log.Scope(t).Close(t)

	testutils.RunTrueAndFalse(t, "keep-refresh-spans", func(t *testing.T, keepRefreshSpans bool) {
		s, _, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
		ctx := context.Background()
		defer s.Stopper().Stop(context.Background())

		if keepRefreshSpans {
			kvcoord.KeepRefreshSpansOnSavepointRollback.Override(ctx, &s.ClusterSettings().SV, true)
		} else {
			kvcoord.KeepRefreshSpansOnSavepointRollback.Override(ctx, &s.ClusterSettings().SV, false)
		}

		keyA := roachpb.Key("a")
		keyB := roachpb.Key("b")
		txn1 := kvDB.NewTxn(ctx, "txn1")
		txn2 := kvDB.NewTxn(ctx, "txn2")

		spt1, err := txn1.CreateSavepoint(ctx)
		require.NoError(t, err)

		_, err = txn1.Get(ctx, keyA)
		require.NoError(t, err)

		err = txn1.RollbackToSavepoint(ctx, spt1)
		require.NoError(t, err)

		_, err = txn2.Get(ctx, keyB)
		require.NoError(t, err)

		err = txn1.Put(ctx, keyB, "bb")
		require.NoError(t, err)

		err = txn2.Put(ctx, keyA, "aa")
		require.NoError(t, err)

		err = txn1.Commit(ctx)
		if keepRefreshSpans {
			require.Regexp(t, ".*RETRY_SERIALIZABLE - failed preemptive refresh due to conflicting locks on \"a\"*", err)
		} else {
			require.NoError(t, err)
		}

		err = txn2.Commit(ctx)
		require.NoError(t, err)
	})
}