Merge #153588 #153592 #153593

153588: restore: set correct number of link workers r=msbutler a=kev-cao

#153065 contains a typo that ensures the number of link workers is always 0 or 1. This patch fixes the bug.

Epic: None

Release note: None

153592: admission: add getterQualification to granter/requester interfaces r=joshimhoff a=sumeerbhola

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

153593: sql: clarify recent change around soft limits r=yuzefovich a=yuzefovich

AI review spotted a "bug" in a recent change around soft limits when that limit is negative. This is impossible by construction, so this commit clarifies the code a bit by switching to `uint64` in non-proto code. We also have "hard limit" that still remains `int64` due to one special case (I was being lazy so only left a TODO to clean this up).

Epic: None
Release note: None

Co-authored-by: Kevin Cao <[email protected]>
Co-authored-by: sumeerbhola <[email protected]>
Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: 4 people

pkg/backup/restore_online.go

@@ -1061,6 +1061,6 @@ func getNumOnlineRestoreLinkWorkers(ctx context.Context, execCtx sql.JobExecCont
 	if err != nil {
 		return 0, err
 	}
-	numNodes := min(len(sqlInstanceIDs), 1)
+	numNodes := max(len(sqlInstanceIDs), 1)
 	return defaultLinkWorkersPerNode * numNodes, nil
 }

pkg/sql/distsql_physical_planner.go

@@ -726,7 +726,7 @@ func checkSupportForPlanNode(
 			var suffix string
 			estimate := n.estimatedRowCount
 			if n.softLimit != 0 && sd.UseSoftLimitForDistributeScan {
-				estimate = uint64(n.softLimit)
+				estimate = n.softLimit
 				suffix = " (using soft limit)"
 			}
 			if estimate >= sd.DistributeScanRowCountThreshold {
@@ -2205,8 +2205,8 @@ func initTableReaderSpecTemplate(
 	var post execinfrapb.PostProcessSpec
 	if n.hardLimit != 0 {
 		post.Limit = uint64(n.hardLimit)
-	} else if n.softLimit != 0 {
-		s.LimitHint = n.softLimit
+	} else if softLimit := int64(n.softLimit); softLimit > 0 {
+		s.LimitHint = softLimit
 	}
 	return s, post, nil
 }

pkg/sql/distsql_spec_exec_factory.go

@@ -337,8 +337,8 @@ func (e *distSQLSpecExecFactory) ConstructScan(
 	post := execinfrapb.PostProcessSpec{}
 	if params.HardLimit != 0 {
 		post.Limit = uint64(params.HardLimit)
-	} else if params.SoftLimit != 0 {
-		trSpec.LimitHint = params.SoftLimit
+	} else if softLimit := int64(params.SoftLimit); softLimit > 0 {
+		trSpec.LimitHint = softLimit
 	}
 
 	err = e.dsp.planTableReaders(

pkg/sql/opt/exec/execbuilder/relational.go

@@ -737,7 +737,7 @@ func (b *Builder) scanParams(
 		sqltelemetry.IncrementPartitioningCounter(sqltelemetry.PartitionConstrainedScan)
 	}
 
-	softLimit := reqProps.LimitHintInt64()
+	softLimit := uint64(reqProps.LimitHintInt64())
 	hardLimit := scan.HardLimit.RowCount()
 	maxResults, maxResultsOk := b.indexConstraintMaxResults(scan, relProps)
 

pkg/sql/opt/exec/explain/emit.go

@@ -1319,7 +1319,7 @@ func (e *emitter) spansStr(table cat.Table, index cat.Index, scanParams exec.Sca
 		if scanParams.HardLimit != 0 {
 			return "LIMITED SCAN"
 		}
-		if scanParams.SoftLimit > 0 {
+		if scanParams.SoftLimit != 0 {
 			return "FULL SCAN (SOFT LIMIT)"
 		}
 		return "FULL SCAN"

pkg/sql/opt/exec/factory.go

@@ -118,12 +118,17 @@ type ScanParams struct {
 	InvertedConstraint inverted.Spans
 
 	// If non-zero, the scan returns this many rows.
+	//
+	// Additionally, in plan-gist decoding path, this will be set to -1 to
+	// indicate presence of a limit, regardless of its value.
+	// TODO(yuzefovich): we could refactor this special case by adding an
+	// additional boolean that would allow us to switch to using uint64.
 	HardLimit int64
 
 	// If non-zero, the scan may still be required to return up to all its rows
 	// (or up to the HardLimit if it is set, but can be optimized under the
 	// assumption that only SoftLimit rows will be needed.
-	SoftLimit int64
+	SoftLimit uint64
 
 	Reverse bool
 

pkg/sql/scan.go

@@ -52,7 +52,7 @@ type scanNode struct {
 	// if non-zero, softLimit is an estimation that only this many rows might be
 	// needed. It is a (potentially optimistic) "hint". If hardLimit is set
 	// (non-zero), softLimit must be unset (zero).
-	softLimit int64
+	softLimit uint64
 
 	// See exec.Factory.ConstructScan.
 	parallelize bool

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 {