go-kv/wal/recovery.go

package wal

import (
	"errors"
	"fmt"
	"math"
)

// BatchReplayer is the interface for replaying recovered WAL entries.
// The recovery process calls the appropriate method for each entry.
type BatchReplayer interface {
	ReplayPut(key, value []byte, sequence uint64)
	ReplayDelete(key []byte, sequence uint64)
}

type replayAction struct {
	opType   uint8
	key      []byte
	value    []byte
	sequence uint64
}

// ReplayBatch validates and replays a decoded WAL batch, returning the next
// expected sequence after the batch.
func ReplayBatch(batch *WalBatch, expectedSequence uint64, replayer BatchReplayer) (nextSequence uint64, err error) {
	if batch == nil {
		return expectedSequence, errors.New("wal: batch is nil")
	}
	if replayer == nil {
		return expectedSequence, errors.New("wal: batch replayer is nil")
	}

	if batch.Flags != 0 {
		return expectedSequence, fmt.Errorf("wal: invalid batch flags %d", batch.Flags)
	}
	if batch.EntryCount == 0 {
		return expectedSequence, errors.New("wal: batch entry count is zero")
	}
	if batch.EntryCount > MaxWalBatchEntryCount {
		return expectedSequence, fmt.Errorf("wal: batch entry count %d exceeds maximum %d", batch.EntryCount, MaxWalBatchEntryCount)
	}
	if batch.EntriesSize != uint32(len(batch.Entries)) {
		return expectedSequence, fmt.Errorf("wal: batch entries size mismatch: header says %d bytes, have %d bytes", batch.EntriesSize, len(batch.Entries))
	}
	if batch.EntriesSize == 0 {
		return expectedSequence, errors.New("wal: batch entries size is zero")
	}
	if batch.EntriesSize > MaxWalBatchEntriesSize {
		return expectedSequence, fmt.Errorf("wal: batch entries size %d exceeds maximum %d", batch.EntriesSize, MaxWalBatchEntriesSize)
	}
	if batch.BaseSequence != expectedSequence {
		return expectedSequence, fmt.Errorf("wal: batch base sequence %d does not match expected sequence %d", batch.BaseSequence, expectedSequence)
	}

	entryCount := uint64(batch.EntryCount)
	if batch.BaseSequence > math.MaxUint64-(entryCount-1) {
		return expectedSequence, fmt.Errorf("wal: batch sequence range overflows uint64: base sequence %d entry count %d", batch.BaseSequence, batch.EntryCount)
	}
	if expectedSequence > math.MaxUint64-entryCount {
		return expectedSequence, fmt.Errorf("wal: next sequence overflows uint64: expected sequence %d entry count %d", expectedSequence, batch.EntryCount)
	}

	actions := make([]replayAction, 0, batch.EntryCount)
	offset := 0
	for i := range entryCount {
		entry, consumed, err := DecodeEntry(batch.Entries[offset:])
		if err != nil {
			return expectedSequence, fmt.Errorf("wal: decode batch entry: %w", err)
		}
		if consumed <= 0 {
			return expectedSequence, errors.New("wal: decoded batch entry consumed no bytes")
		}

		sequence := batch.BaseSequence + i
		actions = append(actions, replayAction{
			opType:   entry.OpType,
			key:      entry.Key,
			value:    entry.Value,
			sequence: sequence,
		})
		offset += consumed
	}

	if offset != len(batch.Entries) {
		return expectedSequence, fmt.Errorf("wal: batch entries contain trailing bytes: parsed %d of %d", offset, len(batch.Entries))
	}

	for _, action := range actions {
		switch action.opType {
		case OpPut:
			replayer.ReplayPut(action.key, action.value, action.sequence)
		case OpDelete:
			replayer.ReplayDelete(action.key, action.sequence)
		default:
			return expectedSequence, fmt.Errorf("wal: invalid op type %d", action.opType)
		}
	}

	return expectedSequence + entryCount, nil
}

// ReplaySegmentFile replays all complete WAL batches from one segment file.
//
// isLastSegment controls how parse errors and CollectingFragments-at-end are
// classified per design §3.2 line 704:
//   - true: tail corruption (TailCorruptionError, truncatable by Recover)
//   - false: hard corruption (plain error, Recover must hard-fail)
func ReplaySegmentFile(filePath string, startSequence uint64, isLastSegment bool, replayer BatchReplayer) (nextSequence uint64, err error) {
	nextSequence = startSequence
	records, parseErr := ParseRecordsFromFile(filePath)
	if parseErr != nil && !IsTailCorruption(parseErr) {
		return nextSequence, fmt.Errorf("wal: parse segment records: %w", parseErr)
	}

	// Attach SegmentPath to parseErr for downstream diagnostics + truncation target.
	if parseErr != nil {
		var tce *TailCorruptionError
		if errors.As(parseErr, &tce) {
			tce.SegmentPath = filePath
		}
	}

	// C4 fix: tail corruption in non-last segment is hard corruption per
	// design §3.2 line 704. Use %v (NOT %w) so IsTailCorruption returns false
	// — otherwise errors.As would still find the underlying *TailCorruptionError
	// through the %w chain and Recover would treat it as truncatable.
	if parseErr != nil && !isLastSegment {
		return nextSequence, fmt.Errorf("wal: corruption in non-last segment %s (hard corruption): %v",
			filePath, parseErr)
	}

	collector := NewFragmentCollector()
	for _, record := range records {
		if err := collector.Append(record.Type, record.Payload); err != nil {
			return nextSequence, fmt.Errorf("wal: collect segment fragments: %w", err)
		}
		if !collector.IsComplete() {
			continue
		}

		batch, err := DecodeWalBatch(collector.BatchData())
		if err != nil {
			return nextSequence, fmt.Errorf("wal: decode recovered batch: %w", err)
		}
		nextSequence, err = ReplayBatch(batch, nextSequence, replayer)
		if err != nil {
			return nextSequence, fmt.Errorf("wal: replay recovered batch: %w", err)
		}
		collector.Reset()
	}

	if parseErr != nil {
		return nextSequence, parseErr
	}
	if collector.State() == FragmentCollecting {
		if isLastSegment {
			return nextSequence, &TailCorruptionError{
				Offset:      0,
				SegmentPath: filePath,
				Err:         errors.New("incomplete fragmented batch at segment tail"),
			}
		}
		// Non-last segment with incomplete fragments = middle corruption.
		// Plain error (no TailCorruptionError) so IsTailCorruption is false.
		return nextSequence, fmt.Errorf("wal: incomplete fragmented batch in non-last segment %s (hard corruption)", filePath)
	}

	return nextSequence, nil
}

// RecoverFromSegments scans and replays WAL segments from recoverySegmentID.
func RecoverFromSegments(dir string, recoverySegmentID uint64, replayer BatchReplayer) (nextSequence uint64, err error) {
	segments, err := ScanSegments(dir, recoverySegmentID)
	if err != nil {
		return 0, fmt.Errorf("wal: scan recovery segments: %w", err)
	}
	if len(segments) == 0 {
		return 0, nil
	}

	nextSequence = segments[0].StartSequence
	for i, segment := range segments {
		if segment.StartSequence != nextSequence {
			return nextSequence, fmt.Errorf("wal: segment start sequence %d does not match expected sequence %d", segment.StartSequence, nextSequence)
		}

		isLastSegment := i == len(segments)-1
		nextSequence, err = ReplaySegmentFile(segment.FilePath, nextSequence, isLastSegment, replayer)
		if err != nil {
			if IsTailCorruption(err) {
				return nextSequence, err
			}
			return nextSequence, fmt.Errorf("wal: replay segment: %w", err)
		}
	}

	return nextSequence, nil
}