Merge #158361

158361: crosscluster/logical: add CPU pacing to batch handling r=dt a=dt

Release note (performance improvement): More of the CPU usage of LDR jobs is subject to background job admission control limits.

Epic: none.

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

Author: craig[bot]

pkg/crosscluster/logical/BUILD.bazel

@@ -92,6 +92,7 @@ go_library(
         "//pkg/sql/syntheticprivilege",
         "//pkg/sql/types",
         "//pkg/storage",
+        "//pkg/util/admission",
         "//pkg/util/admission/admissionpb",
         "//pkg/util/buildutil",
         "//pkg/util/bulk",

pkg/crosscluster/logical/logical_replication_writer_processor.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/crosscluster/streamclient"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
+	kvbulk "github.com/cockroachdb/cockroach/pkg/kv/bulk"
 	"github.com/cockroachdb/cockroach/pkg/repstream/streampb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
@@ -40,6 +41,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlclustersettings"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/bulk"
 	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -148,6 +150,8 @@ type logicalReplicationWriterProcessor struct {
 	dupeCount  int64
 	seenEvery  log.EveryN
 	retryEvery log.EveryN
+
+	pacer *admission.Pacer
 }
 
 var (
@@ -229,6 +233,7 @@ func newLogicalReplicationWriterProcessor(
 		metrics:    flowCtx.Cfg.JobRegistry.MetricsStruct().JobSpecificMetrics[jobspb.TypeLogicalReplication].(*Metrics),
 		seenEvery:  log.Every(1 * time.Minute),
 		retryEvery: log.Every(1 * time.Minute),
+		pacer:      kvbulk.NewCPUPacer(ctx, flowCtx.Cfg.DB.KV(), useLowPriority),
 	}
 	lrw.purgatory = purgatory{
 		deadline:    func() time.Duration { return retryQueueAgeLimit.Get(&flowCtx.Cfg.Settings.SV) },

pkg/crosscluster/logical/lww_kv_processor.go

@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"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/repstream/streampb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -29,6 +30,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
@@ -47,6 +49,8 @@ type kvRowProcessor struct {
 	dstBySrc map[descpb.ID]descpb.ID
 	writers  map[descpb.ID]*kvTableWriter
 
+	pacer *admission.Pacer
+
 	failureInjector
 }
 
@@ -79,6 +83,7 @@ func newKVRowProcessor(
 		dstBySrc: dstBySrc,
 		writers:  make(map[descpb.ID]*kvTableWriter, len(procConfigByDestID)),
 		decoder:  decoder,
+		pacer:    bulk.NewCPUPacer(ctx, cfg.DB.KV(), useLowPriority),
 	}
 	return p, nil
 }
@@ -119,6 +124,10 @@ func (p *kvRowProcessor) HandleBatch(
 	ctx, sp := tracing.ChildSpan(ctx, "kvRowProcessor.HandleBatch")
 	defer sp.Finish()
 
+	if _, err := p.pacer.Pace(ctx); err != nil {
+		return batchStats{}, err
+	}
+
 	if len(batch) == 1 {
 		stats := batchStats{}
 		if p.spec.Discard == jobspb.LogicalReplicationDetails_DiscardAllDeletes && len(batch[0].KeyValue.Value.RawBytes) == 0 {

pkg/crosscluster/logical/sql_crud_writer.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/cdcevent"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/bulk"
 	"github.com/cockroachdb/cockroach/pkg/repstream/streampb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
@@ -18,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/stats"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 )
 
@@ -32,6 +34,7 @@ type sqlCrudWriter struct {
 	handlers map[descpb.ID]*tableHandler
 	settings *cluster.Settings
 	discard  jobspb.LogicalReplicationDetails_Discard
+	pacer    *admission.Pacer
 }
 
 var _ BatchHandler = &sqlCrudWriter{}
@@ -81,6 +84,7 @@ func newCrudSqlWriter(
 		handlers: handlers,
 		settings: evalCtx.Settings,
 		discard:  discard,
+		pacer:    bulk.NewCPUPacer(ctx, cfg.DB.KV(), useLowPriority),
 	}, nil
 }
 
@@ -90,6 +94,10 @@ func (c *sqlCrudWriter) HandleBatch(
 	ctx, sp := tracing.ChildSpan(ctx, "crudBatcher.HandleBatch")
 	defer sp.Finish()
 
+	if _, err := c.pacer.Pace(ctx); err != nil {
+		return batchStats{}, err
+	}
+
 	sortedEvents, err := c.decoder.decodeAndCoalesceEvents(ctx, batch, c.discard)
 	if err != nil {
 		return batchStats{}, err