[compression] NewZstdCompressingWriter

Author: dan-stowell

server/util/compression/BUILD

@@ -7,6 +7,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//server/metrics",
+        "//server/util/bytebufferpool",
         "//server/util/log",
         "@com_github_klauspost_compress//zstd",
         "@com_github_prometheus_client_golang//prometheus",

server/util/compression/compression.go

@@ -7,11 +7,14 @@ import (
 	"sync"
 
 	"github.com/buildbuddy-io/buildbuddy/server/metrics"
+	"github.com/buildbuddy-io/buildbuddy/server/util/bytebufferpool"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/klauspost/compress/zstd"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
+const compressChunkSize = 4 * 1024 * 1024 // 4MB
+
 var (
 	// zstdEncoder can be shared across goroutines to compress chunks of data
 	// using EncodeAll. Streaming functions such as encoder.ReadFrom or io.Copy
@@ -22,6 +25,8 @@ var (
 	// either for streaming decompression using ReadFrom or batch decompression
 	// using DecodeAll. The returned decoders *must not* be closed.
 	zstdDecoderPool = NewZstdDecoderPool()
+
+	compressBufPool = bytebufferpool.FixedSize(compressChunkSize)
 )
 
 func mustGetZstdEncoder() *zstd.Encoder {
@@ -177,6 +182,50 @@ func NewZstdCompressingReader(reader io.ReadCloser, readBuf []byte, compressBuf
 	}, nil
 }
 
+type compressingWriter struct {
+	w               io.Writer
+	compressBuf     []byte
+	poolCompressBuf []byte
+}
+
+func (c *compressingWriter) Write(p []byte) (int, error) {
+	var totalWritten int
+	for len(p) > 0 {
+		chunkSize := min(len(p), cap(c.compressBuf))
+		chunk := p[:chunkSize]
+		c.compressBuf = CompressZstd(c.compressBuf[:0], chunk)
+
+		written, err := c.w.Write(c.compressBuf)
+		if err != nil {
+			return totalWritten, err
+		}
+		if written < len(c.compressBuf) {
+			return totalWritten, io.ErrShortWrite
+		}
+
+		totalWritten += chunkSize
+		p = p[chunkSize:]
+	}
+	return totalWritten, nil
+}
+
+func (c *compressingWriter) Close() error {
+	compressBufPool.Put(c.poolCompressBuf)
+	return nil
+}
+
+// NewZstdCompressingWriter returns a writer that compresses each chunk of the
+// input using zstd and writes the compressed data to the underlying writer.
+// The writer uses a fixed-size 4MB buffer for compression.
+func NewZstdCompressingWriter(w io.Writer) io.WriteCloser {
+	compressBuf := compressBufPool.Get()
+	return &compressingWriter{
+		w:               w,
+		compressBuf:     compressBuf,
+		poolCompressBuf: compressBuf,
+	}
+}
+
 // NewZstdDecompressingReader reads zstd-compressed data from the input
 // reader and makes the decompressed data available on the output reader. The
 // output reader is also an io.WriterTo, which can often prevent allocations

server/util/compression/compression_test.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"fmt"
 	"io"
-	"math"
 	"strconv"
 	"testing"
 
@@ -50,18 +49,39 @@ func TestLossless(t *testing.T) {
 			compress:   compressWithNewZstdCompressingReader,
 			decompress: decompressWithNewZstdDecompressingReader,
 		},
+		{
+			name:       "NewZstdCompressingWriter -> DecompressZstd",
+			compress:   compressWithNewZstdCompressingWriter,
+			decompress: decompressWithDecompressZstd,
+		},
+		{
+			name:       "NewZstdCompressingWriter -> NewZstdDecompressor",
+			compress:   compressWithNewZstdCompressingWriter,
+			decompress: decompressWithNewZstdDecompressor,
+		},
+		{
+			name:       "NewZstdCompressingWriter -> NewZstdDecompressingReader",
+			compress:   compressWithNewZstdCompressingWriter,
+			decompress: decompressWithNewZstdDecompressingReader,
+		},
 	} {
-		for i := 1; i <= 5; i++ {
-			srclen := int(math.Pow10(i))
+		for _, srclen := range []int{9, 99, 999, 1_999_999, 5_999_999} {
 			name := tc.name + "_" + strconv.Itoa(srclen) + "_bytes"
 			t.Run(name, func(t *testing.T) {
 				_, r := testdigest.NewReader(t, int64(srclen))
 				src, err := io.ReadAll(r)
 				require.NoError(t, err)
+				require.Len(t, src, srclen)
 				require.Equal(t, srclen, len(src))
 				compressed := tc.compress(t, src)
 
 				decompressed := tc.decompress(t, len(src), compressed)
+				require.Len(t, decompressed, srclen)
+				if srclen > 1000 {
+					require.Empty(t, cmp.Diff(src[:1000], decompressed[:1000]))
+					require.Empty(t, cmp.Diff(src[len(src)-1000:], decompressed[len(decompressed)-1000:]))
+					return
+				}
 				require.Empty(t, cmp.Diff(src, decompressed))
 			})
 		}
@@ -88,6 +108,17 @@ func compressWithNewZstdCompressingReader(t *testing.T, src []byte) []byte {
 	return compressed
 }
 
+func compressWithNewZstdCompressingWriter(t *testing.T, src []byte) []byte {
+	compressed := &bytes.Buffer{}
+	cw := compression.NewZstdCompressingWriter(compressed)
+	written, err := cw.Write(src)
+	require.NoError(t, err)
+	require.Equal(t, len(src), written)
+	err = cw.Close()
+	require.NoError(t, err)
+	return compressed.Bytes()
+}
+
 func decompressWithDecompressZstd(t *testing.T, srclen int, compressed []byte) []byte {
 	decompressed := make([]byte, srclen)
 	decompressed, err := compression.DecompressZstd(decompressed, compressed)
@@ -111,10 +142,9 @@ func decompressWithNewZstdDecompressingReader(t *testing.T, srclen int, compress
 	rc := io.NopCloser(bytes.NewReader(compressed))
 	d, err := compression.NewZstdDecompressingReader(rc)
 	require.NoError(t, err)
-	buf := make([]byte, srclen)
-	n, err := d.Read(buf)
+	buf, err := io.ReadAll(d)
 	require.NoError(t, err)
-	require.Equal(t, srclen, n)
+	require.Len(t, buf, srclen)
 	err = d.Close()
 	require.NoError(t, err)
 	err = rc.Close()