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.
func ReplaySegmentFile(filePath string, startSequence uint64, 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)
	}

	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 {
		return nextSequence, &TailCorruptionError{
			Offset: 0,
			Err:    errors.New("incomplete fragmented batch at segment tail"),
		}
	}

	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 _, 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)
		}

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

	return nextSequence, nil
}