go-kv/db_e2e_test.go

package go_kv

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/dailz/go-kv/config"
	"github.com/dailz/go-kv/manifest"
	"github.com/dailz/go-kv/wal"
	testifyassert "github.com/stretchr/testify/assert"
	testifyrequire "github.com/stretchr/testify/require"
)

func TestE2EWALRoundtrip(t *testing.T) {
	dir := t.TempDir()
	db := openTestDB(t, dir)

	const entryCount = 1000
	for i := range entryCount {
		key := []byte(e2eKey("roundtrip", i))
		value := []byte(e2eValue("roundtrip", i))
		testifyrequire.NoError(t, db.Put(key, value), "put %q", key)
	}
	closeTestDB(t, db)

	reopened := openTestDB(t, dir)
	for i := range entryCount {
		key := []byte(e2eKey("roundtrip", i))
		got := reopened.Get(key)
		testifyrequire.True(t, got.Found, "get %q after reopen", key)
		testifyassert.Equal(t, []byte(e2eValue("roundtrip", i)), got.Value, "value for %q", key)
	}
}

func TestE2ERecoveryPartialWrite(t *testing.T) {
	dir := t.TempDir()
	db := openTestDB(t, dir)

	testifyrequire.NoError(t, db.Put([]byte("preserved-1"), []byte("value-1")))
	testifyrequire.NoError(t, db.Put([]byte("preserved-2"), []byte("value-2")))
	testifyrequire.NoError(t, db.Put([]byte("truncated-tail"), bytes.Repeat([]byte("x"), 256)))
	closeTestDB(t, db)

	appendWalTail(t, dir, []byte{0xde, 0xad, 0xbe, 0xef})

	reopened := openTestDB(t, dir)
	for _, tc := range []struct {
		name  string
		key   string
		value []byte
	}{
		{name: "first complete batch survives", key: "preserved-1", value: []byte("value-1")},
		{name: "second complete batch survives", key: "preserved-2", value: []byte("value-2")},
	} {
		t.Run(tc.name, func(t *testing.T) {
			got := reopened.Get([]byte(tc.key))
			testifyrequire.True(t, got.Found, "get %q after tail truncation", tc.key)
			testifyassert.Equal(t, tc.value, got.Value)
		})
	}
}

func TestE2ERecoveryMultiSegment(t *testing.T) {
	dir := t.TempDir()
	cfg := e2eSmallSegmentConfig()
	db := openTestDBWithConfig(t, dir, cfg)

	const entryCount = 80
	value := bytes.Repeat([]byte("m"), 512)
	for i := range entryCount {
		key := []byte(e2eKey("multi-segment", i))
		testifyrequire.NoError(t, db.Put(key, e2eValueWithSuffix(value, i)), "put %q", key)
	}
	closeTestDB(t, db)

	testifyassert.Greater(t, len(walSegmentFiles(t, dir)), 1, "test setup should rotate WAL segments")
	prepareAllSegmentsForRecovery(t, dir, cfg)

	reopened := openTestDBWithConfig(t, dir, cfg)
	for i := range entryCount {
		key := []byte(e2eKey("multi-segment", i))
		got := reopened.Get(key)
		testifyrequire.True(t, got.Found, "get %q after multi-segment recovery", key)
		testifyassert.Equal(t, e2eValueWithSuffix(value, i), got.Value, "value for %q", key)
	}
}

func TestE2EWriteStoppedAfterFailure(t *testing.T) {
	t.Skip("write-stopped-after-I/O-failure requires an injectable WAL writer or filesystem fault; current DB architecture has neither")
}

func TestE2EGetSemantics(t *testing.T) {
	db := openTestDB(t, t.TempDir())

	testifyrequire.NoError(t, db.Put([]byte("empty"), []byte{}))
	gotEmpty := db.Get([]byte("empty"))
	testifyrequire.True(t, gotEmpty.Found, "empty value should be found")
	testifyassert.Empty(t, gotEmpty.Value, "empty value should round-trip as empty bytes")

	gotMissing := db.Get([]byte("missing"))
	testifyassert.False(t, gotMissing.Found, "nonexistent key should not be found")
	testifyassert.Empty(t, gotMissing.Value, "nonexistent key should not return a value")

	testifyrequire.NoError(t, db.Delete([]byte("empty")))
	gotDeleted := db.Get([]byte("empty"))
	testifyassert.False(t, gotDeleted.Found, "deleted key should not be found")
}

func TestE2EConcurrentReadWrite(t *testing.T) {
	db := openTestDB(t, t.TempDir())

	const (
		writerCount       = 10
		readerCount       = 10
		entriesPerWriter  = 100
		readsPerReader    = 1000
		totalPossibleKeys = writerCount * entriesPerWriter
	)

	var writersDone atomic.Bool
	errCh := make(chan error, writerCount+readerCount)
	var writerWG sync.WaitGroup
	writerWG.Add(writerCount)
	for writerID := range writerCount {
		go func() {
			defer writerWG.Done()
			for entryID := range entriesPerWriter {
				key := []byte(concurrentKey(writerID, entryID))
				value := []byte(concurrentValue(writerID, entryID))
				if err := db.Put(key, value); err != nil {
					errCh <- fmt.Errorf("writer %d put %q: %w", writerID, key, err)
					return
				}
			}
		}()
	}

	var readerWG sync.WaitGroup
	readerWG.Add(readerCount)
	for readerID := range readerCount {
		go func() {
			defer readerWG.Done()
			for i := range readsPerReader {
				keyIndex := (readerID*readsPerReader + i) % totalPossibleKeys
				writerID := keyIndex / entriesPerWriter
				entryID := keyIndex % entriesPerWriter
				got := db.Get([]byte(concurrentKey(writerID, entryID)))
				if !got.Found {
					if writersDone.Load() {
						errCh <- fmt.Errorf("reader %d missing key after writers done: writer=%d entry=%d", readerID, writerID, entryID)
						return
					}
					continue
				}
				want := []byte(concurrentValue(writerID, entryID))
				if !bytes.Equal(want, got.Value) {
					errCh <- fmt.Errorf("reader %d got value %q for writer=%d entry=%d, want %q", readerID, got.Value, writerID, entryID, want)
					return
				}
			}
		}()
	}

	writerWG.Wait()
	writersDone.Store(true)
	readerWG.Wait()
	close(errCh)

	for err := range errCh {
		testifyrequire.NoError(t, err)
	}
}

func TestE2ELargeBatchSegmentRotation(t *testing.T) {
	dir := t.TempDir()
	cfg := e2eSmallSegmentConfig()
	db := openTestDBWithConfig(t, dir, cfg)

	const entryCount = 120
	value := bytes.Repeat([]byte("b"), 512)
	for i := range entryCount {
		key := []byte(e2eKey("large-rotation", i))
		testifyrequire.NoError(t, db.Put(key, e2eValueWithSuffix(value, i)), "put %q", key)
	}
	closeTestDB(t, db)

	testifyassert.Greater(t, len(walSegmentFiles(t, dir)), 1, "test setup should rotate WAL segments")
	prepareAllSegmentsForRecovery(t, dir, cfg)

	reopened := openTestDBWithConfig(t, dir, cfg)
	for i := range entryCount {
		key := []byte(e2eKey("large-rotation", i))
		got := reopened.Get(key)
		testifyrequire.True(t, got.Found, "get %q after large rotation recovery", key)
		testifyassert.Equal(t, e2eValueWithSuffix(value, i), got.Value, "value for %q", key)
	}
}

func openTestDBWithConfig(t *testing.T, dir string, cfg config.WalConfig) *DB {
	t.Helper()

	db, err := Open(dir, &cfg)
	testifyrequire.NoError(t, err, "Open")
	t.Cleanup(func() {
		if !db.closed.Load() {
			testifyrequire.NoError(t, db.Close(), "Close cleanup")
		}
	})
	return db
}

func e2eSmallSegmentConfig() config.WalConfig {
	cfg := config.Defaults()
	cfg.MaxSegmentSize = 3 * 1024
	cfg.MaxBatchSize = 2 * 1024
	cfg.MaxInlineValue = 1024
	cfg.MemTableSize = 1024 * 1024
	return cfg
}

func e2eKey(prefix string, i int) string {
	return fmt.Sprintf("%s-key-%04d", prefix, i)
}

func e2eValue(prefix string, i int) string {
	return fmt.Sprintf("%s-value-%04d", prefix, i)
}

func e2eValueWithSuffix(prefix []byte, i int) []byte {
	return fmt.Appendf(bytes.Clone(prefix), "-%03d", i)
}

func concurrentKey(writerID int, entryID int) string {
	return fmt.Sprintf("writer-%02d-key-%03d", writerID, entryID)
}

func concurrentValue(writerID int, entryID int) string {
	return fmt.Sprintf("writer-%02d-value-%03d", writerID, entryID)
}

func walSegmentFiles(t *testing.T, dir string) []string {
	t.Helper()

	entries, err := os.ReadDir(dir)
	testifyrequire.NoError(t, err, "ReadDir")
	segments := make([]string, 0)
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasPrefix(entry.Name(), "segment-") || !strings.HasSuffix(entry.Name(), ".wal") {
			continue
		}
		segments = append(segments, filepath.Join(dir, entry.Name()))
	}
	sort.Strings(segments)
	return segments
}

func prepareAllSegmentsForRecovery(t *testing.T, dir string, cfg config.WalConfig) {
	t.Helper()

	testifyrequire.NoError(t, manifest.WriteCurrent(dir, 0), "reset CURRENT to recover from first WAL segment")
	for _, segment := range walSegmentFiles(t, dir) {
		segmentID, ok := wal.ParseSegmentFilename(filepath.Base(segment))
		testifyrequire.True(t, ok, "parse WAL segment filename %q", segment)
		startSequence := firstBatchSequence(t, segment)
		header := wal.EncodeWalHeader(&wal.WalFileHeader{
			BlockSize:     cfg.BlockSize,
			SegmentID:     segmentID,
			StartSequence: startSequence,
		})
		file, err := os.OpenFile(segment, os.O_WRONLY, 0)
		testifyrequire.NoError(t, err, "open WAL segment header for rewrite")
		_, err = file.WriteAt(header[:], 0)
		testifyassert.NoError(t, err, "rewrite WAL segment header")
		testifyassert.NoError(t, file.Sync(), "sync WAL segment header rewrite")
		testifyassert.NoError(t, file.Close(), "close WAL segment header rewrite")
	}
}

func firstBatchSequence(t *testing.T, segment string) uint64 {
	t.Helper()

	records, err := wal.ParseRecordsFromFile(segment)
	testifyrequire.NoError(t, err, "parse WAL records from %s", segment)
	collector := wal.NewFragmentCollector()
	for _, record := range records {
		testifyrequire.NoError(t, collector.Append(record.Type, record.Payload), "collect WAL record fragments")
		if !collector.IsComplete() {
			continue
		}
		batch, err := wal.DecodeWalBatch(collector.BatchData())
		testifyrequire.NoError(t, err, "decode first WAL batch")
		return batch.BaseSequence
	}
	t.Fatalf("segment %s contains no complete WAL batch", segment)
	return 0
}