admission: add getterQualification to granter/requester interfaces

This is in preparation for the CPU time token scheme, where the WorkQueue
will distinguish between burstable and non-burstable tenants.

This is just an interface change, with no behavioral change.

Informs #153591

Epic: none

Release note: None

Author: sumeerbhola

pkg/util/admission/admission.go

@@ -129,13 +129,28 @@ import (
 	"github.com/cockroachdb/pebble"
 )
 
+// burstQualification is an optional behavior of certain WorkQueues (which
+// implement requester), that differentiate between tenants that are qualified
+// to burst (in their token consumption) and those that are not. This is a
+// dynamic attribute of a tenant, based on token consumption history
+// maintained in the WorkQueue. The ordering of tenants is also affected in
+// that burstable tenants are ordered before non-burstable tenants.
+type burstQualification uint8
+
+const (
+	canBurst burstQualification = iota
+	noBurst
+	numBurstQualifications
+)
+
 // requester is an interface implemented by an object that orders admission
 // work for a particular WorkKind. See WorkQueue for the implementation of
 // requester.
 type requester interface {
 	// hasWaitingRequests returns whether there are any waiting/queued requests
-	// of this WorkKind.
-	hasWaitingRequests() bool
+	// of this WorkKind, and when true, the qualification of the highest
+	// importance getter.
+	hasWaitingRequests() (bool, burstQualification)
 	// granted is called by a granter to grant admission to a single queued
 	// request. It returns > 0 if the grant was accepted, else returns 0. A
 	// grant may not be accepted if the grant raced with request cancellation
@@ -155,12 +170,15 @@ type granter interface {
 	grantKind() grantKind
 	// tryGet is used by a requester to get slots/tokens for a piece of work
 	// that has encountered no waiting/queued work. This is the fast path that
-	// avoids queueing in the requester.
+	// avoids queueing in the requester. The optional parameter
+	// burstQualification identifies the qualification of the getter, which is
+	// useful for certain granters.
 	//
 	// REQUIRES: count > 0. count == 1 for slots.
-	tryGet(count int64) (granted bool)
+	tryGet(getterQual burstQualification, count int64) (granted bool)
 	// returnGrant is called for:
 	// - returning slots after use.
+	// - returning tokens after use, if all the granted tokens were not used.
 	// - returning either slots or tokens when the grant raced with the work
 	//   being canceled, and the grantee did not end up doing any work.
 	//
@@ -197,7 +215,9 @@ type granter interface {
 	// the grantee after its goroutine runs and notices that it has been granted
 	// a slot/tokens. This provides a natural throttling that reduces grant
 	// bursts by taking into immediate account the capability of the goroutine
-	// scheduler to schedule such work.
+	// scheduler to schedule such work. Grant chains are only used for the CPU
+	// resource in the hybrid slot/token scheme, where slots are used for KVWork
+	// and tokens for SQLKVResponseWork and SQLSQLResponseWork.
 	//
 	// In an experiment, using such grant chains reduced burstiness of grants by
 	// 5x and shifted ~2s of latency (at p99) from the scheduler into admission

pkg/util/admission/elastic_cpu_granter.go

@@ -146,7 +146,7 @@ func (e *elasticCPUGranter) grantKind() grantKind {
 }
 
 // tryGet implements granter.
-func (e *elasticCPUGranter) tryGet(count int64) (granted bool) {
+func (e *elasticCPUGranter) tryGet(_ burstQualification, count int64) (granted bool) {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
@@ -182,7 +182,7 @@ func (e *elasticCPUGranter) continueGrantChain(grantChainID) {
 
 // tryGrant is used to attempt to grant to waiting requests.
 func (e *elasticCPUGranter) tryGrant() {
-	for e.requester.hasWaitingRequests() && e.tryGet(1) {
+	for e.hasWaitingRequests() && e.tryGet(canBurst /*arbitrary*/, 1) {
 		tokens := e.requester.granted(noGrantChain)
 		if tokens == 0 {
 			e.returnGrantWithoutGrantingElsewhere(1)
@@ -235,7 +235,8 @@ func (e *elasticCPUGranter) getUtilizationLimit() float64 {
 
 // hasWaitingRequests is part of the elasticCPULimiter interface.
 func (e *elasticCPUGranter) hasWaitingRequests() bool {
-	return e.requester.hasWaitingRequests()
+	hasWaiting, _ := e.requester.hasWaitingRequests()
+	return hasWaiting
 }
 
 // computeUtilizationMetric is part of the elasticCPULimiter interface.

pkg/util/admission/elastic_cpu_granter_test.go

@@ -152,7 +152,7 @@ func TestElasticCPUGranter(t *testing.T) {
 				return ""
 
 			case "try-get":
-				granted := elasticCPUGranter.tryGet(duration.Nanoseconds())
+				granted := elasticCPUGranter.tryGet(canBurst /*arbitrary*/, duration.Nanoseconds())
 				return fmt.Sprintf("granted:         %t\n", granted)
 
 			case "return-grant":
@@ -189,13 +189,13 @@ type testElasticCPURequester struct {
 
 var _ requester = &testElasticCPURequester{}
 
-func (t *testElasticCPURequester) hasWaitingRequests() bool {
+func (t *testElasticCPURequester) hasWaitingRequests() (bool, burstQualification) {
 	var padding string
 	if t.buf.Len() > 0 {
 		padding = "                 "
 	}
 	t.buf.WriteString(fmt.Sprintf("%shas-waiting=%t ", padding, t.numWaitingRequests > 0))
-	return t.numWaitingRequests > 0
+	return t.numWaitingRequests > 0, canBurst /*arbitrary*/
 }
 
 func (t *testElasticCPURequester) granted(grantChainID grantChainID) int64 {

pkg/util/admission/elastic_cpu_work_queue_test.go

@@ -130,7 +130,7 @@ func (t *testElasticCPUGranter) grantKind() grantKind {
 	return token
 }
 
-func (t *testElasticCPUGranter) tryGet(count int64) (granted bool) {
+func (t *testElasticCPUGranter) tryGet(_ burstQualification, count int64) (granted bool) {
 	panic("unimplemented")
 }
 
@@ -175,7 +175,7 @@ func (t *testElasticCPUInternalWorkQueue) adjustTenantUsed(
 	}
 }
 
-func (t *testElasticCPUInternalWorkQueue) hasWaitingRequests() bool {
+func (t *testElasticCPUInternalWorkQueue) hasWaitingRequests() (bool, burstQualification) {
 	panic("unimplemented")
 }
 

pkg/util/admission/granter.go

@@ -53,7 +53,7 @@ func (sg *slotGranter) grantKind() grantKind {
 }
 
 // tryGet implements granter.
-func (sg *slotGranter) tryGet(count int64) bool {
+func (sg *slotGranter) tryGet(_ burstQualification, count int64) bool {
 	return sg.coord.tryGet(sg.workKind, count, 0 /*arbitrary*/)
 }
 
@@ -122,7 +122,8 @@ func (sg *slotGranter) continueGrantChain(grantChainID grantChainID) {
 
 // requesterHasWaitingRequests implements granterWithLockedCalls.
 func (sg *slotGranter) requesterHasWaitingRequests() bool {
-	return sg.requester.hasWaitingRequests()
+	hasWaiting, _ := sg.requester.hasWaitingRequests()
+	return hasWaiting
 }
 
 // tryGrantLocked implements granterWithLockedCalls.
@@ -193,7 +194,7 @@ func (tg *tokenGranter) grantKind() grantKind {
 }
 
 // tryGet implements granter.
-func (tg *tokenGranter) tryGet(count int64) bool {
+func (tg *tokenGranter) tryGet(_ burstQualification, count int64) bool {
 	return tg.coord.tryGet(tg.workKind, count, 0 /*arbitrary*/)
 }
 
@@ -239,7 +240,8 @@ func (tg *tokenGranter) continueGrantChain(grantChainID grantChainID) {
 
 // requesterHasWaitingRequests implements granterWithLockedCalls.
 func (tg *tokenGranter) requesterHasWaitingRequests() bool {
-	return tg.requester.hasWaitingRequests()
+	hasWaiting, _ := tg.requester.hasWaitingRequests()
+	return hasWaiting
 }
 
 // tryGrantLocked implements granterWithLockedCalls.
@@ -340,7 +342,7 @@ func (cg *kvStoreTokenChildGranter) grantKind() grantKind {
 }
 
 // tryGet implements granter.
-func (cg *kvStoreTokenChildGranter) tryGet(count int64) bool {
+func (cg *kvStoreTokenChildGranter) tryGet(_ burstQualification, count int64) bool {
 	return cg.parent.tryGet(cg.workType, count)
 }
 
@@ -604,14 +606,6 @@ func (sg *kvStoreTokenGranter) subtractTokensLockedForWorkClass(
 	}
 }
 
-// requesterHasWaitingRequests returns whether some requester associated with
-// the granter has waiting requests. Used by storeGrantCoordinator.
-func (sg *kvStoreTokenGranter) requesterHasWaitingRequests() bool {
-	return sg.regularRequester.hasWaitingRequests() ||
-		sg.elasticRequester.hasWaitingRequests() ||
-		sg.snapshotRequester.hasWaitingRequests()
-}
-
 func (sg *kvStoreTokenGranter) tryGrant() {
 	sg.mu.Lock()
 	defer sg.mu.Unlock()
@@ -620,13 +614,14 @@ func (sg *kvStoreTokenGranter) tryGrant() {
 
 // tryGrantLocked attempts to grant to as many requests as possible.
 func (sg *kvStoreTokenGranter) tryGrantLocked() {
-	for sg.requesterHasWaitingRequests() && sg.tryGrantLockedOne() {
+	for sg.tryGrantLockedOne() {
 	}
 }
 
 // tryGrantLocked is used to attempt to grant to waiting requests. Used by
 // storeGrantCoordinator. It successfully grants to at most one waiting
-// request.
+// request. If there are no waiting requests, or all waiters reject the grant,
+// it returns false.
 func (sg *kvStoreTokenGranter) tryGrantLockedOne() bool {
 	// NB: We grant work in the following priority order: regular, snapshot
 	// ingest, elastic work. Snapshot ingests are a special type of elastic work.
@@ -640,7 +635,8 @@ func (sg *kvStoreTokenGranter) tryGrantLockedOne() bool {
 		} else if wt == admissionpb.SnapshotIngestStoreWorkType {
 			req = sg.snapshotRequester
 		}
-		if req.hasWaitingRequests() {
+		hasWaiting, _ := req.hasWaitingRequests()
+		if hasWaiting {
 			res := sg.tryGetLocked(1, wt)
 			if res {
 				tookTokenCount := req.granted(noGrantChain)

pkg/util/admission/granter_test.go

@@ -512,8 +512,8 @@ type testRequester struct {
 
 var _ requester = &testRequester{}
 
-func (tr *testRequester) hasWaitingRequests() bool {
-	return tr.waitingRequests
+func (tr *testRequester) hasWaitingRequests() (bool, burstQualification) {
+	return tr.waitingRequests, canBurst /*arbitrary*/
 }
 
 func (tr *testRequester) granted(grantChainID grantChainID) int64 {
@@ -527,7 +527,7 @@ func (tr *testRequester) granted(grantChainID grantChainID) int64 {
 func (tr *testRequester) close() {}
 
 func (tr *testRequester) tryGet(count int64) {
-	rv := tr.granter.tryGet(count)
+	rv := tr.granter.tryGet(canBurst /*arbitrary*/, count)
 	fmt.Fprintf(tr.buf, "%s%s: tryGet(%d) returned %t\n", tr.workKind,
 		tr.additionalID, count, rv)
 }

pkg/util/admission/replicated_write_admission_test.go

@@ -384,7 +384,7 @@ func (tg *testReplicatedWriteGranter) grantKind() grantKind {
 	return token
 }
 
-func (tg *testReplicatedWriteGranter) tryGet(count int64) bool {
+func (tg *testReplicatedWriteGranter) tryGet(_ burstQualification, count int64) bool {
 	if count > tg.tokens {
 		tg.buf.printf("[%s] try-get=%s available=%s => insufficient tokens",
 			tg.wc, printTrimmedBytes(count), printTrimmedBytes(tg.tokens))
@@ -414,7 +414,8 @@ func (tg *testReplicatedWriteGranter) grant() {
 		if tg.tokens <= 0 {
 			return // nothing left to do
 		}
-		if !tg.r.hasWaitingRequests() {
+		hasWaiting, _ := tg.r.hasWaitingRequests()
+		if !hasWaiting {
 			return // nothing left to do
 		}
 		_ = tg.r.granted(0 /* unused */)

pkg/util/admission/snapshot_queue.go

@@ -111,10 +111,10 @@ func makeSnapshotQueue(snapshotGranter granter, metrics *SnapshotMetrics) *Snaps
 var _ requester = &SnapshotQueue{}
 var _ snapshotRequester = &SnapshotQueue{}
 
-func (s *SnapshotQueue) hasWaitingRequests() bool {
+func (s *SnapshotQueue) hasWaitingRequests() (bool, burstQualification) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	return !s.mu.q.Empty()
+	return !s.mu.q.Empty(), canBurst /*arbitrary*/
 }
 
 func (s *SnapshotQueue) granted(_ grantChainID) int64 {
@@ -160,7 +160,7 @@ func (s *SnapshotQueue) Admit(ctx context.Context, count int64) error {
 		s.snapshotGranter.returnGrant(count)
 		return nil
 	}
-	if s.snapshotGranter.tryGet(count) {
+	if s.snapshotGranter.tryGet(canBurst /*arbitrary*/, count) {
 		return nil
 	}
 	// We were unable to get tokens for admission, so we queue.

pkg/util/admission/work_queue.go

@@ -639,7 +639,9 @@ func (q *WorkQueue) Admit(ctx context.Context, info WorkInfo) (enabled bool, err
 		// We have unlocked q.mu, so another concurrent request can also do tryGet
 		// and get ahead of this request. We don't need to be fair for such
 		// concurrent requests.
-		if q.granter.tryGet(info.RequestedCount) {
+		//
+		// TODO(sumeer): set a proper burstQualification.
+		if q.granter.tryGet(canBurst /*arbitrary*/, info.RequestedCount) {
 			q.metrics.incAdmitted(info.Priority)
 			if info.ReplicatedWorkInfo.Enabled {
 				// TODO(irfansharif): There's a race here, and could lead to
@@ -871,10 +873,11 @@ func (q *WorkQueue) AdmittedWorkDone(tenantID roachpb.TenantID, cpuTime time.Dur
 	q.granter.returnGrant(1)
 }
 
-func (q *WorkQueue) hasWaitingRequests() bool {
+func (q *WorkQueue) hasWaitingRequests() (bool, burstQualification) {
 	q.mu.Lock()
 	defer q.mu.Unlock()
-	return len(q.mu.tenantHeap) > 0
+	// TODO(sumeer): return a proper burstQualification.
+	return len(q.mu.tenantHeap) > 0, canBurst /*arbitrary*/
 }
 
 func (q *WorkQueue) granted(grantChainID grantChainID) int64 {

pkg/util/admission/work_queue_test.go

@@ -74,7 +74,7 @@ func (tg *testGranter) grantKind() grantKind {
 	return tg.gk
 }
 
-func (tg *testGranter) tryGet(count int64) bool {
+func (tg *testGranter) tryGet(_ burstQualification, count int64) bool {
 	tg.buf.printf("tryGet%s: returning %t", tg.name, tg.returnValueFromTryGet)
 	return tg.returnValueFromTryGet
 }