Merge #153172

153172: colfetcher: elastic AC pace low-pri bulk reads r=dt a=dt

Some bulk low pri reads can read large amounts of data. If KV admission
slots are not contended these scans scan end up doing non-trivial chunks
of client-side work in a tight loop, potentially remaining on-core for
some time. Having background/bulk work remain on-core for long, uninterrupted
blocks of time can push scheduling delays for foreground work higher, so we
typically aim to add scheduling-latency-driven pacer preemption points to these
sorts of loops.

Release note: none.
Epic: none.

Co-authored-by: David Taylor <[email protected]>

Author: craig[bot]

pkg/sql/colfetcher/BUILD.bazel

@@ -56,6 +56,8 @@ go_library(
         "//pkg/sql/sessiondata",
         "//pkg/sql/types",
         "//pkg/storage",
+        "//pkg/util/admission",
+        "//pkg/util/admission/admissionpb",
         "//pkg/util/buildutil",
         "//pkg/util/encoding",
         "//pkg/util/hlc",

pkg/sql/colfetcher/cfetcher.go

@@ -14,6 +14,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/cockroachdb/apd/v3"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
@@ -38,6 +39,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
+	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/intsets"
@@ -218,6 +221,9 @@ type cFetcherArgs struct {
 	alwaysReallocate bool
 	// Txn is the txn for the fetch. It might be nil.
 	txn *kv.Txn
+
+	// tenantID is required in some places for AC/bookkeeping.
+	tenantID roachpb.TenantID
 }
 
 // noOutputColumn is a sentinel value to denote that a system column is not
@@ -279,6 +285,7 @@ type cFetcher struct {
 	// cpuStopWatch tracks the CPU time spent by this cFetcher while fulfilling KV
 	// requests *in the current goroutine*.
 	cpuStopWatch *timeutil.CPUStopWatch
+	pacer        *admission.Pacer
 
 	// machine contains fields that get updated during the run of the fetcher.
 	machine struct {
@@ -544,6 +551,17 @@ func (cf *cFetcher) Init(
 	if cf.cFetcherArgs.collectStats {
 		cf.cpuStopWatch = timeutil.NewCPUStopWatch()
 	}
+	if cf.txn != nil {
+		if pri := admissionpb.WorkPriority(cf.txn.AdmissionHeader().Priority); pri <= admissionpb.BulkNormalPri {
+			cf.pacer = cf.txn.DB().AdmissionPacerFactory.NewPacer(
+				50*time.Millisecond, // Request a realistic per-batch amount.
+				admission.WorkInfo{
+					TenantID: cf.tenantID, Priority: pri, CreateTime: timeutil.Now().UnixNano(),
+				},
+			)
+		}
+	}
+
 	cf.machine.state[0] = stateResetBatch
 	cf.machine.state[1] = stateInitFetch
 
@@ -724,6 +742,9 @@ func (cf *cFetcher) setNextKV(kv roachpb.KeyValue) {
 // rows, the Batch.Length is 0.
 func (cf *cFetcher) NextBatch(ctx context.Context) (coldata.Batch, error) {
 	for {
+		if err := cf.pacer.Pace(ctx); err != nil {
+			return nil, err
+		}
 		if debugState {
 			log.Dev.Infof(ctx, "State %s", cf.machine.state[0])
 		}
@@ -1480,6 +1501,8 @@ func (cf *cFetcher) Release() {
 
 func (cf *cFetcher) Close(ctx context.Context) {
 	if cf != nil {
+		cf.pacer.Close()
+		cf.pacer = nil
 		cf.nextKVer = nil
 		if cf.fetcher != nil {
 			cf.bytesRead = cf.fetcher.GetBytesRead()

pkg/sql/colfetcher/cfetcher_wrapper.go

@@ -223,6 +223,11 @@ func newCFetcherWrapper(
 	const collectStats = false
 	// We cannot reuse batches if we're not serializing the response.
 	alwaysReallocate := !mustSerialize
+
+	tenantID, ok := roachpb.ClientTenantFromContext(ctx)
+	if !ok {
+		tenantID = roachpb.SystemTenantID
+	}
 	// TODO(yuzefovich, 23.1): think through estimatedRowCount (#94850) and
 	// traceKV arguments.
 	fetcher.cFetcherArgs = cFetcherArgs{
@@ -232,7 +237,8 @@ func newCFetcherWrapper(
 		true,  /* singleUse */
 		collectStats,
 		alwaysReallocate,
-		nil, /* txn */
+		nil, /* txn; TODO(dt): this means no AC priority info is passed. */
+		tenantID,
 	}
 
 	// This memory monitor is not connected to the memory accounting system

pkg/sql/colfetcher/colbatch_scan.go

@@ -347,6 +347,7 @@ func NewColBatchScan(
 		shouldCollectStats,
 		false, /* alwaysReallocate */
 		flowCtx.Txn,
+		flowCtx.Codec().TenantID,
 	}
 	if err = fetcher.Init(fetcherAllocator, kvFetcher, tableArgs); err != nil {
 		fetcher.Release()

pkg/sql/colfetcher/index_join.go

@@ -648,6 +648,7 @@ func NewColIndexJoin(
 		shouldCollectStats,
 		false, /* alwaysReallocate */
 		txn,
+		flowCtx.Codec().TenantID,
 	}
 	if err = fetcher.Init(
 		fetcherAllocator, kvFetcher, tableArgs,

pkg/sql/sessiondatapb/local_only_session_data.go

@@ -387,6 +387,8 @@ func ParseQoSLevelFromString(val string) (_ QoSLevel, ok bool) {
 		return UserLow, true
 	case NormalName:
 		return Normal, true
+	case BulkLowName:
+		return BulkLow, true
 	default:
 		return 0, false
 	}
@@ -414,7 +416,7 @@ func ToQoSLevelString(value int32) string {
 // Validate checks for a valid user QoSLevel setting before returning it.
 func (e QoSLevel) Validate() QoSLevel {
 	switch e {
-	case Normal, UserHigh, UserLow:
+	case Normal, UserHigh, UserLow, BulkLow:
 		return e
 	default:
 		panic(errors.AssertionFailedf("use of illegal user QoSLevel: %s", e.String()))