Merge #156775

156775: bulkio: add SST writer and file allocator r=spilchen a=spilchen

Implements bulksst.Writer and FileAllocator to flush sorted KVs into SSTs on local or external storage and track per-file metadata (key range, size, sample key). This is to be used for distributed merge. Much of the code here is taken from: https://github.com/jeffswenson/cockroach/tree/feature-distributed-merge.

Resolves: #156571
Epic: CRDB-48845
Release note (none): none

Co-authored by: `@jeffswenson,` `@fqazi`

Co-authored-by: Matt Spilchen <matt.spilchen@cockroachlabs.com>

Author: craig[bot]

pkg/BUILD.bazel

@@ -396,6 +396,7 @@ ALL_TESTS = [
     "//pkg/sql/appstatspb:appstatspb_test",
     "//pkg/sql/auditlogging:auditlogging_test",
     "//pkg/sql/backfill:backfill_test",
+    "//pkg/sql/bulksst:bulksst_test",
     "//pkg/sql/bulkutil:bulkutil_test",
     "//pkg/sql/cacheutil:cacheutil_test",
     "//pkg/sql/catalog/bootstrap:bootstrap_test",
@@ -1880,6 +1881,8 @@ GO_TARGETS = [
     "//pkg/sql/auditlogging:auditlogging_test",
     "//pkg/sql/backfill:backfill",
     "//pkg/sql/backfill:backfill_test",
+    "//pkg/sql/bulksst:bulksst",
+    "//pkg/sql/bulksst:bulksst_test",
     "//pkg/sql/bulkutil:bulkutil",
     "//pkg/sql/bulkutil:bulkutil_test",
     "//pkg/sql/cacheutil:cacheutil",

pkg/gen/protobuf.bzl

@@ -55,6 +55,7 @@ PROTOBUF_SRCS = [
     "//pkg/server/status/statuspb:statuspb_go_proto",
     "//pkg/settings:settings_go_proto",
     "//pkg/sql/appstatspb:appstatspb_go_proto",
+    "//pkg/sql/bulksst:bulksst_go_proto",
     "//pkg/sql/catalog/catenumpb:catenumpb_go_proto",
     "//pkg/sql/catalog/catpb:catpb_go_proto",
     "//pkg/sql/catalog/descpb:descpb_go_proto",

pkg/sql/bulksst/BUILD.bazel

@@ -0,0 +1,80 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
+
+proto_library(
+    name = "bulksst_proto",
+    srcs = ["sst_info.proto"],
+    strip_import_prefix = "/pkg",
+    visibility = ["//visibility:public"],
+    deps = ["@com_github_gogo_protobuf//gogoproto:gogo_proto"],
+)
+
+go_proto_library(
+    name = "bulksst_go_proto",
+    compilers = ["//pkg/cmd/protoc-gen-gogoroach:protoc-gen-gogoroach_compiler"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/sql/bulksst",
+    proto = ":bulksst_proto",
+    visibility = ["//visibility:public"],
+    deps = ["@com_github_gogo_protobuf//gogoproto"],
+)
+
+go_library(
+    name = "bulksst",
+    srcs = [
+        "sst_file_allocator.go",
+        "sst_writer.go",
+    ],
+    embed = [":bulksst_go_proto"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/sql/bulksst",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/cloud",
+        "//pkg/kv/kvpb",
+        "//pkg/kv/kvserver/kvserverbase",
+        "//pkg/roachpb",
+        "//pkg/settings",
+        "//pkg/settings/cluster",
+        "//pkg/storage",
+        "//pkg/util/hlc",
+        "//pkg/util/log",
+        "//pkg/util/randutil",
+        "//pkg/util/timeutil",
+        "@com_github_cockroachdb_pebble//objstorage",
+        "@com_github_cockroachdb_pebble//objstorage/objstorageprovider",
+    ],
+)
+
+go_test(
+    name = "bulksst_test",
+    srcs = [
+        "main_test.go",
+        "sst_file_allocator_test.go",
+        "sst_writer_test.go",
+    ],
+    embed = [":bulksst"],
+    deps = [
+        "//pkg/base",
+        "//pkg/cloud/cloudpb",
+        "//pkg/cloud/impl:cloudimpl",
+        "//pkg/cloud/nodelocal",
+        "//pkg/kv/kvpb",
+        "//pkg/roachpb",
+        "//pkg/security/securityassets",
+        "//pkg/security/securitytest",
+        "//pkg/server",
+        "//pkg/settings/cluster",
+        "//pkg/storage",
+        "//pkg/testutils/serverutils",
+        "//pkg/testutils/testcluster",
+        "//pkg/util/hlc",
+        "//pkg/util/intsets",
+        "//pkg/util/leaktest",
+        "//pkg/util/log",
+        "//pkg/util/randutil",
+        "@com_github_cockroachdb_pebble//objstorage",
+        "@com_github_cockroachdb_pebble//sstable",
+        "@com_github_cockroachdb_pebble//vfs",
+        "@com_github_stretchr_testify//require",
+    ],
+)

pkg/sql/bulksst/main_test.go

@@ -0,0 +1,33 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package bulksst_test
+
+import (
+	"os"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/security/securityassets"
+	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
+	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+)
+
+//go:generate ../../util/leaktest/add-leaktest.sh *_test.go
+
+func init() {
+	securityassets.SetLoader(securitytest.EmbeddedAssets)
+}
+func TestMain(m *testing.M) {
+	randutil.SeedForTests()
+	serverutils.InitTestServerFactory(server.TestServerFactory)
+	serverutils.InitTestClusterFactory(testcluster.TestClusterFactory)
+
+	code := m.Run()
+
+	os.Exit(code)
+}

pkg/sql/bulksst/sst_file_allocator.go

@@ -0,0 +1,122 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package bulksst
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+
+	"github.com/cockroachdb/cockroach/pkg/cloud"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+	"github.com/cockroachdb/pebble/objstorage"
+	"github.com/cockroachdb/pebble/objstorage/objstorageprovider"
+)
+
+const (
+	maxRowSamples = 1024
+)
+
+// FileAllocator is used to allocate new files for SSTs ingested via the Writer.
+type FileAllocator interface {
+	// AddFile creates a new file and returns the objstorage.Writable and URI for
+	// tracking. The allocator generates a unique filename internally. The caller
+	// is responsible for closing the returned objstorage.Writable.
+	AddFile(ctx context.Context) (objstorage.Writable, string, error)
+
+	// CommitFile records metadata for a successfully written SST file.
+	CommitFile(uri string, span roachpb.Span, rowSample roachpb.Key, fileSize uint64)
+
+	// GetFileList gets all the files created by this file allocator.
+	GetFileList() *SSTFiles
+}
+
+// fileAllocatorBase helps track metadata for created SST files.
+type fileAllocatorBase struct {
+	fileInfo        SSTFiles
+	rowSampleRand   *rand.Rand
+	totalRowSamples int
+}
+
+// GetFileList gets all the files created by this file allocator.
+func (f *fileAllocatorBase) GetFileList() *SSTFiles {
+	return &f.fileInfo
+}
+
+// addFile helps track metadata for created SST files.
+func (f *fileAllocatorBase) addFile(
+	uri string, span roachpb.Span, rowSample roachpb.Key, fileSize uint64,
+) {
+	f.fileInfo.SST = append(f.fileInfo.SST, &SSTFileInfo{
+		URI:      uri,
+		StartKey: span.Key,
+		EndKey:   span.EndKey,
+		FileSize: fileSize,
+	})
+	f.fileInfo.TotalSize += fileSize
+	f.recordRowSample(rowSample)
+}
+
+func (f *fileAllocatorBase) recordRowSample(rowSample roachpb.Key) {
+	if len(rowSample) == 0 {
+		return
+	}
+	f.totalRowSamples++
+	if len(f.fileInfo.RowSamples) < maxRowSamples {
+		f.fileInfo.RowSamples = append(f.fileInfo.RowSamples, string(rowSample))
+		return
+	}
+	if f.rowSampleRand == nil {
+		rng, _ := randutil.NewLockedPseudoRand()
+		f.rowSampleRand = rng
+	}
+	// Reservoir sampling: replace an existing sample with probability maxRowSamples / totalRowSamples.
+	idx := f.rowSampleRand.Intn(f.totalRowSamples)
+	if idx < maxRowSamples {
+		f.fileInfo.RowSamples[idx] = string(rowSample)
+	}
+}
+
+// ExternalFileAllocator allocates external files for SSTs.
+type ExternalFileAllocator struct {
+	es      cloud.ExternalStorage
+	baseURI string
+	clock   *hlc.Clock
+	fileAllocatorBase
+}
+
+func NewExternalFileAllocator(
+	es cloud.ExternalStorage, baseURI string, clock *hlc.Clock,
+) FileAllocator {
+	return &ExternalFileAllocator{
+		es:                es,
+		baseURI:           baseURI,
+		clock:             clock,
+		fileAllocatorBase: fileAllocatorBase{},
+	}
+}
+
+// AddFile creates a new file with an HLC timestamp-based unique name.
+func (e *ExternalFileAllocator) AddFile(ctx context.Context) (objstorage.Writable, string, error) {
+	// Use HLC timestamp for unique filename generation.
+	ts := e.clock.Now()
+	fileName := fmt.Sprintf("%d-%d.sst", ts.WallTime, ts.Logical)
+	writer, err := e.es.Writer(ctx, fileName)
+	if err != nil {
+		return nil, "", err
+	}
+	remoteWritable := objstorageprovider.NewRemoteWritable(writer)
+	return remoteWritable, e.baseURI + fileName, nil
+}
+
+// CommitFile records metadata for a successfully written SST file.
+func (e *ExternalFileAllocator) CommitFile(
+	uri string, span roachpb.Span, rowSample roachpb.Key, fileSize uint64,
+) {
+	e.fileAllocatorBase.addFile(uri, span, rowSample, fileSize)
+}

pkg/sql/bulksst/sst_file_allocator_test.go

@@ -0,0 +1,115 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package bulksst
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/pebble/objstorage"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRowSampleReservoirCapsSamples(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	var base fileAllocatorBase
+	expected := make(map[string]struct{})
+	total := maxRowSamples*2 + 17
+	for i := 0; i < total; i++ {
+		key := roachpb.Key(fmt.Sprintf("key-%04d", i))
+		expected[string(key)] = struct{}{}
+		span := roachpb.Span{Key: key, EndKey: key.Next()}
+		base.addFile(fmt.Sprintf("uri-%d", i), span, key, 1)
+	}
+
+	require.Equal(t, total, base.totalRowSamples)
+	require.Len(t, base.fileInfo.RowSamples, maxRowSamples)
+
+	for _, sample := range base.fileInfo.RowSamples {
+		_, ok := expected[sample]
+		require.True(t, ok, "sample %q must correspond to ingested keys", sample)
+	}
+}
+
+func TestRowSampleSkipsEmptyKeys(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	var base fileAllocatorBase
+	key := roachpb.Key("some-key")
+	span := roachpb.Span{Key: key, EndKey: key.Next()}
+
+	base.addFile("uri-1", span, key, 1)
+	base.addFile("uri-2", span, nil, 1)
+
+	require.Len(t, base.fileInfo.RowSamples, 1)
+	require.Equal(t, string(key), base.fileInfo.RowSamples[0])
+}
+
+func TestWriterFlushFailureDoesNotPersistMetadata(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	st := cluster.MakeTestingClusterSettings()
+	failErr := errors.New("finish failed")
+	allocator := &mockFailingAllocator{finishErr: failErr}
+
+	writer := NewUnsortedSSTBatcher(st, allocator)
+	ts := hlc.Timestamp{WallTime: 1}
+	require.NoError(t, writer.AddMVCCKey(ctx, storage.MVCCKey{Key: roachpb.Key("k"), Timestamp: ts}, []byte("v")))
+
+	err := writer.Flush(ctx)
+	require.ErrorIs(t, err, failErr)
+
+	fileList := allocator.GetFileList()
+	require.Empty(t, fileList.SST, "no files should be tracked after a failed flush")
+	require.Zero(t, fileList.TotalSize)
+	require.Empty(t, fileList.RowSamples)
+}
+
+type mockFailingAllocator struct {
+	fileAllocatorBase
+	finishErr error
+	nextID    int
+}
+
+func (m *mockFailingAllocator) AddFile(ctx context.Context) (objstorage.Writable, string, error) {
+	uri := fmt.Sprintf("mock://%d", m.nextID)
+	m.nextID++
+	w := &failingWritable{finishErr: m.finishErr}
+	return w, uri, nil
+}
+
+func (m *mockFailingAllocator) CommitFile(
+	uri string, span roachpb.Span, rowSample roachpb.Key, fileSize uint64,
+) {
+	m.fileAllocatorBase.addFile(uri, span, rowSample, fileSize)
+}
+
+type failingWritable struct {
+	finishErr error
+}
+
+func (f *failingWritable) Write(p []byte) error {
+	return nil
+}
+
+func (f *failingWritable) Finish() error {
+	return f.finishErr
+}
+
+func (f *failingWritable) Abort() {}