go-kv/wal/batch_test.go

package wal

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

// helper: make a Put+Inline entry.
func makePutEntry(key, value string) *WalEntry {
	return &WalEntry{
		OpType:    OpPut,
		ValueKind: VKInline,
		Key:       []byte(key),
		Value:     []byte(value),
	}
}

// helper: make a Delete entry.
func makeDeleteEntry(key string) *WalEntry {
	return &WalEntry{
		OpType:    OpDelete,
		ValueKind: VKNone,
		Key:       []byte(key),
		Value:     nil,
	}
}

func TestBatchRoundtrip(t *testing.T) {
	// Build 10 entries: 5 Put+Inline, 3 Delete, 2 Put+Inline with empty value.
	entries := []*WalEntry{
		makePutEntry("key1", "val1"),
		makePutEntry("key2", "val2"),
		makePutEntry("key3", "val3"),
		makePutEntry("key4", "val4"),
		makePutEntry("key5", "val5"),
		makeDeleteEntry("key6"),
		makeDeleteEntry("key7"),
		makeDeleteEntry("key8"),
		makePutEntry("key9", ""),  // empty value
		makePutEntry("key10", ""), // empty value
	}

	const baseSeq uint64 = 42

	// Encode.
	encoded, err := EncodeWalBatch(baseSeq, entries)
	if err != nil {
		t.Fatalf("EncodeWalBatch: %v", err)
	}

	// Decode.
	batch, err := DecodeWalBatch(encoded)
	if err != nil {
		t.Fatalf("DecodeWalBatch: %v", err)
	}

	// Verify header fields.
	if batch.Flags != 0 {
		t.Errorf("Flags = %d, want 0", batch.Flags)
	}
	if batch.BaseSequence != baseSeq {
		t.Errorf("BaseSequence = %d, want %d", batch.BaseSequence, baseSeq)
	}
	if batch.EntryCount != 10 {
		t.Errorf("EntryCount = %d, want 10", batch.EntryCount)
	}
	if batch.EntriesSize != uint32(len(batch.Entries)) {
		t.Errorf("EntriesSize = %d, len(Entries) = %d", batch.EntriesSize, len(batch.Entries))
	}
	if uint32(len(encoded)) != WalBatchHeaderSize+batch.EntriesSize {
		t.Errorf("total encoded = %d, want %d", len(encoded), WalBatchHeaderSize+batch.EntriesSize)
	}

	// Parse entries from Entries bytes.
	offset := 0
	for i, want := range entries {
		got, consumed, err := DecodeEntry(batch.Entries[offset:])
		if err != nil {
			t.Fatalf("DecodeEntry[%d]: %v", i, err)
		}
		if got.OpType != want.OpType {
			t.Errorf("entry[%d].OpType = %d, want %d", i, got.OpType, want.OpType)
		}
		if got.ValueKind != want.ValueKind {
			t.Errorf("entry[%d].ValueKind = %d, want %d", i, got.ValueKind, want.ValueKind)
		}
		if !bytes.Equal(got.Key, want.Key) {
			t.Errorf("entry[%d].Key = %q, want %q", i, got.Key, want.Key)
		}
		if !bytes.Equal(got.Value, want.Value) {
			t.Errorf("entry[%d].Value = %q, want %q", i, got.Value, want.Value)
		}
		offset += consumed
	}

	if offset != len(batch.Entries) {
		t.Errorf("parsed %d bytes, Entries region is %d bytes", offset, len(batch.Entries))
	}
}

func TestFragmentCollector(t *testing.T) {
	// Create a batch large enough to span multiple records (~100 KB).
	var entries []*WalEntry
	for i := 0; i < 5000; i++ {
		key := fmt.Sprintf("key_%06d", i)
		val := fmt.Sprintf("val_%06d_%050s", i, "x") // ~57 bytes value
		entries = append(entries, makePutEntry(key, val))
	}

	encoded, err := EncodeWalBatch(1, entries)
	if err != nil {
		t.Fatalf("EncodeWalBatch: %v", err)
	}

	// Split into records.
	records := SplitIntoRecords(encoded)
	if len(records) < 2 {
		t.Fatalf("expected multiple records for ~100KB batch, got %d", len(records))
	}

	// Decode each physical record to get (type, payload).
	collector := NewFragmentCollector()

	for i, rec := range records {
		decoded, _, err := DecodePhysicalRecord(rec)
		if err != nil {
			t.Fatalf("DecodePhysicalRecord[%d]: %v", i, err)
		}

		// Verify state transitions.
		switch {
		case i == 0:
			if decoded.Type != RecFirst {
				t.Errorf("record[0] type = %d, want RecFirst(%d)", decoded.Type, RecFirst)
			}
			if collector.State() != FragmentIdle {
				t.Errorf("before append[0] state = %d, want FragmentIdle", collector.State())
			}
		case i == len(records)-1:
			if decoded.Type != RecLast {
				t.Errorf("record[%d] type = %d, want RecLast(%d)", i, decoded.Type, RecLast)
			}
			if collector.State() != FragmentCollecting {
				t.Errorf("before append[%d] state = %d, want FragmentCollecting", i, collector.State())
			}
		default:
			if decoded.Type != RecMiddle {
				t.Errorf("record[%d] type = %d, want RecMiddle(%d)", i, decoded.Type, RecMiddle)
			}
			if collector.State() != FragmentCollecting {
				t.Errorf("before append[%d] state = %d, want FragmentCollecting", i, collector.State())
			}
		}

		if err := collector.Append(decoded.Type, decoded.Payload); err != nil {
			t.Fatalf("Append[%d]: %v", i, err)
		}
	}

	// After all fragments, batch should be complete.
	if !collector.IsComplete() {
		t.Fatal("expected IsComplete() after all fragments")
	}

	// Collected data should match original encoded batch.
	if !bytes.Equal(collector.BatchData(), encoded) {
		t.Errorf("collected data length = %d, want %d", len(collector.BatchData()), len(encoded))
	}

	// Verify the batch can be decoded.
	batch, err := DecodeWalBatch(collector.BatchData())
	if err != nil {
		t.Fatalf("DecodeWalBatch from collected: %v", err)
	}
	if batch.EntryCount != uint32(len(entries)) {
		t.Errorf("EntryCount = %d, want %d", batch.EntryCount, len(entries))
	}
}

func TestFragmentCollectorSingleRecord(t *testing.T) {
	// Small batch fits in one record.
	entries := []*WalEntry{makePutEntry("k", "v")}
	encoded, err := EncodeWalBatch(1, entries)
	if err != nil {
		t.Fatalf("EncodeWalBatch: %v", err)
	}

	records := SplitIntoRecords(encoded)
	if len(records) != 1 {
		t.Fatalf("expected 1 record, got %d", len(records))
	}

	decoded, _, err := DecodePhysicalRecord(records[0])
	if err != nil {
		t.Fatalf("DecodePhysicalRecord: %v", err)
	}
	if decoded.Type != RecFull {
		t.Errorf("type = %d, want RecFull(%d)", decoded.Type, RecFull)
	}

	collector := NewFragmentCollector()
	if err := collector.Append(decoded.Type, decoded.Payload); err != nil {
		t.Fatalf("Append: %v", err)
	}
	if !collector.IsComplete() {
		t.Fatal("expected IsComplete()")
	}
	if !bytes.Equal(collector.BatchData(), encoded) {
		t.Error("collected data mismatch")
	}
}

func TestFragmentIllegalTransitions(t *testing.T) {
	tests := []struct {
		name     string
		prepFunc func(fc *FragmentCollector) // set up initial state
		recType  uint8
	}{
		{
			name:     "Idle+RecMiddle",
			prepFunc: func(fc *FragmentCollector) {},
			recType:  RecMiddle,
		},
		{
			name:     "Idle+RecLast",
			prepFunc: func(fc *FragmentCollector) {},
			recType:  RecLast,
		},
		{
			name: "Collecting+RecFull",
			prepFunc: func(fc *FragmentCollector) {
				_ = fc.Append(RecFirst, []byte("data"))
			},
			recType: RecFull,
		},
		{
			name: "Collecting+RecFirst",
			prepFunc: func(fc *FragmentCollector) {
				_ = fc.Append(RecFirst, []byte("data"))
			},
			recType: RecFirst,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			fc := NewFragmentCollector()
			tt.prepFunc(fc)
			err := fc.Append(tt.recType, []byte("payload"))
			if err == nil {
				t.Errorf("expected error for %s, got nil", tt.name)
			}
		})
	}
}

func TestFragmentBufferSizeLimit(t *testing.T) {
	// Exceed the maximum fragment buffer size.
	fc := NewFragmentCollector()

	// Start collecting.
	if err := fc.Append(RecFirst, make([]byte, maxFragmentBufferSize-1)); err != nil {
		t.Fatalf("Append First: %v", err)
	}
	if fc.State() != FragmentCollecting {
		t.Fatal("expected FragmentCollecting state")
	}

	// This should exceed the limit.
	err := fc.Append(RecMiddle, make([]byte, 10))
	if err == nil {
		t.Fatal("expected error for exceeding buffer size limit")
	}
}

func TestFragmentReset(t *testing.T) {
	fc := NewFragmentCollector()
	if err := fc.Append(RecFirst, []byte("data")); err != nil {
		t.Fatalf("Append: %v", err)
	}
	if fc.State() != FragmentCollecting {
		t.Fatal("expected FragmentCollecting")
	}

	fc.Reset()
	if fc.State() != FragmentIdle {
		t.Errorf("state after Reset = %d, want FragmentIdle", fc.State())
	}
	if fc.IsComplete() {
		t.Error("IsComplete() should be false after Reset")
	}
	if fc.buf.Len() != 0 {
		t.Errorf("buffer len after Reset = %d, want 0", fc.buf.Len())
	}
}

func TestBatchValidationErrors(t *testing.T) {
	t.Run("EntryCountZero", func(t *testing.T) {
		// Manually craft a batch with EntryCount=0.
		buf := make([]byte, WalBatchHeaderSize)
		binary.LittleEndian.PutUint16(buf[0:2], 0)
		binary.LittleEndian.PutUint64(buf[2:10], 1)
		binary.LittleEndian.PutUint32(buf[10:14], 0) // entryCount = 0
		binary.LittleEndian.PutUint32(buf[14:18], 0)
		_, err := DecodeWalBatch(buf)
		if err == nil {
			t.Fatal("expected error for entryCount=0")
		}
	})

	t.Run("EntryCountExceedsMax", func(t *testing.T) {
		buf := make([]byte, WalBatchHeaderSize)
		binary.LittleEndian.PutUint16(buf[0:2], 0)
		binary.LittleEndian.PutUint64(buf[2:10], 1)
		binary.LittleEndian.PutUint32(buf[10:14], MaxWalBatchEntryCount+1)
		binary.LittleEndian.PutUint32(buf[14:18], 1)
		_, err := DecodeWalBatch(buf)
		if err == nil {
			t.Fatalf("expected error for entryCount > max")
		}
	})

	t.Run("EntriesSizeMismatch", func(t *testing.T) {
		buf := make([]byte, WalBatchHeaderSize)
		binary.LittleEndian.PutUint16(buf[0:2], 0)
		binary.LittleEndian.PutUint64(buf[2:10], 1)
		binary.LittleEndian.PutUint32(buf[10:14], 1)
		binary.LittleEndian.PutUint32(buf[14:18], 100) // entriesSize=100 but no data after header
		_, err := DecodeWalBatch(buf)
		if err == nil {
			t.Fatal("expected error for entriesSize mismatch")
		}
	})

	t.Run("DataTooShort", func(t *testing.T) {
		_, err := DecodeWalBatch([]byte{1, 2, 3})
		if err == nil {
			t.Fatal("expected error for data too short")
		}
	})

	t.Run("EmptyEntriesSlice", func(t *testing.T) {
		_, err := EncodeWalBatch(1, nil)
		if err == nil {
			t.Fatal("expected error for empty entries slice")
		}
	})
}