feat(wal): implement recovery orchestrator with tail truncation

- wal/recover.go: top-level Recover() function that orchestrates scan + replay + truncation Reads MANIFEST/CURRENT for recovery segment ID, handles tail corruption by truncating - Comprehensive tests for empty dir, full flow, tail corruption, and MANIFEST updates - All tests pass with -race
2026-06-12 14:02:19 +08:00
parent fe2d4fc5f0
commit 294ab9181f
4 changed files with 392 additions and 3 deletions
@@ -0,0 +1,206 @@
+package wal
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/dailz/go-kv/manifest"
+)
+
+// RecoveryResult holds the outcome of a WAL recovery pass.
+type RecoveryResult struct {
+	NextSequence    uint64
+	NextSegmentID   uint64
+	ReplayedEntries int
+	Truncated       bool
+	TruncateError   error // non-nil if tail corruption was found
+}
+
+// Recover performs a full WAL recovery: reads the recovery checkpoint from
+// MANIFEST (or CURRENT), scans segments, replays entries, and handles tail
+// truncation. On success the MANIFEST is updated with the new recovery state.
+func Recover(dir string, replayer BatchReplayer) (*RecoveryResult, error) {
+	if replayer == nil {
+		return nil, fmt.Errorf("wal: recover: replayer is nil")
+	}
+
+	// Step 1: Determine recovery segment ID from MANIFEST or CURRENT.
+	recoverySegmentID, err := resolveRecoverySegmentID(dir)
+	if err != nil {
+		return nil, fmt.Errorf("wal: recover: resolve segment id: %w", err)
+	}
+
+	// Step 2: Scan and replay segments.
+	nextSequence, err := RecoverFromSegments(dir, recoverySegmentID, replayer)
+	if err != nil {
+		if !IsTailCorruption(err) {
+			return nil, fmt.Errorf("wal: recover: %w", err)
+		}
+
+		// Step 3: Tail corruption — truncate the last segment and accept
+		// partial data loss for Phase 1.
+		result := &RecoveryResult{
+			NextSequence:  nextSequence,
+			Truncated:     true,
+			TruncateError: err,
+		}
+
+		// Determine nextSegmentID from scanned segments.
+		segments, scanErr := ScanSegments(dir, recoverySegmentID)
+		if scanErr != nil {
+			return nil, fmt.Errorf("wal: recover: scan after tail corruption: %w", scanErr)
+		}
+		if len(segments) > 0 {
+			result.NextSegmentID = segments[len(segments)-1].SegmentID + 1
+		} else {
+			result.NextSegmentID = recoverySegmentID
+		}
+
+		// Truncate the last segment file to remove corrupted tail.
+		if len(segments) > 0 {
+			lastSeg := segments[len(segments)-1]
+			validOffset, truncErr := findValidOffset(lastSeg.FilePath)
+			if truncErr != nil {
+				// Best-effort: record the truncation error but don't fail recovery.
+				result.TruncateError = fmt.Errorf("%w (find valid offset: %v)", err, truncErr)
+			} else if truncErr := truncateSegment(lastSeg.FilePath, validOffset); truncErr != nil {
+				result.TruncateError = fmt.Errorf("%w (truncate: %v)", err, truncErr)
+			}
+		}
+
+		// Count replayed entries by re-scanning the replayer state.
+		// For Phase 1 we accept that ReplayedEntries may be approximate;
+		// the replayer interface doesn't expose a count.
+		result.ReplayedEntries = 0 // caller can inspect replayer directly
+
+		// Update MANIFEST with new recovery state.
+		if saveErr := manifest.Save(dir, result.NextSegmentID); saveErr != nil {
+			return nil, fmt.Errorf("wal: recover: save manifest after truncation: %w", saveErr)
+		}
+
+		return result, nil
+	}
+
+	// Step 4: Successful recovery — compute result.
+	segments, scanErr := ScanSegments(dir, recoverySegmentID)
+	if scanErr != nil {
+		return nil, fmt.Errorf("wal: recover: scan after replay: %w", scanErr)
+	}
+
+	result := &RecoveryResult{
+		NextSequence:  nextSequence,
+		NextSegmentID: recoverySegmentID,
+		Truncated:     false,
+	}
+	if len(segments) > 0 {
+		result.NextSegmentID = segments[len(segments)-1].SegmentID + 1
+	}
+
+	// Update MANIFEST with new recovery state.
+	if saveErr := manifest.Save(dir, result.NextSegmentID); saveErr != nil {
+		return nil, fmt.Errorf("wal: recover: save manifest: %w", saveErr)
+	}
+
+	return result, nil
+}
+
+// resolveRecoverySegmentID determines the starting segment ID for recovery.
+// It tries MANIFEST first, then falls back to CURRENT, then defaults to 0.
+func resolveRecoverySegmentID(dir string) (uint64, error) {
+	mf, err := manifest.Load(dir)
+	if err != nil {
+		return 0, fmt.Errorf("load manifest: %w", err)
+	}
+	if mf.RecoverySegmentID > 0 {
+		return mf.RecoverySegmentID, nil
+	}
+
+	// MANIFEST had 0 (fresh DB or not yet written). Try CURRENT.
+	if segID, ok := manifest.ReadCurrent(dir); ok {
+		return segID, nil
+	}
+
+	return 0, nil
+}
+
+// truncateSegment truncates the file at filePath to validOffset bytes,
+// removing any corrupted data after that point.
+func truncateSegment(filePath string, validOffset int64) error {
+	if validOffset < 0 {
+		return fmt.Errorf("wal: truncate: invalid offset %d", validOffset)
+	}
+	return os.Truncate(filePath, validOffset)
+}
+
+// findValidOffset parses a segment file and returns the byte offset of the
+// last valid record boundary. The offset includes the file header size.
+func findValidOffset(filePath string) (int64, error) {
+	// Re-parse the file to find where valid records end.
+	// We need to track the byte offset as we parse.
+	f, err := os.Open(filePath)
+	if err != nil {
+		return 0, fmt.Errorf("open for offset scan: %w", err)
+	}
+	defer f.Close()
+
+	if _, err := f.Seek(WalFileHeaderSize, 0); err != nil {
+		return 0, fmt.Errorf("seek past header: %w", err)
+	}
+
+	validOffset := int64(WalFileHeaderSize)
+	buf := make([]byte, WalBlockSize)
+
+	for {
+		n, readErr := f.Read(buf)
+		if readErr != nil {
+			break
+		}
+		if n == 0 {
+			break
+		}
+
+		blockData := buf[:n]
+		blockStartOffset := validOffset
+		pos := 0
+
+		for pos < len(blockData) {
+			remaining := len(blockData) - pos
+
+			if remaining < PhysicalRecordHeaderSize {
+				// Check if remaining bytes are zero-padding.
+				if isAllZeros(blockData[pos:]) {
+					// Valid padding — update offset to end of last valid record.
+					validOffset = blockStartOffset + int64(pos)
+				}
+				// Either way, we're done with this block.
+				break
+			}
+
+			if isAllZeros(blockData[pos : pos+PhysicalRecordHeaderSize]) {
+				if isAllZeros(blockData[pos:]) {
+					validOffset = blockStartOffset + int64(pos)
+				}
+				break
+			}
+
+			rec, consumed, err := DecodePhysicalRecord(blockData[pos:])
+			if err != nil {
+				// Corruption starts here — offset is up to last valid record.
+				validOffset = blockStartOffset + int64(pos)
+				return validOffset, nil
+			}
+
+			// Valid record found.
+			validOffset = blockStartOffset + int64(pos+consumed)
+			// Keep the rec reference alive so the compiler doesn't optimize it away.
+			_ = rec
+			pos += consumed
+		}
+
+		if n < WalBlockSize {
+			break
+		}
+	}
+
+	return validOffset, nil
+}
@@ -0,0 +1,157 @@
+package wal
+
+import (
+	"os"
+	"path/filepath"
+	"reflect"
+	"testing"
+
+	"github.com/dailz/go-kv/manifest"
+)
+
+func TestRecoverFromEmptyDir(t *testing.T) {
+	dir := t.TempDir()
+
+	replayer := &mockReplayer{}
+	result, err := Recover(dir, replayer)
+	if err != nil {
+		t.Fatalf("Recover empty dir: %v", err)
+	}
+	if result.NextSequence != 0 {
+		t.Errorf("NextSequence = %d, want 0", result.NextSequence)
+	}
+	if result.Truncated {
+		t.Error("Truncated = true, want false")
+	}
+	if result.TruncateError != nil {
+		t.Errorf("TruncateError = %v, want nil", result.TruncateError)
+	}
+}
+
+func TestRecoverFullFlow(t *testing.T) {
+	dir := t.TempDir()
+
+	// Write test data: 2 batches across 1 segment.
+	writeTestSegment(t, dir, 0, 0, [][]*WalEntry{
+		{makePutEntry("key1", "val1"), makeDeleteEntry("key2")},
+		{makePutEntry("key3", "val3")},
+	})
+
+	replayer := &mockReplayer{}
+	result, err := Recover(dir, replayer)
+	if err != nil {
+		t.Fatalf("Recover: %v", err)
+	}
+	if result.NextSequence != 3 {
+		t.Errorf("NextSequence = %d, want 3", result.NextSequence)
+	}
+	if result.NextSegmentID != 1 {
+		t.Errorf("NextSegmentID = %d, want 1", result.NextSegmentID)
+	}
+	if result.Truncated {
+		t.Error("Truncated = true, want false")
+	}
+
+	wantPuts := []replayPut{
+		{key: "key1", value: "val1", seq: 0},
+		{key: "key3", value: "val3", seq: 2},
+	}
+	if !reflect.DeepEqual(replayer.puts, wantPuts) {
+		t.Errorf("puts = %#v, want %#v", replayer.puts, wantPuts)
+	}
+	wantDeletes := []replayDelete{
+		{key: "key2", seq: 1},
+	}
+	if !reflect.DeepEqual(replayer.deletes, wantDeletes) {
+		t.Errorf("deletes = %#v, want %#v", replayer.deletes, wantDeletes)
+	}
+}
+
+func TestRecoverWithTailCorruption(t *testing.T) {
+	dir := t.TempDir()
+
+	// Write valid batches, then corrupt the tail.
+	filePath := writeTestSegment(t, dir, 0, 100, [][]*WalEntry{
+		{makePutEntry("good1", "before-corruption")},
+		{makePutEntry("good2", "also-before")},
+	})
+	appendFileBytes(t, filePath, []byte{0xDE, 0xAD, 0xBE, 0xEF})
+
+	replayer := &mockReplayer{}
+	result, err := Recover(dir, replayer)
+	if err != nil {
+		t.Fatalf("Recover with tail corruption: %v", err)
+	}
+	if !result.Truncated {
+		t.Fatal("Truncated = false, want true")
+	}
+	if result.TruncateError == nil {
+		t.Fatal("TruncateError = nil, want non-nil")
+	}
+	if result.NextSequence != 102 {
+		t.Errorf("NextSequence = %d, want 102", result.NextSequence)
+	}
+	if result.NextSegmentID != 1 {
+		t.Errorf("NextSegmentID = %d, want 1", result.NextSegmentID)
+	}
+
+	// Verify only the valid entries were replayed.
+	wantPuts := []replayPut{
+		{key: "good1", value: "before-corruption", seq: 100},
+		{key: "good2", value: "also-before", seq: 101},
+	}
+	if !reflect.DeepEqual(replayer.puts, wantPuts) {
+		t.Errorf("puts = %#v, want %#v", replayer.puts, wantPuts)
+	}
+
+	// Verify the file was truncated — it should be smaller than before.
+	fi, fiErr := os.Stat(filePath)
+	if fiErr != nil {
+		t.Fatalf("Stat truncated file: %v", fiErr)
+	}
+	if fi.Size() == 0 {
+		t.Error("truncated file is empty")
+	}
+
+	// Verify the truncated file still parses cleanly.
+	replayer2 := &mockReplayer{}
+	_, parseErr := ReplaySegmentFile(filePath, 100, replayer2)
+	if parseErr != nil {
+		t.Fatalf("replay after truncation: %v", parseErr)
+	}
+	if !reflect.DeepEqual(replayer2.puts, wantPuts) {
+		t.Errorf("replay after truncation puts = %#v, want %#v", replayer2.puts, wantPuts)
+	}
+}
+
+func TestRecoverUpdatesManifest(t *testing.T) {
+	dir := t.TempDir()
+
+	// Write test data.
+	writeTestSegment(t, dir, 0, 50, [][]*WalEntry{
+		{makePutEntry("a", "b")},
+		{makePutEntry("c", "d")},
+	})
+
+	replayer := &mockReplayer{}
+	result, err := Recover(dir, replayer)
+	if err != nil {
+		t.Fatalf("Recover: %v", err)
+	}
+
+	// Verify MANIFEST file exists and has correct content.
+	mf, err := manifest.Load(dir)
+	if err != nil {
+		t.Fatalf("manifest.Load after recover: %v", err)
+	}
+	if mf.RecoverySegmentID != result.NextSegmentID {
+		t.Errorf("MANIFEST RecoverySegmentID = %d, want %d",
+			mf.RecoverySegmentID, result.NextSegmentID)
+	}
+
+	// Verify MANIFEST file path.
+	manifestPath := filepath.Join(dir, "MANIFEST")
+	if _, err := os.Stat(manifestPath); os.IsNotExist(err) {
+		t.Error("MANIFEST file should exist after Recover")
+	}
+}