go-kv/wal/segment_manager_test.go

package wal

import (
	"os"
	"path/filepath"
	"testing"

	"github.com/dailz/go-kv/config"
	"github.com/dailz/go-kv/manifest"
)

// tinyWalConfig returns a config with a very small MaxSegmentSize to force
// quick rotation. The minimum is derived from config.Validate: we need enough
// room for the file header, batch header, physical record overhead, and block
// padding. We use 256 bytes which is well above the minimum for default
// block/batch settings.
func tinyWalConfig() *config.WalConfig {
	cfg := config.Defaults()
	// Use a small segment size to force rotation quickly.
	// MaxSegmentSize must be > WalFileHeaderSize (32) and pass Validate().
	// With defaults, minimum is around 4MB+overhead, so we must also reduce
	// MaxBatchSize and BlockSize to make a small segment valid.
	cfg.BlockSize = 512
	cfg.MaxBatchSize = 64   // very small batches
	cfg.MaxBatchEntries = 5
	cfg.MaxKeyBytes = 16
	cfg.MaxInlineValue = 16
	cfg.MaxSegmentSize = 512 // small enough to trigger rotation with a few writes
	return &cfg
}

func TestSegmentManagerCreation(t *testing.T) {
	dir := t.TempDir()
	cfg := testWalConfig()

	sm, err := NewSegmentManager(dir, 0, 1, cfg)
	if err != nil {
		t.Fatalf("NewSegmentManager: %v", err)
	}
	defer sm.Close()

	// Verify segment-0.wal exists.
	expected := filepath.Join(dir, "segment-0.wal")
	if _, err := os.Stat(expected); err != nil {
		t.Errorf("segment file %q should exist: %v", expected, err)
	}

	if sm.ActiveSegmentID() != 0 {
		t.Errorf("ActiveSegmentID = %d, want 0", sm.ActiveSegmentID())
	}

	// Verify CURRENT file points to segment-0.
	segID, ok := manifest.ReadCurrent(dir)
	if !ok {
		t.Fatal("ReadCurrent: expected CURRENT file to exist")
	}
	if segID != 0 {
		t.Errorf("CURRENT segment ID = %d, want 0", segID)
	}
}

func TestSegmentManagerRotation(t *testing.T) {
	dir := t.TempDir()
	cfg := tinyWalConfig()

	sm, err := NewSegmentManager(dir, 0, 1, cfg)
	if err != nil {
		t.Fatalf("NewSegmentManager: %v", err)
	}
	defer sm.Close()

	// Write small batches until rotation occurs.
	// Each batch is a minimal encoded WAL batch: just a small payload.
	// We'll write enough to exhaust the tiny segment.
	batch := make([]byte, 32) // 32-byte dummy batch
	for i := range batch {
		batch[i] = byte(i)
	}

	// Write until we rotate past segment 0.
	for i := 0; i < 20; i++ {
		if err := sm.AppendBatch(batch); err != nil {
			t.Fatalf("AppendBatch %d: %v", i, err)
		}
	}

	// After many writes, we should have rotated to a higher segment.
	if sm.ActiveSegmentID() == 0 {
		t.Error("expected segment rotation, but still on segment 0")
	}

	// Verify that segment-1.wal (or higher) exists on disk.
	segment1Path := filepath.Join(dir, "segment-1.wal")
	if _, err := os.Stat(segment1Path); err != nil {
		t.Errorf("segment-1.wal should exist after rotation: %v", err)
	}
}

func TestSegmentManagerBatchNotSplit(t *testing.T) {
	dir := t.TempDir()
	cfg := tinyWalConfig()

	sm, err := NewSegmentManager(dir, 0, 1, cfg)
	if err != nil {
		t.Fatalf("NewSegmentManager: %v", err)
	}
	defer sm.Close()

	// Fill segment 0 until it's nearly full.
	smallBatch := make([]byte, 16)
	for i := range smallBatch {
		smallBatch[i] = byte(i)
	}

	// Write until we're close to rotation threshold.
	for sm.RemainingPayload() > 256 {
		if err := sm.AppendBatch(smallBatch); err != nil {
			t.Fatalf("AppendBatch small: %v", err)
		}
	}

	// Now write a batch that triggers rotation.
	// This batch must go entirely into the new segment.
	triggerBatch := make([]byte, 128)
	for i := range triggerBatch {
		triggerBatch[i] = 0xAA
	}

	segIDBefore := sm.ActiveSegmentID()
	if err := sm.AppendBatch(triggerBatch); err != nil {
		t.Fatalf("AppendBatch trigger: %v", err)
	}
	segIDAfter := sm.ActiveSegmentID()

	// The trigger batch should have caused rotation (or the segment was big enough).
	// If rotation happened, verify the batch is in the new segment.
	if segIDAfter != segIDBefore {
		// Rotation occurred — the batch should be in the new segment.
		// Read the new segment file and verify it contains our trigger data.
		newSegPath := filepath.Join(dir, fmtSegName(segIDAfter))
		data, err := os.ReadFile(newSegPath)
		if err != nil {
			t.Fatalf("read new segment: %v", err)
		}
		// The trigger batch bytes should appear somewhere after the file header.
		found := false
		for i := WalFileHeaderSize; i <= len(data)-len(triggerBatch); i++ {
			if data[i] == 0xAA {
				found = true
				break
			}
		}
		if !found {
			t.Error("trigger batch data not found in new segment after rotation")
		}
	}
}

func TestSegmentManagerCurrentFile(t *testing.T) {
	dir := t.TempDir()
	cfg := tinyWalConfig()

	sm, err := NewSegmentManager(dir, 0, 1, cfg)
	if err != nil {
		t.Fatalf("NewSegmentManager: %v", err)
	}
	defer sm.Close()

	// Initial CURRENT should point to segment 0.
	segID, ok := manifest.ReadCurrent(dir)
	if !ok || segID != 0 {
		t.Fatalf("initial CURRENT: got segment %d, ok=%v, want 0", segID, ok)
	}

	// Write enough to force rotation.
	batch := make([]byte, 32)
	for i := 0; i < 20; i++ {
		if err := sm.AppendBatch(batch); err != nil {
			t.Fatalf("AppendBatch %d: %v", i, err)
		}
	}

	// CURRENT should now point to the active segment.
	currentSegID, ok := manifest.ReadCurrent(dir)
	if !ok {
		t.Fatal("ReadCurrent after rotation: expected CURRENT file to exist")
	}
	if currentSegID != sm.ActiveSegmentID() {
		t.Errorf("CURRENT segment ID = %d, want %d", currentSegID, sm.ActiveSegmentID())
	}
}

func TestSegmentManagerSync(t *testing.T) {
	dir := t.TempDir()
	cfg := testWalConfig()

	sm, err := NewSegmentManager(dir, 0, 1, cfg)
	if err != nil {
		t.Fatalf("NewSegmentManager: %v", err)
	}
	defer sm.Close()

	if err := sm.Sync(); err != nil {
		t.Errorf("Sync: %v", err)
	}
}

func TestSegmentManagerRemainingPayload(t *testing.T) {
	dir := t.TempDir()
	cfg := testWalConfig()

	sm, err := NewSegmentManager(dir, 0, 1, cfg)
	if err != nil {
		t.Fatalf("NewSegmentManager: %v", err)
	}
	defer sm.Close()

	expected := cfg.MaxSegmentSize - WalFileHeaderSize
	if got := sm.RemainingPayload(); got != expected {
		t.Errorf("RemainingPayload = %d, want %d", got, expected)
	}
}

// fmtSegName formats a segment filename.
func fmtSegName(segID uint64) string {
	return filepath.Join("", "segment-"+itoa(segID)+".wal")
}

func itoa(n uint64) string {
	if n == 0 {
		return "0"
	}
	var buf [20]byte
	i := len(buf)
	for n > 0 {
		i--
		buf[i] = byte('0' + n%10)
		n /= 10
	}
	return string(buf[i:])
}