kv/bulk: pace bulk SST generation/addition with elastic AC

Various bulk operations end up using SST Batcher to construct an SST,
which is can be a CPU-bound operation when a large buffer of keys each
need to be compared, compressed and flushed in a tight loop.

This adds AC pacing to the method that is called per key so that it can
pace these operations as needed.

Release note: none.
Epic: none.

Author: dt

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