Merge #63891

63891: changefeedccl: Correctly account for memory when closing gzip files. r=miretskiy a=miretskiy

Correctly account for used memory when closing gzip compressed cloud storage
file.

Fixes #63888 

Release Notes: Bug fix; correctly account for used memory when closing
compressed files.

Co-authored-by: Yevgeniy Miretskiy <[email protected]>

pkg/ccl/changefeedccl/sink_cloudstorage.go

@@ -512,8 +512,15 @@ func (s *cloudStorageSink) flushFile(ctx context.Context, file *cloudStorageSink
 		return nil
 	}
 
-	// If the file is written via compression codec, close the codec to ensure it
-	// has flushed to the underlying buffer.
+	// Release memory allocated for this file.  Note, closing codec
+	// below may as well write more data to our buffer (and that may cause buffer
+	// to grow due to reallocation).  But we don't account for that additional memory
+	// because a) we don't know if buffer will be resized (nor by how much), and
+	// b) if we're out of memory we'd OOMed when trying to close codec anyway.
+	defer func(delta int) {
+		s.mem.Shrink(ctx, int64(delta))
+	}(file.buf.Cap())
+
 	if file.codec != nil {
 		if err := file.codec.Close(); err != nil {
 			return err
@@ -538,7 +545,6 @@ func (s *cloudStorageSink) flushFile(ctx context.Context, file *cloudStorageSink
 	if err := s.es.WriteFile(ctx, filepath.Join(s.dataFilePartition, filename), bytes.NewReader(file.buf.Bytes())); err != nil {
 		return err
 	}
-	s.mem.Shrink(ctx, int64(file.buf.Cap()))
 	return nil
 }
 

pkg/ccl/changefeedccl/sink_cloudstorage_test.go

@@ -35,6 +35,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/span"
 	"github.com/stretchr/testify/require"
 )
@@ -575,4 +576,43 @@ func TestCloudStorageSink(t *testing.T) {
 		}
 		require.Regexp(t, "memory budget exceeded", err)
 	})
+	t.Run(`memory-accounting`, func(t *testing.T) {
+		before := opts[changefeedbase.OptCompression]
+		// Compression codecs include buffering that interferes with other tests,
+		// e.g. the bucketing test that configures very small flush sizes.
+		defer func() {
+			opts[changefeedbase.OptCompression] = before
+		}()
+
+		// A bit of magic constant: we're using bytes.Buffer internally, which
+		// allocates "small" buffer (64 bytes) initially.  We will try to target
+		// our file size to be less than that value; but we will write
+		// larger amount of data (thus, hopefully causing multiple bytes.Buffer
+		// reallocs, plus a file flush).
+		const targetFileSize = 63
+
+		rnd, _ := randutil.NewPseudoRand()
+		for _, compression := range []string{"", "gzip"} {
+			opts[changefeedbase.OptCompression] = compression
+			t.Run("compress="+compression, func(t *testing.T) {
+				t1 := makeTopic(`t1`)
+				testSpan := roachpb.Span{Key: []byte("a"), EndKey: []byte("b")}
+				sf := span.MakeFrontier(testSpan)
+				timestampOracle := &changeAggregatorLowerBoundOracle{sf: sf}
+				sinkDir := `memory-accounting`
+				s, err := makeCloudStorageSink(
+					ctx, `nodelocal://0/`+sinkDir, 1, targetFileSize,
+					settings, opts, timestampOracle, externalStorageFromURI, user, memAcc,
+				)
+				require.NoError(t, err)
+				defer func() { require.NoError(t, s.Close()) }()
+				s.(*cloudStorageSink).sinkID = 7 // Force a deterministic sinkID.
+
+				data := randutil.RandBytes(rnd, 1+rnd.Intn(targetFileSize))
+				require.NoError(t, s.EmitRow(ctx, t1, noKey, data, ts(0)))
+
+				require.NoError(t, s.Flush(ctx))
+			})
+		}
+	})
 }