go-kv/wal/recover.go

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
}