feat(wal): implement batch codec, validation, and Arena allocator

- wal/batch.go: WalBatch encode/decode with FragmentCollector state machine - wal/validate.go: ValidateBatchLimits with checked arithmetic - memtable/arena.go: Arena allocator with 8-byte alignment and mutex - Comprehensive tests for all modules, all pass with -race
2026-06-12 13:27:17 +08:00
parent cf913b1d52
commit fede839eb6
8 changed files with 1014 additions and 5 deletions
@@ -0,0 +1,216 @@
+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
+}