Merge cockroachdb#162285

162285: kvserver: remove kv.Sender implementation from Store, Stores, Replica r=tbg a=sumeerbhola

Previously, Store, Stores, and Replica all implemented kv.Sender via their Send methods. This allowed callers to use them interchangeably with other Senders, but obscured an important detail: these Send methods were called from code that did not integrate with admission control.

This commit removes the Send methods from these types. They now implement only SenderWithWriteBytes, which returns the StoreWriteBytes that callers that use admission control must handle appropriately.

For test code that needs a kv.Sender, ToSenderForTesting provides an explicit wrapper that discards the StoreWriteBytes. The "ForTesting" suffix makes it clear that this wrapper bypasses admission control accounting and should only be used in tests.

This change makes admission control bypass explicit: any code path that converts these types to kv.Sender must now visibly call ToSenderForTesting, making it easy to audit which paths skip admission control.

Release note: None

Epic: None

Co-authored-by: sumeerbhola <sumeer@cockroachlabs.com>

Author: craig[bot]

pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go

@@ -183,7 +183,7 @@ func TestTxnCoordSenderHeartbeat(t *testing.T) {
 		kvcoord.NewDistSenderForLocalTestCluster(
 			ctx,
 			s.Cfg.Settings, &roachpb.NodeDescriptor{NodeID: 1},
-			ambient.Tracer, s.Clock, s.Latency, s.Stores, s.Stopper(), s.Gossip,
+			ambient.Tracer, s.Clock, s.Latency, kvserver.ToSenderForTesting(s.Stores), s.Stopper(), s.Gossip,
 		),
 	)
 	quickHeartbeatDB := kv.NewDB(ambient, tsf, s.Clock, s.Stopper())
@@ -287,7 +287,7 @@ func TestDB_PrepareForRetryAfterHeartbeatFailure(t *testing.T) {
 		kvcoord.NewDistSenderForLocalTestCluster(
 			ctx,
 			s.Cfg.Settings, &roachpb.NodeDescriptor{NodeID: 1},
-			ambient.Tracer, s.Clock, s.Latency, s.Stores, s.Stopper(), s.Gossip,
+			ambient.Tracer, s.Clock, s.Latency, kvserver.ToSenderForTesting(s.Stores), s.Stopper(), s.Gossip,
 		),
 	)
 	db := kv.NewDB(ambient, tsf, s.Clock, s.Stopper())

pkg/kv/kvserver/client_merge_test.go

@@ -393,7 +393,7 @@ func mergeWithData(t *testing.T, retries int64) {
 	verify(rhsDesc.StartKey.Next().AsRawKey(), rhsRepl.RangeID, newContent)
 
 	gArgs := getArgs(lhsDesc.StartKey.Next().AsRawKey())
-	if _, pErr := kv.SendWrappedWith(ctx, store, kvpb.Header{
+	if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store), kvpb.Header{
 		RangeID: rhsDesc.RangeID,
 	}, gArgs); !testutils.IsPError(
 		pErr, `was not found on s`,
@@ -561,7 +561,7 @@ func mergeCheckingTimestampCaches(
 
 	// Write a key to the RHS.
 	rhsKey := scratchKey("c")
-	if _, pErr := kv.SendWrappedWith(ctx, rhsStore, kvpb.Header{
+	if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(rhsStore), kvpb.Header{
 		RangeID: rhsDesc.RangeID,
 	}, incrementArgs(rhsKey, 1)); pErr != nil {
 		t.Fatal(pErr)
@@ -577,7 +577,7 @@ func mergeCheckingTimestampCaches(
 	ba.Timestamp = readTS
 	ba.RangeID = rhsDesc.RangeID
 	ba.Add(getArgs(rhsKey))
-	if br, pErr := rhsStore.Send(ctx, ba); pErr != nil {
+	if br, pErr := kvserver.ToSenderForTesting(rhsStore).Send(ctx, ba); pErr != nil {
 		t.Fatal(pErr)
 	} else if v, err := br.Responses[0].GetGet().Value.GetInt(); err != nil {
 		t.Fatal(err)
@@ -596,7 +596,7 @@ func mergeCheckingTimestampCaches(
 	ba.Timestamp = readTS.Next()
 	ba.RangeID = rhsDesc.RangeID
 	ba.Add(pushTxnArgs(&pusher, &pushee, kvpb.PUSH_ABORT))
-	if br, pErr := rhsStore.Send(ctx, ba); pErr != nil {
+	if br, pErr := kvserver.ToSenderForTesting(rhsStore).Send(ctx, ba); pErr != nil {
 		t.Fatal(pErr)
 	} else if txn := br.Responses[0].GetPushTxn().PusheeTxn; txn.Status != roachpb.ABORTED {
 		t.Fatalf("expected aborted pushee, but got %v", txn)
@@ -719,7 +719,7 @@ func mergeCheckingTimestampCaches(
 			// be the only replica that does not apply the proposal.
 			go func() {
 				incArgs := incrementArgs(lhsKey, 4)
-				_, pErr := kv.SendWrappedWith(ctx, lhsStore, kvpb.Header{RangeID: lhsDesc.RangeID}, incArgs)
+				_, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(lhsStore), kvpb.Header{RangeID: lhsDesc.RangeID}, incArgs)
 				incChan <- pErr
 			}()
 			// NB: the operation won't complete, so peek below Raft and wait for
@@ -749,7 +749,7 @@ func mergeCheckingTimestampCaches(
 			go func() {
 				truncArgs := truncateLogArgs(truncIndex, lhsDesc.RangeID)
 				truncArgs.Key = lhsKey
-				_, pErr := kv.SendWrappedWith(ctx, lhsStore, kvpb.Header{RangeID: lhsDesc.RangeID}, truncArgs)
+				_, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(lhsStore), kvpb.Header{RangeID: lhsDesc.RangeID}, truncArgs)
 				truncChan <- pErr
 			}()
 			// NB: the operation won't complete, so peek below Raft and wait for
@@ -858,7 +858,7 @@ func mergeCheckingTimestampCaches(
 	ba.Timestamp = readTS
 	ba.RangeID = lhsDesc.RangeID
 	ba.Add(incrementArgs(rhsKey, 1))
-	if br, pErr := lhsStore.Send(ctx, ba); pErr != nil {
+	if br, pErr := kvserver.ToSenderForTesting(lhsStore).Send(ctx, ba); pErr != nil {
 		t.Fatal(pErr)
 	} else if br.Timestamp.LessEq(readTS) {
 		t.Fatalf("expected write to execute after %v, but executed at %v", readTS, br.Timestamp)
@@ -883,7 +883,7 @@ func mergeCheckingTimestampCaches(
 	} else {
 		expReason = kvpb.ABORT_REASON_ABORTED_RECORD_FOUND
 	}
-	if _, pErr := lhsStore.Send(ctx, ba); pErr == nil {
+	if _, pErr := kvserver.ToSenderForTesting(lhsStore).Send(ctx, ba); pErr == nil {
 		t.Fatalf("expected TransactionAbortedError(%s) but got %v", expReason, pErr)
 	} else if abortErr, ok := pErr.GetDetail().(*kvpb.TransactionAbortedError); !ok {
 		t.Fatalf("expected TransactionAbortedError(%s) but got %v", expReason, pErr)
@@ -945,7 +945,7 @@ func TestStoreRangeMergeTimestampCacheCausality(t *testing.T) {
 				gba.Timestamp = ba.Timestamp.Add(42 /* wallTime */, 0 /* logical */)
 				gba.Add(getArgs(rhsKey))
 				store := tc.GetFirstStoreFromServer(t, int(ba.Header.Replica.NodeID-1))
-				gbr, pErr := store.Send(ctx, gba)
+				gbr, pErr := kvserver.ToSenderForTesting(store).Send(ctx, gba)
 				if pErr != nil {
 					t.Error(pErr) // different goroutine, so can't use t.Fatal
 					return pErr
@@ -1011,7 +1011,7 @@ func TestStoreRangeMergeTimestampCacheCausality(t *testing.T) {
 		// Merge [a, b) and [b, Max). Our request filter above will intercept the
 		// merge and execute a read with a large timestamp immediately before the
 		// Subsume request executes.
-		if _, pErr := kv.SendWrappedWith(ctx, tc.GetFirstStoreFromServer(t, 2), kvpb.Header{
+		if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(tc.GetFirstStoreFromServer(t, 2)), kvpb.Header{
 			RangeID: lhsRangeDesc.RangeID,
 		}, adminMergeArgs(scratchKey("a"))); pErr != nil {
 			t.Fatal(pErr)
@@ -1044,7 +1044,7 @@ func TestStoreRangeMergeTimestampCacheCausality(t *testing.T) {
 		ba.Timestamp = readTS
 		ba.RangeID = lhsRangeDesc.RangeID
 		ba.Add(incrementArgs(rhsKey, 1))
-		if br, pErr := tc.GetFirstStoreFromServer(t, 1).Send(ctx, ba); pErr != nil {
+		if br, pErr := kvserver.ToSenderForTesting(tc.GetFirstStoreFromServer(t, 1)).Send(ctx, ba); pErr != nil {
 			t.Fatal(pErr)
 		} else if br.Timestamp.LessEq(readTS) {
 			t.Fatalf("expected write to execute after %v, but executed at %v", readTS, br.Timestamp)

pkg/kv/kvserver/client_raft_test.go

@@ -2623,7 +2623,7 @@ func TestQuotaPool(t *testing.T) {
 		if err := ba.SetActiveTimestamp(tc.Servers[0].Clock()); err != nil {
 			t.Fatal(err)
 		}
-		if _, pErr := leaderRepl.Send(ctx, ba); pErr != nil {
+		if _, pErr := kvserver.ToSenderForTesting(leaderRepl).Send(ctx, ba); pErr != nil {
 			t.Fatal(pErr)
 		}
 
@@ -2645,7 +2645,7 @@ func TestQuotaPool(t *testing.T) {
 				ch <- kvpb.NewError(err)
 				return
 			}
-			_, pErr := leaderRepl.Send(ctx, ba)
+			_, pErr := kvserver.ToSenderForTesting(leaderRepl).Send(ctx, ba)
 			ch <- pErr
 		}()
 	}()
@@ -2780,7 +2780,7 @@ func TestWedgedReplicaDetection(t *testing.T) {
 		if err := ba.SetActiveTimestamp(leaderClock); err != nil {
 			t.Fatal(err)
 		}
-		if _, pErr := leaderRepl.Send(ctx, ba); pErr != nil {
+		if _, pErr := kvserver.ToSenderForTesting(leaderRepl).Send(ctx, ba); pErr != nil {
 			t.Fatal(pErr)
 		}
 
@@ -2803,7 +2803,7 @@ func TestWedgedReplicaDetection(t *testing.T) {
 
 			// Send another request to the leader replica. followerRepl is locked
 			// so it will not respond.
-			if _, pErr := leaderRepl.Send(ctx, ba); pErr != nil {
+			if _, pErr := kvserver.ToSenderForTesting(leaderRepl).Send(ctx, ba); pErr != nil {
 				t.Fatal(pErr)
 			}
 
@@ -3105,7 +3105,7 @@ func TestReplicaRemovalCampaign(t *testing.T) {
 			replica2 := store0.LookupReplica(roachpb.RKey(key2))
 
 			rg2 := func(s *kvserver.Store) kv.Sender {
-				return kv.Wrap(s, func(ba *kvpb.BatchRequest) *kvpb.BatchRequest {
+				return kv.Wrap(kvserver.ToSenderForTesting(s), func(ba *kvpb.BatchRequest) *kvpb.BatchRequest {
 					if ba.RangeID == 0 {
 						ba.RangeID = replica2.RangeID
 					}
@@ -4227,7 +4227,7 @@ func TestRemovedReplicaError(t *testing.T) {
 	// start seeing the RangeNotFoundError after a little bit of time has passed.
 	getArgs := getArgs(key)
 	testutils.SucceedsSoon(t, func() error {
-		_, pErr := kv.SendWrappedWith(ctx, store, kvpb.Header{}, getArgs)
+		_, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store), kvpb.Header{}, getArgs)
 		switch pErr.GetDetail().(type) {
 		case *kvpb.AmbiguousResultError:
 			return pErr.GoError()
@@ -4283,7 +4283,7 @@ func TestTransferRaftLeadership(t *testing.T) {
 			require.NoError(t, err)
 			require.Equal(t, roachpb.VOTER_FULL, rd1.Type)
 
-			_, pErr := kv.SendWrappedWith(ctx, store0, kvpb.Header{RangeID: repl0.RangeID}, getArgs(key))
+			_, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store0), kvpb.Header{RangeID: repl0.RangeID}, getArgs(key))
 			require.NoError(t, pErr.GoError())
 
 			status := repl0.RaftStatus()
@@ -4294,7 +4294,7 @@ func TestTransferRaftLeadership(t *testing.T) {
 			// Transfer the lease. We'll then check that the leadership follows
 			// automatically.
 			transferLeaseArgs := adminTransferLeaseArgs(key, store1.StoreID())
-			_, pErr = kv.SendWrappedWith(ctx, store0, kvpb.Header{RangeID: repl0.RangeID}, transferLeaseArgs)
+			_, pErr = kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store0), kvpb.Header{RangeID: repl0.RangeID}, transferLeaseArgs)
 			require.NoError(t, pErr.GoError())
 
 			// Verify leadership is transferred.
@@ -6398,7 +6398,7 @@ func TestInvalidConfChangeRejection(t *testing.T) {
 		},
 	})
 
-	_, pErr := repl.Send(ctx, &ba)
+	_, pErr := kvserver.ToSenderForTesting(repl).Send(ctx, &ba)
 	// Verify that we see the configuration change below raft, where we rejected it
 	// (since it would've otherwise blow up the Replica: after all, we intentionally
 	// proposed an invalid configuration change.

pkg/kv/kvserver/client_replica_circuit_breaker_test.go

@@ -1315,7 +1315,7 @@ func (cbt *circuitBreakerTest) SendCtxTS(
 	// going to leak memory.
 	ctx = context.WithValue(ctx, req, struct{}{})
 
-	_, pErr := repl.Send(ctx, ba)
+	_, pErr := kvserver.ToSenderForTesting(repl.Replica).Send(ctx, ba)
 	// If our context got canceled, return an opaque error regardless of presence or
 	// absence of actual error. This makes sure we don't accidentally pass tests as
 	// a result of our context cancellation.

pkg/kv/kvserver/client_replica_test.go

@@ -1251,7 +1251,7 @@ func TestNonTxnReadWithinUncertaintyIntervalAfterLeaseTransfer(t *testing.T) {
 		ba := &kvpb.BatchRequest{}
 		ba.RangeID = desc.RangeID
 		ba.Add(getArgs(key))
-		br, pErr := tc.GetFirstStoreFromServer(t, 1).Send(ctx, ba)
+		br, pErr := kvserver.ToSenderForTesting(tc.GetFirstStoreFromServer(t, 1)).Send(ctx, ba)
 		nonTxnRespC <- resp{br, pErr}
 	})
 
@@ -1290,7 +1290,7 @@ func TestNonTxnReadWithinUncertaintyIntervalAfterLeaseTransfer(t *testing.T) {
 	ba := &kvpb.BatchRequest{}
 	ba.RangeID = desc.RangeID
 	ba.Add(putArgs(key, []byte("val")))
-	br, pErr := tc.GetFirstStoreFromServer(t, 0).Send(ctx, ba)
+	br, pErr := kvserver.ToSenderForTesting(tc.GetFirstStoreFromServer(t, 0)).Send(ctx, ba)
 	require.Nil(t, pErr)
 	writeTs := br.Timestamp
 	require.True(t, nonTxnOrigTs.Less(writeTs), "nonTxnOrigTs: %v, writeTs: %v", nonTxnOrigTs, writeTs)
@@ -2589,7 +2589,7 @@ func TestConsistencyQueueDelaysProcessingNewRanges(t *testing.T) {
 		rngID := store.LookupReplica(roachpb.RKey(key)).RangeID
 		h := kvpb.Header{RangeID: rngID}
 		args := adminSplitArgs(key)
-		if _, pErr := kv.SendWrappedWith(ctx, store, h, args); pErr != nil {
+		if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store), h, args); pErr != nil {
 			return pErr.GoError()
 		}
 		return nil
@@ -2600,7 +2600,7 @@ func TestConsistencyQueueDelaysProcessingNewRanges(t *testing.T) {
 		rngID := store.LookupReplica(roachpb.RKey(key)).RangeID
 		h := kvpb.Header{RangeID: rngID}
 		args := adminMergeArgs(key)
-		if _, pErr := kv.SendWrappedWith(ctx, store, h, args); pErr != nil {
+		if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store), h, args); pErr != nil {
 			return pErr.GoError()
 		}
 		return nil
@@ -2744,7 +2744,7 @@ func TestLeaseInfoRequest(t *testing.T) {
 		},
 	}
 	reply, pErr := kv.SendWrappedWith(
-		context.Background(), s, kvpb.Header{
+		context.Background(), kvserver.ToSenderForTesting(s), kvpb.Header{
 			RangeID:         rangeDesc.RangeID,
 			ReadConsistency: kvpb.INCONSISTENT,
 		}, leaseInfoReq)
@@ -2870,7 +2870,7 @@ func TestRangeInfoAfterSplit(t *testing.T) {
 				},
 			}
 			ba.Add(gArgs)
-			br, pErr := store.Send(ctx, ba)
+			br, pErr := kvserver.ToSenderForTesting(store).Send(ctx, ba)
 			require.NoError(t, pErr.GoError())
 			descs := make([]roachpb.RangeDescriptor, len(br.RangeInfos))
 			for i, ri := range br.RangeInfos {
@@ -3051,7 +3051,7 @@ func TestLossQuorumCauseLeaderlessWatcherToSignalUnavailable(t *testing.T) {
 		ba.RangeID = desc.RangeID
 		ba.Timestamp = repl.Clock().Now()
 		ba.Add(putArgs(key, []byte("foo")))
-		_, pErr := repl.Send(ctx, ba)
+		_, pErr := kvserver.ToSenderForTesting(repl).Send(ctx, ba)
 		return pErr, ctx.Err()
 	}
 

pkg/kv/kvserver/client_split_test.go

@@ -570,14 +570,14 @@ func TestStoreRangeSplitAtRangeBounds(t *testing.T) {
 	rngID := store.LookupReplica(roachpb.RKey(key)).RangeID
 	h := kvpb.Header{RangeID: rngID}
 	args := adminSplitArgs(key)
-	if _, pErr := kv.SendWrappedWith(ctx, store, h, args); pErr != nil {
+	if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store), h, args); pErr != nil {
 		t.Fatal(pErr)
 	}
 	replCount := store.ReplicaCount()
 
 	// An AdminSplit request sent to the end of the old range should be re-routed
 	// to the start of the new range, succeeding but without creating any new ranges.
-	if _, pErr := kv.SendWrappedWith(ctx, store, h, args); pErr != nil {
+	if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store), h, args); pErr != nil {
 		t.Fatal(pErr)
 	}
 
@@ -591,7 +591,7 @@ func TestStoreRangeSplitAtRangeBounds(t *testing.T) {
 	// should succeed but no new ranges should be created.
 	newRng := store.LookupReplica(roachpb.RKey(key))
 	h.RangeID = newRng.RangeID
-	if _, pErr := kv.SendWrappedWith(ctx, store, h, args); pErr != nil {
+	if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store), h, args); pErr != nil {
 		t.Fatal(pErr)
 	}
 
@@ -1401,7 +1401,7 @@ func fillRange(
 		}
 		val := randutil.RandBytes(src, 200000)
 		pArgs := putArgs(key, val)
-		_, pErr := kv.SendWrappedWith(context.Background(), store, kvpb.Header{
+		_, pErr := kv.SendWrappedWith(context.Background(), kvserver.ToSenderForTesting(store), kvpb.Header{
 			RangeID: rangeID,
 		}, pArgs)
 		// When the split occurs in the background, our writes may start failing.
@@ -2891,7 +2891,7 @@ func TestDistributedTxnCleanup(t *testing.T) {
 	// Split at "b".
 	rhsKey := roachpb.Key("b")
 	args = adminSplitArgs(rhsKey)
-	if _, pErr := kv.SendWrappedWith(context.Background(), store, kvpb.Header{
+	if _, pErr := kv.SendWrappedWith(context.Background(), kvserver.ToSenderForTesting(store), kvpb.Header{
 		RangeID: lhs.RangeID,
 	}, args); pErr != nil {
 		t.Fatalf("split at %q: %s", rhsKey, pErr)
@@ -2932,7 +2932,7 @@ func TestDistributedTxnCleanup(t *testing.T) {
 						PushType:  kvpb.PUSH_ABORT,
 						Force:     true,
 					})
-					_, pErr := store.Send(ctx, ba)
+					_, pErr := kvserver.ToSenderForTesting(store).Send(ctx, ba)
 					if pErr != nil {
 						t.Fatalf("failed to abort the txn: %s", pErr)
 					}
@@ -3186,7 +3186,7 @@ func TestTxnWaitQueueDependencyCycleWithRangeSplit(t *testing.T) {
 
 		// Split at "b".
 		args = adminSplitArgs(rhsKey)
-		if _, pErr := kv.SendWrappedWith(ctx, store, kvpb.Header{
+		if _, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(store), kvpb.Header{
 			RangeID: lhs.RangeID,
 		}, args); pErr != nil {
 			t.Fatalf("split at %q: %s", rhsKey, pErr)
@@ -3533,7 +3533,7 @@ func TestStoreSplitRangeLookupRace(t *testing.T) {
 
 		// Don't use s.DistSenderI().(kv.Sender) so that we don't disturb the RangeDescriptorCache.
 		rangeID := store.LookupReplica(roachpb.RKey(splitKey)).RangeID
-		_, pErr := kv.SendWrappedWith(context.Background(), store, kvpb.Header{
+		_, pErr := kv.SendWrappedWith(context.Background(), kvserver.ToSenderForTesting(store), kvpb.Header{
 			RangeID: rangeID,
 		}, args)
 		if pErr != nil {
@@ -3978,7 +3978,7 @@ func TestSplitBlocksReadsToRHS(t *testing.T) {
 			g.GoCtx(func(ctx context.Context) error {
 				// Send directly to repl to avoid racing with the
 				// split and routing requests to the post-split RHS.
-				_, pErr := kv.SendWrappedWith(ctx, repl, h, args)
+				_, pErr := kv.SendWrappedWith(ctx, kvserver.ToSenderForTesting(repl), h, args)
 				errCh <- pErr.GoError()
 				return nil
 			})