Merge #149499

149499: kvserver: clean up prepareSnapApply r=arulajmani a=pav-kv

This PR introduces a builder type which prepares a set of writes when applying a snapshot. This removes extensive lists of parameters dragged across function calls, and wraps repetitive SST creation code into a helper method.

It also fixes a TODO in `TestPrepareSnapApply`: all the writes are now printed in the original order, thanks to the `storage.Writer` abstraction instead of direct SST generation code.

The new `storage.Writer`-based interface is more flexible, and e.g. allows eliminating the `ConvertFilesToBatchAndCommit` [quirk](https://github.com/cockroachdb/cockroach/blob/255f60e14a03060b9b2ca81e38f7188c33173a7a/pkg/kv/kvserver/replica_raftstorage.go#L640) as well as the "cleared spans" tracking. We can know when the snapshot is ["simple/small"](https://github.com/cockroachdb/cockroach/blob/255f60e14a03060b9b2ca81e38f7188c33173a7a/pkg/kv/kvserver/replica_raftstorage.go#L624-L627) before `MultiSSTWriter` starts writing files (e.g. from the handshake), so we can generate a batch directly and skip the files writing phase; this batch can then be extended with `prepareSnapApply` data and committed. This would make small snapshot ingestions faster.

Epic: CRDB-49111

Co-authored-by: Pavel Kalinnikov <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/BUILD.bazel

@@ -450,6 +450,7 @@ go_test(
         "//pkg/kv/kvserver/load",
         "//pkg/kv/kvserver/lockspanset",
         "//pkg/kv/kvserver/logstore",
+        "//pkg/kv/kvserver/print",
         "//pkg/kv/kvserver/protectedts",
         "//pkg/kv/kvserver/protectedts/ptpb",
         "//pkg/kv/kvserver/protectedts/ptstorage",

pkg/kv/kvserver/kvstorage/snaprecv/sst_snapshot_storage.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"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/fs"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -53,14 +54,15 @@ func NewSSTSnapshotStorage(engine storage.Engine, limiter *rate.Limiter) SSTSnap
 // NewScratchSpace creates a new storage scratch space for SSTs for a specific
 // snapshot.
 func (s *SSTSnapshotStorage) NewScratchSpace(
-	rangeID roachpb.RangeID, snapUUID uuid.UUID,
+	rangeID roachpb.RangeID, snapUUID uuid.UUID, st *cluster.Settings,
 ) *SSTSnapshotStorageScratch {
 	s.mu.Lock()
 	s.mu.rangeRefCount[rangeID]++
 	s.mu.Unlock()
 	snapDir := filepath.Join(s.dir, strconv.Itoa(int(rangeID)), snapUUID.String())
 	return &SSTSnapshotStorageScratch{
 		storage: s,
+		st:      st,
 		rangeID: rangeID,
 		snapDir: snapDir,
 	}
@@ -98,6 +100,7 @@ func (s *SSTSnapshotStorage) scratchClosed(rangeID roachpb.RangeID) {
 // snapshot.
 type SSTSnapshotStorageScratch struct {
 	storage    *SSTSnapshotStorage
+	st         *cluster.Settings
 	rangeID    roachpb.RangeID
 	ssts       []string
 	snapDir    string
@@ -138,21 +141,34 @@ func (s *SSTSnapshotStorageScratch) NewFile(
 	return f, nil
 }
 
-// WriteSST writes SST data to a file. The method closes
-// the provided SST when it is finished using it. If the provided SST is empty,
-// then no file will be created and nothing will be written.
-func (s *SSTSnapshotStorageScratch) WriteSST(ctx context.Context, data []byte) error {
+// WriteSST creates an SST populated with the given write function, and writes
+// it to a file. Does nothing if no data is written.
+func (s *SSTSnapshotStorageScratch) WriteSST(
+	ctx context.Context, write func(context.Context, storage.Writer) error,
+) error {
 	if s.closed {
 		return errors.AssertionFailedf("SSTSnapshotStorageScratch closed")
 	}
-	if len(data) == 0 {
+
+	// TODO(itsbilal): Write to SST directly rather than buffer in a MemObject.
+	sstFile := &storage.MemObject{}
+	w := storage.MakeIngestionSSTWriter(ctx, s.st, sstFile)
+	defer w.Close()
+	if err := write(ctx, &w); err != nil {
+		return err
+	}
+	if err := w.Finish(); err != nil {
+		return err
+	}
+	if w.DataSize == 0 {
 		return nil
 	}
+
 	f, err := s.NewFile(ctx, 512<<10 /* 512 KB */)
 	if err != nil {
 		return err
 	}
-	if err := f.Write(data); err != nil {
+	if err := f.Write(sstFile.Data()); err != nil {
 		f.Abort()
 		return err
 	}

pkg/kv/kvserver/kvstorage/snaprecv/sst_snapshot_storage_test.go

@@ -57,7 +57,7 @@ func TestSSTSnapshotStorage(t *testing.T) {
 	defer eng.Close()
 
 	sstSnapshotStorage := NewSSTSnapshotStorage(eng, testLimiter)
-	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID)
+	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID, nil)
 
 	// Check that the storage lazily creates the directories on first write.
 	_, err := eng.Env().Stat(scratch.snapDir)
@@ -139,7 +139,7 @@ func TestSSTSnapshotStorageConcurrentRange(t *testing.T) {
 	sstSnapshotStorage := NewSSTSnapshotStorage(eng, testLimiter)
 
 	runForSnap := func(snapUUID uuid.UUID) error {
-		scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, snapUUID)
+		scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, snapUUID, nil)
 
 		// Check that the storage lazily creates the directories on first write.
 		_, err := eng.Env().Stat(scratch.snapDir)
@@ -242,7 +242,7 @@ func TestSSTSnapshotStorageContextCancellation(t *testing.T) {
 	defer eng.Close()
 
 	sstSnapshotStorage := NewSSTSnapshotStorage(eng, testLimiter)
-	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID)
+	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID, nil)
 
 	var cancel func()
 	ctx, cancel = context.WithCancel(ctx)
@@ -283,7 +283,7 @@ func testMultiSSTWriterInitSSTInner(t *testing.T, interesting bool) {
 	defer eng.Close()
 
 	sstSnapshotStorage := NewSSTSnapshotStorage(eng, testLimiter)
-	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID)
+	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID, nil)
 	desc := roachpb.RangeDescriptor{
 		RangeID:  100,
 		StartKey: roachpb.RKey("d"),
@@ -415,8 +415,8 @@ func TestMultiSSTWriterSize(t *testing.T) {
 	defer eng.Close()
 
 	sstSnapshotStorage := NewSSTSnapshotStorage(eng, testLimiter)
-	ref := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID)
-	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID)
+	ref := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID, nil)
+	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID, nil)
 	settings := cluster.MakeTestingClusterSettings()
 
 	desc := roachpb.RangeDescriptor{
@@ -533,7 +533,7 @@ func TestMultiSSTWriterAddLastSpan(t *testing.T) {
 	defer eng.Close()
 
 	sstSnapshotStorage := NewSSTSnapshotStorage(eng, testLimiter)
-	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID)
+	scratch := sstSnapshotStorage.NewScratchSpace(testRangeID, testSnapUUID, nil)
 	desc := roachpb.RangeDescriptor{
 		StartKey: roachpb.RKey("d"),
 		EndKey:   roachpb.RKeyMax,

pkg/kv/kvserver/replica_raftstorage.go

@@ -561,7 +561,6 @@ func (r *Replica) applySnapshotRaftMuLocked(
 		Index: kvpb.RaftIndex(nonemptySnap.Metadata.Index),
 		Term:  kvpb.RaftTerm(nonemptySnap.Metadata.Term),
 	}
-	clearedSpans := inSnap.clearedSpans
 
 	subsumedDescs := make([]*roachpb.RangeDescriptor, 0, len(subsumedRepls))
 	for _, sr := range subsumedRepls {
@@ -583,11 +582,9 @@ func (r *Replica) applySnapshotRaftMuLocked(
 		subsumedDescs = append(subsumedDescs, sr.Desc())
 	}
 
-	st := r.ClusterSettings()
-	prepInput := prepareSnapApplyInput{
+	sb := snapWriteBuilder{
 		id: r.ID(),
 
-		st:       st,
 		todoEng:  r.store.TODOEngine(),
 		sl:       r.raftMu.stateLoader,
 		writeSST: inSnap.SSTStorageScratch.WriteSST,
@@ -596,15 +593,13 @@ func (r *Replica) applySnapshotRaftMuLocked(
 		hardState:     hs,
 		desc:          desc,
 		subsumedDescs: subsumedDescs,
-	}
 
-	_ = applySnapshotTODO
-	clearedUnreplicatedSpan, clearedSubsumedSpans, err := prepareSnapApply(ctx, prepInput)
-	if err != nil {
+		cleared: inSnap.clearedSpans,
+	}
+	_ = applySnapshotTODO // 2.4 is written, the rest is handled below
+	if err := sb.prepareSnapApply(ctx); err != nil {
 		return err
 	}
-	clearedSpans = append(clearedSpans, clearedUnreplicatedSpan)
-	clearedSpans = append(clearedSpans, clearedSubsumedSpans...)
 
 	ls := r.asLogStorage()
 
@@ -622,7 +617,7 @@ func (r *Replica) applySnapshotRaftMuLocked(
 	}
 
 	if len(inSnap.externalSSTs)+len(inSnap.sharedSSTs) == 0 && /* simple */
-		inSnap.SSTSize <= snapshotIngestAsWriteThreshold.Get(&st.SV) /* small */ {
+		inSnap.SSTSize <= snapshotIngestAsWriteThreshold.Get(&r.ClusterSettings().SV) /* small */ {
 		applyAsIngest = false
 	}
 
@@ -638,7 +633,7 @@ func (r *Replica) applySnapshotRaftMuLocked(
 	} else {
 		_ = applySnapshotTODO // all atomic
 		err := r.store.TODOEngine().ConvertFilesToBatchAndCommit(
-			ctx, inSnap.SSTStorageScratch.SSTs(), clearedSpans)
+			ctx, inSnap.SSTStorageScratch.SSTs(), sb.cleared)
 		if err != nil {
 			return errors.Wrapf(err, "while applying as batch %s", inSnap.SSTStorageScratch.SSTs())
 		}