kv/bulk: ensure in-flight requests end on batcher reset

pkg/ccl/backupccl/restore_data_processor.go

@@ -767,7 +767,7 @@ 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) error
+	Reset(ctx context.Context)
 	Flush(ctx context.Context) error
 	Close(ctx context.Context)
 	GetSummary() kvpb.BulkOpSummary
@@ -786,9 +786,7 @@ func (b *sstBatcherNoop) AddMVCCKey(ctx context.Context, key storage.MVCCKey, va
 }
 
 // Reset resets the counter
-func (b *sstBatcherNoop) Reset(ctx context.Context) error {
-	return nil
-}
+func (b *sstBatcherNoop) Reset(ctx context.Context) {}
 
 // Flush noops.
 func (b *sstBatcherNoop) Flush(ctx context.Context) error {

pkg/ccl/streamingccl/streamingest/stream_ingestion_processor.go

@@ -1155,6 +1155,8 @@ type flushableBuffer struct {
 func (sip *streamIngestionProcessor) flushBuffer(b flushableBuffer) (*jobspb.ResolvedSpans, 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)
 
 	// First process the point KVs.
 	//
@@ -1188,10 +1190,6 @@ func (sip *streamIngestionProcessor) flushBuffer(b flushableBuffer) (*jobspb.Res
 	sip.metrics.IngestedEvents.Inc(int64(len(b.buffer.curKVBatch)))
 	sip.metrics.IngestedEvents.Inc(int64(len(b.buffer.curRangeKVBatch)))
 
-	if err := sip.batcher.Reset(ctx); err != nil {
-		return b.checkpoint, err
-	}
-
 	releaseBuffer(b.buffer)
 
 	return b.checkpoint, nil

pkg/kv/bulk/buffering_adder.go

@@ -256,9 +256,7 @@ func (b *BufferingAdder) doFlush(ctx context.Context, forSize bool) error {
 		b.curBufSummary.Reset()
 		return nil
 	}
-	if err := b.sink.Reset(ctx); err != nil {
-		return err
-	}
+	b.sink.Reset(ctx)
 	b.sink.currentStats.BufferFlushes++
 
 	var before *bulkpb.IngestionPerformanceStats

pkg/kv/bulk/sst_batcher.go

@@ -232,8 +232,8 @@ func MakeSSTBatcher(
 	}
 	b.mu.lastFlush = timeutil.Now()
 	b.mu.tracingSpan = tracing.SpanFromContext(ctx)
-	err := b.Reset(ctx)
-	return b, err
+	b.Reset(ctx)
+	return b, nil
 }
 
 // MakeStreamSSTBatcher creates a batcher configured to ingest duplicate keys
@@ -258,8 +258,8 @@ func MakeStreamSSTBatcher(
 	b.mu.lastFlush = timeutil.Now()
 	b.mu.tracingSpan = tracing.SpanFromContext(ctx)
 	b.SetOnFlush(onFlush)
-	err := b.Reset(ctx)
-	return b, err
+	b.Reset(ctx)
+	return b, nil
 }
 
 // MakeTestingSSTBatcher creates a batcher for testing, allowing setting options
@@ -281,8 +281,8 @@ func MakeTestingSSTBatcher(
 		mem:            mem,
 		limiter:        sendLimiter,
 	}
-	err := b.Reset(ctx)
-	return b, err
+	b.Reset(ctx)
+	return b, nil
 }
 
 func (b *SSTBatcher) updateMVCCStats(key storage.MVCCKey, value []byte) {
@@ -373,8 +373,14 @@ func (b *SSTBatcher) AddMVCCKey(ctx context.Context, key storage.MVCCKey, value
 }
 
 // Reset clears all state in the batcher and prepares it for reuse.
-func (b *SSTBatcher) Reset(ctx context.Context) error {
+func (b *SSTBatcher) Reset(ctx context.Context) {
+	if err := b.asyncAddSSTs.Wait(); err != nil {
+		log.Warningf(ctx, "closing with flushes in-progress encountered an error: %v", err)
+	}
+	b.asyncAddSSTs = ctxgroup.Group{}
+
 	b.sstWriter.Close()
+
 	b.sstFile = &storage.MemObject{}
 	// Create sstables intended for ingestion using the newest format that all
 	// nodes can support. MakeIngestionSSTWriter will handle cluster version
@@ -402,8 +408,6 @@ func (b *SSTBatcher) Reset(ctx context.Context) error {
 	if b.mu.totalStats.SendWaitByStore == nil {
 		b.mu.totalStats.SendWaitByStore = make(map[roachpb.StoreID]time.Duration)
 	}
-
-	return nil
 }
 
 const (
@@ -438,7 +442,8 @@ func (b *SSTBatcher) flushIfNeeded(ctx context.Context, nextKey roachpb.Key) err
 		if err := b.doFlush(ctx, rangeFlush); err != nil {
 			return err
 		}
-		return b.Reset(ctx)
+		b.Reset(ctx)
+		return nil
 	}
 
 	if b.sstWriter.DataSize >= ingestFileSize(b.settings) {
@@ -462,7 +467,8 @@ func (b *SSTBatcher) flushIfNeeded(ctx context.Context, nextKey roachpb.Key) err
 		if err := b.doFlush(ctx, sizeFlush); err != nil {
 			return err
 		}
-		return b.Reset(ctx)
+		b.Reset(ctx)
+		return nil
 	}
 	return nil
 }