Merge #148395

148395: kv/bulk: enable async flushing r=jeffswenson a=jeffswenson

This change re-enables async flushing in the SST batcher. Async flushing
is a throughput win for large (1 Tib+) out of order writers. This occurs
during index builds, out of order imports, and PCR.

Previously, async flushing was a special case in the batcher. This
change reworks async flushing so that all flushes are async. This
ensures that context handling and error handling are consistent. This
should reduce the probability of an async only bug in the sst writer
that is hard to exercise.

Fixes: #146828

Co-authored-by: Jeff Swenson <[email protected]>

Author: craig[bot]

pkg/backup/restore_data_processor.go

@@ -514,6 +514,7 @@ func (rd *restoreDataProcessor) processRestoreSpanEntry(
 			// tests to fail.
 			rd.FlowCtx.Cfg.BackupMonitor.MakeConcurrentBoundAccount(),
 			rd.FlowCtx.Cfg.BulkSenderLimiter,
+			nil,
 		)
 		if err != nil {
 			return summary, err
@@ -717,7 +718,6 @@ func reserveRestoreWorkerMemory(
 // implement a mock SSTBatcher used purely for job progress tracking.
 type SSTBatcherExecutor interface {
 	AddMVCCKey(ctx context.Context, key storage.MVCCKey, value []byte) error
-	Reset(ctx context.Context)
 	Flush(ctx context.Context) error
 	Close(ctx context.Context)
 	GetSummary() kvpb.BulkOpSummary
@@ -735,9 +735,6 @@ func (b *sstBatcherNoop) AddMVCCKey(ctx context.Context, key storage.MVCCKey, va
 	return b.totalRows.Count(key.Key)
 }
 
-// Reset resets the counter
-func (b *sstBatcherNoop) Reset(ctx context.Context) {}
-
 // Flush noops.
 func (b *sstBatcherNoop) Flush(ctx context.Context) error {
 	return nil

pkg/crosscluster/logical/offline_initial_scan_processor.go

@@ -141,7 +141,9 @@ func (o *offlineInitialScanProcessor) setup(ctx context.Context) error {
 		true,            /* writeAtBatchTs */
 		true,            /* splitAndScatterRanges */
 		o.FlowCtx.Cfg.BackupMonitor.MakeConcurrentBoundAccount(),
-		o.FlowCtx.Cfg.BulkSenderLimiter)
+		o.FlowCtx.Cfg.BulkSenderLimiter,
+		nil,
+	)
 	if err != nil {
 		return err
 	}
@@ -377,7 +379,9 @@ func (o *offlineInitialScanProcessor) checkpoint(
 	if err := o.flushBatch(ctx); err != nil {
 		return errors.Wrap(err, "flushing batcher on checkpoint")
 	}
-	o.batcher.Reset(ctx)
+	if err := o.batcher.Reset(ctx); err != nil {
+		return errors.Wrap(err, "resetting batcher on checkpoint")
+	}
 
 	select {
 	case o.checkpointCh <- offlineCheckpoint{
@@ -409,7 +413,9 @@ func (o *offlineInitialScanProcessor) flushBatch(ctx context.Context) error {
 	if err := o.batcher.Flush(ctx); err != nil {
 		return err
 	}
-	o.batcher.Reset(ctx)
+	if err := o.batcher.Reset(ctx); err != nil {
+		return errors.Wrap(err, "resetting batcher after flush")
+	}
 	o.lastKeyAdded = roachpb.Key{}
 	return nil
 }

pkg/crosscluster/physical/stream_ingestion_processor.go

@@ -1273,11 +1273,15 @@ type flushableBuffer struct {
 
 // flushBuffer flushes the given streamIngestionBuffer via the SST
 // batchers and returns the underlying streamIngestionBuffer to the pool.
-func (sip *streamIngestionProcessor) flushBuffer(b flushableBuffer) (*jobspb.ResolvedSpans, error) {
+func (sip *streamIngestionProcessor) flushBuffer(
+	b flushableBuffer,
+) (_ *jobspb.ResolvedSpans, err error) {
 	ctx, sp := tracing.ChildSpan(sip.Ctx(), "stream-ingestion-flush")
 	defer sp.Finish()
-	// Ensure the batcher is always reset, even on early error returns.
-	defer sip.batcher.Reset(ctx)
+	defer func() {
+		// Ensure the batcher is always reset, even on early error returns.
+		err = errors.CombineErrors(err, sip.batcher.Reset(ctx))
+	}()
 
 	// First process the point KVs.
 	//

pkg/kv/bulk/BUILD.bazel

@@ -25,6 +25,7 @@ go_library(
         "//pkg/storage",
         "//pkg/storage/enginepb",
         "//pkg/util/admission/admissionpb",
+        "//pkg/util/buildutil",
         "//pkg/util/ctxgroup",
         "//pkg/util/hlc",
         "//pkg/util/humanizeutil",
@@ -57,15 +58,23 @@ go_test(
         "//pkg/kv/kvclient/kvcoord",
         "//pkg/kv/kvclient/rangecache",
         "//pkg/kv/kvpb",
+        "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/kvserverbase",
         "//pkg/roachpb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/server",
         "//pkg/settings/cluster",
+        "//pkg/sql/catalog",
+        "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/desctestutils",
+        "//pkg/sql/catalog/tabledesc",
+        "//pkg/sql/rowenc",
+        "//pkg/sql/sem/tree",
         "//pkg/storage",
         "//pkg/storage/enginepb",
         "//pkg/testutils/serverutils",
+        "//pkg/testutils/sqlutils",
         "//pkg/testutils/storageutils",
         "//pkg/testutils/testcluster",
         "//pkg/util/encoding",
@@ -77,6 +86,8 @@ go_test(
         "//pkg/util/randutil",
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
+        "@com_github_cockroachdb_apd_v3//:apd",
+        "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/kv/bulk/buffering_adder.go

@@ -145,6 +145,7 @@ func MakeBulkAdder(
 		lastFlush:      timeutil.Now(),
 		curBufSummary:  kvpb.BulkOpSummary{},
 	}
+	b.sink.init(ctx)
 
 	// Register a callback with the underlying sink to accumulate the summary for
 	// the current buffered KVs. The curBufSummary is reset when the buffering
@@ -282,7 +283,9 @@ func (b *BufferingAdder) doFlush(ctx context.Context, forSize bool) error {
 		b.curBufSummary.Reset()
 		return nil
 	}
-	b.sink.Reset(ctx)
+	if err := b.sink.Reset(ctx); err != nil {
+		return errors.Wrapf(err, "failed to reset %s adder before flush", b.name)
+	}
 	b.sink.batch.stats.BufferFlushes++
 
 	var before *bulkpb.IngestionPerformanceStats