kvclient: serve reads from the buffer on the response path

Previously, we were serving reads from the buffer when transforming
a batch. Because of the last commit, this meant that any writes from
earlier in the batch that needed to be added to the buffer would be
added at response time, so the reads served on the request path are
no longer correct. This patch fixes the issue by serving reads from
the buffer on the response path.

Currently, we're not being smart about omitting Get requests for keys
that are reading keys that were written to earlier in the batch. We
could change this behaviour in the future if we think its worthwhile.
SQL never constructs batches that read and write to the same key, so
this is likely not a big deal for now.

Epic: none

Release note: None

Author: arulajmani

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -424,7 +424,9 @@ func (twb *txnWriteBuffer) adjustError(
 		baIdx := int32(0)
 		for i := range numOriginalRequests {
 			if len(ts) > 0 && ts[0].index == i {
-				if ts[0].stripped {
+				curTs := ts[0]
+				ts = ts[1:]
+				if curTs.stripped {
 					numStripped++
 				} else {
 					// This is a transformed request (for example a LockingGet that was
@@ -433,14 +435,19 @@ func (twb *txnWriteBuffer) adjustError(
 					// exactly the request the user sent.
 					//
 					// For now, we handle this by logging and removing the error index.
-					if baIdx == pErr.Index.Index {
+					//
+					// [1] Get requests are always collected as transformations, but
+					// they're never transformed. Attributing an error to them shouldn't
+					// confuse the client.
+					if baIdx == pErr.Index.Index && curTs.origRequest.Method() != kvpb.Get {
 						log.Warningf(ctx, "error index %d is part of a transformed request", pErr.Index.Index)
 						pErr.Index = nil
 						return pErr
 					}
 				}
-				ts = ts[1:]
-				continue
+				if curTs.origRequest.Method() != kvpb.Get {
+					continue
+				}
 			}
 			if baIdx == pErr.Index.Index {
 				break
@@ -713,42 +720,43 @@ func (twb *txnWriteBuffer) applyTransformations(
 
 		case *kvpb.GetRequest:
 			// If the key is in the buffer, we must serve the read from the buffer.
-			val, served := twb.maybeServeRead(t.Key, t.Sequence)
+			// The actual serving of the read will happen on the response path though.
+			stripped := false
+			_, 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 resp kvpb.ResponseUnion
-				getResp := &kvpb.GetResponse{}
-				if val.IsPresent() {
-					getResp.Value = val
-				}
-				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.
+					// request from the remote 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
+					log.VEventf(ctx, 2, "locking %s on key %s must be sent to the server", t.Method(), t.Key)
 					baRemote.Requests = append(baRemote.Requests, ru)
+				} else {
+					// We'll synthesize the response from the buffer on the response path;
+					// eschew sending the request to the KV layer as we don't need to
+					// acquire a lock.
+					stripped = true
+					log.VEventf(
+						ctx, 2, "non-locking %s on key %s can be fully served by the client; not sending to KV", t.Method(), t.Key,
+					)
 				}
-
-				ts = append(ts, transformation{
-					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.
-				continue
+			} else {
+				// Wasn't served locally; send the request to the KV layer.
+				baRemote.Requests = append(baRemote.Requests, ru)
 			}
-			// Wasn't served locally; send the request to the KV layer.
-			baRemote.Requests = append(baRemote.Requests, ru)
+			// Even if the request wasn't served from the buffer here, we still track
+			// a transformation for it. That's because we haven't buffered any writes
+			// from our current batch in the buffer yet, so checking the buffer above
+			// isn't sufficient to determine whether the request needs to serve a read
+			// from the buffer before returning a response or not.
+			ts = append(ts, transformation{
+				stripped:    stripped,
+				index:       i,
+				origRequest: req,
+			})
 
 		case *kvpb.ScanRequest:
 			overlaps := twb.scanOverlaps(t.Key, t.EndKey)
@@ -1054,12 +1062,6 @@ type transformation struct {
 	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). Must be set if stripped is true, but the converse doesn't
-	// hold.
-	resp kvpb.ResponseUnion
 }
 
 // toResp returns the response that should be added to the batch response as
@@ -1142,13 +1144,20 @@ func (t transformation) toResp(
 		twb.addToBuffer(req.Key, roachpb.Value{}, req.Sequence)
 
 	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.
-		assertTrue(t.stripped == (req.KeyLockingStrength == lock.None),
-			"Get requests should either be stripped or be locking")
-		ru = t.resp
+		val, served := twb.maybeServeRead(req.Key, req.Sequence)
+		if served {
+			getResp := &kvpb.GetResponse{}
+			if val.IsPresent() {
+				getResp.Value = val
+			}
+			ru.MustSetInner(getResp)
+			log.VEventf(ctx, 2, "serving %s on key %s from the buffer", req.Method(), req.Key)
+		} else {
+			// The request wasn't served from the buffer; return the response from the
+			// KV layer.
+			assertTrue(!t.stripped, "we shouldn't be stripping requests that aren't served from the buffer")
+			ru = br
+		}
 
 	case *kvpb.ScanRequest:
 		scanResp, err := twb.mergeWithScanResp(

pkg/kv/kvclient/kvcoord/txn_test.go

@@ -42,6 +42,14 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func checkGetResults(t *testing.T, expected map[string][]byte, results ...kv.Result) {
+	for _, result := range results {
+		require.Equal(t, 1, len(result.Rows))
+		require.Equal(t, expected[string(result.Rows[0].Key)], result.Rows[0].ValueBytes())
+	}
+	require.Len(t, expected, len(results))
+}
+
 // TestTxnDBBasics verifies that a simple transaction can be run and
 // either committed or aborted. On commit, mutations are visible; on
 // abort, mutations are never visible. During the txn, verify that
@@ -2161,3 +2169,85 @@ func TestTxnBufferedWriteRetriesCorrectly(t *testing.T) {
 		require.Equal(t, 1, len(keys))
 	})
 }
+
+// TestTxnBufferedWriteReadYourOwnWrites tests that read-your-own-writes are
+// served correctly from the transaction's buffer. We test two cases:
+// 1. The write that needs to be observed was part of an earlier batch.
+// 2. The write that needs to be observed was part of the same batch.
+func TestTxnBufferedWriteReadYourOwnWrites(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+
+	s := createTestDB(t)
+	defer s.Stop()
+
+	value1 := []byte("value1")
+	value21 := []byte("value21")
+	value22 := []byte("value22")
+	value3 := []byte("value3")
+
+	keyA := []byte("keyA")
+	keyB := []byte("keyB")
+	keyC := []byte("keyC")
+
+	// Before the test begins, write a value to keyC.
+	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
+	require.NoError(t, txn.Put(ctx, keyC, value3))
+	require.NoError(t, txn.Commit(ctx))
+
+	err := s.DB.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+		txn.SetBufferedWritesEnabled(true)
+
+		// Put transactional value at keyA.
+		if err := txn.Put(ctx, keyA, value1); err != nil {
+			return err
+		}
+
+		// Construct a batch that contains two Gets -- one on keyA, which will be
+		// served from the buffer, and another on keyC, which will be served by
+		// the server.
+		b := txn.NewBatch()
+		b.Get(keyA)
+		b.Get(keyC)
+		if err := txn.Run(ctx, b); err != nil {
+			return err
+		}
+		expected := map[string][]byte{
+			"keyA": value1,
+			"keyC": value3,
+		}
+		checkGetResults(t, expected, b.Results...)
+
+		// Next, construct a batch that contains both Puts and Gets to keyB. The Get
+		// should see the value written by the Put preceding it in the batch.
+		b = txn.NewBatch()
+		b.Get(keyB)
+		b.Put(keyB, value21)
+		b.Get(keyB)
+		b.Put(keyB, value22)
+		b.Get(keyB)
+
+		if err := txn.Run(ctx, b); err != nil {
+			return err
+		}
+		checkGetResults(t, map[string][]byte{
+			"keyB": nil,
+		},
+			b.Results[0],
+		)
+		checkGetResults(t, map[string][]byte{
+			"keyB": value21,
+		},
+			b.Results[2],
+		)
+		checkGetResults(t, map[string][]byte{
+			"keyB": value22,
+		},
+			b.Results[4],
+		)
+
+		return nil
+	})
+	require.NoError(t, err)
+}