Merge #153056

153056: cloud: create and use an objstorage.Writable writer wrapper r=jeffswenson a=jeffswenson

Previously, CRDB code interacting with the pebble SST writer used a wrapper that no-op'd Finish/Abort calls on the cloud object. This caused us to write invalid SSTs to cloud storage when we would have otherwise aborted the write.

Implementing Finish/Abort also lets us clean up some of the lifetime handling. Creating an SST writer normally takes ownership of the object. But since we no-op'd the cleanup methods, we had to hold onto the object and manually Close it to trigger the flush.

Issue: #153055
Epic: CRDB-53946
Release note: none

Co-authored-by: Jeff Swenson <[email protected]>

Author: craig[bot]

pkg/backup/BUILD.bazel

@@ -343,6 +343,7 @@ go_test(
         "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_errors//oserror",
+        "@com_github_cockroachdb_pebble//objstorage",
         "@com_github_cockroachdb_pebble//sstable",
         "@com_github_cockroachdb_pebble//vfs",
         "@com_github_cockroachdb_redact//:redact",

pkg/backup/backupinfo/BUILD.bazel

@@ -51,6 +51,7 @@ go_library(
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_pebble//objstorage",
         "@com_github_klauspost_compress//gzip",
     ],
 )

pkg/backup/backupinfo/desc_sst.go

@@ -49,19 +49,15 @@ func WriteDescsSST(
 	if err != nil {
 		return err
 	}
-	defer w.Close()
 	descSST := storage.MakeTransportSSTWriter(ctx, dest.Settings(), w)
 	defer descSST.Close()
 
-	if err := writeDescsToMetadata(ctx, descSST, m); err != nil {
-		return err
-	}
-
-	if err := descSST.Finish(); err != nil {
+	err = writeDescsToMetadata(ctx, descSST, m)
+	if err != nil {
 		return err
 	}
 
-	return w.Close()
+	return descSST.Finish()
 }
 
 func DescChangesLess(

pkg/backup/backupinfo/external_sst_util.go

@@ -8,7 +8,6 @@ package backupinfo
 import (
 	"bytes"
 	"context"
-	"io"
 
 	"github.com/cockroachdb/cockroach/pkg/backup/backupencryption"
 	"github.com/cockroachdb/cockroach/pkg/ccl/storageccl"
@@ -17,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/pebble/objstorage"
 )
 
 func makeWriter(
@@ -25,8 +25,9 @@ func makeWriter(
 	filename string,
 	enc *jobspb.BackupEncryptionOptions,
 	kmsEnv cloud.KMSEnv,
-) (io.WriteCloser, error) {
-	w, err := dest.Writer(ctx, filename)
+) (objstorage.Writable, error) {
+	var w objstorage.Writable
+	w, err := cloud.OpenAbortableWriter(ctx, dest, filename)
 	if err != nil {
 		return nil, err
 	}

pkg/backup/backupinfo/file_sst.go

@@ -55,7 +55,6 @@ func writeFilesSST(
 	if err != nil {
 		return err
 	}
-	defer w.Close()
 	fileSST := storage.MakeTransportSSTWriter(ctx, dest.Settings(), w)
 	defer fileSST.Close()
 
@@ -75,11 +74,7 @@ func writeFilesSST(
 		}
 	}
 
-	err = fileSST.Finish()
-	if err != nil {
-		return err
-	}
-	return w.Close()
+	return fileSST.Finish()
 }
 
 func encodeFileSSTKey(spanStart roachpb.Key, filename string) roachpb.Key {

pkg/backup/backupsink/BUILD.bazel

@@ -25,8 +25,8 @@ go_library(
         "//pkg/util/log",
         "//pkg/util/unique",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_pebble//objstorage",
         "@com_github_gogo_protobuf//types",
-        "@com_github_kr_pretty//:pretty",
     ],
 )
 
@@ -56,6 +56,7 @@ go_test(
         "//pkg/util/leaktest",
         "//pkg/util/log",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_pebble//objstorage",
         "@com_github_gogo_protobuf//types",
         "@com_github_stretchr_testify//require",
     ],

pkg/backup/backupsink/file_sst_sink.go

@@ -8,7 +8,6 @@ package backupsink
 import (
 	"bytes"
 	"context"
-	io "io"
 
 	"github.com/cockroachdb/cockroach/pkg/backup/backuppb"
 	"github.com/cockroachdb/cockroach/pkg/base"
@@ -24,8 +23,8 @@ import (
 	hlc "github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble/objstorage"
 	gogotypes "github.com/gogo/protobuf/types"
-	"github.com/kr/pretty"
 )
 
 var (
@@ -58,10 +57,11 @@ type FileSSTSink struct {
 	conf  SSTSinkConf
 	pacer *admission.Pacer
 
-	sst     storage.SSTWriter
+	isOpen bool
+	sst    storage.SSTWriter
+
 	ctx     context.Context
 	cancel  func()
-	out     io.WriteCloser
 	outName string
 
 	flushedFiles []backuppb.BackupManifest_File
@@ -133,7 +133,7 @@ func (s *FileSSTSink) Write(ctx context.Context, resp ExportedSpan) (roachpb.Key
 	}
 
 	// Initialize the writer if needed.
-	if s.out == nil {
+	if !s.isOpen {
 		if err := s.open(ctx); err != nil {
 			return nil, err
 		}
@@ -232,16 +232,14 @@ func (s *FileSSTSink) Close() error {
 		s.cancel()
 	}
 
-	var err error
-	if s.out != nil {
-		err = s.out.Close()
-	}
 	s.sst.Close()
-	return err
+	s.isOpen = false
+
+	return nil
 }
 
 func (s *FileSSTSink) Flush(ctx context.Context) error {
-	if s.out == nil {
+	if !s.isOpen {
 		// If the writer was not initialized but the sink has reported completed
 		// spans then there were empty ExportRequests that were processed by the
 		// owner of this sink. These still need to reported to the coordinator as
@@ -279,13 +277,10 @@ func (s *FileSSTSink) Flush(ctx context.Context) error {
 	if err := s.sst.Finish(); err != nil {
 		return err
 	}
-	if err := s.out.Close(); err != nil {
-		log.Dev.Warningf(ctx, "failed to close write in FileSSTSink: % #v", pretty.Formatter(err))
-		return errors.Wrap(err, "writing SST")
-	}
+
 	wroteSize := s.sst.Meta.Size
 	s.outName = ""
-	s.out = nil
+	s.isOpen = false
 
 	for i := range s.flushedFiles {
 		s.flushedFiles[i].BackingFileSize = wroteSize
@@ -322,30 +317,30 @@ func (s *FileSSTSink) open(ctx context.Context) error {
 	if s.ctx == nil {
 		s.ctx, s.cancel = context.WithCancel(ctx)
 	}
-	w, err := s.dest.Writer(s.ctx, s.outName)
+	var w objstorage.Writable
+	w, err := cloud.OpenAbortableWriter(s.ctx, s.dest, s.outName)
 	if err != nil {
 		return err
 	}
-	s.out = w
 	if s.conf.Enc != nil {
-		e, err := storageccl.EncryptingWriter(w, s.conf.Enc.Key)
+		w, err = storageccl.EncryptingWriter(w, s.conf.Enc.Key)
 		if err != nil {
 			return err
 		}
-		s.out = e
 	}
-	// TODO(dt): make ExternalStorage.Writer return objstorage.Writable.
-	//
+	// NOTE: the sst writer takes ownership of the object writer. So we don't
+	// hold on to it to close it.
 	// Value blocks are disabled since such SSTs can be huge (e.g. 750MB in the
 	// mixed_version_backup.go roachtest), which can cause OOMs due to value
 	// block buffering.
 	s.sst = storage.MakeIngestionSSTWriterWithOverrides(
-		ctx, s.dest.Settings(), storage.NoopFinishAbortWritable(s.out),
+		ctx, s.dest.Settings(), w,
 		storage.WithValueBlocksDisabled,
 		storage.WithCompressionFromClusterSetting(
 			ctx, s.dest.Settings(), storage.CompressionAlgorithmBackupStorage,
 		),
 	)
+	s.isOpen = true
 
 	return nil
 }

pkg/backup/backupsink/file_sst_sink_test.go

@@ -6,7 +6,6 @@
 package backupsink
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"reflect"
@@ -29,6 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble/objstorage"
 	"github.com/gogo/protobuf/types"
 	"github.com/stretchr/testify/require"
 )
@@ -43,13 +43,13 @@ func TestFileSSTSinkExtendOneFile(t *testing.T) {
 	ctx := context.Background()
 
 	getKeys := func(prefix string, n int) []byte {
-		var b bytes.Buffer
+		var b objstorage.MemObj
 		sst := storage.MakeTransportSSTWriter(ctx, cluster.MakeTestingClusterSettings(), &b)
 		for i := 0; i < n; i++ {
 			require.NoError(t, sst.PutUnversioned([]byte(fmt.Sprintf("%s%08d", prefix, i)), nil))
 		}
 		sst.Close()
-		return b.Bytes()
+		return b.Data()
 	}
 
 	exportResponse1 := ExportedSpan{
@@ -661,8 +661,8 @@ func TestFileSSTSinkCopyPointKeys(t *testing.T) {
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
-			buf := &bytes.Buffer{}
-			sst := storage.MakeTransportSSTWriter(ctx, settings, buf)
+			buf := objstorage.MemObj{}
+			sst := storage.MakeTransportSSTWriter(ctx, settings, &buf)
 			sink := FileSSTSink{sst: sst}
 			compareSST := true
 
@@ -713,7 +713,7 @@ func TestFileSSTSinkCopyPointKeys(t *testing.T) {
 				UpperBound: keys.MaxKey,
 			}
 
-			iter, err := storage.NewMemSSTIterator(buf.Bytes(), false, iterOpts)
+			iter, err := storage.NewMemSSTIterator(buf.Data(), false, iterOpts)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -840,8 +840,8 @@ func TestFileSSTSinkCopyRangeKeys(t *testing.T) {
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
-			buf := &bytes.Buffer{}
-			sst := storage.MakeTransportSSTWriter(ctx, settings, buf)
+			buf := objstorage.MemObj{}
+			sst := storage.MakeTransportSSTWriter(ctx, settings, &buf)
 			sink := FileSSTSink{sst: sst}
 			compareSST := true
 
@@ -887,7 +887,7 @@ func TestFileSSTSinkCopyRangeKeys(t *testing.T) {
 				UpperBound: keys.MaxKey,
 			}
 
-			iter, err := storage.NewMemSSTIterator(buf.Bytes(), false, iterOpts)
+			iter, err := storage.NewMemSSTIterator(buf.Data(), false, iterOpts)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -1029,8 +1029,8 @@ func (b *exportedSpanBuilder) build() ExportedSpan {
 func (b *exportedSpanBuilder) buildWithEncoding(stringToKey func(string) roachpb.Key) ExportedSpan {
 	ctx := context.Background()
 	settings := cluster.MakeTestingClusterSettings()
-	buf := &bytes.Buffer{}
-	sst := storage.MakeTransportSSTWriter(ctx, settings, buf)
+	buf := objstorage.MemObj{}
+	sst := storage.MakeTransportSSTWriter(ctx, settings, &buf)
 	for _, d := range b.keyValues {
 		v := roachpb.Value{}
 		v.SetBytes(d.value)
@@ -1058,7 +1058,7 @@ func (b *exportedSpanBuilder) buildWithEncoding(stringToKey func(string) roachpb
 
 	sst.Close()
 
-	b.es.DataSST = buf.Bytes()
+	b.es.DataSST = buf.Data()
 
 	return *b.es
 }

pkg/backup/backupsink/sst_sink_key_writer.go

@@ -135,7 +135,7 @@ func (s *SSTSinkKeyWriter) Reset(ctx context.Context, newSpan roachpb.Span) erro
 			return err
 		}
 	}
-	if s.out == nil {
+	if !s.isOpen {
 		if err := s.open(ctx); err != nil {
 			return err
 		}

pkg/backup/bench_test.go

@@ -8,7 +8,6 @@ package backup
 import (
 	"context"
 	"fmt"
-	"io"
 	"math/rand"
 	"testing"
 
@@ -29,6 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/pebble/objstorage"
 	"github.com/stretchr/testify/require"
 )
 
@@ -59,24 +59,22 @@ func BenchmarkIteratorMemory(b *testing.B) {
 	makeWriter := func(
 		store cloud.ExternalStorage,
 		filename string,
-		enc *jobspb.BackupEncryptionOptions) (io.WriteCloser, error) {
-		w, err := store.Writer(ctx, filename)
+		enc *jobspb.BackupEncryptionOptions) (objstorage.Writable, error) {
+		w, err := cloud.OpenAbortableWriter(ctx, store, filename)
 		if err != nil {
 			return nil, err
 		}
 
-		if enc != nil {
-			key, err := backupencryption.GetEncryptionKey(ctx, enc, nil)
-			if err != nil {
-				return nil, err
-			}
-			encW, err := storageccl.EncryptingWriter(w, key)
-			if err != nil {
-				return nil, err
-			}
-			w = encW
+		if enc == nil {
+			return w, nil
 		}
-		return w, nil
+
+		key, err := backupencryption.GetEncryptionKey(ctx, enc, nil)
+		if err != nil {
+			return nil, err
+		}
+
+		return storageccl.EncryptingWriter(w, key)
 	}
 
 	getRandomPayload := func(buf []byte) {
@@ -87,7 +85,7 @@ func BenchmarkIteratorMemory(b *testing.B) {
 		}
 	}
 
-	writeSST := func(w io.WriteCloser, store cloud.ExternalStorage, payloadSize int, numKeys int) {
+	writeSST := func(w objstorage.Writable, store cloud.ExternalStorage, payloadSize int, numKeys int) {
 		sst := storage.MakeTransportSSTWriter(ctx, store.Settings(), w)
 
 		buf := make([]byte, payloadSize)
@@ -156,7 +154,7 @@ func BenchmarkIteratorMemory(b *testing.B) {
 							require.NoError(b, err)
 
 							writeSST(w, store, 100, rows)
-							require.NoError(b, w.Close())
+							require.NoError(b, w.Finish())
 
 							sz, err := store.Size(ctx, filename)
 							require.NoError(b, err)