go-kv/wal/recovery_test.go

package wal

import (
	"errors"
	"math"
	"os"
	"reflect"
	"strings"
	"testing"

	"github.com/dailz/go-kv/config"
)

type mockReplayer struct {
	puts    []replayPut
	deletes []replayDelete
}

type replayPut struct {
	key   string
	value string
	seq   uint64
}

type replayDelete struct {
	key string
	seq uint64
}

func (m *mockReplayer) ReplayPut(key, value []byte, sequence uint64) {
	m.puts = append(m.puts, replayPut{key: string(key), value: string(value), seq: sequence})
}

func (m *mockReplayer) ReplayDelete(key []byte, sequence uint64) {
	m.deletes = append(m.deletes, replayDelete{key: string(key), seq: sequence})
}

func TestReplayBatchValid(t *testing.T) {
	encoded, err := EncodeWalBatch(100, []*WalEntry{
		makePutEntry("alpha", "one"),
		makeDeleteEntry("beta"),
	})
	if err != nil {
		t.Fatalf("EncodeWalBatch: %v", err)
	}
	batch, err := DecodeWalBatch(encoded)
	if err != nil {
		t.Fatalf("DecodeWalBatch: %v", err)
	}

	replayer := &mockReplayer{}
	next, err := ReplayBatch(batch, 100, replayer)
	if err != nil {
		t.Fatalf("ReplayBatch: %v", err)
	}
	if next != 102 {
		t.Fatalf("nextSequence = %d, want 102", next)
	}

	wantPuts := []replayPut{{key: "alpha", value: "one", seq: 100}}
	if !reflect.DeepEqual(replayer.puts, wantPuts) {
		t.Fatalf("puts = %#v, want %#v", replayer.puts, wantPuts)
	}
	wantDeletes := []replayDelete{{key: "beta", seq: 101}}
	if !reflect.DeepEqual(replayer.deletes, wantDeletes) {
		t.Fatalf("deletes = %#v, want %#v", replayer.deletes, wantDeletes)
	}
}

func TestReplayBatchSequenceMismatch(t *testing.T) {
	batch := mustDecodeTestBatch(t, 10, []*WalEntry{makePutEntry("k", "v")})

	_, err := ReplayBatch(batch, 11, &mockReplayer{})
	if err == nil {
		t.Fatal("ReplayBatch succeeded, want sequence mismatch error")
	}
	if !strings.Contains(err.Error(), "base sequence") {
		t.Fatalf("error = %v, want base sequence context", err)
	}
}

func TestReplayBatchZeroEntries(t *testing.T) {
	batch := &WalBatch{
		Flags:        0,
		BaseSequence: 10,
		EntryCount:   0,
		EntriesSize:  0,
		Entries:      nil,
	}

	_, err := ReplayBatch(batch, 10, &mockReplayer{})
	if err == nil {
		t.Fatal("ReplayBatch succeeded, want zero entries error")
	}
}

func TestReplayBatchOverflowCheck(t *testing.T) {
	entryBytes, err := EncodeEntry(makePutEntry("k", "v"))
	if err != nil {
		t.Fatalf("EncodeEntry: %v", err)
	}
	batch := &WalBatch{
		Flags:        0,
		BaseSequence: math.MaxUint64 - 1,
		EntryCount:   3,
		EntriesSize:  uint32(len(entryBytes) * 3),
		Entries:      append(append(append([]byte{}, entryBytes...), entryBytes...), entryBytes...),
	}

	_, err = ReplayBatch(batch, math.MaxUint64-1, &mockReplayer{})
	if err == nil {
		t.Fatal("ReplayBatch succeeded, want overflow error")
	}
}

func TestReplaySegmentFileFull(t *testing.T) {
	dir := t.TempDir()
	filePath := writeTestSegment(t, dir, 0, 50, [][]*WalEntry{
		{makePutEntry("a", "1"), makeDeleteEntry("b")},
		{makePutEntry("c", "3")},
	})

	replayer := &mockReplayer{}
	next, err := ReplaySegmentFile(filePath, 50, replayer)
	if err != nil {
		t.Fatalf("ReplaySegmentFile: %v", err)
	}
	if next != 53 {
		t.Fatalf("nextSequence = %d, want 53", next)
	}

	wantPuts := []replayPut{{key: "a", value: "1", seq: 50}, {key: "c", value: "3", seq: 52}}
	if !reflect.DeepEqual(replayer.puts, wantPuts) {
		t.Fatalf("puts = %#v, want %#v", replayer.puts, wantPuts)
	}
	wantDeletes := []replayDelete{{key: "b", seq: 51}}
	if !reflect.DeepEqual(replayer.deletes, wantDeletes) {
		t.Fatalf("deletes = %#v, want %#v", replayer.deletes, wantDeletes)
	}
}

func TestReplaySegmentFileTailCorruption(t *testing.T) {
	dir := t.TempDir()
	filePath := writeTestSegment(t, dir, 0, 70, [][]*WalEntry{
		{makePutEntry("ok", "before-corruption")},
	})
	appendFileBytes(t, filePath, []byte{0x01, 0x02, 0x03})

	replayer := &mockReplayer{}
	next, err := ReplaySegmentFile(filePath, 70, replayer)
	if err == nil {
		t.Fatal("ReplaySegmentFile succeeded, want tail corruption error")
	}
	if !IsTailCorruption(err) {
		t.Fatalf("error = %v, want tail corruption", err)
	}
	if next != 71 {
		t.Fatalf("nextSequence = %d, want 71", next)
	}
	wantPuts := []replayPut{{key: "ok", value: "before-corruption", seq: 70}}
	if !reflect.DeepEqual(replayer.puts, wantPuts) {
		t.Fatalf("puts = %#v, want %#v", replayer.puts, wantPuts)
	}
}

func TestRecoverFromSegmentsMultiple(t *testing.T) {
	dir := t.TempDir()
	writeTestSegment(t, dir, 0, 10, [][]*WalEntry{
		{makePutEntry("s0-a", "a"), makeDeleteEntry("s0-b")},
	})
	writeTestSegment(t, dir, 1, 12, [][]*WalEntry{
		{makePutEntry("s1-c", "c")},
		{makeDeleteEntry("s1-d")},
	})

	replayer := &mockReplayer{}
	next, err := RecoverFromSegments(dir, 0, replayer)
	if err != nil {
		t.Fatalf("RecoverFromSegments: %v", err)
	}
	if next != 14 {
		t.Fatalf("nextSequence = %d, want 14", next)
	}

	wantPuts := []replayPut{{key: "s0-a", value: "a", seq: 10}, {key: "s1-c", value: "c", seq: 12}}
	if !reflect.DeepEqual(replayer.puts, wantPuts) {
		t.Fatalf("puts = %#v, want %#v", replayer.puts, wantPuts)
	}
	wantDeletes := []replayDelete{{key: "s0-b", seq: 11}, {key: "s1-d", seq: 13}}
	if !reflect.DeepEqual(replayer.deletes, wantDeletes) {
		t.Fatalf("deletes = %#v, want %#v", replayer.deletes, wantDeletes)
	}
}

func mustDecodeTestBatch(t *testing.T, baseSequence uint64, entries []*WalEntry) *WalBatch {
	t.Helper()
	encoded, err := EncodeWalBatch(baseSequence, entries)
	if err != nil {
		t.Fatalf("EncodeWalBatch: %v", err)
	}
	batch, err := DecodeWalBatch(encoded)
	if err != nil {
		t.Fatalf("DecodeWalBatch: %v", err)
	}
	return batch
}

func writeTestSegment(t *testing.T, dir string, segmentID, startSequence uint64, batches [][]*WalEntry) string {
	t.Helper()
	cfg := config.Defaults()
	sw, err := NewSegmentWriter(dir, segmentID, startSequence, &cfg)
	if err != nil {
		t.Fatalf("NewSegmentWriter: %v", err)
	}
	closed := false
	defer func() {
		if !closed {
			if closeErr := sw.Close(); closeErr != nil && !errors.Is(closeErr, os.ErrClosed) {
				t.Fatalf("SegmentWriter.Close cleanup: %v", closeErr)
			}
		}
	}()

	nextSequence := startSequence
	for i, entries := range batches {
		encoded, err := EncodeWalBatch(nextSequence, entries)
		if err != nil {
			t.Fatalf("EncodeWalBatch[%d]: %v", i, err)
		}
		if err := sw.AppendBatch(encoded); err != nil {
			t.Fatalf("AppendBatch[%d]: %v", i, err)
		}
		nextSequence += uint64(len(entries))
	}
	if err := sw.Close(); err != nil {
		t.Fatalf("SegmentWriter.Close: %v", err)
	}
	closed = true
	return sw.SegmentPath()
}

func appendFileBytes(t *testing.T, filePath string, data []byte) {
	t.Helper()
	f, err := os.OpenFile(filePath, os.O_WRONLY|os.O_APPEND, 0)
	if err != nil {
		t.Fatalf("OpenFile append: %v", err)
	}
	defer f.Close()
	if _, err := f.Write(data); err != nil {
		t.Fatalf("Write corruption bytes: %v", err)
	}
}