Merge #148571

148571: kvclient: change lock elision logic r=stevendanna a=arulajmani

Previously, when write requests had MustAcquireExclusiveLock set on them, we would elide a transformed locking Get request regardless of the strength with which a lock was known to exist on the key. This patch cleans things up to check the strength with which the known to exist lock is held, only eliding the locking Get if the lock is held with Exclusive lock strength.

I've modified the TestTxnWriteBufferElidesUnnecessaryLockingRequests to account for buffered shared locks in the future. This isn't testing anything meaningful yet, though, because no lock knowledge is added to the buffer for locking reads currently.

Epic: none

Release note: None

Co-authored-by: Arul Ajmani <[email protected]>

Author: craig[bot]

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -743,7 +743,7 @@ func (twb *txnWriteBuffer) applyTransformations(
 			_, lockStr, isServed := twb.maybeServeRead(t.Key, t.Sequence)
 			// To elide the locking request, we must have both a value (to evaluate
 			// the condition) and a lock.
-			if isServed && lockExcludesWrites(lockStr) {
+			if isServed && lockStr == lock.Exclusive {
 				record.stripped = true
 			} else {
 				record.transformed = true
@@ -767,7 +767,7 @@ func (twb *txnWriteBuffer) applyTransformations(
 			// If the MustAcquireExclusiveLock flag is set then we need to add a
 			// locking Get to the BatchRequest, including if the key doesn't exist. We
 			// can elide this locking request when we already have an existing lock.
-			lockRequired := t.MustAcquireExclusiveLock && !lockExcludesWrites(lockStr)
+			lockRequired := t.MustAcquireExclusiveLock && lockStr != lock.Exclusive
 			if lockRequired {
 				var getReqU kvpb.RequestUnion
 				getReqU.MustSetInner(&kvpb.GetRequest{
@@ -787,9 +787,9 @@ func (twb *txnWriteBuffer) applyTransformations(
 			_, lockStr, served := twb.maybeServeRead(t.Key, t.Sequence)
 			// If MustAcquireExclusiveLock flag is set on the DeleteRequest, then we
 			// need to add a locking Get to the BatchRequest, including if the key
-			// doesn't exist. We can elide this locking request when we have both a
-			// value (to populate the FoundKey field in the response) and a lock.
-			lockRequired := t.MustAcquireExclusiveLock && !(served && lockExcludesWrites(lockStr))
+			// doesn't exist. We can only elide this locking request when we have both
+			// a value (to populate the FoundKey field in the response) and a lock.
+			lockRequired := t.MustAcquireExclusiveLock && !(served && lockStr == lock.Exclusive)
 			if lockRequired {
 				var getReqU kvpb.RequestUnion
 				getReqU.MustSetInner(&kvpb.GetRequest{
@@ -1799,10 +1799,6 @@ func (li *lockedKeyInfo) rollbackSequence(seq enginepb.TxnSeq) bool {
 	return stillHeld
 }
 
-func lockExcludesWrites(str lock.Strength) bool {
-	return str >= lock.Shared
-}
-
 // getKey reads the key for the next KV from a slice of BatchResponses field of
 // {,Reverse}ScanResponse. The KV is encoded in the following format:
 //

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -2865,12 +2865,18 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 		// it locks. Every place where this is false represents a missed opportunity
 		// to avoid a locking request.
 		buffersLock bool
+		// bufferedLockStr indicates the strength of the lock that is buffered. Used
+		// in conjunction with buffersLock above.
+		bufferedLockStr lock.Strength
 		// isFullyCovered, when true, indicates that this request can be served
-		// completely from the buffer if a previous request locked the value. This
-		// is different from requiresValue as it indicates that requests that aren't
-		// fully covered must always be sent to KV since there is no way to know if
-		// the buffer has all required required values.
-		isFullyCovered bool
+		// completely from the buffer if a previous request locked the value at an
+		// equal or higher lock strength. If the request is never fully covered, use
+		// lock.None.
+		//
+		// This is different from requiresValue as it
+		// indicates that requests that aren't fully covered must always be sent to
+		// KV since there is no way to know if the buffer has all required values.
+		isFullyCoveredByLockStrength lock.Strength
 
 		generateRequest             func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction)
 		optionalGenSecondRequest    func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction)
@@ -2884,24 +2890,26 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 	keyC := roachpb.Key("c")
 
 	// A number of the requests expect a single locking get to be produced. This
-	// validate function is shared across them.
-	validateLockingGet := func(t *testing.T, ba *kvpb.BatchRequest) *kvpb.BatchResponse {
-		getReq := ba.Requests[0].GetGet()
-		require.NotNil(t, getReq)
-		require.Equal(t, lock.Exclusive, getReq.KeyLockingStrength)
-		require.Equal(t, lock.Unreplicated, getReq.KeyLockingDurability)
-		resp := ba.CreateReply()
-		resp.Txn = ba.Txn
-		return resp
+	// validator function is shared across them.
+	validateLockingGetOfStr := func(str lock.Strength) func(t *testing.T, ba *kvpb.BatchRequest) *kvpb.BatchResponse {
+		return func(t *testing.T, ba *kvpb.BatchRequest) *kvpb.BatchResponse {
+			getReq := ba.Requests[0].GetGet()
+			require.NotNil(t, getReq)
+			require.Equal(t, str, getReq.KeyLockingStrength)
+			require.Equal(t, lock.Unreplicated, getReq.KeyLockingDurability) // NB: always Unreplicated
+			resp := ba.CreateReply()
+			resp.Txn = ba.Txn
+			return resp
+		}
 	}
 
 	reqs := []lockingRequests{
 		{
-			name:           "ReplicatedLockingScanScanFormat=KEY_VALUES",
-			buffersLock:    false,
-			buffersValue:   false,
-			requiresValue:  true,
-			isFullyCovered: false,
+			name:                         "ReplicatedLockingScanScanFormat=KEY_VALUES",
+			buffersLock:                  false,
+			buffersValue:                 false,
+			requiresValue:                true,
+			isFullyCoveredByLockStrength: lock.None, // Scans are never covered
 			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
 				ba.Add(&kvpb.ScanRequest{
 					KeyLockingStrength:   lock.Exclusive,
@@ -2923,11 +2931,11 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 			},
 		},
 		{
-			name:           "ReplicatedLockingScanScanFormat=BATCH_RESPONSE",
-			buffersLock:    false,
-			buffersValue:   false,
-			requiresValue:  true,
-			isFullyCovered: false,
+			name:                         "ReplicatedLockingScanScanFormat=BATCH_RESPONSE",
+			buffersLock:                  false,
+			buffersValue:                 false,
+			requiresValue:                true,
+			isFullyCoveredByLockStrength: lock.None, // Scans are never covered
 			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
 				ba.Add(&kvpb.ScanRequest{
 					KeyLockingStrength:   lock.Exclusive,
@@ -2953,11 +2961,11 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 			},
 		},
 		{
-			name:           "ReplicatedLockingReverseScan",
-			buffersLock:    false,
-			buffersValue:   false,
-			requiresValue:  true,
-			isFullyCovered: false,
+			name:                         "ReplicatedLockingReverseScan",
+			buffersLock:                  false,
+			buffersValue:                 false,
+			requiresValue:                true,
+			isFullyCoveredByLockStrength: lock.None, // ReverseScans are never covered
 			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
 				ba.Add(&kvpb.ReverseScanRequest{
 					KeyLockingStrength:   lock.Exclusive,
@@ -2980,11 +2988,11 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 			},
 		},
 		{
-			name:           "ReplicatedLockingGet",
-			buffersLock:    false,
-			buffersValue:   false,
-			requiresValue:  true,
-			isFullyCovered: true,
+			name:                         "ReplicatedLockingGet(Strength=Exclusive)",
+			buffersLock:                  false,
+			buffersValue:                 false,
+			requiresValue:                true,
+			isFullyCoveredByLockStrength: lock.Exclusive,
 			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
 				getReq := &kvpb.GetRequest{
 					RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
@@ -2993,42 +3001,77 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 				getReq.KeyLockingStrength = lock.Exclusive
 				ba.Add(getReq)
 			},
-			validateAndGenerateResponse: validateLockingGet,
+			validateAndGenerateResponse: validateLockingGetOfStr(lock.Exclusive),
+		},
+		{
+			name:                         "UnReplicatedLockingGet(Strength=Exclusive)",
+			buffersLock:                  false,
+			buffersValue:                 false,
+			requiresValue:                true,
+			isFullyCoveredByLockStrength: lock.Exclusive,
+			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
+				getReq := &kvpb.GetRequest{
+					RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+				}
+				getReq.KeyLockingDurability = lock.Unreplicated
+				getReq.KeyLockingStrength = lock.Exclusive
+				ba.Add(getReq)
+			},
+			validateAndGenerateResponse: validateLockingGetOfStr(lock.Exclusive),
+		},
+		{
+			name:                         "ReplicatedLockingGet(Strength=Shared)",
+			buffersLock:                  false,
+			buffersValue:                 false,
+			requiresValue:                true,
+			isFullyCoveredByLockStrength: lock.Shared,
+			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
+				getReq := &kvpb.GetRequest{
+					RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+				}
+				getReq.KeyLockingDurability = lock.Replicated
+				getReq.KeyLockingStrength = lock.Shared
+				ba.Add(getReq)
+			},
+			validateAndGenerateResponse: validateLockingGetOfStr(lock.Shared),
 		},
 		{
-			name:           "PutMustAcquireExclusive",
-			buffersLock:    true,
-			buffersValue:   true,
-			requiresValue:  false,
-			isFullyCovered: true,
+			name:                         "PutMustAcquireExclusive",
+			buffersLock:                  true,
+			bufferedLockStr:              lock.Exclusive,
+			buffersValue:                 true,
+			requiresValue:                false,
+			isFullyCoveredByLockStrength: lock.Exclusive,
 			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
 				putReq := putArgs(keyA, valueAStr, txn.Sequence)
 				putReq.MustAcquireExclusiveLock = true
 				ba.Add(putReq)
 			},
-			validateAndGenerateResponse: validateLockingGet,
+			validateAndGenerateResponse: validateLockingGetOfStr(lock.Exclusive),
 		},
 		{
-			name:         "DeleteMustAcquireExclusive",
-			buffersLock:  true,
-			buffersValue: true,
+			name:            "DeleteMustAcquireExclusive",
+			buffersLock:     true,
+			bufferedLockStr: lock.Exclusive,
+			buffersValue:    true,
 			// Delete requires a value to correctly return the number of deleted row.
-			requiresValue:  true,
-			isFullyCovered: true,
+			requiresValue:                true,
+			isFullyCoveredByLockStrength: lock.Exclusive,
 			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
 				delReq := delArgs(keyA, txn.Sequence)
 				delReq.MustAcquireExclusiveLock = true
 				ba.Add(delReq)
 			},
-			validateAndGenerateResponse: validateLockingGet,
+			validateAndGenerateResponse: validateLockingGetOfStr(lock.Exclusive),
 		},
 		{
-			name:         "ConditionalPut",
-			buffersLock:  true,
-			buffersValue: true,
+			name:            "ConditionalPut",
+			buffersLock:     true,
+			bufferedLockStr: lock.Exclusive,
+			buffersValue:    true,
 			// ConditionalPut requires a value to evaluate its expected bytes condition.
-			requiresValue:  true,
-			isFullyCovered: true,
+			requiresValue:                true,
+			isFullyCoveredByLockStrength: lock.Exclusive,
 			generateRequest: func(t *testing.T, ba *kvpb.BatchRequest, txn *roachpb.Transaction) {
 				ba.Add(cputArgs(keyA, valueAStr, "", txn.Sequence))
 			},
@@ -3039,7 +3082,7 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 				cput.AllowIfDoesNotExist = true
 				ba.Add(cput)
 			},
-			validateAndGenerateResponse: validateLockingGet,
+			validateAndGenerateResponse: validateLockingGetOfStr(lock.Exclusive),
 		},
 	}
 
@@ -3080,8 +3123,12 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 					secondReq.generateRequest(t, ba, &txn)
 				}
 
+				// The second request is fully covered if the first request has already
+				// acquired a lock at a sufficient lock strength.
+				isFullyCovered := secondReq.isFullyCoveredByLockStrength != lock.None &&
+					firstReq.bufferedLockStr >= secondReq.isFullyCoveredByLockStrength
 				mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
-					if secondReq.isFullyCovered {
+					if isFullyCovered {
 						resp := ba.CreateReply()
 						resp.Txn = ba.Txn
 						return resp, nil
@@ -3092,7 +3139,7 @@ func TestTxnWriteBufferElidesUnnecessaryLockingRequests(t *testing.T) {
 				})
 
 				expectedCalls := 0
-				if !secondReq.isFullyCovered {
+				if !isFullyCovered {
 					expectedCalls = 1
 				}
 				numCalled := mockSender.NumCalled()