kv/bulk: extract sst add logic from sst batcher

jeffswenson · jeffswenson · commit 3641703f6164 · 2025-06-16T08:17:24.000-04:00
This extracts the logic related to flushing a single sst from the logic that controls batching and flushing. Release note: none Part of: #146828
diff --git a/pkg/kv/bulk/buffering_adder.go b/pkg/kv/bulk/buffering_adder.go
@@ -111,17 +111,23 @@ func MakeBulkAdder(
 		name:        opts.Name,
 		importEpoch: opts.ImportEpoch,
 		sink: SSTBatcher{
-			name:                   opts.Name,
-			db:                     db,
-			rc:                     rangeCache,
+			name: opts.Name,
+			db:   db,
+			rc:   rangeCache,
+			adder: newSSTAdder(
+				db,
+				settings,
+				opts.WriteAtBatchTimestamp,
+				opts.DisallowShadowingBelow,
+				admissionpb.BulkNormalPri,
+			),
 			settings:               settings,
 			skipDuplicates:         opts.SkipDuplicates,
 			disallowShadowingBelow: opts.DisallowShadowingBelow,
 			batchTS:                opts.BatchTimestamp,
 			writeAtBatchTS:         opts.WriteAtBatchTimestamp,
 			mem:                    bulkMon.MakeConcurrentBoundAccount(),
 			limiter:                sendLimiter,
-			priority:               admissionpb.BulkNormalPri,
 		},
 		timestamp:      timestamp,
 		maxBufferLimit: opts.MaxBufferSize,
diff --git a/pkg/kv/bulk/sst_adder.go b/pkg/kv/bulk/sst_adder.go
@@ -9,12 +9,14 @@ import (
 	"context"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/bulk/bulkpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
+	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
@@ -23,30 +25,71 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
+type sstAdder struct {
+	settings *cluster.Settings
+	db       *kv.DB
+
+	// disallowShadowingBelow is described on kvpb.AddSSTableRequest.
+	disallowShadowingBelow hlc.Timestamp
+
+	// priority is the admission priority used for AddSSTable
+	// requests.
+	priority admissionpb.WorkPriority
+
+	// writeAtBatchTS is true if the SST should be written at the
+	// batch timestamp. If this is set to true, then all the keys in
+	//the sst must have the same timestamp and they must be equal to
+	// the batch timestamp.
+	// TODO(jeffswenson): remove this from `sstAdder` in a stand alone PR.
+	writeAtBatchTS bool
+}
+
+func newSSTAdder(
+	db *kv.DB,
+	settings *cluster.Settings,
+	writeAtBatchTS bool,
+	disallowShadowingBelow hlc.Timestamp,
+	priority admissionpb.WorkPriority,
+) *sstAdder {
+	return &sstAdder{
+		db:                     db,
+		disallowShadowingBelow: disallowShadowingBelow,
+		priority:               priority,
+		settings:               settings,
+		writeAtBatchTS:         writeAtBatchTS,
+	}
+}
+
 type sstSpan struct {
 	start, end roachpb.Key // [inclusive, exclusive)
 	sstBytes   []byte
 	stats      enginepb.MVCCStats
 }
 
+type addSSTResult struct {
+	timestamp                        hlc.Timestamp
+	availableBytes                   int64
+	followingLikelyNonEmptySpanStart roachpb.Key
+	rangeSpan                        roachpb.Span
+}
+
 // addSSTable retries db.AddSSTable if retryable errors occur, including if the
 // SST spans a split, in which case it is iterated and split into two SSTs, one
 // for each side of the split in the error, and each are retried.
-func (b *SSTBatcher) addSSTable(
+func (a *sstAdder) AddSSTable(
 	ctx context.Context,
 	batchTS hlc.Timestamp,
 	start, end roachpb.Key,
 	sstBytes []byte,
 	stats enginepb.MVCCStats,
-	updatesLastRange bool,
 	ingestionPerformanceStats *bulkpb.IngestionPerformanceStats,
-) error {
+) ([]addSSTResult, error) {
 	ctx, sp := tracing.ChildSpan(ctx, "*SSTBatcher.addSSTable")
 	defer sp.Finish()
 
 	sendStart := timeutil.Now()
 	if ingestionPerformanceStats == nil {
-		return errors.AssertionFailedf("ingestionPerformanceStats should not be nil")
+		return nil, errors.AssertionFailedf("ingestionPerformanceStats should not be nil")
 	}
 
 	// Currently, the SSTBatcher cannot ingest range keys, so it is safe to
@@ -58,21 +101,25 @@ func (b *SSTBatcher) addSSTable(
 	}
 	iter, err := storage.NewMemSSTIterator(sstBytes, true, iterOpts)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	defer iter.Close()
 
-	if (stats == enginepb.MVCCStats{}) {
+	if stats == (enginepb.MVCCStats{}) {
+		// TODO(jeffswenson): Audit AddSST callers to see if they generate
+		// server side stats now. Accurately computing stats in the face of replays
+		// requires the server to do it.
 		iter.SeekGE(storage.MVCCKey{Key: start})
 		// NB: even though this ComputeStatsForIter call exhausts the iterator, we
 		// can reuse/re-seek on the iterator, as part of the MVCCIterator contract.
 		stats, err = storage.ComputeStatsForIter(iter, sendStart.UnixNano())
 		if err != nil {
-			return errors.Wrapf(err, "computing stats for SST [%s, %s)", start, end)
+			return nil, errors.Wrapf(err, "computing stats for SST [%s, %s)", start, end)
 		}
 	}
 
 	work := []*sstSpan{{start: start, end: end, sstBytes: sstBytes, stats: stats}}
+	var results []addSSTResult
 	var files int
 	for len(work) > 0 {
 		item := work[0]
@@ -97,7 +144,7 @@ func (b *SSTBatcher) addSSTable(
 				// detection - making it is simpler to just always use the same API
 				// and just switch how it writes its result.
 				ingestAsWriteBatch := false
-				if b.settings != nil && int64(len(item.sstBytes)) < tooSmallSSTSize.Get(&b.settings.SV) {
+				if a.settings != nil && int64(len(item.sstBytes)) < tooSmallSSTSize.Get(&a.settings.SV) {
 					log.VEventf(ctx, 3, "ingest data is too small (%d keys/%d bytes) for SSTable, adding via regular batch", item.stats.KeyCount, len(item.sstBytes))
 					ingestAsWriteBatch = true
 					ingestionPerformanceStats.AsWrites++
@@ -106,29 +153,30 @@ func (b *SSTBatcher) addSSTable(
 				req := &kvpb.AddSSTableRequest{
 					RequestHeader:                          kvpb.RequestHeader{Key: item.start, EndKey: item.end},
 					Data:                                   item.sstBytes,
-					DisallowShadowingBelow:                 b.disallowShadowingBelow,
+					DisallowShadowingBelow:                 a.disallowShadowingBelow,
 					MVCCStats:                              &item.stats,
 					IngestAsWrites:                         ingestAsWriteBatch,
 					ReturnFollowingLikelyNonEmptySpanStart: true,
 				}
-				if b.writeAtBatchTS {
+				if a.writeAtBatchTS {
 					req.SSTTimestampToRequestTimestamp = batchTS
 				}
 
 				ba := &kvpb.BatchRequest{
 					Header: kvpb.Header{Timestamp: batchTS, ClientRangeInfo: roachpb.ClientRangeInfo{ExplicitlyRequested: true}},
 					AdmissionHeader: kvpb.AdmissionHeader{
-						Priority:                 int32(b.priority),
+						Priority:                 int32(a.priority),
 						CreateTime:               timeutil.Now().UnixNano(),
 						Source:                   kvpb.AdmissionHeader_FROM_SQL,
 						NoMemoryReservedAtSource: true,
 					},
 				}
+
 				ba.Add(req)
 				beforeSend := timeutil.Now()
 
 				sendCtx, sendSp := tracing.ChildSpan(ctx, "*SSTBatcher.addSSTable/Send")
-				br, pErr := b.db.NonTransactionalSender().Send(sendCtx, ba)
+				br, pErr := a.db.NonTransactionalSender().Send(sendCtx, ba)
 				sendSp.Finish()
 
 				sendTime := timeutil.Since(beforeSend)
@@ -147,23 +195,12 @@ func (b *SSTBatcher) addSSTable(
 
 				if pErr == nil {
 					resp := br.Responses[0].GetInner().(*kvpb.AddSSTableResponse)
-					b.mu.Lock()
-					if b.writeAtBatchTS {
-						b.mu.maxWriteTS.Forward(br.Timestamp)
-					}
-					b.mu.Unlock()
-					// If this was sent async then, by the time the reply gets back, it
-					// might not be the last range anymore. We can just discard the last
-					// range reply in this case though because async sends are only used
-					// for SSTs sent due to range boundaries, i.e. when we are done with
-					// with that range anyway.
-					if updatesLastRange {
-						b.lastRange.span = resp.RangeSpan
-						if resp.RangeSpan.Valid() {
-							b.lastRange.remaining = sz(resp.AvailableBytes)
-							b.lastRange.nextExistingKey = resp.FollowingLikelyNonEmptySpanStart
-						}
-					}
+					results = append(results, addSSTResult{
+						timestamp:                        br.Timestamp,
+						availableBytes:                   resp.AvailableBytes,
+						followingLikelyNonEmptySpanStart: resp.FollowingLikelyNonEmptySpanStart,
+						rangeSpan:                        resp.RangeSpan,
+					})
 					files++
 					log.VEventf(ctx, 3, "adding %s AddSSTable [%s,%s) took %v", sz(len(item.sstBytes)), item.start, item.end, sendTime)
 					return nil
@@ -185,7 +222,7 @@ func (b *SSTBatcher) addSSTable(
 					}
 					split := mr.Desc.EndKey.AsRawKey()
 					log.Infof(ctx, "SSTable cannot be added spanning range bounds %v, retrying...", split)
-					left, right, err := createSplitSSTable(ctx, item.start, split, iter, b.settings)
+					left, right, err := createSplitSSTable(ctx, item.start, split, iter, a.settings)
 					if err != nil {
 						return err
 					}
@@ -199,7 +236,7 @@ func (b *SSTBatcher) addSSTable(
 			}
 			return err
 		}(); err != nil {
-			return errors.Wrapf(err, "addsstable [%s,%s)", item.start, item.end)
+			return nil, errors.Wrapf(err, "addsstable [%s,%s)", item.start, item.end)
 		}
 		// explicitly deallocate SST. This will not deallocate the
 		// top level SST which is kept around to iterate over.
@@ -208,7 +245,7 @@ func (b *SSTBatcher) addSSTable(
 	ingestionPerformanceStats.SplitRetries += int64(files - 1)
 
 	log.VEventf(ctx, 3, "AddSSTable [%v, %v) added %d files and took %v", start, end, files, timeutil.Since(sendStart))
-	return nil
+	return results, nil
 }
 
 // createSplitSSTable is a helper for splitting up SSTs. The iterator
diff --git a/pkg/kv/bulk/sst_batcher.go b/pkg/kv/bulk/sst_batcher.go
@@ -94,15 +94,12 @@ func MakeAndRegisterConcurrencyLimiter(sv *settings.Values) limit.ConcurrentRequ
 type SSTBatcher struct {
 	name     string
 	db       *kv.DB
+	adder    *sstAdder
 	rc       *rangecache.RangeCache
 	settings *cluster.Settings
 	mem      *mon.ConcurrentBoundAccount
 	limiter  limit.ConcurrentRequestLimiter
 
-	// priority is the admission priority used for AddSSTable
-	// requests.
-	priority admissionpb.WorkPriority
-
 	// disallowShadowingBelow is described on kvpb.AddSSTableRequest.
 	disallowShadowingBelow hlc.Timestamp
 
@@ -222,13 +219,13 @@ func MakeSSTBatcher(
 	b := &SSTBatcher{
 		name:                   name,
 		db:                     db,
+		adder:                  newSSTAdder(db, settings, writeAtBatchTs, disallowShadowingBelow, admissionpb.BulkNormalPri),
 		settings:               settings,
 		disallowShadowingBelow: disallowShadowingBelow,
 		writeAtBatchTS:         writeAtBatchTs,
 		disableScatters:        !scatterSplitRanges,
 		mem:                    mem,
 		limiter:                sendLimiter,
-		priority:               admissionpb.BulkNormalPri,
 	}
 	b.mu.lastFlush = timeutil.Now()
 	b.mu.tracingSpan = tracing.SpanFromContext(ctx)
@@ -248,8 +245,13 @@ func MakeStreamSSTBatcher(
 	onFlush func(summary kvpb.BulkOpSummary),
 ) (*SSTBatcher, error) {
 	b := &SSTBatcher{
-		db:        db,
-		rc:        rc,
+		db: db,
+		rc: rc,
+		// We use NormalPri since anything lower than normal priority is assumed to
+		// be able to handle reduced throughput. We are OK with his for now since
+		// the consuming cluster of a replication stream does not have a latency
+		// sensitive workload running against it.
+		adder:     newSSTAdder(db, settings, false /*writeAtBatchTS*/, hlc.Timestamp{}, admissionpb.BulkNormalPri),
 		settings:  settings,
 		ingestAll: true,
 		mem:       mem,
@@ -263,11 +265,6 @@ func MakeStreamSSTBatcher(
 		// does not however make sense to scatter that range as the RHS maybe
 		// non-empty.
 		disableScatters: true,
-		// We use NormalPri since anything lower than normal priority is assumed to
-		// be able to handle reduced throughput. We are OK with his for now since
-		// the consuming cluster of a replication stream does not have a latency
-		// sensitive workload running against it.
-		priority: admissionpb.NormalPri,
 	}
 	b.mu.lastFlush = timeutil.Now()
 	b.mu.tracingSpan = tracing.SpanFromContext(ctx)
@@ -289,12 +286,12 @@ func MakeTestingSSTBatcher(
 ) (*SSTBatcher, error) {
 	b := &SSTBatcher{
 		db:             db,
+		adder:          newSSTAdder(db, settings, false, hlc.Timestamp{}, admissionpb.BulkNormalPri),
 		settings:       settings,
 		skipDuplicates: skipDuplicates,
 		ingestAll:      ingestAll,
 		mem:            mem,
 		limiter:        sendLimiter,
-		priority:       admissionpb.BulkNormalPri,
 	}
 	b.Reset(ctx)
 	return b, nil
@@ -724,15 +721,34 @@ func (b *SSTBatcher) doFlush(ctx context.Context, reason int) error {
 	fn := func(ctx context.Context) error {
 		defer res.Release()
 		defer b.mem.Shrink(ctx, reserved)
-		if err := b.addSSTable(ctx, batchTS, start, end, data, stats, !flushAsync, currentBatchStatsCopy); err != nil {
+		results, err := b.adder.AddSSTable(ctx, batchTS, start, end, data, stats, currentBatchStatsCopy)
+		if err != nil {
 			return err
 		}
 
 		// Now that we have completed ingesting the SSTables we take a lock and
-		// update the statistics on the SSTBatcher.
+		// process the flush results.
 		b.mu.Lock()
 		defer b.mu.Unlock()
 
+		for _, addResult := range results {
+			if b.writeAtBatchTS {
+				b.mu.maxWriteTS.Forward(addResult.timestamp)
+			}
+			if !flushAsync {
+				// If this was sent async then, by the time the reply gets back, it
+				// might not be the last range anymore. We can just discard the last
+				// range reply in this case though because async sends are only used
+				// for SSTs sent due to range boundaries, i.e. when we are done with
+				// with that range anyway.
+				b.lastRange.span = addResult.rangeSpan
+				if addResult.rangeSpan.Valid() {
+					b.lastRange.remaining = sz(addResult.availableBytes)
+					b.lastRange.nextExistingKey = addResult.followingLikelyNonEmptySpanStart
+				}
+			}
+		}
+
 		// Update the statistics associated with the current batch. We do this on
 		// our captured copy of the currentBatchSummary instead of the
 		// b.batchRowCounter since the caller may have reset the batcher by the time
