Add Snappy streaming compression to WAL segments (#70)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: nakabonne

disk_wal.go

@@ -13,6 +13,8 @@ import (
 	"strconv"
 	"sync"
 	"sync/atomic"
+
+	"github.com/golang/snappy"
 )
 
 // diskWAL contains multiple segment files. One segment is responsible for one partition.
@@ -28,6 +30,8 @@ type diskWAL struct {
 	bufferedSize int
 	// Buffered-writer to the active segment
 	w *bufio.Writer
+	// Snappy streaming writer
+	sw *snappy.Writer
 	// File descriptor to the active segment
 	fd    *os.File
 	index uint32
@@ -46,12 +50,29 @@ func newDiskWAL(dir string, bufferedSize int) (wal, error) {
 	if err != nil {
 		return nil, err
 	}
-	w.fd = f
-	w.w = bufio.NewWriterSize(f, bufferedSize)
+	w.initWriter(f)
 
 	return w, nil
 }
 
+// initWriter sets up the writer chain: bufio.Writer → snappy.Writer → os.File.
+func (w *diskWAL) initWriter(f *os.File) {
+	w.fd = f
+	w.sw = snappy.NewBufferedWriter(f)
+	w.w = bufio.NewWriterSize(w.sw, w.bufferedSize)
+}
+
+// closeWriter flushes and closes the entire writer chain.
+func (w *diskWAL) closeWriter() error {
+	if err := w.w.Flush(); err != nil {
+		return fmt.Errorf("failed to flush buffered-data into the underlying WAL file: %w", err)
+	}
+	if err := w.sw.Close(); err != nil {
+		return fmt.Errorf("failed to close snappy writer: %w", err)
+	}
+	return w.fd.Close()
+}
+
 // append appends the given entry to the end of a file via the file descriptor it has.
 func (w *diskWAL) append(op walOperation, rows []Row) error {
 	w.mu.Lock()
@@ -103,25 +124,24 @@ func (w *diskWAL) flush() error {
 	if err := w.w.Flush(); err != nil {
 		return fmt.Errorf("failed to flush buffered-data into the underlying WAL file: %w", err)
 	}
+	if err := w.sw.Flush(); err != nil {
+		return fmt.Errorf("failed to flush snappy writer: %w", err)
+	}
 	return nil
 }
 
 // punctuate set boundary and creates a new segment.
 func (w *diskWAL) punctuate() error {
 	w.mu.Lock()
 	defer w.mu.Unlock()
-	if err := w.flush(); err != nil {
-		return err
-	}
-	if err := w.fd.Close(); err != nil {
+	if err := w.closeWriter(); err != nil {
 		return err
 	}
 	f, err := w.createSegmentFile(w.dir)
 	if err != nil {
 		return err
 	}
-	w.fd = f
-	w.w = bufio.NewWriterSize(f, w.bufferedSize)
+	w.initWriter(f)
 	return nil
 }
 
@@ -143,6 +163,9 @@ func (w *diskWAL) removeOldest() error {
 func (w *diskWAL) removeAll() error {
 	w.mu.Lock()
 	defer w.mu.Unlock()
+	// Best-effort flush/close of the writer chain before removing files.
+	_ = w.w.Flush()
+	_ = w.sw.Close()
 	if err := w.fd.Close(); err != nil {
 		return err
 	}
@@ -164,8 +187,7 @@ func (w *diskWAL) refresh() error {
 	if err != nil {
 		return err
 	}
-	w.fd = f
-	w.w = bufio.NewWriterSize(f, w.bufferedSize)
+	w.initWriter(f)
 	return nil
 }
 
@@ -216,7 +238,7 @@ func (f *diskWALReader) readAll() error {
 		}
 		segment := &segment{
 			file: fd,
-			r:    bufio.NewReader(fd),
+			r:    bufio.NewReader(snappy.NewReader(fd)),
 		}
 		for segment.next() {
 			rec := segment.record()
@@ -230,7 +252,7 @@ func (f *diskWALReader) readAll() error {
 		}
 
 		err = segment.error()
-		if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) {
+		if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) || errors.Is(err, snappy.ErrCorrupt) {
 			// It is not unusual for a line to be invalid, as it may well terminate in the middle of writing to the WAL.
 			return nil
 		}

disk_wal_test.go

@@ -51,6 +51,37 @@ func Test_diskWAL_append_read(t *testing.T) {
 	assert.Equal(t, rows, got)
 }
 
+func Test_diskWAL_snappy_compressed(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "tstorage-test")
+	defer os.RemoveAll(tmpDir)
+	require.NoError(t, err)
+	path := filepath.Join(tmpDir, "wal")
+
+	wal, err := newDiskWAL(path, 4096)
+	require.NoError(t, err)
+
+	rows := []Row{
+		{Metric: "metric-1", DataPoint: DataPoint{Value: 0.1, Timestamp: 1600000000}},
+	}
+	err = wal.append(operationInsert, rows)
+	require.NoError(t, err)
+	err = wal.flush()
+	require.NoError(t, err)
+
+	// Read the segment file and check for snappy stream magic bytes.
+	files, err := os.ReadDir(path)
+	require.NoError(t, err)
+	require.NotEmpty(t, files)
+
+	data, err := os.ReadFile(filepath.Join(path, files[0].Name()))
+	require.NoError(t, err)
+
+	// Snappy stream format magic: 0xff 0x06 0x00 0x00 s N a P p Y
+	snappyMagic := []byte{0xff, 0x06, 0x00, 0x00, 0x73, 0x4e, 0x61, 0x50, 0x70, 0x59}
+	require.True(t, len(data) >= len(snappyMagic), "segment file too small")
+	assert.Equal(t, snappyMagic, data[:len(snappyMagic)], "segment file should start with snappy stream magic bytes")
+}
+
 func Test_diskWAL_removeOldest(t *testing.T) {
 	tmpDir, err := os.MkdirTemp("", "tstorage-test")
 	require.NoError(t, err)

go.mod

@@ -2,9 +2,11 @@ module github.com/nakabonne/tstorage
 
 go 1.20
 
+require github.com/stretchr/testify v1.7.0
+
 require (
 	github.com/davecgh/go-spew v1.1.0 // indirect
+	github.com/golang/snappy v1.0.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/stretchr/testify v1.7.0 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
 )

go.sum

@@ -1,5 +1,7 @@
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/golang/snappy v1.0.0 h1:Oy607GVXHs7RtbggtPBnr2RmDArIsAefDwvrdWvRhGs=
+github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=