kvserver,storage: add ExpectExclusionSince to Get

This adds ExpectExclusionSince to Get requests. This option semantically works
the same as the ExpectExclusionSince option on Put and Del requests. The
implementation, however, is slightly different. It adjusts the read timestamp
used by storage.MVCCGet and then configures the scanner to return an
ExclusionViolationError rather than a WriteTooOld error.

Informs #142977

Release note: None

Author: stevendanna

pkg/kv/kvpb/api.go

@@ -2530,6 +2530,13 @@ func (r *ConditionalPutRequest) Validate(_ Header) error {
 	return nil
 }
 
+func (r *GetRequest) Validate(_ Header) error {
+	if !r.ExpectExclusionSince.IsEmpty() && r.KeyLockingStrength == lock.None {
+		return errors.AssertionFailedf("invalid GetRequest: ExpectExclusionSince is non-empty for non-locking request")
+	}
+	return nil
+}
+
 func (r *PutRequest) Validate(bh Header) error {
 	if err := validateExclusionTimestampForBatch(r.ExpectExclusionSince, bh); err != nil {
 		return errors.NewAssertionErrorWithWrappedErrf(err, "invalid PutRequest")

pkg/kv/kvpb/api.proto

@@ -223,6 +223,13 @@ message GetRequest {
   // LockNonExisting indicates whether the Get request should acquire a lock
   // on keys that don't exist.
   bool lock_non_existing = 5;
+
+  // ExpectExclusionSince, if set, indicates that the request should return an
+  // error if this key has been written to at a timestamp at or after the given
+  // timestamp. This is different from the ReadTimestamp of the transaction in
+  // that it allows a writer to verify an unreplicated lock acquired at an
+  // earlier timestamp has provided the required protection.
+  util.hlc.Timestamp expect_exclusion_since = 6 [(gogoproto.nullable) = false];
 }
 
 // A GetResponse is the return value from the Get() method.
@@ -2959,7 +2966,7 @@ message Header {
 
   // HasBufferedAllPrecedingWrites, if set, indicates that the batch belongs to
   // a transaction that has buffered all of its writes (from preceding batches)
-  // on the client. 
+  // on the client.
   //
   // The server may use this field to omit checking the AbortSpan. Transactions
   // use the AbortSpan to check whether they've been aborted or not. If they
@@ -2974,7 +2981,7 @@ message Header {
   // read-your-own-writes. As such, they can eschew checking the AbortSpan on
   // the server.
   bool has_buffered_all_preceding_writes = 37;
-  
+
   reserved 7, 10, 12, 14, 20;
 
   // Next ID: 38

pkg/kv/kvpb/api_test.go

@@ -370,6 +370,20 @@ func TestFlagCombinations(t *testing.T) {
 	}
 }
 
+func TestGetValidate(t *testing.T) {
+	t.Run("ExpectExclusionSinceOnNonLockingGet", func(t *testing.T) {
+		getReq := &GetRequest{ExpectExclusionSince: hlc.Timestamp{WallTime: 1}}
+		require.Error(t, getReq.Validate(Header{}))
+	})
+	t.Run("ExpectExclusionSinceOnLockingGet", func(t *testing.T) {
+		getReq := &GetRequest{
+			ExpectExclusionSince: hlc.Timestamp{WallTime: 1},
+			KeyLockingStrength:   lock.Exclusive,
+		}
+		require.NoError(t, getReq.Validate(Header{}))
+	})
+}
+
 func TestRequestHeaderRoundTrip(t *testing.T) {
 	var seq kvnemesisutil.Container
 	seq.Set(123)

pkg/kv/kvpb/batch.go

@@ -108,6 +108,15 @@ func (ba *BatchRequest) EarliestActiveTimestamp() hlc.Timestamp {
 				// NB: StartTime.Next() because StartTime is exclusive.
 				ts.Backward(t.StartTime.Next())
 			}
+		case *GetRequest:
+			// A GetRequest with ExpectExclusionSince set need to be able to observe
+			// MVCC versions from the specified time to correctly detect isolation
+			// violations.
+			//
+			// See the example in RefreshRequest for more details.
+			if !t.ExpectExclusionSince.IsEmpty() {
+				ts.Backward(t.ExpectExclusionSince.Next())
+			}
 		case *PutRequest:
 			// A PutRequest with ExpectExclusionSince set need to be able to observe MVCC
 			// versions from the specified time to correctly detect isolation

pkg/kv/kvserver/batcheval/cmd_get.go

@@ -29,6 +29,10 @@ func Get(
 	h := cArgs.Header
 	reply := resp.(*kvpb.GetResponse)
 
+	if err := args.Validate(h); err != nil {
+		return result.Result{}, err
+	}
+
 	var lockTableForSkipLocked storage.LockTableView
 	if h.WaitPolicy == lock.WaitPolicy_SkipLocked {
 		lockTableForSkipLocked = newRequestBoundLockTableView(
@@ -37,11 +41,25 @@ func Get(
 		defer lockTableForSkipLocked.Close()
 	}
 
-	getRes, err := storage.MVCCGet(ctx, readWriter, args.Key, h.Timestamp, storage.MVCCGetOptions{
+	readTimestamp := h.Timestamp
+	moreRecentPolicy := storage.IgnoreMoreRecent
+
+	// If ExpectExclusionSince is set, use it as the read timestamp and set the
+	// MoreRecentPolicy to ExclusionViolationErrorOnMoreRecent to ensure that an
+	// exclusion violation error is returned if a write has occurred since the
+	// exclusion timestamp.
+	if !args.ExpectExclusionSince.IsEmpty() {
+		readTimestamp = args.ExpectExclusionSince
+		moreRecentPolicy = storage.ExclusionViolationErrorOnMoreRecent
+	} else if args.KeyLockingStrength != lock.None {
+		moreRecentPolicy = storage.WriteTooOldErrorOnMoreRecent
+	}
+
+	getRes, err := storage.MVCCGet(ctx, readWriter, args.Key, readTimestamp, storage.MVCCGetOptions{
 		Inconsistent:          h.ReadConsistency != kvpb.CONSISTENT,
 		SkipLocked:            h.WaitPolicy == lock.WaitPolicy_SkipLocked,
 		Txn:                   h.Txn,
-		FailOnMoreRecent:      args.KeyLockingStrength != lock.None,
+		MoreRecentPolicy:      moreRecentPolicy,
 		ScanStats:             cArgs.ScanStats,
 		Uncertainty:           cArgs.Uncertainty,
 		MemoryAccount:         cArgs.EvalCtx.GetResponseMemoryAccount(),

pkg/kv/kvserver/batcheval/cmd_get_test.go

@@ -11,10 +11,14 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/isolation"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
 )
 
@@ -106,3 +110,251 @@ func TestGetResumeSpan(t *testing.T) {
 		})
 	}
 }
+
+// TestExpectExclusionSince tests evaluation of the ExpectExclusionSince option
+// on all commands that support it.
+func TestExpectExclusionSince(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+
+	var (
+		key   = roachpb.Key([]byte{'a'})
+		value = roachpb.MakeValueFromString("woohoo")
+
+		clock    = hlc.NewClockForTesting(nil)
+		settings = cluster.MakeTestingClusterSettings()
+		evalCtx  = (&MockEvalCtx{Clock: clock, ClusterSettings: settings}).EvalContext()
+		writeTS  = clock.Now()
+	)
+
+	putWithTxn := func(key roachpb.Key, value roachpb.Value, exclusionTS hlc.Timestamp, db storage.Engine, txn *roachpb.Transaction) error {
+		putResp := kvpb.PutResponse{}
+		ts := writeTS
+		if txn != nil {
+			ts = txn.ReadTimestamp
+		}
+		_, err := Put(ctx, db, CommandArgs{
+			EvalCtx: evalCtx,
+			Header: kvpb.Header{
+				Txn:       txn,
+				Timestamp: ts,
+			},
+			Args: &kvpb.PutRequest{
+				ExpectExclusionSince: exclusionTS,
+				RequestHeader:        kvpb.RequestHeader{Key: key},
+				Value:                value,
+			},
+		}, &putResp)
+		return err
+	}
+
+	getWithTxn := func(key roachpb.Key, exclusionTS hlc.Timestamp,
+		str lock.Strength, dur lock.Durability, db storage.Engine, txn *roachpb.Transaction) error {
+		getResp := kvpb.GetResponse{}
+		ts := writeTS
+		if txn != nil {
+			ts = txn.ReadTimestamp
+		}
+		_, err := Get(ctx, db, CommandArgs{
+			EvalCtx: evalCtx,
+			Header: kvpb.Header{
+				Txn:       txn,
+				Timestamp: ts,
+			},
+			Args: &kvpb.GetRequest{
+				LockNonExisting:      true,
+				ExpectExclusionSince: exclusionTS,
+				KeyLockingStrength:   str,
+				KeyLockingDurability: dur,
+				RequestHeader:        kvpb.RequestHeader{Key: key},
+			},
+		}, &getResp)
+		return err
+	}
+
+	// We'll setup each test in a new engine. Placing a write or lock on `key` at
+	// `writeTS`.
+	type existingWriteType string
+	const (
+		intentWrite    existingWriteType = "intent"
+		committedWrite existingWriteType = "committed"
+		replicatedLock existingWriteType = "replicated-lock"
+	)
+
+	var (
+		beforeExistingTS = writeTS.Prev()
+		equalExistingTS  = writeTS
+		afterExistingTS  = writeTS.Next()
+	)
+
+	type testCase struct {
+		writeType                existingWriteType
+		exclusionTS              hlc.Timestamp
+		expectExclusionViolation bool
+		expectLockConflictError  bool
+	}
+
+	setup := func(t *testing.T, writeType existingWriteType) (storage.Engine, *roachpb.Transaction) {
+		db := storage.NewDefaultInMemForTesting()
+
+		var txn *roachpb.Transaction
+		if writeType == intentWrite || writeType == replicatedLock {
+			txn1 := roachpb.MakeTransaction("test", nil, /* baseKey */
+				isolation.Serializable,
+				roachpb.NormalUserPriority, writeTS, 0, 0, 0, false /* omitInRangefeeds */)
+			txn = &txn1
+		}
+		if writeType == replicatedLock {
+			require.NoError(t, getWithTxn(key, hlc.Timestamp{}, lock.Exclusive, lock.Replicated, db, txn))
+		} else {
+			require.NoError(t, putWithTxn(key, value, hlc.Timestamp{}, db, txn))
+		}
+		return db, txn
+	}
+
+	// ops are the operations that support sending a ExpectExclusionSince
+	// timestamp.
+	ops := []struct {
+		name    string
+		request func(key roachpb.Key, exclusionTS hlc.Timestamp, db storage.Engine, txn *roachpb.Transaction) error
+	}{
+		{
+			name: "Get",
+			request: func(key roachpb.Key, exclusionTS hlc.Timestamp, db storage.Engine, txn *roachpb.Transaction) error {
+				return getWithTxn(key, exclusionTS, lock.Exclusive, lock.Replicated, db, txn)
+			},
+		},
+		{
+			name: "Put",
+			request: func(key roachpb.Key, exclusionTS hlc.Timestamp, db storage.Engine, txn *roachpb.Transaction) error {
+				return putWithTxn(key, value, exclusionTS, db, txn)
+			},
+		},
+		{
+			name: "Delete",
+			request: func(key roachpb.Key, exclusionTS hlc.Timestamp, db storage.Engine, txn *roachpb.Transaction) error {
+				delResp := kvpb.DeleteResponse{}
+				_, err := Delete(ctx, db, CommandArgs{
+					EvalCtx: evalCtx,
+					Header: kvpb.Header{
+						Txn:       txn,
+						Timestamp: txn.ReadTimestamp,
+					},
+					Args: &kvpb.DeleteRequest{
+						ExpectExclusionSince: exclusionTS,
+						RequestHeader:        kvpb.RequestHeader{Key: key},
+					},
+				}, &delResp)
+				return err
+			},
+		},
+	}
+
+	testCases := []testCase{
+		// If an intent write exists, it doesn't matter the timestamp it is at. We
+		// expect a lock conflict error and then wait on whoever violated our write
+		// exclusion.
+		{
+			writeType:               intentWrite,
+			exclusionTS:             beforeExistingTS,
+			expectLockConflictError: true,
+		},
+		{
+			writeType:               intentWrite,
+			exclusionTS:             equalExistingTS,
+			expectLockConflictError: true,
+		},
+		{
+			writeType:               intentWrite,
+			exclusionTS:             afterExistingTS,
+			expectLockConflictError: true,
+		},
+
+		// Committed writes are where we expect to see a write exclusion violation
+		// error.
+		{
+			writeType:                committedWrite,
+			exclusionTS:              beforeExistingTS,
+			expectExclusionViolation: true,
+		},
+		{
+			writeType:                committedWrite,
+			exclusionTS:              equalExistingTS,
+			expectExclusionViolation: true,
+		},
+		{
+			writeType:   committedWrite,
+			exclusionTS: afterExistingTS,
+		},
+
+		// For replicatedLocks, we get a lock conflict error in all case just like
+		// an intent write.
+		{
+			writeType:               replicatedLock,
+			exclusionTS:             beforeExistingTS,
+			expectLockConflictError: true,
+		},
+		{
+			writeType:               replicatedLock,
+			exclusionTS:             equalExistingTS,
+			expectLockConflictError: true,
+		},
+		{
+			writeType:               replicatedLock,
+			exclusionTS:             afterExistingTS,
+			expectLockConflictError: true,
+		},
+	}
+
+	for _, op := range ops {
+		for _, tc := range testCases {
+			exclusionTSString := "unknown"
+			if tc.exclusionTS.Equal(equalExistingTS) {
+				exclusionTSString = "equal"
+			} else if tc.exclusionTS.Equal(beforeExistingTS) {
+				exclusionTSString = "before"
+			} else if tc.exclusionTS.Equal(afterExistingTS) {
+				exclusionTSString = "after"
+			}
+
+			for _, sameTxn := range []bool{true, false} {
+				// Committed writes don't have an associated transaction so the sameTxn
+				// variant doesn't make sense here.
+				if sameTxn && tc.writeType == committedWrite {
+					continue
+				}
+
+				name := fmt.Sprintf("%s/%s/exclusion_ts=%s/sameTxn=%v",
+					op.name, tc.writeType, exclusionTSString, sameTxn)
+
+				t.Run(name, func(t *testing.T) {
+					db, txn := setup(t, tc.writeType)
+					defer db.Close()
+
+					if !sameTxn {
+						txn1 := roachpb.MakeTransaction("test", nil, /* baseKey */
+							isolation.Serializable,
+							roachpb.NormalUserPriority, clock.Now(), 0, 0, 0, false /* omitInRangefeeds */)
+						txn = &txn1
+					} else {
+						txn.Sequence++
+						txn.BumpReadTimestamp(clock.Now())
+					}
+
+					err := op.request(key, tc.exclusionTS, db, txn)
+					if sameTxn {
+						require.NoError(t, err, "expected no error for write in the same txn")
+					} else if tc.expectExclusionViolation {
+						require.ErrorContains(t, err, "write exclusion on key")
+					} else if tc.expectLockConflictError {
+						require.ErrorContains(t, err, "conflicting locks on")
+					} else {
+						require.NoError(t, err)
+					}
+				})
+			}
+		}
+	}
+}