go-kv/wal/batch.go

package wal

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
)

// WalBatch represents a single WAL batch — the physical persistence unit.
// A batch consists of an 18-byte header followed by a contiguous Entries region.
type WalBatch struct {
	Flags        uint16
	BaseSequence uint64
	EntryCount   uint32
	EntriesSize  uint32
	Entries      []byte
}

// EncodeWalBatch encodes a WAL batch from the given base sequence and entries.
// The returned byte slice is: BatchHeader(18) + encoded entries.
func EncodeWalBatch(baseSequence uint64, entries []*WalEntry) ([]byte, error) {
	if len(entries) == 0 {
		return nil, errors.New("wal: batch requires at least one entry")
	}
	if uint32(len(entries)) > MaxWalBatchEntryCount {
		return nil, fmt.Errorf("wal: entry count %d exceeds maximum %d", len(entries), MaxWalBatchEntryCount)
	}

	// Encode all entries.
	var entriesBuf bytes.Buffer
	for i, e := range entries {
		encoded, err := EncodeEntry(e)
		if err != nil {
			return nil, fmt.Errorf("wal: encoding entry %d: %w", i, err)
		}
		entriesBuf.Write(encoded)
	}

	entriesSize := uint32(entriesBuf.Len())
	if entriesSize > MaxWalBatchEntriesSize {
		return nil, fmt.Errorf("wal: entries size %d exceeds maximum %d", entriesSize, MaxWalBatchEntriesSize)
	}

	// Build batch: header(18) + entries.
	buf := make([]byte, WalBatchHeaderSize+entriesSize)
	binary.LittleEndian.PutUint16(buf[0:2], 0)          // flags
	binary.LittleEndian.PutUint64(buf[2:10], baseSequence)
	binary.LittleEndian.PutUint32(buf[10:14], uint32(len(entries)))
	binary.LittleEndian.PutUint32(buf[14:18], entriesSize)
	copy(buf[WalBatchHeaderSize:], entriesBuf.Bytes())

	return buf, nil
}

// DecodeWalBatch decodes a WAL batch from raw bytes.
// The data must contain the full batch (header + entries).
func DecodeWalBatch(data []byte) (*WalBatch, error) {
	if len(data) < WalBatchHeaderSize {
		return nil, fmt.Errorf("wal: batch data too short: %d < %d", len(data), WalBatchHeaderSize)
	}

	flags := binary.LittleEndian.Uint16(data[0:2])
	baseSeq := binary.LittleEndian.Uint64(data[2:10])
	entryCount := binary.LittleEndian.Uint32(data[10:14])
	entriesSize := binary.LittleEndian.Uint32(data[14:18])

	if entryCount == 0 {
		return nil, errors.New("wal: batch entry count is zero")
	}
	if entryCount > MaxWalBatchEntryCount {
		return nil, fmt.Errorf("wal: entry count %d exceeds maximum %d", entryCount, MaxWalBatchEntryCount)
	}
	if entriesSize > MaxWalBatchEntriesSize {
		return nil, fmt.Errorf("wal: entries size %d exceeds maximum %d", entriesSize, MaxWalBatchEntriesSize)
	}

	expectedLen := WalBatchHeaderSize + entriesSize
	if uint32(len(data)) < expectedLen {
		return nil, fmt.Errorf("wal: entries size mismatch: header says %d bytes, have %d bytes after header",
			entriesSize, len(data)-WalBatchHeaderSize)
	}

	entries := make([]byte, entriesSize)
	copy(entries, data[WalBatchHeaderSize:WalBatchHeaderSize+entriesSize])

	return &WalBatch{
		Flags:        flags,
		BaseSequence: baseSeq,
		EntryCount:   entryCount,
		EntriesSize:  entriesSize,
		Entries:      entries,
	}, nil
}

// FragmentState represents the state of the fragment collector state machine.
type FragmentState uint8

const (
	// FragmentIdle means no fragments are being collected.
	FragmentIdle FragmentState = 0
	// FragmentCollecting means fragments are being accumulated.
	FragmentCollecting FragmentState = 1
)

// maxFragmentBufferSize is the maximum total bytes the fragment buffer can hold:
// Batch Header size + max entries size.
var maxFragmentBufferSize = uint32(WalBatchHeaderSize) + MaxWalBatchEntriesSize

// FragmentCollector reassembles WAL batches from physical record fragments.
// The state machine transitions between Idle and Collecting based on the
// record type (Full, First, Middle, Last).
type FragmentCollector struct {
	state FragmentState
	buf   bytes.Buffer
}

// NewFragmentCollector creates a new FragmentCollector in Idle state.
func NewFragmentCollector() *FragmentCollector {
	return &FragmentCollector{
		state: FragmentIdle,
	}
}

// Reset clears all collected data and returns the collector to Idle state.
func (fc *FragmentCollector) Reset() {
	fc.state = FragmentIdle
	fc.buf.Reset()
}

// Append feeds a physical record fragment to the collector.
// The state machine enforces valid transitions:
//   - Idle + Full   → collect payload, stay Idle (complete)
//   - Idle + First  → collect payload, → Collecting
//   - Idle + Middle → error
//   - Idle + Last   → error
//   - Collecting + Middle → collect payload
//   - Collecting + Last   → collect payload, → Idle (complete)
//   - Collecting + Full   → error
//   - Collecting + First  → error
func (fc *FragmentCollector) Append(recType uint8, payload []byte) error {
	switch fc.state {
	case FragmentIdle:
		switch recType {
		case RecFull:
			// Complete batch in one record.
			if err := fc.checkBufferCapacity(len(payload)); err != nil {
				return err
			}
			fc.buf.Write(payload)
			// State stays Idle — batch is complete.
			return nil
		case RecFirst:
			// Start collecting fragments.
			if err := fc.checkBufferCapacity(len(payload)); err != nil {
				return err
			}
			fc.buf.Write(payload)
			fc.state = FragmentCollecting
			return nil
		case RecMiddle, RecLast:
			return fmt.Errorf("wal: unexpected fragment type %d in Idle state", recType)
		default:
			return fmt.Errorf("wal: invalid fragment type %d", recType)
		}

	case FragmentCollecting:
		switch recType {
		case RecMiddle:
			if err := fc.checkBufferCapacity(len(payload)); err != nil {
				return err
			}
			fc.buf.Write(payload)
			return nil
		case RecLast:
			if err := fc.checkBufferCapacity(len(payload)); err != nil {
				return err
			}
			fc.buf.Write(payload)
			fc.state = FragmentIdle
			return nil
		case RecFull, RecFirst:
			return fmt.Errorf("wal: unexpected fragment type %d in Collecting state", recType)
		default:
			return fmt.Errorf("wal: invalid fragment type %d", recType)
		}

	default:
		return fmt.Errorf("wal: invalid fragment collector state %d", fc.state)
	}
}

// IsComplete reports whether a complete batch has been collected.
// A batch is complete when the state returns to Idle after collecting data.
func (fc *FragmentCollector) IsComplete() bool {
	return fc.state == FragmentIdle && fc.buf.Len() > 0
}

// BatchData returns the collected batch bytes. Only valid when IsComplete() is true.
func (fc *FragmentCollector) BatchData() []byte {
	return fc.buf.Bytes()
}

// State returns the current fragment collector state.
func (fc *FragmentCollector) State() FragmentState {
	return fc.state
}

// checkBufferCapacity ensures the total collected bytes do not exceed the limit.
func (fc *FragmentCollector) checkBufferCapacity(additional int) error {
	newSize := uint32(fc.buf.Len() + additional)
	if newSize > maxFragmentBufferSize {
		return fmt.Errorf("wal: fragment buffer size %d exceeds maximum %d", newSize, maxFragmentBufferSize)
	}
	return nil
}