kvcoord: minor txnWriteBuffer refactor

This makes some future changes we plan to make a bit easier.

Release note: None

Epic: none

Author: stevendanna

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -267,17 +267,18 @@ func (twb *txnWriteBuffer) SendLocked(
 		return twb.flushBufferAndSendBatch(ctx, ba)
 	}
 
-	if twb.batchRequiresFlush(ctx, ba) {
+	// We check if scan transforms are enabled once and use that answer until the
+	// end of SendLocked.
+	transformScans := bufferedWritesScanTransformEnabled.Get(&twb.st.SV)
+
+	if twb.batchRequiresFlush(ctx, ba, transformScans) {
 		return twb.flushBufferAndSendBatch(ctx, ba)
 	}
 
 	// Check if buffering writes from the supplied batch will run us over
 	// budget. If it will, we shouldn't buffer writes from the current batch,
 	// and flush the buffer.
 	maxSize := bufferedWritesMaxBufferSize.Get(&twb.st.SV)
-	// We check if scan transforms are enabled once and use that answer until the
-	// end of SendLocked.
-	transformScans := bufferedWritesScanTransformEnabled.Get(&twb.st.SV)
 	bufSize := twb.estimateSize(ba, transformScans) + twb.bufferSize
 
 	// NB: if bufferedWritesMaxBufferSize is set to 0 then we effectively disable
@@ -331,7 +332,9 @@ func (twb *txnWriteBuffer) SendLocked(
 	return twb.mergeResponseWithRequestRecords(ctx, rr, br)
 }
 
-func (twb *txnWriteBuffer) batchRequiresFlush(ctx context.Context, ba *kvpb.BatchRequest) bool {
+func (twb *txnWriteBuffer) batchRequiresFlush(
+	ctx context.Context, ba *kvpb.BatchRequest, transformScans bool,
+) bool {
 	for _, ru := range ba.Requests {
 		req := ru.GetInner()
 		switch req.(type) {
@@ -477,7 +480,7 @@ func (twb *txnWriteBuffer) estimateSize(ba *kvpb.BatchRequest, transformScans bo
 			estimate += scratch.size()
 			estimate += lockKeyInfoSize
 		case *kvpb.GetRequest:
-			if t.KeyLockingDurability == lock.Replicated {
+			if IsReplicatedLockingRequest(t) {
 				scratch.key = t.Key
 				estimate += scratch.size()
 				estimate += lockKeyInfoSize
@@ -514,9 +517,7 @@ func (twb *txnWriteBuffer) estimateSize(ba *kvpb.BatchRequest, transformScans bo
 			// the buffer. Here, we assume at least 1 key will be returned that is
 			// about the size of the scan start boundary. We try to protect from large
 			// buffer overflows by transforming the batch's MaxSpanRequestKeys later.
-			shouldTransform := t.KeyLockingStrength > lock.None && t.KeyLockingDurability == lock.Replicated
-			shouldTransform = shouldTransform && transformScans
-			if shouldTransform {
+			if IsReplicatedLockingRequest(t) && transformScans {
 				scratch.key = t.Key
 				scratch.vals[0] = bufferedValue{
 					seq: t.Sequence,
@@ -525,9 +526,7 @@ func (twb *txnWriteBuffer) estimateSize(ba *kvpb.BatchRequest, transformScans bo
 			}
 		case *kvpb.ReverseScanRequest:
 			// See the comment on the ScanRequest case for more details.
-			shouldTransform := t.KeyLockingStrength > lock.None && t.KeyLockingDurability == lock.Replicated
-			shouldTransform = shouldTransform && transformScans
-			if shouldTransform {
+			if IsReplicatedLockingRequest(t) && transformScans {
 				scratch.key = t.Key
 				scratch.vals[0] = bufferedValue{
 					seq: t.Sequence,
@@ -932,7 +931,7 @@ func (twb *txnWriteBuffer) applyTransformations(
 			_, lockStr, served := twb.maybeServeRead(t.Key, t.Sequence)
 
 			requiresAdditionalLocking := t.KeyLockingStrength > lockStr
-			requiresLockTransform := t.KeyLockingStrength != lock.None && t.KeyLockingDurability == lock.Replicated
+			requiresLockTransform := IsReplicatedLockingRequest(t)
 			requestRequired := requiresAdditionalLocking || !served
 
 			if requestRequired && requiresLockTransform {
@@ -953,9 +952,7 @@ func (twb *txnWriteBuffer) applyTransformations(
 			// Regardless of whether the scan overlaps with any writes in the buffer
 			// or not, we must send the request to the KV layer. We can't know for
 			// sure that there's nothing else to read.
-			shouldTransform := t.KeyLockingStrength > lock.None && t.KeyLockingDurability == lock.Replicated
-			shouldTransform = shouldTransform && transformScans
-			if shouldTransform {
+			if IsReplicatedLockingRequest(t) && transformScans {
 				var scanReqU kvpb.RequestUnion
 				scanReq := t.ShallowCopy().(*kvpb.ScanRequest)
 				scanReq.KeyLockingDurability = lock.Unreplicated
@@ -973,9 +970,7 @@ func (twb *txnWriteBuffer) applyTransformations(
 			// Regardless of whether the reverse scan overlaps with any writes in the
 			// buffer or not, we must send the request to the KV layer. We can't know
 			// for sure that there's nothing else to read.
-			shouldTransform := t.KeyLockingStrength > lock.None && t.KeyLockingDurability == lock.Replicated
-			shouldTransform = shouldTransform && transformScans
-			if shouldTransform {
+			if IsReplicatedLockingRequest(t) && transformScans {
 				var rScanReqU kvpb.RequestUnion
 				rScanReq := t.ShallowCopy().(*kvpb.ReverseScanRequest)
 				rScanReq.KeyLockingDurability = lock.Unreplicated
@@ -1635,6 +1630,19 @@ func (twb *txnWriteBuffer) addToBuffer(
 	}
 }
 
+func IsReplicatedLockingRequest(req kvpb.Request) bool {
+	switch r := req.(type) {
+	case *kvpb.ScanRequest:
+		return r.KeyLockingStrength > lock.None && r.KeyLockingDurability == lock.Replicated
+	case *kvpb.ReverseScanRequest:
+		return r.KeyLockingStrength > lock.None && r.KeyLockingDurability == lock.Replicated
+	case *kvpb.GetRequest:
+		return r.KeyLockingStrength > lock.None && r.KeyLockingDurability == lock.Replicated
+	default:
+		return false
+	}
+}
+
 // addDurableLockedReadToBuffer adds a locking read to the given buffer.
 //
 // TODO(ssd): Determine if we need to track the kvnemesis sequence number for