Add a basic WAL

The WAL will serialize upserts and deletions for later replay.

Author: antw

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/spf13/cobra v1.2.1
 	github.com/spf13/viper v1.9.0
 	github.com/stretchr/testify v1.7.0
+	google.golang.org/api v0.56.0
 	google.golang.org/genproto v0.0.0-20211129164237-f09f9a12af12
 	google.golang.org/grpc v1.42.0
 	google.golang.org/protobuf v1.27.1

go.sum

@@ -630,6 +630,7 @@ google.golang.org/api v0.48.0/go.mod h1:71Pr1vy+TAZRPkPs/xlCf5SsU8WjuAWv1Pfjbtuk
 google.golang.org/api v0.50.0/go.mod h1:4bNT5pAuq5ji4SRZm+5QIkjny9JAyVD/3gaSihNefaw=
 google.golang.org/api v0.51.0/go.mod h1:t4HdrdoNgyN5cbEfm7Lum0lcLDLiise1F8qDKX00sOU=
 google.golang.org/api v0.54.0/go.mod h1:7C4bFFOvVDGXjfDTAsgGwDgAxRDeQ4X8NvUedIt6z3k=
+google.golang.org/api v0.56.0 h1:08F9XVYTLOGeSQb3xI9C0gXMuQanhdGed0cWFhDozbI=
 google.golang.org/api v0.56.0/go.mod h1:38yMfeP1kfjsl8isn0tliTjIb1rJXcQi4UXlbqivdVE=
 google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
 google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=

internal/wal/iterator.go

@@ -0,0 +1,114 @@
+package wal
+
+import (
+	"bufio"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+
+	"google.golang.org/api/iterator"
+	"google.golang.org/protobuf/proto"
+)
+
+var (
+	ErrInvalidSkip = errors.New("wal: cannot skip, already at a later txid")
+	ErrNoSuchTxid  = errors.New("wal: no such txid")
+)
+
+type Iterator struct {
+	source  *bufio.Reader
+	peeked  *Record
+	started bool
+}
+
+// NewIterator creates a new iterator for the given reader.
+func NewIterator(source io.Reader) *Iterator {
+	return &Iterator{source: bufio.NewReader(source)}
+}
+
+// Next returns the next record in the WAL and advances the iterator.
+func (r *Iterator) Next() (*Record, error) {
+	if !r.started {
+		r.started = true
+		if err := verifySchemaVersion(r.source); err != nil {
+			return nil, err
+		}
+	}
+
+	if r.peeked != nil {
+		rec := r.peeked
+		r.peeked = nil
+		return rec, nil
+	}
+
+	rec, err := r.nextRecord()
+	if err == io.EOF {
+		return nil, iterator.Done
+	}
+
+	return rec, nil
+}
+
+// SkipTo advances the iterator to the first record whose transaction ID is equal to or higher than
+// txid.
+func (r *Iterator) SkipTo(txid uint64) error {
+	for {
+		rec, err := r.Next()
+		if err != nil {
+			if errors.Is(err, iterator.Done) {
+				return ErrNoSuchTxid
+			}
+			return err
+		}
+
+		if rec.GetTxid() > txid {
+			return ErrInvalidSkip
+		}
+
+		if rec.GetTxid() == txid {
+			r.peeked = rec
+			return nil
+		}
+	}
+}
+
+// nextRecord fetches the next record from the source.
+func (r *Iterator) nextRecord() (*Record, error) {
+	var recordLen uint32
+	if err := binary.Read(r.source, binary.LittleEndian, &recordLen); err != nil {
+		return nil, err
+	}
+
+	recordBytes := make([]byte, recordLen)
+	if _, err := r.source.Read(recordBytes); err != nil {
+		return nil, err
+	}
+
+	var record Record
+	if err := proto.Unmarshal(recordBytes, &record); err != nil {
+		return nil, err
+	}
+
+	return &record, nil
+}
+
+// verifySchemaVersion checks the schema version of the SSTable index file and returns an error if
+// the version is not supported.
+func verifySchemaVersion(r io.Reader) error {
+	var version uint32
+
+	if err := binary.Read(r, binary.LittleEndian, &version); err != nil {
+		if errors.Is(err, io.EOF) {
+			// Allow completely empty indices to simplify testing.
+			return nil
+		}
+		return err
+	}
+
+	if version != schemaVersion {
+		return fmt.Errorf("unsupported WAL schema version: %d", version)
+	}
+
+	return nil
+}

internal/wal/iterator_test.go

@@ -0,0 +1,89 @@
+package wal
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/api/iterator"
+)
+
+func TestIterator_Next(t *testing.T) {
+	w, teardown := createNewWAL(t, "wal_iterator_next")
+	defer teardown()
+
+	err := w.Upsert(1, "foo", []byte("foo value"))
+	require.NoError(t, err)
+
+	err = w.Delete(2, "bar")
+	require.NoError(t, err)
+
+	err = w.Upsert(3, "bar", []byte("bar value"))
+	require.NoError(t, err)
+
+	require.NoError(t, w.Close())
+
+	it := createIterator(t, w.file.Name())
+
+	rec, err := it.Next()
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), rec.GetTxid())
+	require.Equal(t, "foo", rec.GetUpsert().GetKey())
+	require.Equal(t, "foo value", string(rec.GetUpsert().GetValue()))
+
+	rec, err = it.Next()
+	require.NoError(t, err)
+	require.Equal(t, uint64(2), rec.GetTxid())
+	require.Equal(t, "bar", rec.GetDelete().GetKey())
+
+	rec, err = it.Next()
+	require.NoError(t, err)
+	require.Equal(t, uint64(3), rec.GetTxid())
+	require.Equal(t, "bar", rec.GetUpsert().GetKey())
+	require.Equal(t, "bar value", string(rec.GetUpsert().GetValue()))
+
+	_, err = it.Next()
+	require.ErrorIs(t, iterator.Done, err)
+}
+
+func TestIterator_SkipTo(t *testing.T) {
+	w, teardown := createNewWAL(t, "wal_iterator_peek")
+	defer teardown()
+
+	err := w.Upsert(1, "foo", []byte("foo value"))
+	require.NoError(t, err)
+
+	err = w.Delete(2, "bar")
+	require.NoError(t, err)
+
+	err = w.Upsert(3, "bar", []byte("bar value"))
+	require.NoError(t, err)
+
+	require.NoError(t, w.Close())
+
+	it := createIterator(t, w.file.Name())
+
+	require.NoError(t, it.SkipTo(2))
+
+	rec, err := it.Next()
+	require.NoError(t, err)
+	require.Equal(t, uint64(2), rec.GetTxid())
+	require.Equal(t, "bar", rec.GetDelete().GetKey())
+
+	// Can't skip backwards.
+	err = it.SkipTo(1)
+	require.ErrorIs(t, ErrInvalidSkip, err)
+
+	// Can't skip to missing ID.
+	err = it.SkipTo(4)
+	require.ErrorIs(t, ErrNoSuchTxid, err)
+}
+
+// -------------------------------------------------------------------------------------------------
+
+func createIterator(t *testing.T, path string) *Iterator {
+	f, err := os.Open(path)
+	require.NoError(t, err)
+
+	return NewIterator(f)
+}