requestbatcher: batch requests per-AC priority

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

Author: arulajmani

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)
+	}
+}