kvcoord: prevent concurrent EndTxn requests

`TxnCoordSender` generally operates synchronously (i.e. the client waits
for the previous response before sending the next request). However, the
`txnHeartbeater` sends asynchronous `EndTxn(commit=false)` rollbacks
when it discovers an aborted transaction record. Unfortunately, some code
assumes synchrony, which caused race conditions with txn rollbacks.

In particular, the `txnPipeliner` attaches lock spans and in-flight
writes to the `EndTxn` request for e.g. intent cleanup, but it only
records this information when it receives responses. Thus, if an
`EndTxn(commit=false)` is sent concurrently with a write request, the
lock spans and in-flight writes of that write request will not get
attached to the `EndTxn` request and the intents will not get cleaned
up.

This patch makes the `txnHeartbeater` wait for any in-flight requests to
complete before sending asynchronous rollbacks.

Release note (bug fix): Fixed a race condition where transaction cleanup
would fail to take into account ongoing writes and clean up their
intents.

Author: erikgrinaker

pkg/kv/kvclient/kvcoord/BUILD.bazel

@@ -47,6 +47,7 @@ go_library(
         "//pkg/settings/cluster",
         "//pkg/storage/enginepb",
         "//pkg/util",
+        "//pkg/util/contextutil",
         "//pkg/util/ctxgroup",
         "//pkg/util/errorutil/unimplemented",
         "//pkg/util/grpcutil",

pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go

@@ -16,12 +16,21 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/contextutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 )
 
+// abortTxnAsyncTimeout is the context timeout for abortTxnAsyncLocked()
+// rollbacks. If the intent resolver has spare async task capacity, this timeout
+// only needs to be long enough for the EndTxn request to make it through Raft,
+// but if the cleanup task is synchronous (to backpressure clients) then cleanup
+// will be abandoned when the timeout expires. We generally want to clean up if
+// possible, but not at any cost, so we set it high at 1 minute.
+const abortTxnAsyncTimeout = time.Minute
+
 // txnHeartbeater is a txnInterceptor in charge of a transaction's heartbeat
 // loop. Transaction coordinators heartbeat their transaction record
 // periodically to indicate the liveness of their transaction. Other actors like
@@ -104,9 +113,37 @@ type txnHeartbeater struct {
 		// future requests sent though it (which indicates that the heartbeat
 		// loop did not race with an EndTxn request).
 		finalObservedStatus roachpb.TransactionStatus
+
+		// ifReqs tracks the number of in-flight requests. This is expected to
+		// be either 0 or 1, but we let the txnLockGatekeeper enforce that.
+		//
+		// This is used to make sure we don't send EndTxn(commit=false) from
+		// abortTxnAsyncLocked() concurrently with another in-flight request.
+		// The TxnCoordSender assumes synchronous operation; in particular,
+		// the txnPipeliner must update its lock spans with pending responses
+		// before attaching the final lock spans to the EndTxn request.
+		ifReqs uint8
+
+		// abortTxnAsyncPending, if true, signals that an abortTxnAsyncLocked()
+		// call is waiting for in-flight requests to complete. Once the last
+		// request returns (setting ifReqs=0), it calls abortTxnAsyncLocked().
+		abortTxnAsyncPending bool
+
+		// abortTxnAsyncResultC is non-nil when an abortTxnAsyncLocked()
+		// rollback is in-flight. If a client rollback arrives concurrently, it
+		// will wait for the result on this channel, collapsing the requests.
+		// This prevents violating the txnLockGatekeeper concurrency assertion.
+		// Only EndTxn(commit=false) requests can arrive during rollback, the
+		// TxnCoordSender blocks any others due to finalObservedStatus.
+		abortTxnAsyncResultC chan abortTxnAsyncResult
 	}
 }
 
+type abortTxnAsyncResult struct {
+	br   *roachpb.BatchResponse
+	pErr *roachpb.Error
+}
+
 // init initializes the txnHeartbeater. This method exists instead of a
 // constructor because txnHeartbeaters live in a pool in the TxnCoordSender.
 func (h *txnHeartbeater) init(
@@ -175,9 +212,38 @@ func (h *txnHeartbeater) SendLocked(
 		// with the next epoch.
 		if !et.Commit {
 			h.cancelHeartbeatLoopLocked()
+
+			// If an abortTxnAsyncLocked() rollback is in flight, we'll wait
+			// for the result here to avoid violating the txnLockGatekeeper
+			// concurrency assertion.
+			if resultC := h.mu.abortTxnAsyncResultC; resultC != nil {
+				h.mu.Unlock()
+				defer h.mu.Lock()
+				select {
+				case res := <-resultC:
+					return res.br, res.pErr
+				case <-ctx.Done():
+					return nil, roachpb.NewError(ctx.Err())
+				}
+			}
 		}
 	}
 
+	// Record the in-flight request and its response.
+	h.mu.ifReqs++
+	defer func() {
+		h.mu.ifReqs--
+
+		// If an abortTxnAsyncLocked() call is waiting for this in-flight
+		// request to complete, call it. At this point, finalObservedStatus has
+		// already been set, so we don't have to worry about additional incoming
+		// requests (except rollbacks), the TxnCoordSender will block them.
+		if h.mu.abortTxnAsyncPending && h.mu.ifReqs == 0 {
+			h.abortTxnAsyncLocked(ctx)
+			h.mu.abortTxnAsyncPending = false
+		}
+	}()
+
 	// Forward the batch through the wrapped lockedSender.
 	return h.wrapped.SendLocked(ctx, ba)
 }
@@ -379,11 +445,26 @@ func (h *txnHeartbeater) heartbeat(ctx context.Context) bool {
 // The purpose of the async cleanup is to resolve transaction intents as soon
 // as possible when a transaction coordinator observes an ABORTED transaction.
 func (h *txnHeartbeater) abortTxnAsyncLocked(ctx context.Context) {
-	log.VEventf(ctx, 1, "Heartbeat detected aborted txn. Cleaning up.")
 
-	// NB: We use context.Background() here because we don't want a canceled
-	// context to interrupt the aborting.
-	ctx = h.AnnotateCtx(context.Background())
+	// If a request is in flight, we must wait for it to complete first such
+	// that txnPipeliner can record its lock spans and attach them to the EndTxn
+	// request we'll send.
+	if h.mu.ifReqs > 0 {
+		h.mu.abortTxnAsyncPending = true
+		log.VEventf(ctx, 1,
+			"Heartbeat detected aborted txn. Waiting for in-flight request before cleaning up.")
+		return
+	}
+
+	// If another abortTxnAsyncLocked attempt is already in flight, bail out.
+	// We check this since the TxnCoordSender also calls abortTxnAsyncLocked()
+	// independently of the heartbeat loop.
+	if h.mu.abortTxnAsyncResultC != nil {
+		log.VEventf(ctx, 1, "Heartbeat detected aborted txn, but found concurrent abort attempt.")
+		return
+	}
+
+	log.VEventf(ctx, 1, "Heartbeat detected aborted txn. Cleaning up.")
 
 	// Construct a batch with an EndTxn request.
 	txn := h.mu.txn.Clone()
@@ -397,21 +478,39 @@ func (h *txnHeartbeater) abortTxnAsyncLocked(ctx context.Context) {
 		TxnHeartbeating: true,
 	})
 
+	// Set up a result channel to signal to an incoming client rollback that
+	// an async rollback is already in progress and pass it the result. The
+	// buffer allows storing the result even when no client rollback arrives.
+	h.mu.abortTxnAsyncResultC = make(chan abortTxnAsyncResult, 1)
+
 	log.VEventf(ctx, 2, "async abort for txn: %s", txn)
-	if err := h.stopper.RunAsyncTask(
-		ctx, "txnHeartbeater: aborting txn", func(ctx context.Context) {
-			// Send the abort request through the interceptor stack. This is
-			// important because we need the txnPipeliner to append lock spans
-			// to the EndTxn request.
-			h.mu.Lock()
-			defer h.mu.Unlock()
-			_, pErr := h.wrapped.SendLocked(ctx, ba)
-			if pErr != nil {
-				log.VErrEventf(ctx, 1, "async abort failed for %s: %s ", txn, pErr)
+	const taskName = "txnHeartbeater: aborting txn"
+	if err := h.stopper.RunAsyncTask(h.AnnotateCtx(context.Background()), taskName,
+		func(ctx context.Context) {
+			if err := contextutil.RunWithTimeout(ctx, taskName, abortTxnAsyncTimeout,
+				func(ctx context.Context) error {
+					// Send the abort request through the interceptor stack. This is
+					// important because we need the txnPipeliner to append lock spans
+					// to the EndTxn request.
+					h.mu.Lock()
+					defer h.mu.Unlock()
+					br, pErr := h.wrapped.SendLocked(ctx, ba)
+					if pErr != nil {
+						log.VErrEventf(ctx, 1, "async abort failed for %s: %s ", txn, pErr)
+					}
+					// Pass the result to a waiting client rollback, if any, and
+					// remove the channel since we're no longer in flight.
+					h.mu.abortTxnAsyncResultC <- abortTxnAsyncResult{br: br, pErr: pErr}
+					h.mu.abortTxnAsyncResultC = nil
+					return nil
+				},
+			); err != nil {
+				log.VEventf(ctx, 1, "async abort failed for %s: %s", txn, err)
 			}
 		},
 	); err != nil {
 		log.Warningf(ctx, "%v", err)
+		h.mu.abortTxnAsyncResultC = nil // task wasn't started after all
 	}
 }
 

pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater_test.go

@@ -417,6 +417,203 @@ func TestTxnHeartbeaterAsyncAbort(t *testing.T) {
 	})
 }
 
+// TestTxnHeartbeaterAsyncAbortWaitsForInFlight tests that the txnHeartbeater
+// will wait for an in-flight request to complete before sending the
+// EndTxn rollback request.
+func TestTxnHeartbeaterAsyncAbortWaitsForInFlight(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	txn := makeTxnProto()
+	th, mockSender, mockGatekeeper := makeMockTxnHeartbeater(&txn)
+	defer th.stopper.Stop(ctx)
+
+	// Mock the heartbeat request, which should wait for an in-flight put via
+	// putReady then return an aborted txn and signal hbAborted.
+	putReady := make(chan struct{})
+	hbAborted := make(chan struct{})
+	mockGatekeeper.MockSend(func(ba roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
+		<-putReady
+		defer close(hbAborted)
+
+		require.Len(t, ba.Requests, 1)
+		require.IsType(t, &roachpb.HeartbeatTxnRequest{}, ba.Requests[0].GetInner())
+
+		br := ba.CreateReply()
+		br.Txn = ba.Txn
+		br.Txn.Status = roachpb.ABORTED
+		return br, nil
+	})
+
+	putResume := make(chan struct{})
+	rollbackSent := make(chan struct{})
+	mockSender.ChainMockSend(
+		// Mock a Put, which signals putReady and then waits for putResume
+		// before returning a response.
+		func(ba roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
+			th.mu.Unlock() // without txnLockGatekeeper, we must unlock manually
+			defer th.mu.Lock()
+			close(putReady)
+			require.Len(t, ba.Requests, 1)
+			require.IsType(t, &roachpb.PutRequest{}, ba.Requests[0].GetInner())
+
+			<-putResume
+
+			br := ba.CreateReply()
+			br.Txn = ba.Txn
+			return br, nil
+		},
+		// Mock an EndTxn, which signals rollbackSent.
+		func(ba roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
+			defer close(rollbackSent)
+			require.Len(t, ba.Requests, 1)
+			require.IsType(t, &roachpb.EndTxnRequest{}, ba.Requests[0].GetInner())
+
+			etReq := ba.Requests[0].GetInner().(*roachpb.EndTxnRequest)
+			require.Equal(t, &txn, ba.Txn)
+			require.False(t, etReq.Commit)
+			require.True(t, etReq.Poison)
+			require.True(t, etReq.TxnHeartbeating)
+
+			br := ba.CreateReply()
+			br.Txn = ba.Txn
+			br.Txn.Status = roachpb.ABORTED
+			return br, nil
+		},
+	)
+
+	// Spawn a goroutine to send the Put.
+	require.NoError(t, th.stopper.RunAsyncTask(ctx, "put", func(ctx context.Context) {
+		var ba roachpb.BatchRequest
+		ba.Header = roachpb.Header{Txn: txn.Clone()}
+		ba.Add(&roachpb.PutRequest{RequestHeader: roachpb.RequestHeader{Key: roachpb.Key("a")}})
+
+		th.mu.Lock() // without TxnCoordSender, we must lock manually
+		defer th.mu.Unlock()
+		br, pErr := th.SendLocked(ctx, ba)
+		require.Nil(t, pErr)
+		require.NotNil(t, br)
+	}))
+
+	<-putReady  // wait for put
+	<-hbAborted // wait for heartbeat abort
+	select {
+	case <-rollbackSent: // we don't expect a rollback yet
+		require.Fail(t, "received unexpected EndTxn")
+	case <-time.After(20 * time.Millisecond):
+	}
+	close(putResume) // make put return
+	<-rollbackSent   // we now expect the rollback
+
+	// The heartbeat loop should eventually close.
+	waitForHeartbeatLoopToStop(t, &th)
+}
+
+// TestTxnHeartbeaterAsyncAbortCollapsesRequests tests that when the
+// txnHeartbeater has an async abort rollback in flight, any client
+// rollbacks will wait for the async rollback to complete and return
+// its result.
+func TestTxnHeartbeaterAsyncAbortCollapsesRequests(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	txn := makeTxnProto()
+	th, mockSender, mockGatekeeper := makeMockTxnHeartbeater(&txn)
+	defer th.stopper.Stop(ctx)
+
+	// Mock the heartbeat request, which simply aborts and signals hbAborted.
+	hbAborted := make(chan struct{})
+	mockGatekeeper.MockSend(func(ba roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
+		defer close(hbAborted)
+
+		require.Len(t, ba.Requests, 1)
+		require.IsType(t, &roachpb.HeartbeatTxnRequest{}, ba.Requests[0].GetInner())
+
+		br := ba.CreateReply()
+		br.Txn = ba.Txn
+		br.Txn.Status = roachpb.ABORTED
+		return br, nil
+	})
+
+	// Mock an EndTxn response, which signals rollbackReady and blocks
+	// until rollbackUnblock is closed.
+	rollbackReady := make(chan struct{})
+	rollbackUnblock := make(chan struct{})
+	mockSender.ChainMockSend(
+		// The first Put request is expected and should just return.
+		func(ba roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
+			require.Len(t, ba.Requests, 1)
+			require.IsType(t, &roachpb.PutRequest{}, ba.Requests[0].GetInner())
+
+			br := ba.CreateReply()
+			br.Txn = ba.Txn
+			return br, nil
+		},
+		// The first EndTxn request from the heartbeater is expected, so block and return.
+		func(ba roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
+			th.mu.Unlock() // manually unlock for concurrency, no txnLockGatekeeper
+			defer th.mu.Lock()
+			close(rollbackReady)
+			require.Len(t, ba.Requests, 1)
+			require.IsType(t, &roachpb.EndTxnRequest{}, ba.Requests[0].GetInner())
+
+			<-rollbackUnblock
+
+			etReq := ba.Requests[0].GetInner().(*roachpb.EndTxnRequest)
+			require.Equal(t, &txn, ba.Txn)
+			require.False(t, etReq.Commit)
+			require.True(t, etReq.Poison)
+			require.True(t, etReq.TxnHeartbeating)
+
+			br := ba.CreateReply()
+			br.Txn = ba.Txn
+			br.Txn.Status = roachpb.ABORTED
+			return br, nil
+		},
+		// The second EndTxn request from the client is unexpected, so
+		// return an error response.
+		func(ba roachpb.BatchRequest) (*roachpb.BatchResponse, *roachpb.Error) {
+			return nil, roachpb.NewError(errors.Errorf("unexpected request: %v", ba))
+		},
+	)
+
+	// Kick off the heartbeat loop.
+	var ba roachpb.BatchRequest
+	ba.Header = roachpb.Header{Txn: txn.Clone()}
+	ba.Add(&roachpb.PutRequest{RequestHeader: roachpb.RequestHeader{Key: roachpb.Key("a")}})
+
+	th.mu.Lock() // manually lock, there's no TxnCoordSender
+	br, pErr := th.SendLocked(ctx, ba)
+	th.mu.Unlock()
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+
+	// Wait for the heartbeater to abort and send an EndTxn.
+	<-hbAborted
+	<-rollbackReady
+
+	// Send a rollback from the client. This should be collapsed together
+	// with the heartbeat abort, and block until it returns. We spawn
+	// a goroutine to unblock the rollback.
+	require.NoError(t, th.stopper.RunAsyncTask(ctx, "put", func(ctx context.Context) {
+		time.Sleep(100 * time.Millisecond)
+		close(rollbackUnblock)
+	}))
+
+	ba = roachpb.BatchRequest{}
+	ba.Header = roachpb.Header{Txn: txn.Clone()}
+	ba.Add(&roachpb.EndTxnRequest{Commit: false})
+
+	th.mu.Lock() // manually lock, there's no TxnCoordSender
+	br, pErr = th.SendLocked(ctx, ba)
+	th.mu.Unlock()
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+
+	// The heartbeat loop should eventually close.
+	waitForHeartbeatLoopToStop(t, &th)
+}
+
 func heartbeaterRunning(th *txnHeartbeater) bool {
 	th.mu.Lock()
 	defer th.mu.Unlock()