kvcoord: correctly initialize txnWriteBuffer

Previously we were initializing the pipelineEnabler on a copy of the
txnWriteBuffer that would later be overwritten. As a result, we were
never re-enabling pipelining after a mid-transaction flush. This is OK
from a correctness perspective since this was only intended as a
potential performance improvement for some workloads under buffered
writes.

There is a potentially strong argument that we shouldn't do this at all
and that if someone sees a regression of their workload under buffered
writes they should turn it off rather than relying on this fallback.

Epic: None
Release note: None

Author: stevendanna

pkg/kv/kvclient/kvcoord/txn_coord_sender.go

@@ -281,9 +281,6 @@ func newRootTxnCoordSender(
 		timeSource: timeutil.DefaultTimeSource{},
 		txn:        &tcs.mu.txn,
 	}
-	tcs.interceptorAlloc.txnWriteBuffer.init(
-		&tcs.interceptorAlloc.txnPipeliner,
-	)
 
 	tcs.initCommonInterceptors(tcf, txn, kv.RootTxn)
 
@@ -363,6 +360,7 @@ func (tc *TxnCoordSender) initCommonInterceptors(
 		allowConcurrentRequests: typ == kv.LeafTxn,
 	}
 	tc.interceptorAlloc.txnSeqNumAllocator.writeSeq = txn.Sequence
+	tc.interceptorAlloc.txnWriteBuffer.init(&tc.interceptorAlloc.txnPipeliner)
 }
 
 func (tc *TxnCoordSender) connectInterceptors() {

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_client_test.go

@@ -92,3 +92,71 @@ func TestTxnCoordSenderWriteBufferingDisablesPipelining(t *testing.T) {
 		kvpb.Scan, kvpb.Put, kvpb.EndTxn,
 	}, calls)
 }
+
+// TestTxnCoordSenderWriteBufferingReEnablesPipelining verifies that pipelining
+// is re-enabled after a mid-transaction flush.
+func TestTxnCoordSenderWriteBufferingReEnablesPipelining(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	s := serverutils.StartServerOnly(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(ctx)
+
+	distSender := s.DistSenderI().(*DistSender)
+	batchCount := 0
+	var calls []kvpb.Method
+	var senderFn kv.SenderFunc = func(
+		ctx context.Context, ba *kvpb.BatchRequest,
+	) (*kvpb.BatchResponse, *kvpb.Error) {
+		batchCount++
+		t.Logf("Batch: %#+v", ba.Methods())
+		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)
+	}
+
+	st := s.ClusterSettings()
+	BufferedWritesMaxBufferSize.Override(ctx, &st.SV, defaultBufferSize)
+
+	tsf := NewTxnCoordSenderFactory(TxnCoordSenderFactoryConfig{
+		AmbientCtx: s.AmbientCtx(),
+		Settings:   st,
+		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())
+
+	require.NoError(t, db.Put(ctx, "test-key-a", "hello"))
+	require.NoError(t, db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+		txn.SetBufferedWritesEnabled(true)
+		if err := txn.Put(ctx, "test-key-c", "hello"); err != nil {
+			return err
+		}
+		if _, err := txn.DelRange(ctx, "test-key", "test-key-d", true); err != nil {
+			return err
+		}
+		if err := txn.Put(ctx, "test-key-a", "hello"); err != nil {
+			return err
+		}
+		return nil
+	}))
+
+	require.Equal(t, 4, batchCount)
+	require.Equal(t, []kvpb.Method{
+		// The initial setup
+		kvpb.Put,
+		// The first (buffered) Put and the DeleteRange that flushes the buffer.
+		kvpb.Put, kvpb.DeleteRange,
+		// The second (pipelined) Put
+		kvpb.Put,
+		// EndTxn with the QueryIntent because pipelining was turned back on.
+		kvpb.QueryIntent, kvpb.EndTxn,
+	}, calls)
+}

pkg/sql/opt/exec/execbuilder/testdata/show_trace

@@ -398,6 +398,7 @@ SELECT operation, message FROM [SHOW KV TRACE FOR SESSION]
 WHERE message     LIKE '%r$rangeid: sending batch%'
   AND message NOT LIKE '%PushTxn%'
   AND message NOT LIKE '%QueryTxn%'
+  AND message NOT LIKE '%QueryIntent%'
 ----
 dist sender send  r77: sending batch 6 CPut to (n1,s1):1
 dist sender send  r77: sending batch 6 CPut to (n1,s1):1