Merge #150703

150703: bulk,importer,rowexec: add CPU AC to bulk ops r=dt a=dt

This adds the option to enable a CPU-based pacer to a few loops that have the potential to become tight, CPU-bound loops during bulk operations, namely:
* SST Batcher's Add method, which is called in a tight loop after sorting a buffered pool of KVs
* index backfiller's key generation loop, which is a tight loop over a buffer of read rows
* IMPORT's row -> datum reader loop and datum -> KV loops.

Each of these are controlled by a setting, which are default off for now pending more production-scale testing and benchmarking. 

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

Author: craig[bot]

pkg/kv/bulk/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "buffering_adder.go",
         "bulk_metrics.go",
+        "cpu_pacer.go",
         "kv_buf.go",
         "setting.go",
         "sst_adder.go",
@@ -24,6 +25,7 @@ go_library(
         "//pkg/settings/cluster",
         "//pkg/storage",
         "//pkg/storage/enginepb",
+        "//pkg/util/admission",
         "//pkg/util/admission/admissionpb",
         "//pkg/util/buildutil",
         "//pkg/util/ctxgroup",

pkg/kv/bulk/buffering_adder.go

@@ -153,6 +153,7 @@ func MakeBulkAdder(
 			writeAtBatchTS: opts.WriteAtBatchTimestamp,
 			mem:            bulkMon.MakeConcurrentBoundAccount(),
 			limiter:        sendLimiter,
+			pacer:          NewCPUPacer(ctx, db, sstBatcherElasticCPUControlEnabled),
 		},
 		timestamp:      timestamp,
 		maxBufferLimit: opts.MaxBufferSize,

pkg/kv/bulk/cpu_pacer.go

@@ -0,0 +1,71 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package bulk
+
+import (
+	"context"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/settings"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
+	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+)
+
+// CPUPacer wraps an admission.Pacer for use in bulk operations where any errors
+// pacing can just be logged (at most once a minute).
+type CPUPacer struct {
+	pacer    *admission.Pacer
+	logEvery *log.EveryN
+}
+
+func (p *CPUPacer) Pace(ctx context.Context) {
+	if err := p.pacer.Pace(ctx); err != nil {
+		if p.logEvery.ShouldLog() {
+			// TODO(dt): consider just returning this error/eliminating this wrapper,
+			// and making callers bubble up this error if it is only ever a ctx cancel
+			// error that the caller should be bubbling up anyway.
+			log.Warningf(ctx, "failed to pace SST batcher: %v", err)
+		}
+	}
+}
+func (p *CPUPacer) Close() {
+	p.pacer.Close()
+}
+
+var cpuPacerRequestDuration = settings.RegisterDurationSetting(
+	settings.ApplicationLevel,
+	"bulkio.elastic_cpu_control.request_duration",
+	"exeuction time unit to request when pacing CPU requests during various bulk operations",
+	50*time.Millisecond,
+)
+
+// NewCPUPacer creates a new AC pacer for SST batcher. It may return an empty
+// Pacer which noops if pacing is disabled or its arguments are nil.
+func NewCPUPacer(ctx context.Context, db *kv.DB, setting *settings.BoolSetting) CPUPacer {
+	if db == nil || db.AdmissionPacerFactory == nil || !setting.Get(db.SettingsValues()) {
+		log.Infof(ctx, "admission control is not configured to pace bulk ingestion")
+		return CPUPacer{}
+	}
+	tenantID, ok := roachpb.ClientTenantFromContext(ctx)
+	if !ok {
+		tenantID = roachpb.SystemTenantID
+	}
+	every := log.Every(time.Minute)
+	return CPUPacer{pacer: db.AdmissionPacerFactory.NewPacer(
+		cpuPacerRequestDuration.Get(db.SettingsValues()),
+		admission.WorkInfo{
+			TenantID:        tenantID,
+			Priority:        admissionpb.BulkNormalPri,
+			CreateTime:      timeutil.Now().UnixNano(),
+			BypassAdmission: false,
+		}),
+		logEvery: &every,
+	}
+}

pkg/kv/bulk/sst_batcher.go

@@ -71,6 +71,13 @@ var (
 		"if set, kvserver will compute an accurate stats diff for every addsstable request",
 		metamorphic.ConstantWithTestBool("computeStatsDiffInStreamBatcher", true),
 	)
+
+	sstBatcherElasticCPUControlEnabled = settings.RegisterBoolSetting(
+		settings.ApplicationLevel,
+		"bulkio.ingest.sst_batcher_elastic_control.enabled",
+		"determines whether the sst batcher integrates with elastic CPU control",
+		false, // TODO(dt): enable this by default.
+	)
 )
 
 // MakeAndRegisterConcurrencyLimiter makes a concurrency limiter and registers it
@@ -212,6 +219,9 @@ type SSTBatcher struct {
 	// disallowShadowingBelow is described on kvpb.AddSSTableRequest.
 	disallowShadowingBelow hlc.Timestamp
 
+	// pacer for admission control during SST ingestion
+	pacer CPUPacer
+
 	// skips duplicate keys (iff they are buffered together). This is true when
 	// used to backfill an inverted index. An array in JSONB with multiple values
 	// which are the same, will all correspond to the same kv in the inverted
@@ -313,6 +323,7 @@ func MakeSSTBatcher(
 		mem:                    mem,
 		limiter:                sendLimiter,
 		rc:                     rc,
+		pacer:                  NewCPUPacer(ctx, db, sstBatcherElasticCPUControlEnabled),
 	}
 	b.mu.lastFlush = timeutil.Now()
 	b.mu.tracingSpan = tracing.SpanFromContext(ctx)
@@ -352,6 +363,7 @@ func MakeStreamSSTBatcher(
 		// does not however make sense to scatter that range as the RHS maybe
 		// non-empty.
 		disableScatters: true,
+		pacer:           NewCPUPacer(ctx, db, sstBatcherElasticCPUControlEnabled),
 	}
 	b.mu.lastFlush = timeutil.Now()
 	b.mu.tracingSpan = tracing.SpanFromContext(ctx)
@@ -379,6 +391,7 @@ func MakeTestingSSTBatcher(
 		ingestAll:      ingestAll,
 		mem:            mem,
 		limiter:        sendLimiter,
+		pacer:          NewCPUPacer(ctx, db, sstBatcherElasticCPUControlEnabled),
 	}
 	b.init(ctx)
 	return b, nil
@@ -394,7 +407,6 @@ func (b *SSTBatcher) SetOnFlush(onFlush func(summary kvpb.BulkOpSummary)) {
 func (b *SSTBatcher) AddMVCCKeyWithImportEpoch(
 	ctx context.Context, key storage.MVCCKey, value []byte, importEpoch uint32,
 ) error {
-
 	mvccVal, err := storage.DecodeMVCCValue(value)
 	if err != nil {
 		return err
@@ -411,7 +423,6 @@ func (b *SSTBatcher) AddMVCCKeyWithImportEpoch(
 }
 
 func (b *SSTBatcher) AddMVCCKeyLDR(ctx context.Context, key storage.MVCCKey, value []byte) error {
-
 	mvccVal, err := storage.DecodeMVCCValue(value)
 	if err != nil {
 		return err
@@ -432,6 +443,9 @@ func (b *SSTBatcher) AddMVCCKeyLDR(ctx context.Context, key storage.MVCCKey, val
 // keys -- like RESTORE where we want the restored data to look like the backup.
 // Keys must be added in order.
 func (b *SSTBatcher) AddMVCCKey(ctx context.Context, key storage.MVCCKey, value []byte) error {
+	// Pace based on admission control before adding the key.
+	b.pacer.Pace(ctx)
+
 	if len(b.batch.endKey) > 0 && bytes.Equal(b.batch.endKey, key.Key) {
 		if b.ingestAll && key.Timestamp.Equal(b.batch.endTimestamp) {
 			if bytes.Equal(b.batch.endValue, value) {
@@ -924,6 +938,7 @@ func (b *SSTBatcher) Close(ctx context.Context) {
 	if err := b.syncFlush(); err != nil {
 		log.Warningf(ctx, "closing with flushes in-progress encountered an error: %v", err)
 	}
+	b.pacer.Close()
 	b.mem.Close(ctx)
 }
 

pkg/sql/importer/BUILD.bazel

@@ -35,6 +35,7 @@ go_library(
         "//pkg/jobs/jobspb",
         "//pkg/jobs/jobsprofiler",
         "//pkg/kv",
+        "//pkg/kv/bulk",
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/kvserverbase",
         "//pkg/roachpb",

pkg/sql/importer/read_import_base.go

@@ -21,9 +21,11 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cloud"
 	"github.com/cockroachdb/cockroach/pkg/crosscluster"
 	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/kv/bulk"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
+	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
@@ -44,6 +46,13 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
+var importElasticCPUControlEnabled = settings.RegisterBoolSetting(
+	settings.ApplicationLevel,
+	"bulkio.import.elastic_control.enabled",
+	"determines whether import operations integrate with elastic CPU control",
+	false, // TODO(dt): enable this by default after more benchmarking.
+)
+
 func runImport(
 	ctx context.Context,
 	flowCtx *execinfra.FlowCtx,
@@ -567,9 +576,15 @@ func runParallelImport(
 		var span *tracing.Span
 		ctx, span = tracing.ChildSpan(ctx, "import-file-to-rows")
 		defer span.Finish()
+
+		// Create a pacer for admission control for the producer.
+		pacer := bulk.NewCPUPacer(ctx, importCtx.db, importElasticCPUControlEnabled)
+		defer pacer.Close()
+
 		var numSkipped int64
 		var count int64
 		for producer.Scan() {
+			pacer.Pace(ctx)
 			// Skip rows if needed.
 			count++
 			if count <= fileCtx.skip {
@@ -660,6 +675,10 @@ func (p *parallelImporter) importWorker(
 	fileCtx *importFileContext,
 	minEmitted []int64,
 ) error {
+	// Create a pacer for admission control for this worker.
+	pacer := bulk.NewCPUPacer(ctx, importCtx.db, importElasticCPUControlEnabled)
+	defer pacer.Close()
+
 	conv, err := makeDatumConverter(ctx, importCtx, fileCtx, importCtx.db)
 	if err != nil {
 		return err
@@ -680,6 +699,8 @@ func (p *parallelImporter) importWorker(
 		conv.KvBatch.Progress = batch.progress
 		for batchIdx, record := range batch.data {
 			rowNum = batch.startPos + int64(batchIdx)
+			// Pace the admission control before processing each row.
+			pacer.Pace(ctx)
 			if err := consumer.FillDatums(ctx, record, rowNum, conv); err != nil {
 				if err = handleCorruptRow(ctx, fileCtx, err); err != nil {
 					return err

pkg/sql/importer/read_import_workload.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cloud"
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
 	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/kv/bulk"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
@@ -245,6 +246,12 @@ func NewWorkloadKVConverter(
 func (w *WorkloadKVConverter) Worker(
 	ctx context.Context, evalCtx *eval.Context, semaCtx *tree.SemaContext,
 ) error {
+	// Workload needs to pace itself explicitly since it manages its own workers
+	// and loops rather than using the "runParallelImport" helper which the other
+	// formats use and which has pacing built-in.
+	pacer := bulk.NewCPUPacer(ctx, w.db, importElasticCPUControlEnabled)
+	defer pacer.Close()
+
 	conv, err := row.NewDatumRowConverter(
 		ctx, semaCtx, w.tableDesc, nil, /* targetColNames */
 		evalCtx, w.kvCh, nil /* seqChunkProvider */, nil /* metrics */, w.db,
@@ -268,6 +275,7 @@ func (w *WorkloadKVConverter) Worker(
 		a = a.Truncate()
 		w.rows.FillBatch(batchIdx, cb, &a)
 		for rowIdx, numRows := 0, cb.Length(); rowIdx < numRows; rowIdx++ {
+			pacer.Pace(ctx)
 			for colIdx, col := range cb.ColVecs() {
 				// TODO(dan): This does a type switch once per-datum. Reduce this to
 				// a one-time switch per column.

pkg/sql/rowexec/BUILD.bazel

@@ -46,6 +46,7 @@ go_library(
         "//pkg/jobs/jobspb",
         "//pkg/keys",
         "//pkg/kv",
+        "//pkg/kv/bulk",
         "//pkg/kv/kvclient/kvstreamer",
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/kvserverbase",

pkg/sql/rowexec/indexbackfiller.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/kv/bulk"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -78,6 +79,13 @@ var indexBackfillIngestConcurrency = settings.RegisterIntSetting(
 	settings.PositiveInt, /* validateFn */
 )
 
+var indexBackfillElasticCPUControlEnabled = settings.RegisterBoolSetting(
+	settings.ApplicationLevel,
+	"bulkio.index_backfill.elastic_control.enabled",
+	"determines whether index backfill operations integrate with elastic CPU control",
+	false, // TODO(dt): enable this by default after more benchmarking.
+)
+
 func newIndexBackfiller(
 	ctx context.Context,
 	flowCtx *execinfra.FlowCtx,
@@ -329,9 +337,16 @@ func (ib *indexBackfiller) ingestIndexEntries(
 	g.GoCtx(func(ctx context.Context) error {
 		defer close(stopProgress)
 
+		// Create a pacer for admission control for index entry processing.
+		pacer := bulk.NewCPUPacer(ctx, ib.flowCtx.Cfg.DB.KV(), indexBackfillElasticCPUControlEnabled)
+		defer pacer.Close()
+
 		var vectorInputEntry rowenc.IndexEntry
 		for indexBatch := range indexEntryCh {
 			for _, indexEntry := range indexBatch.indexEntries {
+				// Pace the admission control before processing each index entry.
+				pacer.Pace(ctx)
+
 				// If there is at least one vector index being written, we need to check to see
 				// if this IndexEntry is going to a vector index and then re-encode it for that
 				// index if so.