kvclient: ensure locking Get requests are always sent to KV

A Get request to a key that's been previously buffered on the client
must always be served locally. However, if the Get request is locking
in nature, we must also acquire a lock at the leaseholder. This patch
addresses a TODO to do exactly this, while ensuring the response from
KV is discarded and served locally instead.

Informs #139054

Release note: None

Author: arulajmani

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -9,6 +9,7 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
@@ -289,9 +290,10 @@ func (twb *txnWriteBuffer) applyTransformations(
 			// TODO(yuzefovich): ensure that we elide the lock acquisition
 			// whenever possible (e.g. blind UPSERT in an implicit txn).
 			ts = append(ts, transformation{
-				stripped: true,
-				index:    i,
-				resp:     ru,
+				stripped:    true,
+				index:       i,
+				origRequest: req,
+				resp:        ru,
 			})
 			twb.addToBuffer(t.Key, t.Value, t.Sequence)
 
@@ -304,9 +306,10 @@ func (twb *txnWriteBuffer) applyTransformations(
 				FoundKey: false,
 			})
 			ts = append(ts, transformation{
-				stripped: true,
-				index:    i,
-				resp:     ru,
+				stripped:    true,
+				index:       i,
+				origRequest: req,
+				resp:        ru,
 			})
 			twb.addToBuffer(t.Key, roachpb.Value{}, t.Sequence)
 
@@ -315,25 +318,36 @@ func (twb *txnWriteBuffer) applyTransformations(
 			val, served := twb.maybeServeRead(t.Key, t.Sequence)
 			if served {
 				log.VEventf(ctx, 2, "serving %s on key %s from the buffer", t.Method(), t.Key)
-				var ru kvpb.ResponseUnion
+				var resp kvpb.ResponseUnion
 				getResp := &kvpb.GetResponse{}
 				if val.IsPresent() {
 					getResp.Value = val
 				}
-				ru.MustSetInner(getResp)
+				resp.MustSetInner(getResp)
+
+				stripped := true
+				if t.KeyLockingStrength != lock.None {
+					// Even though the Get request must be served from the buffer, as the
+					// transaction performed a previous write to the key, we still need to
+					// acquire a lock at the leaseholder. As a result, we can't strip the
+					// request from the batch.
+					//
+					// TODO(arul): we could eschew sending this request if we knew there
+					// was a sufficiently strong lock already present on the key.
+					stripped = false
+					baRemote.Requests = append(baRemote.Requests, ru)
+				}
+
 				ts = append(ts, transformation{
-					stripped: true,
-					index:    i,
-					resp:     ru,
+					stripped:    stripped,
+					index:       i,
+					origRequest: req,
+					resp:        resp,
 				})
 				// We've constructed a response that we'll stitch together with the
 				// result on the response path; eschew sending the request to the KV
 				// layer.
 				//
-				// TODO(arul): if this is a locking Get request, we can't omit the KV
-				// request. At least, not unless we know a lock is already present on
-				// the key.
-				//
 				// TODO(arul): if the ReturnRawMVCCValues flag is set, we'll need to
 				// flush the buffer.
 				continue
@@ -579,14 +593,14 @@ type transformation struct {
 	// index of the request in the original batch to which the transformation
 	// applies.
 	index int
+	// origRequest is the original request that was transformed.
+	origRequest kvpb.Request
 	// resp is locally produced response that needs to be merged with any
 	// responses returned by the KV layer. This is set for requests that can be
 	// evaluated locally (e.g. blind writes, reads that can be served entirely
-	// from the buffer). If non-empty, stripped must also be true.
+	// from the buffer). Must be set if stripped is true, but the converse doesn't
+	// hold.
 	resp kvpb.ResponseUnion
-	// origRequest is the original request that was transformed. Set iff stripped
-	// is false.
-	origRequest kvpb.Request
 }
 
 // toResp returns the response that should be added to the batch response as
@@ -599,8 +613,21 @@ func (t transformation) toResp(
 	}
 
 	var ru kvpb.ResponseUnion
-	switch br.GetInner().(type) {
-	case *kvpb.ScanResponse:
+	switch t.origRequest.(type) {
+	case *kvpb.PutRequest:
+		ru = t.resp
+	case *kvpb.DeleteRequest:
+		ru = t.resp
+	case *kvpb.GetRequest:
+		// Get requests must be served from the local buffer if a transaction
+		// performed a previous write to the key being read. However, Get requests
+		// must be sent to the KV layer (i.e. not be stripped) iff they are locking
+		// in nature.
+		req := t.origRequest.(*kvpb.GetRequest)
+		assertTrue(t.stripped == (req.KeyLockingStrength == lock.None),
+			"Get requests should either be stripped or be locking")
+		ru = t.resp
+	case *kvpb.ScanRequest:
 		scanResp, err := twb.mergeWithScanResp(
 			t.origRequest.(*kvpb.ScanRequest), br.GetInner().(*kvpb.ScanResponse),
 		)

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -11,8 +11,10 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
@@ -825,3 +827,108 @@ func TestTxnWriteBufferServesOverlappingReadsCorrectly(t *testing.T) {
 	require.Len(t, br.Responses, 1)
 	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
 }
+
+// TestTxnWriteBufferLockingGetRequests ensures that locking get requests are
+// handled appropriately -- they're sent to KV, to acquire a lock, but the
+// read is served from the buffer (upholding read-your-own-write semantics).
+func TestTxnWriteBufferLockingGetRequests(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	twb, mockSender := makeMockTxnWriteBuffer()
+
+	txn := makeTxnProto()
+	txn.Sequence = 10
+	keyA := roachpb.Key("a")
+	valA := "val"
+
+	// Blindly write to keys A.
+	ba := &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	putA := putArgs(keyA, valA, txn.Sequence)
+	ba.Add(putA)
+
+	numCalled := mockSender.NumCalled()
+	br, pErr := twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+	// All the requests should be buffered and not make it past the
+	// txnWriteBuffer.
+	require.Equal(t, numCalled, mockSender.NumCalled())
+	// Even though the txnWriteBuffer did not send any Put requests to the KV
+	// layer above, the responses should still be populated.
+	require.Len(t, br.Responses, 1)
+	require.Equal(t, br.Responses[0].GetInner(), &kvpb.PutResponse{})
+
+	// Verify the writes were buffered correctly.
+	expBufferedWrites := []bufferedWrite{
+		makeBufferedWrite(keyA, makeBufferedValue(valA, 10)),
+	}
+	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
+
+	// Perform a locking read on keyA. Ensure a request is sent to the KV layer,
+	// but the response is served from the buffer.
+	testutils.RunValues(t, "str", []lock.Strength{lock.None, lock.Shared, lock.Exclusive, lock.Update}, func(t *testing.T, strength lock.Strength) {
+		txn.Sequence = 11
+		ba = &kvpb.BatchRequest{}
+		ba.Header = kvpb.Header{Txn: &txn}
+		getA := &kvpb.GetRequest{
+			RequestHeader:      kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+			KeyLockingStrength: strength,
+		}
+		ba.Add(getA)
+		numCalled = mockSender.NumCalled()
+
+		mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+			require.Len(t, ba.Requests, 1)
+			require.IsType(t, &kvpb.GetRequest{}, ba.Requests[0].GetInner())
+			require.Equal(t, strength, ba.Requests[0].GetInner().(*kvpb.GetRequest).KeyLockingStrength)
+			require.True(t, strength != lock.None) // non-locking Gets aren't sent to KV
+
+			br = ba.CreateReply()
+			br.Txn = ba.Txn
+			return br, nil
+		})
+
+		br, pErr = twb.SendLocked(ctx, ba)
+		require.Nil(t, pErr)
+		require.NotNil(t, br)
+		require.Len(t, br.Responses, 1)
+		require.Equal(t, roachpb.MakeValueFromString(valA), *br.Responses[0].GetInner().(*kvpb.GetResponse).Value)
+
+		var expNumCalled int
+		if strength == lock.None {
+			expNumCalled = numCalled // nothing should be sent to KV
+		} else {
+			expNumCalled = numCalled + 1 // a locking request should still be sent to KV
+		}
+		require.Equal(t, expNumCalled, mockSender.NumCalled())
+	})
+
+	// Lastly, for completeness, commit the transaction and ensure that the buffer
+	// is correctly flushed.
+	ba = &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	ba.Add(&kvpb.EndTxnRequest{Commit: true})
+
+	mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+		require.Len(t, ba.Requests, 2)
+
+		// We now expect the buffer to be flushed along with the commit.
+		require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+		require.IsType(t, &kvpb.EndTxnRequest{}, ba.Requests[1].GetInner())
+
+		br = ba.CreateReply()
+		br.Txn = ba.Txn
+		return br, nil
+	})
+
+	br, pErr = twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+
+	// Even though we flushed the buffer, responses from the blind writes should
+	// not be returned.
+	require.Len(t, br.Responses, 1)
+	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
+}