Merge #145417

145417: kvserver,snaprecv: move MultiSSTWriter r=tbg a=tbg

This series of commits moves the multiSSTWriter to newly created package `kvserver/kvstorage/snaprecv`. The plan is that it will be joined by the code extracted in #145328, so that we can build a datadriven test that exercises the snapshot receive path's storage logic end to end.

Epic: CRDB-46488
Release note: none

Co-authored-by: Tobias Grieger <[email protected]>

Author: craig[bot]

pkg/BUILD.bazel

@@ -266,6 +266,7 @@ ALL_TESTS = [
     "//pkg/kv/kvserver/kvflowcontrol/node_rac2:node_rac2_test",
     "//pkg/kv/kvserver/kvflowcontrol/rac2:rac2_test",
     "//pkg/kv/kvserver/kvflowcontrol/replica_rac2:replica_rac2_test",
+    "//pkg/kv/kvserver/kvstorage/snaprecv:snaprecv_test",
     "//pkg/kv/kvserver/kvstorage:kvstorage_test",
     "//pkg/kv/kvserver/leases:leases_test",
     "//pkg/kv/kvserver/liveness:liveness_test",
@@ -1542,6 +1543,8 @@ GO_TARGETS = [
     "//pkg/kv/kvserver/kvflowcontrol:kvflowcontrol",
     "//pkg/kv/kvserver/kvserverbase:kvserverbase",
     "//pkg/kv/kvserver/kvserverpb:kvserverpb",
+    "//pkg/kv/kvserver/kvstorage/snaprecv:snaprecv",
+    "//pkg/kv/kvserver/kvstorage/snaprecv:snaprecv_test",
     "//pkg/kv/kvserver/kvstorage:kvstorage",
     "//pkg/kv/kvserver/kvstorage:kvstorage_test",
     "//pkg/kv/kvserver/leases:leases",

pkg/kv/kvserver/BUILD.bazel

@@ -22,7 +22,6 @@ go_library(
         "merge_queue.go",
         "metric_rules.go",
         "metrics.go",
-        "multi_sst_writer.go",
         "mvcc_gc_queue.go",
         "queue.go",
         "queue_helpers_testutil.go",
@@ -77,7 +76,6 @@ go_library(
         "replica_read.go",
         "replica_send.go",
         "replica_split_load.go",
-        "replica_sst_snapshot_storage.go",
         "replica_store_liveness.go",
         "replica_store_liveness_sleep.go",
         "replica_tscache.go",
@@ -163,6 +161,7 @@ go_library(
         "//pkg/kv/kvserver/kvserverbase",
         "//pkg/kv/kvserver/kvserverpb",
         "//pkg/kv/kvserver/kvstorage",
+        "//pkg/kv/kvserver/kvstorage/snaprecv",
         "//pkg/kv/kvserver/leases",
         "//pkg/kv/kvserver/liveness",
         "//pkg/kv/kvserver/liveness/livenesspb",
@@ -255,7 +254,6 @@ go_library(
         "@com_github_cockroachdb_pebble//:pebble",
         "@com_github_cockroachdb_pebble//objstorage",
         "@com_github_cockroachdb_pebble//objstorage/remote",
-        "@com_github_cockroachdb_pebble//rangedel",
         "@com_github_cockroachdb_pebble//rangekey",
         "@com_github_cockroachdb_pebble//sstable/block",
         "@com_github_cockroachdb_pebble//vfs",
@@ -363,7 +361,6 @@ go_test(
         "replica_rankings_test.go",
         "replica_sideload_test.go",
         "replica_split_load_test.go",
-        "replica_sst_snapshot_storage_test.go",
         "replica_store_liveness_test.go",
         "replica_test.go",
         "replica_tscache_test.go",
@@ -567,8 +564,6 @@ go_test(
         "@com_github_cockroachdb_errors//oserror",
         "@com_github_cockroachdb_logtags//:logtags",
         "@com_github_cockroachdb_pebble//:pebble",
-        "@com_github_cockroachdb_pebble//rangekey",
-        "@com_github_cockroachdb_pebble//sstable",
         "@com_github_cockroachdb_pebble//vfs",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_dustin_go_humanize//:go-humanize",
@@ -585,7 +580,6 @@ go_test(
         "@org_golang_google_grpc//metadata",
         "@org_golang_x_sync//errgroup",
         "@org_golang_x_sync//syncmap",
-        "@org_golang_x_time//rate",
     ],
 )
 

pkg/kv/kvserver/client_merge_test.go

@@ -3910,7 +3910,7 @@ func TestStoreRangeMergeRaftSnapshot(t *testing.T) {
 			writer := storage.MakeIngestionSSTWriter(ctx, st, file)
 			if i < len(keySpans)-1 {
 				// The last span is the MVCC span, and is always cleared via Excise.
-				// See multiSSTWriter.
+				// See MultiSSTWriter.
 				if err := writer.ClearRawRange(span.Key, span.EndKey, true /* pointKeys */, true /* rangeKeys */); err != nil {
 					return err
 				}

pkg/kv/kvserver/kv_snapshot_strategy.go

@@ -9,6 +9,7 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/snaprecv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/rditer"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -43,7 +44,7 @@ type kvBatchSnapshotStrategy struct {
 	// before flushing to disk. Only used on the receiver side.
 	sstChunkSize int64
 	// Only used on the receiver side.
-	scratch   *SSTSnapshotStorageScratch
+	scratch   *snaprecv.SSTSnapshotStorageScratch
 	st        *cluster.Settings
 	clusterID uuid.UUID
 }
@@ -104,21 +105,24 @@ func (kvSS *kvBatchSnapshotStrategy) Receive(
 		return noSnap, sendSnapshotError(ctx, s, stream, errors.New("cannot accept shared sstables"))
 	}
 
-	// We rely on the last keyRange passed into multiSSTWriter being the user key
+	// We rely on the last keyRange passed into MultiSSTWriter being the user key
 	// span. If the sender signals that it can no longer do shared replication
 	// (with a TransitionFromSharedToRegularReplicate = true), we will have to
-	// switch to adding a rangedel for that span. Since multiSSTWriter acts on an
+	// switch to adding a rangedel for that span. Since MultiSSTWriter acts on an
 	// opaque slice of keyRanges, we just tell it to add a rangedel for the last
 	// span. To avoid bugs, assert on the last span in keyRanges actually being
 	// equal to the user key span.
 	if !keyRanges[len(keyRanges)-1].Equal(header.State.Desc.KeySpan().AsRawSpanWithNoLocals()) {
-		return noSnap, errors.AssertionFailedf("last span in multiSSTWriter did not equal the user key span: %s", keyRanges[len(keyRanges)-1].String())
+		return noSnap, errors.AssertionFailedf("last span in MultiSSTWriter did not equal the user key span: %s", keyRanges[len(keyRanges)-1].String())
 	}
 
 	// The last key range is the user key span.
 	localRanges := keyRanges[:len(keyRanges)-1]
 	mvccRange := keyRanges[len(keyRanges)-1]
-	msstw, err := newMultiSSTWriter(ctx, kvSS.st, kvSS.scratch, localRanges, mvccRange, kvSS.sstChunkSize)
+	msstw, err := snaprecv.NewMultiSSTWriter(ctx, kvSS.st, kvSS.scratch, localRanges, mvccRange, snaprecv.MultiSSTWriterOptions{
+		SSTChunkSize: kvSS.sstChunkSize,
+		MaxSSTSize:   MaxSnapshotSSTableSize.Get(&kvSS.st.SV),
+	})
 	if err != nil {
 		return noSnap, err
 	}
@@ -168,9 +172,9 @@ func (kvSS *kvBatchSnapshotStrategy) Receive(
 			// All batch operations are guaranteed to be point key or range key puts.
 			for batchReader.Next() {
 				// TODO(lyang24): maybe avoid decoding engine key twice.
-				// msstw calls (i.e. PutInternalPointKey) can use the decoded engine key here as input.
+				// msstw calls (i.e. putInternalPointKey) can use the decoded engine key here as input.
 
-				bytesEstimate := msstw.estimatedDataSize()
+				bytesEstimate := msstw.EstimatedDataSize()
 				delta := bytesEstimate - prevBytesEstimate
 				// Calling nil pacer is a noop.
 				if err := pacer.Pace(ctx, delta, false /* final */); err != nil {
@@ -189,7 +193,7 @@ func (kvSS *kvBatchSnapshotStrategy) Receive(
 					}
 				}
 
-				if err := kvSS.readOneToBatch(ctx, ek, header.SharedReplicate, batchReader, msstw); err != nil {
+				if err := msstw.ReadOne(ctx, ek, header.SharedReplicate, batchReader); err != nil {
 					return noSnap, err
 				}
 			}
@@ -238,8 +242,7 @@ func (kvSS *kvBatchSnapshotStrategy) Receive(
 			// we must still construct SSTs with range deletion tombstones to remove
 			// the data.
 			timingTag.start("sst")
-			dataSize, err := msstw.Finish(ctx)
-			sstSize := msstw.sstSize
+			dataSize, sstSize, err := msstw.Finish(ctx)
 			if err != nil {
 				return noSnap, errors.Wrapf(err, "finishing sst for raft snapshot")
 			}
@@ -286,79 +289,6 @@ func (kvSS *kvBatchSnapshotStrategy) Receive(
 	}
 }
 
-func (kvSS *kvBatchSnapshotStrategy) readOneToBatch(
-	ctx context.Context,
-	ek storage.EngineKey,
-	shared bool, // may receive shared SSTs
-	batchReader *storage.BatchReader,
-	msstw *multiSSTWriter,
-) error {
-	switch batchReader.KeyKind() {
-	case pebble.InternalKeyKindSet, pebble.InternalKeyKindSetWithDelete:
-		if err := msstw.Put(ctx, ek, batchReader.Value()); err != nil {
-			return errors.Wrapf(err, "writing sst for raft snapshot")
-		}
-	case pebble.InternalKeyKindDelete, pebble.InternalKeyKindDeleteSized:
-		if !shared {
-			return errors.AssertionFailedf("unexpected batch entry key kind %d", batchReader.KeyKind())
-		}
-		if err := msstw.PutInternalPointKey(ctx, batchReader.Key(), batchReader.KeyKind(), nil); err != nil {
-			return errors.Wrapf(err, "writing sst for raft snapshot")
-		}
-	case pebble.InternalKeyKindRangeDelete:
-		if !shared {
-			return errors.AssertionFailedf("unexpected batch entry key kind %d", batchReader.KeyKind())
-		}
-		start := batchReader.Key()
-		end, err := batchReader.EndKey()
-		if err != nil {
-			return err
-		}
-		if err := msstw.PutInternalRangeDelete(ctx, start, end); err != nil {
-			return errors.Wrapf(err, "writing sst for raft snapshot")
-		}
-
-	case pebble.InternalKeyKindRangeKeyUnset, pebble.InternalKeyKindRangeKeyDelete:
-		if !shared {
-			return errors.AssertionFailedf("unexpected batch entry key kind %d", batchReader.KeyKind())
-		}
-		start := batchReader.Key()
-		end, err := batchReader.EndKey()
-		if err != nil {
-			return err
-		}
-		rangeKeys, err := batchReader.RawRangeKeys()
-		if err != nil {
-			return err
-		}
-		for _, rkv := range rangeKeys {
-			err := msstw.PutInternalRangeKey(ctx, start, end, rkv)
-			if err != nil {
-				return errors.Wrapf(err, "writing sst for raft snapshot")
-			}
-		}
-	case pebble.InternalKeyKindRangeKeySet:
-		start := ek
-		end, err := batchReader.EngineEndKey()
-		if err != nil {
-			return err
-		}
-		rangeKeys, err := batchReader.EngineRangeKeys()
-		if err != nil {
-			return err
-		}
-		for _, rkv := range rangeKeys {
-			err := msstw.PutRangeKey(ctx, start.Key, end.Key, rkv.Version, rkv.Value)
-			if err != nil {
-				return errors.Wrapf(err, "writing sst for raft snapshot")
-			}
-		}
-	default:
-		return errors.AssertionFailedf("unexpected batch entry key kind %d", batchReader.KeyKind())
-	}
-	return nil
-}
-
 // Send implements the snapshotStrategy interface.
 func (kvSS *kvBatchSnapshotStrategy) Send(
 	ctx context.Context,

pkg/kv/kvserver/kvstorage/snaprecv/BUILD.bazel

@@ -0,0 +1,61 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "snaprecv",
+    srcs = [
+        "multi_sst_writer.go",
+        "sst_snapshot_storage.go",
+    ],
+    importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/snaprecv",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/kv/kvserver/kvserverbase",
+        "//pkg/roachpb",
+        "//pkg/settings/cluster",
+        "//pkg/storage",
+        "//pkg/storage/fs",
+        "//pkg/util/syncutil",
+        "//pkg/util/uuid",
+        "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_pebble//:pebble",
+        "@com_github_cockroachdb_pebble//objstorage",
+        "@com_github_cockroachdb_pebble//rangedel",
+        "@com_github_cockroachdb_pebble//rangekey",
+        "@com_github_cockroachdb_pebble//vfs",
+        "@org_golang_x_time//rate",
+    ],
+)
+
+go_test(
+    name = "snaprecv_test",
+    srcs = ["sst_snapshot_storage_test.go"],
+    data = glob(["testdata/**"]),
+    embed = [":snaprecv"],
+    deps = [
+        "//pkg/keys",
+        "//pkg/kv/kvserver/rditer",
+        "//pkg/roachpb",
+        "//pkg/settings/cluster",
+        "//pkg/storage",
+        "//pkg/storage/fs",
+        "//pkg/storage/mvccencoding",
+        "//pkg/testutils",
+        "//pkg/testutils/datapathutils",
+        "//pkg/testutils/echotest",
+        "//pkg/testutils/storageutils",
+        "//pkg/util/hlc",
+        "//pkg/util/leaktest",
+        "//pkg/util/log",
+        "//pkg/util/timeutil",
+        "//pkg/util/uuid",
+        "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_errors//oserror",
+        "@com_github_cockroachdb_pebble//:pebble",
+        "@com_github_cockroachdb_pebble//rangekey",
+        "@com_github_cockroachdb_pebble//sstable",
+        "@com_github_cockroachdb_pebble//vfs",
+        "@com_github_cockroachdb_redact//:redact",
+        "@com_github_stretchr_testify//require",
+        "@org_golang_x_time//rate",
+    ],
+)