go-kv/bench_test.go

package go_kv

import (
	"fmt"
	"sync"
	"testing"
)

// openBenchDB opens a fresh database in b.TempDir() with default config.
// Returns the DB handle; caller must Close().
func openBenchDB(b *testing.B) *DB {
	b.Helper()
	db, err := Open(b.TempDir(), nil)
	if err != nil {
		b.Fatalf("Open: %v", err)
	}
	return db
}

// BenchmarkSinglePut measures single-threaded sequential Put throughput.
func BenchmarkSinglePut(b *testing.B) {
	db := openBenchDB(b)
	defer func() {
		if err := db.Close(); err != nil {
			b.Fatalf("Close: %v", err)
		}
	}()

	b.ReportAllocs()
	b.ResetTimer()

	i := 0
	for b.Loop() {
		key := fmt.Appendf(nil, "key-%08d", i)
		val := fmt.Appendf(nil, "value-%08d", i)
		if err := db.Put(key, val); err != nil {
			b.Fatalf("Put: %v", err)
		}
		i++
	}
}

// BenchmarkConcurrentPut measures multi-goroutine Put throughput.
// Uses b.N directly because b.Loop() cannot be called from spawned goroutines.
func BenchmarkConcurrentPut(b *testing.B) {
	for _, workers := range []int{2, 4, 8} {
		b.Run(fmt.Sprintf("workers=%d", workers), func(b *testing.B) {
			db := openBenchDB(b)
			defer func() {
				if err := db.Close(); err != nil {
					b.Fatalf("Close: %v", err)
				}
			}()

			b.ReportAllocs()
			b.ResetTimer()

			var wg sync.WaitGroup
			opsPerWorker := b.N / workers
			for w := range workers {
				wg.Add(1)
				go func(workerID int) {
					defer wg.Done()
					for i := range opsPerWorker {
						key := fmt.Appendf(nil, "key-%d-%08d", workerID, i)
						val := fmt.Appendf(nil, "value-%d-%08d", workerID, i)
						if err := db.Put(key, val); err != nil {
							b.Errorf("Put: %v", err)
							return
						}
					}
				}(w)
			}
			wg.Wait()
		})
	}
}

// BenchmarkGet measures Get latency after pre-loading 1000 keys.
func BenchmarkGet(b *testing.B) {
	const numKeys = 1000

	db := openBenchDB(b)
	defer func() {
		if err := db.Close(); err != nil {
			b.Fatalf("Close: %v", err)
		}
	}()

	// Pre-load keys.
	for i := range numKeys {
		key := fmt.Appendf(nil, "key-%08d", i)
		val := fmt.Appendf(nil, "value-%08d", i)
		if err := db.Put(key, val); err != nil {
			b.Fatalf("Put preload: %v", err)
		}
	}

	b.ReportAllocs()
	b.ResetTimer()

	i := 0
	for b.Loop() {
		key := fmt.Appendf(nil, "key-%08d", i%numKeys)
		res := db.Get(key)
		if !res.Found {
			b.Fatalf("Get key %s: not found", key)
		}
		i++
	}
}

// BenchmarkConcurrentGet measures parallel Get after pre-loading keys.
func BenchmarkConcurrentGet(b *testing.B) {
	const numKeys = 1000

	db := openBenchDB(b)
	defer func() {
		if err := db.Close(); err != nil {
			b.Fatalf("Close: %v", err)
		}
	}()

	// Pre-load keys.
	for i := range numKeys {
		key := fmt.Appendf(nil, "key-%08d", i)
		val := fmt.Appendf(nil, "value-%08d", i)
		if err := db.Put(key, val); err != nil {
			b.Fatalf("Put preload: %v", err)
		}
	}

	b.ReportAllocs()
	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		i := 0
		for pb.Next() {
			key := fmt.Appendf(nil, "key-%08d", i%numKeys)
			res := db.Get(key)
			if !res.Found {
				b.Fatalf("Get key %s: not found", key)
			}
			i++
		}
	})
}

// BenchmarkMixedReadWrite measures a 50% Put + 50% Get workload mix.
func BenchmarkMixedReadWrite(b *testing.B) {
	const numKeys = 1000

	db := openBenchDB(b)
	defer func() {
		if err := db.Close(); err != nil {
			b.Fatalf("Close: %v", err)
		}
	}()

	// Pre-load half the keys so Gets are not all misses.
	for i := range numKeys / 2 {
		key := fmt.Appendf(nil, "key-%08d", i)
		val := fmt.Appendf(nil, "value-%08d", i)
		if err := db.Put(key, val); err != nil {
			b.Fatalf("Put preload: %v", err)
		}
	}

	b.ReportAllocs()
	b.ResetTimer()

	i := 0
	for b.Loop() {
		key := fmt.Appendf(nil, "key-%08d", i%numKeys)
		val := fmt.Appendf(nil, "value-%08d", i%numKeys)
		if i%2 == 0 {
			if err := db.Put(key, val); err != nil {
				b.Fatalf("Put: %v", err)
			}
		} else {
			_ = db.Get(key)
		}
		i++
	}
}

// BenchmarkWALRecovery measures recovery (Open) time for different data sizes.
func BenchmarkWALRecovery(b *testing.B) {
	for _, nKeys := range []int{100, 1000, 10000} {
		b.Run(fmt.Sprintf("keys=%d", nKeys), func(b *testing.B) {
			dir := b.TempDir()

			// Write N keys to populate WAL segments.
			func() {
				db, err := Open(dir, nil)
				if err != nil {
					b.Fatalf("Open for write: %v", err)
				}
				defer db.Close()

				for i := range nKeys {
					key := fmt.Appendf(nil, "key-%08d", i)
					val := fmt.Appendf(nil, "value-%08d", i)
					if err := db.Put(key, val); err != nil {
						b.Fatalf("Put: %v", err)
					}
				}
			}()

			// Now benchmark only the recovery (Open) phase.
			b.ReportAllocs()
			b.ResetTimer()

			for b.Loop() {
				db, err := Open(dir, nil)
				if err != nil {
					b.Fatalf("Open recovery: %v", err)
				}
				if err := db.Close(); err != nil {
					b.Fatalf("Close: %v", err)
				}
			}
		})
	}
}