kvcoord: allow write buffer to re-enable pipelining

This is an attempt to allow the pipeliner to re-enable pipelining when
it is flushing its buffer.

Epic: none
Release note: None

Author: stevendanna

pkg/kv/kvclient/kvcoord/txn_coord_sender.go

@@ -281,6 +281,10 @@ func newRootTxnCoordSender(
 		timeSource: timeutil.DefaultTimeSource{},
 		txn:        &tcs.mu.txn,
 	}
+	tcs.interceptorAlloc.txnWriteBuffer.init(
+		&tcs.interceptorAlloc.txnPipeliner,
+	)
+
 	tcs.initCommonInterceptors(tcf, txn, kv.RootTxn)
 
 	// Once the interceptors are initialized, piece them all together in the
@@ -446,6 +450,7 @@ func (tc *TxnCoordSender) DisablePipelining() error {
 	if tc.mu.active {
 		return errors.Errorf("cannot disable pipelining on a running transaction")
 	}
+	tc.interceptorAlloc.txnPipeliner.disabledExplicitly = true
 	tc.interceptorAlloc.txnPipeliner.disabled = true
 	return nil
 }

pkg/kv/kvclient/kvcoord/txn_interceptor_pipeliner.go

@@ -226,9 +226,14 @@ var rejectTxnMaxCount = settings.RegisterIntSetting(
 // attached to any end transaction request that is passed through the pipeliner
 // to ensure that they the locks within them are released.
 type txnPipeliner struct {
-	st                       *cluster.Settings
-	riGen                    rangeIteratorFactory // used to condense lock spans, if provided
-	wrapped                  lockedSender
+	st      *cluster.Settings
+	riGen   rangeIteratorFactory // used to condense lock spans, if provided
+	wrapped lockedSender
+	// disabledExplicitly tracks whether the user called txn.DisablePipelining().
+	// This is separate from disabled so that the txnWriteBuffer can enable
+	// pipelining if it has flushed its buffer iff pipelining wasn't previously
+	// explicitly disabled.
+	disabledExplicitly       bool
 	disabled                 bool
 	txnMetrics               *TxnMetrics
 	condensedIntentsEveryN   *log.EveryN
@@ -587,6 +592,12 @@ func (tp *txnPipeliner) canUseAsyncConsensus(ctx context.Context, ba *kvpb.Batch
 	return true
 }
 
+// enableImplicitPipelining enables pipelining unless pipelining was explicitly
+// disabled previously.
+func (tp *txnPipeliner) enableImplicitPipelining() {
+	tp.disabled = tp.disabledExplicitly
+}
+
 // chainToInFlightWrites ensures that we "chain" on to any in-flight writes that
 // overlap the keys we're trying to read/write. We do this by prepending
 // QueryIntent requests with the ErrorIfMissing option before each request that

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -183,6 +183,8 @@ type txnWriteBuffer struct {
 	// `flushed` tracks whether the buffer has been previously flushed.
 	flushed bool
 
+	pipelineEnabler pipelineEnabler
+
 	// flushOnNextBatch, if set, indicates that write buffering has just been
 	// disabled, and the interceptor should flush any buffered writes when it
 	// sees the next BatchRequest.
@@ -215,6 +217,14 @@ type txnWriteBuffer struct {
 	testingOverrideCPutEvalFn func(expBytes []byte, actVal *roachpb.Value, actValPresent bool, allowNoExisting bool) *kvpb.ConditionFailedError
 }
 
+type pipelineEnabler interface {
+	enableImplicitPipelining()
+}
+
+func (twb *txnWriteBuffer) init(pe pipelineEnabler) {
+	twb.pipelineEnabler = pe
+}
+
 func (twb *txnWriteBuffer) setEnabled(enabled bool) {
 	if !enabled && twb.buffer.Len() > 0 {
 		// When disabling write buffering, if we evaluated any requests, we need
@@ -1678,7 +1688,14 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 	// Once we've flushed the buffer, we disable write buffering going forward. We
 	// do this even if the buffer is empty since once we've called this function,
 	// our buffer no longer represents all of the writes in the transaction.
-	log.VEventf(ctx, 2, "disabling write buffering for this epoch")
+	log.VEventf(ctx, 2, "disabling write buffering")
+	if twb.pipelineEnabler != nil {
+		// We enable pipelining, but only after this request returns.
+		//
+		// TODO(ssd): Consider enabling pipelining for this batch as well once we
+		// have more discipline around the invariants of the batches we are sending.
+		defer twb.pipelineEnabler.enableImplicitPipelining()
+	}
 	twb.flushed = true
 
 	numKeysBuffered := twb.buffer.Len()