Merge #139612 #147469

139612: requestbatcher: batch requests per-AC priority r=arulajmani a=arulajmani

Previously, the batcher would batch requests to a given range together, and merge AC headers to pick the most permissive. This meant that lower priority AC work could get bumped to higher priority. This patch prevents this by batching per-range, per-AC priority.

There's some ugliness here around requests that bypass AC entirely. They do so by setting the Source field; to ensure these are batched correctly, we bucket them separately.

Closes #136029

Release note: None

147469: stateloader: untangle SynthesizeRaftState r=pav-kv a=pav-kv

This PR starts untangling the opaque code in `SynthesizeHardState` and `SynthesizeRaftState`, with the goals of (a) making it simple, (b) decoupling log and state machine reads/writes. Right now, these funcs are used in exactly 2 cases:

1. during cluster bootstrap, for creating the initial set of ranges
2. when applying a split trigger, for initializing the RHS (which is uninitialized and only possibly can have a `HardState`)

In both cases, the replica's state machine and log are initialized at `RaftInitialLogIndex/Term`. This means that everything these funcs generate is mostly based on constants (modulo the `HardState` in case 2). But it is really hard to tell by reading it. After this PR, case (1) is decoupled into its own `WriteInitialRaftState` func, and directly produces the right output. The "synthesize" funcs are now only used by splits, and can be further simplified in follow-up PRs.

The first half of this PR adds a datadriven test for `WriteInitialRangeState` showing the state of a typical replica created at cluster bootstrap. The printout only includes `RangeID`-local state, and no user space keys. The test helps ensuring that subsequent commits are no-ops.

Epic: CRDB-49111

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

Author: 3 people

pkg/internal/client/requestbatcher/BUILD.bazel

@@ -9,6 +9,8 @@ go_library(
         "//pkg/kv",
         "//pkg/kv/kvpb",
         "//pkg/roachpb",
+        "//pkg/util/admission/admissionpb",
+        "//pkg/util/buildutil",
         "//pkg/util/log",
         "//pkg/util/stop",
         "//pkg/util/timeutil",
@@ -25,7 +27,9 @@ go_test(
         "//pkg/kv/kvpb",
         "//pkg/roachpb",
         "//pkg/testutils",
+        "//pkg/util/admission/admissionpb",
         "//pkg/util/leaktest",
+        "//pkg/util/log",
         "//pkg/util/stop",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",

pkg/internal/client/requestbatcher/batcher.go

@@ -12,9 +12,8 @@
 // this library.
 //
 // Batching assumes that data with the same key can be sent in a single batch.
-// The initial implementation uses rangeID as the key explicitly to avoid
-// creating an overly general solution without motivation but interested readers
-// should recognize that it would be easy to extend this package to accept an
+// The implementation here uses rangeID and admission priority to construct
+// batches. This may be extended, or generalized, in the future by accepting a
 // arbitrary comparable key.
 package requestbatcher
 
@@ -27,6 +26,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -162,8 +163,8 @@ type Config struct {
 	// RequestBatcher -- a BatchRequest on that range will timeout, which will
 	// timeout and fail the entire batch.
 	//
-	// TODO(sumeer): once intent resolution is subject to admission control, we
-	// could have timeouts even though a range is available. Is that desirable?
+	// TODO(sumeer): we could have timeouts even though a range is available. Is
+	// that desirable?
 	MaxTimeout time.Duration
 
 	// InFlightBackpressureLimit is the number of batches in flight above which
@@ -437,6 +438,11 @@ func (b *RequestBatcher) sendResponse(req *request, resp Response) {
 }
 
 func addRequestToBatch(cfg *Config, now time.Time, ba *batch, r *request) (shouldSend bool) {
+	testingAssert(ba.empty() || admissionPriority(ba.admissionHeader()) == admissionPriority(r.header),
+		"requests with different admission headers shouldn't be added to the same batch")
+	testingAssert(ba.empty() || ba.rangeID() == r.rangeID,
+		"requests with different range IDs shouldn't be added to the same batch")
+
 	// Update the deadline for the batch if this requests's deadline is later
 	// than the current latest.
 	rDeadline, rHasDeadline := r.ctx.Deadline()
@@ -521,7 +527,7 @@ func (b *RequestBatcher) run(ctx context.Context) {
 		}
 		handleRequest = func(req *request) {
 			now := b.now()
-			ba, existsInQueue := b.batches.get(req.rangeID)
+			ba, existsInQueue := b.batches.get(req.rangeID, req.header)
 			if !existsInQueue {
 				ba = b.pool.newBatch(now)
 			}
@@ -614,27 +620,49 @@ func (b *batch) rangeID() roachpb.RangeID {
 	return b.reqs[0].rangeID
 }
 
-func (b *batch) batchRequest(cfg *Config) *kvpb.BatchRequest {
-	req := &kvpb.BatchRequest{
-		// Preallocate the Requests slice.
-		Requests: make([]kvpb.RequestUnion, 0, len(b.reqs)),
+// admissionPriority returns the priority with which to bucket requests with
+// the supplied header.
+func admissionPriority(header kvpb.AdmissionHeader) int32 {
+	if header.Source == kvpb.AdmissionHeader_OTHER {
+		// AdmissionHeader_OTHER bypass admission control, so bucket them separately
+		// and treat them as the highest priority.
+		return int32(admissionpb.OneAboveHighPri)
 	}
+	return header.Priority
+}
+
+func (b *batch) admissionHeader() kvpb.AdmissionHeader {
+	testingAssert(len(b.reqs) != 0, "admission header should not be called on an empty batch")
 	var admissionHeader kvpb.AdmissionHeader
 	for i, r := range b.reqs {
-		req.Add(r.req)
 		if i == 0 {
 			admissionHeader = r.header
 		} else {
 			admissionHeader = kv.MergeAdmissionHeaderForBatch(admissionHeader, r.header)
 		}
 	}
+	return admissionHeader
+}
+
+func (b *batch) empty() bool {
+	return len(b.reqs) == 0
+}
+
+func (b *batch) batchRequest(cfg *Config) *kvpb.BatchRequest {
+	req := &kvpb.BatchRequest{
+		// Preallocate the Requests slice.
+		Requests: make([]kvpb.RequestUnion, 0, len(b.reqs)),
+	}
+	for _, r := range b.reqs {
+		req.Add(r.req)
+	}
 	if cfg.MaxKeysPerBatchReq > 0 {
 		req.MaxSpanRequestKeys = int64(cfg.MaxKeysPerBatchReq)
 	}
 	if cfg.TargetBytesPerBatchReq > 0 {
 		req.TargetBytes = cfg.TargetBytesPerBatchReq
 	}
-	req.AdmissionHeader = admissionHeader
+	req.AdmissionHeader = b.admissionHeader()
 	return req
 }
 
@@ -705,6 +733,23 @@ func (p *pool) putBatch(b *batch) {
 	p.batchPool.Put(b)
 }
 
+// rangePriorityTuple is a container for a RangeID and admission priority pair.
+// It's intended to allow the batchQueue to build per-range, per-priority
+// batches.
+type rangePriorityPair struct {
+	rangeID  roachpb.RangeID
+	priority int32
+}
+
+// makeRangePriorityPair returns a new  rangePriorityPair for the supplied
+// rangeID and admission header.
+func makeRangePriorityPair(rangeID roachpb.RangeID, header kvpb.AdmissionHeader) rangePriorityPair {
+	return rangePriorityPair{
+		rangeID:  rangeID,
+		priority: admissionPriority(header),
+	}
+}
+
 // batchQueue is a container for batch objects which offers O(1) get based on
 // rangeID and peekFront as well as O(log(n)) upsert, removal, popFront.
 // Batch structs are heap ordered inside of the batches slice based on their
@@ -717,14 +762,14 @@ func (p *pool) putBatch(b *batch) {
 // per RequestBatcher.
 type batchQueue struct {
 	batches []*batch
-	byRange map[roachpb.RangeID]*batch
+	byRange map[rangePriorityPair]*batch
 }
 
 var _ heap.Interface = (*batchQueue)(nil)
 
 func makeBatchQueue() batchQueue {
 	return batchQueue{
-		byRange: map[roachpb.RangeID]*batch{},
+		byRange: map[rangePriorityPair]*batch{},
 	}
 }
 
@@ -742,8 +787,8 @@ func (q *batchQueue) popFront() *batch {
 	return heap.Pop(q).(*batch)
 }
 
-func (q *batchQueue) get(id roachpb.RangeID) (*batch, bool) {
-	b, exists := q.byRange[id]
+func (q *batchQueue) get(id roachpb.RangeID, header kvpb.AdmissionHeader) (*batch, bool) {
+	b, exists := q.byRange[makeRangePriorityPair(id, header)]
 	return b, exists
 }
 
@@ -771,24 +816,40 @@ func (q *batchQueue) Swap(i, j int) {
 
 func (q *batchQueue) Less(i, j int) bool {
 	idl, jdl := q.batches[i].deadline, q.batches[j].deadline
-	if before := idl.Before(jdl); before || !idl.Equal(jdl) {
-		return before
+	if !idl.Equal(jdl) {
+		return idl.Before(jdl)
+	}
+	iPri, jPri := admissionPriority(q.batches[i].admissionHeader()), admissionPriority(q.batches[j].admissionHeader())
+	if iPri != jPri {
+		// NB: We've got a min-heap, so we want to prefer higher AC priorities. In
+		// practice, this doesn't matter, because the batcher sends out all batches
+		// with the same deadline in parallel. See RequestBatcher.run.
+		return iPri > jPri
 	}
+	// Equal AC priorities; arbitrarily sort by rangeID.
 	return q.batches[i].rangeID() < q.batches[j].rangeID()
 }
 
 func (q *batchQueue) Push(v interface{}) {
 	ba := v.(*batch)
 	ba.idx = len(q.batches)
-	q.byRange[ba.rangeID()] = ba
+	q.byRange[makeRangePriorityPair(ba.rangeID(), ba.admissionHeader())] = ba
 	q.batches = append(q.batches, ba)
 }
 
 func (q *batchQueue) Pop() interface{} {
 	ba := q.batches[len(q.batches)-1]
 	q.batches[len(q.batches)-1] = nil // for GC
 	q.batches = q.batches[:len(q.batches)-1]
-	delete(q.byRange, ba.rangeID())
+	delete(q.byRange, makeRangePriorityPair(ba.rangeID(), ba.admissionHeader()))
 	ba.idx = -1
 	return ba
 }
+
+// testingAssert panics with the supplied message if the conditional doesn't
+// hold.
+func testingAssert(cond bool, msg string) {
+	if buildutil.CrdbTestBuild && !cond {
+		panic(msg)
+	}
+}