go-kv/wal/segment_manager_test.go

package wal

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"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.
	currentSeq := uint64(1) // matches NewSegmentManager startSequence above
	for i := 0; i < 20; i++ {
		if err := sm.AppendBatch(batch, currentSeq); err != nil {
			t.Fatalf("AppendBatch %d: %v", i, err)
		}
		currentSeq++
	}

	// 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.
	currentSeq := uint64(1) // matches NewSegmentManager startSequence above
	for sm.RemainingPayload() > 256 {
		if err := sm.AppendBatch(smallBatch, currentSeq); err != nil {
			t.Fatalf("AppendBatch small: %v", err)
		}
		currentSeq++
	}

	// 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, currentSeq); 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)
	currentSeq := uint64(1) // matches NewSegmentManager startSequence above
	for i := 0; i < 20; i++ {
		if err := sm.AppendBatch(batch, currentSeq); err != nil {
			t.Fatalf("AppendBatch %d: %v", i, err)
		}
		currentSeq++
	}

	// 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:])
}

// Regression guard for C6: NewSegmentManager must propagate dir fsync
// failure from initial segment creation. This is the DB.Open failure path.
func TestNewSegmentManagerFailsOnDirFsyncFailure(t *testing.T) {
	dir := t.TempDir()
	cfg := tinyWalConfig()

	orig := dirFsyncFn
	dirFsyncFn = func(string) error { return errors.New("simulated dir fsync failure") }
	t.Cleanup(func() { dirFsyncFn = orig })

	sm, err := NewSegmentManager(dir, 0, 0, cfg)
	if err == nil {
		if sm != nil {
			sm.Close()
		}
		t.Fatal("NewSegmentManager: expected error on dir fsync failure, got nil")
	}
	if !strings.Contains(err.Error(), "create initial segment") {
		t.Errorf("error should be wrapped as 'create initial segment', got: %v", err)
	}
}

// Regression guard for C6: SegmentManager.AppendBatch must propagate
// rotation failure (which now includes dir fsync failure) as error.
//
// Note: C8 (segment_manager.go:64-66 passes byte offset as startSequence)
// makes multi-segment recovery broken, but this test only verifies error
// propagation through AppendBatch; it does not exercise recovery.
func TestSegmentManagerRotateFailsOnDirFsyncFailure(t *testing.T) {
	dir := t.TempDir()
	cfg := tinyWalConfig()

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

	encoded, err := EncodeWalBatch(0, []*WalEntry{
		{OpType: OpPut, ValueKind: VKInline, Key: []byte("k"), Value: []byte("v")},
	})
	if err != nil {
		t.Fatalf("EncodeWalBatch: %v", err)
	}

	// Fill the active segment until next AppendBatch would trigger rotation.
	// segment_manager.go:62 triggers rotate when
	//   RemainingPayload() < len(encoded) + 2*PhysicalRecordHeaderSize
	worstCaseSize := uint64(len(encoded)) + 2*uint64(PhysicalRecordHeaderSize)
	currentSeq := uint64(0) // matches NewSegmentManager startSequence above
	for sm.RemainingPayload() >= worstCaseSize {
		if err := sm.AppendBatch(encoded, currentSeq); err != nil {
			t.Fatalf("fill AppendBatch: %v", err)
		}
		currentSeq++
	}

	orig := dirFsyncFn
	dirFsyncFn = func(string) error { return errors.New("simulated dir fsync failure") }
	t.Cleanup(func() { dirFsyncFn = orig })

	if err := sm.AppendBatch(encoded, currentSeq); err == nil {
		t.Fatal("AppendBatch: expected rotation failure, got nil")
	}
}

// -------- C8 regression guards --------

// Regression guard for C8: after rotation, the new segment's header
// startSequence must equal the rotated batch's baseSequence, NOT the
// previous segment's byte offset.
func TestSegmentManagerRotationWritesCorrectStartSequence(t *testing.T) {
	dir := t.TempDir()
	cfg := tinyWalConfig()

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

	var firstRotationSeq uint64
	var oldOffset uint64
	var hadRotation bool

	for i := 0; ; i++ {
		// Snapshot offset + segment ID BEFORE append to capture pre-rotation state.
		offsetBefore := sm.active.CurrentOffset()
		segIDBefore := sm.ActiveSegmentID()

		encoded, err := EncodeWalBatch(uint64(i), []*WalEntry{
			makePutEntry(fmt.Sprintf("k%d", i), "v"),
		})
		if err != nil {
			t.Fatalf("EncodeWalBatch %d: %v", i, err)
		}
		if err := sm.AppendBatch(encoded, uint64(i)); err != nil {
			t.Fatalf("AppendBatch %d: %v", i, err)
		}

		if !hadRotation && sm.ActiveSegmentID() != segIDBefore {
			hadRotation = true
			firstRotationSeq = uint64(i)
			oldOffset = offsetBefore
			break
		}
	}

	if !hadRotation {
		t.Skip("no rotation occurred with tinyWalConfig; test setup needs adjustment")
	}

	seg1Path := filepath.Join(dir, "segment-1.wal")
	if _, err := os.Stat(seg1Path); err != nil {
		t.Fatalf("segment-1 should exist after rotation: %v", err)
	}

	seg1Header := readSegmentHeader(t, seg1Path)
	if seg1Header.StartSequence != firstRotationSeq {
		t.Errorf("segment-1 startSequence = %d, want %d (first rotation sequence)",
			seg1Header.StartSequence, firstRotationSeq)
	}
	if seg1Header.StartSequence == oldOffset {
		t.Errorf("segment-1 startSequence = %d matches old byte offset (C8 regression)",
			seg1Header.StartSequence)
	}
}

// Regression guard for C8: writing across multiple segments, then closing
// and reopening, must successfully recover ALL data. Before C8 fix, the
// second segment's startSequence was a byte offset, causing recovery's
// continuity check to fail.
func TestSegmentManagerMultiSegmentRecoveryRoundTrip(t *testing.T) {
	dir := t.TempDir()
	cfg := tinyWalConfig()

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

	const totalBatches = 50
	for i := 0; i < totalBatches; i++ {
		encoded, err := EncodeWalBatch(uint64(i), []*WalEntry{
			makePutEntry(fmt.Sprintf("k%d", i), fmt.Sprintf("v%d", i)),
		})
		if err != nil {
			t.Fatalf("EncodeWalBatch %d: %v", i, err)
		}
		if err := sm.AppendBatch(encoded, uint64(i)); err != nil {
			t.Fatalf("AppendBatch %d: %v", i, err)
		}
	}
	if err := sm.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	entries, _ := os.ReadDir(dir)
	segCount := 0
	for _, e := range entries {
		if _, ok := ParseSegmentFilename(e.Name()); ok {
			segCount++
		}
	}
	if segCount < 2 {
		t.Fatalf("expected at least 2 segments after rotation, got %d", segCount)
	}

	replayer := &mockReplayer{}
	result, err := Recover(dir, replayer)
	if err != nil {
		t.Fatalf("Recover failed (C8 regression): %v", err)
	}
	if result.NextSequence != totalBatches {
		t.Errorf("NextSequence = %d, want %d", result.NextSequence, totalBatches)
	}
	if len(replayer.puts) != totalBatches {
		t.Errorf("replayed puts = %d, want %d", len(replayer.puts), totalBatches)
	}
}

func readSegmentHeader(t *testing.T, path string) *WalFileHeader {
	t.Helper()
	f, err := os.Open(path)
	if err != nil {
		t.Fatalf("Open %s: %v", path, err)
	}
	defer f.Close()
	hdrBuf := make([]byte, WalFileHeaderSize)
	if _, err := io.ReadFull(f, hdrBuf); err != nil {
		t.Fatalf("read header: %v", err)
	}
	hdr, err := DecodeWalHeader(hdrBuf)
	if err != nil {
		t.Fatalf("DecodeWalHeader: %v", err)
	}
	return hdr
}