Merge #144724

144724: backup: fix compaction issue with attempting to write keys outside of span r=msbutler a=kev-cao

After the fix applied by #144652, the test failures outlined in #144216 resurfaced as job failures due to attempts by the `SSTSinkKeyWriter` to write keys outside of the `BackupManifest_File` range. The behavior was unexpected, because for some key `/Table/1/1/5`, we would see an attempt to write it to a span starting with `/Table/1/1/5/1`, despite the the last `SSTSinkKeyWriter.Reset` call being made on a span starting with `/Table/1/1/1`.

This was ultimately determined to be caused by the fact that in the compaction processor's `compactSpanEntry`, we re-use the same underlying memory for the key that is being passed to `SSTSinkKeyWriter.WriteKey` to save on memory allocations. However, if the sink performed a flush due to size constraints in `maybeDoSizeFlush`, the start key of the span used to reset the sink referenced the same memory location as key that was passed in. So for subsequent keys in that span that were written, the re-use of that underlying memory in `compactSpanEntry` would consistently mutate the span referenced by that `BackupManifest_File`, causing corruption. This would usually result in job failures as eventually a key may be written outside of the span, but occasionally result in insidious job successes and the resulting `BackupManifest_File` would report a far smaller span than it actually covered.

One solution was to perform a clone of the key in `maybeDoSizeFlush` when creating the span to reset the sink with, but ultimately decided to instead ensure that anytime `SSTSinkKeyWriter.Reset` is called, we pass a clone of the span to the `BackupManifest_File`. This ensures that once `Reset` is called, the caller is free to reuse the underlying memory of the span however they wish. The same holds true for the key passed to `WriteKey`, as in any instances in which the passed in key is actually persisted, we always persist a copy.

Fixes: #144216, #144339

Release note: None

Co-authored-by: Kevin Cao <[email protected]>

Author: craig[bot]

pkg/backup/backupsink/sst_sink_key_writer.go

@@ -45,10 +45,12 @@ func MakeSSTSinkKeyWriter(conf SSTSinkConf, dest cloud.ExternalStorage) (*SSTSin
 // including the prefix. Reset needs to be called prior to WriteKey whenever
 // writing keys from a new span. Keys must also be written in order.
 //
-// After writing the last key for a span, AssumeNotMidRow must be called to
-// enforce the invariant that BackupManifest_File spans do not end mid-row.
+// Once a key has been written, the caller may safely reuse the underlying
+// memory for the passed in key.
 //
-// Flush should be called before the sink is closed to ensure the SST
+// NOTE: After writing the last key for a span, AssumeNotMidRow must be called
+// to enforce the invariant that BackupManifest_File spans do not end mid-row.
+// Flush should also be called before the sink is closed to ensure the SST
 // is written to the destination.
 func (s *SSTSinkKeyWriter) WriteKey(ctx context.Context, key storage.MVCCKey, value []byte) error {
 	if len(s.flushedFiles) == 0 {
@@ -117,6 +119,9 @@ func (s *SSTSinkKeyWriter) AssumeNotMidRow() {
 // and calling AssumeNotMidRow before resetting to enforce this invariant.
 // Any time a new span is being written, Reset MUST be called prior to any
 // WriteKey calls.
+//
+// Once Reset has been called, the caller may safely reuse the underlying memory
+// of the passed in span.
 func (s *SSTSinkKeyWriter) Reset(ctx context.Context, newSpan roachpb.Span) error {
 	log.VEventf(ctx, 2, "resetting sink to span %s", newSpan)
 	if s.midRow {
@@ -143,7 +148,8 @@ func (s *SSTSinkKeyWriter) Reset(ctx context.Context, newSpan roachpb.Span) erro
 			lastFile.EntryCounts.DataSize < fileSpanByteLimit {
 			log.VEventf(ctx, 2, "extending span %s to %s", lastFile.Span, newSpan)
 			s.stats.spanGrows++
-			lastFile.Span.EndKey = newSpan.EndKey
+			// See reason for Clone() below.
+			lastFile.Span.EndKey = newSpan.EndKey.Clone()
 			return nil
 		}
 	}
@@ -156,7 +162,11 @@ func (s *SSTSinkKeyWriter) Reset(ctx context.Context, newSpan roachpb.Span) erro
 	s.flushedFiles = append(
 		s.flushedFiles,
 		backuppb.BackupManifest_File{
-			Span: newSpan,
+			// Because there are situations where the underlying memory for keys of
+			// the span is reused to optimize memory usage, we need to clone the keys
+			// to ensure that the BackupManifest_File's span is not unintentionally
+			// mutated outside of the SSTSinkKeyWriter.
+			Span: newSpan.Clone(),
 			Path: s.outName,
 		},
 	)

pkg/backup/backupsink/sst_sink_key_writer_test.go

@@ -481,6 +481,48 @@ func TestEnforceFileSSTSinkAssumeNotMidRow(t *testing.T) {
 	})
 }
 
+func TestSSTSinkWriterSafeAgainstKeyMutation(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	st := cluster.MakeTestingClusterSettings()
+	sink, _ := sstSinkKeyWriterTestSetup(t, st, execinfrapb.ElidePrefix_TenantAndTable)
+	defer func() {
+		require.NoError(t, sink.Close())
+	}()
+
+	mutateSpan := func(span *roachpb.Span) {
+		span.Key = append(span.Key[:0], make([]byte, len(span.Key))...)
+		span.EndKey = append(span.EndKey[:0], make([]byte, len(span.EndKey))...)
+	}
+
+	t.Run("safe on reset on new span", func(t *testing.T) {
+		keySet := newMVCCKeySet("a", "c").withKVs([]kvAndTS{{key: "a", timestamp: 10}})
+		require.NoError(t, sink.Reset(ctx, keySet.span))
+		require.Len(t, sink.flushedFiles, 1)
+		originalSpan := sink.flushedFiles[0].Span.Clone()
+		mutateSpan(&keySet.span)
+		require.Equal(t, originalSpan, sink.flushedFiles[0].Span)
+		require.NoError(t, sink.Flush(ctx))
+	})
+
+	t.Run("safe on reset on extending span", func(t *testing.T) {
+		keySet := newMVCCKeySet("a", "c").withKVs([]kvAndTS{{key: "a", timestamp: 10}})
+		extendingSet := newMVCCKeySet("c", "e").withKVs([]kvAndTS{{key: "c", timestamp: 10}})
+		require.NoError(t, sink.Reset(ctx, keySet.span))
+		require.NoError(t, sink.WriteKey(ctx, keySet.kvs[0].key, keySet.kvs[0].value))
+		sink.AssumeNotMidRow()
+
+		require.NoError(t, sink.Reset(ctx, extendingSet.span))
+		require.Len(t, sink.flushedFiles, 1)
+		originalSpan := sink.flushedFiles[0].Span.Clone()
+		mutateSpan(&extendingSet.span)
+		require.Equal(t, originalSpan, sink.flushedFiles[0].Span)
+		require.NoError(t, sink.Flush(ctx))
+	})
+}
+
 func sstSinkKeyWriterTestSetup(
 	t *testing.T, settings *cluster.Settings, elideMode execinfrapb.ElidePrefix,
 ) (*SSTSinkKeyWriter, cloud.ExternalStorage) {

pkg/backup/compaction_processor.go

@@ -395,6 +395,9 @@ func compactSpanEntry(
 ) error {
 	defer sstIter.cleanup()
 	entry := sstIter.entry
+	if err := assertCommonPrefix(entry.Span, entry.ElidedPrefix); err != nil {
+		return err
+	}
 	prefix, err := backupsink.ElidedPrefix(entry.Span.Key, entry.ElidedPrefix)
 	if err != nil {
 		return err