Implement a miniature LSM-tree storage engine in Go in this directory as package lsm: m…

grokgrok-4.5 · 67 steps · 4m 54s · 10 filessession 5728d376-d04@patrick-toulme
Implement a miniature LSM-tree storage engine in Go in this directory as package lsm: memtable (sorted structure of your choice), write-ahead log with replay on open, SSTable flush with a sparse index and bloom filter, manual compaction via Compact() that merges overlapping SSTables, Get/Put/Delete with tombstones, and a range iterator that merges memtable and SSTables. Include lsm_test.go with table-driven tests covering WAL crash-recovery (reopen after a simulated crash), overwrite/delete semantics, compaction correctness, and a 10k-operation randomized test against a map[string]string oracle. go vet and go test must pass using only the standard library. Also write a brief README.md describing the design.
step 67 / 67
lsm_test.gochange 1/1
package lsm
 
import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"sort"
	"testing"
)
 
func testOpts() Options {
	return Options{
		MemtableBytes:       512, // small so flushes happen often
		SparseIndexInterval: 4,
	}
}
 
func openTestDB(t *testing.T, opts Options) (*DB, string) {
	t.Helper()
	dir := t.TempDir()
	db, err := Open(dir, opts)
	if err != nil {
		t.Fatalf("Open: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })
	return db, dir
}
 
func mustPut(t *testing.T, db *DB, k, v string) {
	t.Helper()
	if err := db.Put([]byte(k), []byte(v)); err != nil {
		t.Fatalf("Put(%q): %v", k, err)
	}
}
 
func mustDelete(t *testing.T, db *DB, k string) {
	t.Helper()
	if err := db.Delete([]byte(k)); err != nil {
		t.Fatalf("Delete(%q): %v", k, err)
	}
}
 
func mustGet(t *testing.T, db *DB, k, want string, wantOK bool) {
	t.Helper()
	v, ok, err := db.Get([]byte(k))
	if err != nil {
		t.Fatalf("Get(%q): %v", k, err)
	}
	if ok != wantOK {
		t.Fatalf("Get(%q): ok=%v want %v (value=%q)", k, ok, wantOK, v)
	}
	if wantOK && string(v) != want {
		t.Fatalf("Get(%q): got %q want %q", k, v, want)
	}
}
 
func TestPutGetBasic(t *testing.T) {
	db, _ := openTestDB(t, testOpts())
	mustPut(t, db, "a", "1")
	mustPut(t, db, "b", "2")
	mustGet(t, db, "a", "1", true)
	mustGet(t, db, "b", "2", true)
	mustGet(t, db, "c", "", false)
}
 
func TestOverwriteDeleteSemantics(t *testing.T) {
	tests := []struct {
		name  string
		ops   []func(t *testing.T, db *DB)
		check []struct {
			key string
			val string
			ok  bool
		}
	}{
		{
			name: "overwrite in memtable",
			ops: []func(t *testing.T, db *DB){
				func(t *testing.T, db *DB) { mustPut(t, db, "k", "v1") },
				func(t *testing.T, db *DB) { mustPut(t, db, "k", "v2") },
			},
			check: []struct {
				key, val string
				ok       bool
			}{
				{"k", "v2", true},
			},
		},
		{
			name: "delete in memtable",
			ops: []func(t *testing.T, db *DB){
				func(t *testing.T, db *DB) { mustPut(t, db, "k", "v1") },
				func(t *testing.T, db *DB) { mustDelete(t, db, "k") },
			},
			check: []struct {
				key, val string
				ok       bool
			}{
				{"k", "", false},
			},
		},
		{
			name: "delete then put",
			ops: []func(t *testing.T, db *DB){
				func(t *testing.T, db *DB) { mustPut(t, db, "k", "v1") },
				func(t *testing.T, db *DB) { mustDelete(t, db, "k") },
				func(t *testing.T, db *DB) { mustPut(t, db, "k", "v2") },
			},
			check: []struct {
				key, val string
				ok       bool
			}{
				{"k", "v2", true},
			},
		},
		{
			name: "overwrite after flush",
			ops: []func(t *testing.T, db *DB){
				func(t *testing.T, db *DB) { mustPut(t, db, "k", "v1") },
				func(t *testing.T, db *DB) {
					if err := db.Flush(); err != nil {
						t.Fatal(err)
					}
				},
				func(t *testing.T, db *DB) { mustPut(t, db, "k", "v2") },
			},
			check: []struct {
				key, val string
				ok       bool
			}{
				{"k", "v2", true},
			},
		},
		{
			name: "delete after flush shadows sstable",
			ops: []func(t *testing.T, db *DB){
				func(t *testing.T, db *DB) { mustPut(t, db, "k", "v1") },
				func(t *testing.T, db *DB) {
					if err := db.Flush(); err != nil {
						t.Fatal(err)
					}
				},
				func(t *testing.T, db *DB) { mustDelete(t, db, "k") },
			},
			check: []struct {
				key, val string
				ok       bool
			}{
				{"k", "", false},
			},
		},
		{
			name: "delete flushed then flush tombstone",
			ops: []func(t *testing.T, db *DB){
				func(t *testing.T, db *DB) { mustPut(t, db, "a", "1") },
				func(t *testing.T, db *DB) {
					if err := db.Flush(); err != nil {
						t.Fatal(err)
					}
				},
				func(t *testing.T, db *DB) { mustDelete(t, db, "a") },
				func(t *testing.T, db *DB) {
					if err := db.Flush(); err != nil {
						t.Fatal(err)
					}
				},
			},
			check: []struct {
				key, val string
				ok       bool
			}{
				{"a", "", false},
			},
		},
		{
			name: "multiple keys overwrite mix",
			ops: []func(t *testing.T, db *DB){
				func(t *testing.T, db *DB) {
					mustPut(t, db, "a", "1")
					mustPut(t, db, "b", "2")
					mustPut(t, db, "c", "3")
				},
				func(t *testing.T, db *DB) {
					if err := db.Flush(); err != nil {
						t.Fatal(err)
					}
				},
				func(t *testing.T, db *DB) {
					mustPut(t, db, "b", "22")
					mustDelete(t, db, "c")
					mustPut(t, db, "d", "4")
				},
			},
			check: []struct {
				key, val string
				ok       bool
			}{
				{"a", "1", true},
				{"b", "22", true},
				{"c", "", false},
				{"d", "4", true},
			},
		},
	}
 
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			db, _ := openTestDB(t, testOpts())
			for _, op := range tc.ops {
				op(t, db)
			}
			for _, c := range tc.check {
				mustGet(t, db, c.key, c.val, c.ok)
			}
		})
	}
}
 
func TestWALCrashRecovery(t *testing.T) {
	tests := []struct {
		name string
		ops  func(t *testing.T, db *DB)
		want map[string]string // key -> value; missing keys should be absent
	}{
		{
			name: "puts survive reopen without close flush",
			ops: func(t *testing.T, db *DB) {
				mustPut(t, db, "x", "1")
				mustPut(t, db, "y", "2")
			},
			want: map[string]string{"x": "1", "y": "2"},
		},
		{
			name: "overwrite and delete in WAL",
			ops: func(t *testing.T, db *DB) {
				mustPut(t, db, "a", "old")
				mustPut(t, db, "a", "new")
				mustPut(t, db, "b", "bye")
				mustDelete(t, db, "b")
				mustPut(t, db, "c", "keep")
			},
			want: map[string]string{"a": "new", "c": "keep"},
		},
		{
			name: "recovery after flush then more WAL",
			ops: func(t *testing.T, db *DB) {
				mustPut(t, db, "f", "flushed")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
				mustPut(t, db, "w", "wal-only")
				mustPut(t, db, "f", "updated")
			},
			want: map[string]string{"f": "updated", "w": "wal-only"},
		},
	}
 
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			dir := t.TempDir()
			db, err := Open(dir, testOpts())
			if err != nil {
				t.Fatal(err)
			}
			tc.ops(t, db)
 
			// Simulated crash: close file handles without orderly Close()
			// (no final flush). WAL has been synced on each write.
			db.mu.Lock()
			if db.wal != nil {
				_ = db.wal.f.Sync()
				_ = db.wal.close()
			}
			for _, tab := range db.tables {
				_ = tab.close()
			}
			db.closed = true
			db.mu.Unlock()
 
			// Reopen — WAL should replay.
			db2, err := Open(dir, testOpts())
			if err != nil {
				t.Fatalf("reopen: %v", err)
			}
			defer db2.Close()
 
			for k, v := range tc.want {
				mustGet(t, db2, k, v, true)
			}
			// keys not in want that were deleted
			if _, ok := tc.want["b"]; !ok && tc.name == "overwrite and delete in WAL" {
				mustGet(t, db2, "b", "", false)
			}
		})
	}
}
 
func TestWALPartialRecordIgnored(t *testing.T) {
	dir := t.TempDir()
	db, err := Open(dir, testOpts())
	if err != nil {
		t.Fatal(err)
	}
	mustPut(t, db, "ok", "yes")
	// Append a truncated record to the WAL file to simulate crash mid-write.
	db.mu.Lock()
	walPath := db.wal.path
	_ = db.wal.f.Sync()
	_ = db.wal.close()
	for _, tab := range db.tables {
		_ = tab.close()
	}
	db.closed = true
	db.mu.Unlock()
 
	f, err := os.OpenFile(walPath, os.O_APPEND|os.O_WRONLY, 0o644)
	if err != nil {
		t.Fatal(err)
	}
	// partial: op + incomplete key length
	if _, err := f.Write([]byte{walOpPut, 0x10}); err != nil {
		t.Fatal(err)
	}
	_ = f.Close()
 
	db2, err := Open(dir, testOpts())
	if err != nil {
		t.Fatalf("reopen after partial WAL: %v", err)
	}
	defer db2.Close()
	mustGet(t, db2, "ok", "yes", true)
}
 
func TestCompactionCorrectness(t *testing.T) {
	tests := []struct {
		name string
		run  func(t *testing.T, db *DB)
	}{
		{
			name: "merge overlapping tables keeps newest",
			run: func(t *testing.T, db *DB) {
				// Force multiple overlapping SSTables.
				mustPut(t, db, "a", "1")
				mustPut(t, db, "b", "1")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
				mustPut(t, db, "a", "2")
				mustPut(t, db, "c", "1")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
				mustPut(t, db, "b", "2")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
 
				_, nBefore := db.Stats()
				if nBefore < 2 {
					t.Fatalf("expected multiple tables, got %d", nBefore)
				}
				if err := db.Compact(); err != nil {
					t.Fatal(err)
				}
				_, nAfter := db.Stats()
				if nAfter >= nBefore {
					t.Fatalf("compaction should reduce tables: before=%d after=%d", nBefore, nAfter)
				}
				mustGet(t, db, "a", "2", true)
				mustGet(t, db, "b", "2", true)
				mustGet(t, db, "c", "1", true)
			},
		},
		{
			name: "compaction drops deleted keys",
			run: func(t *testing.T, db *DB) {
				mustPut(t, db, "x", "1")
				mustPut(t, db, "y", "1")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
				mustDelete(t, db, "x")
				mustPut(t, db, "y", "2")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
				if err := db.Compact(); err != nil {
					t.Fatal(err)
				}
				mustGet(t, db, "x", "", false)
				mustGet(t, db, "y", "2", true)
			},
		},
		{
			name: "non-overlapping tables stay separate",
			run: func(t *testing.T, db *DB) {
				// Keys a,b then keys y,z — ranges [a,b] and [y,z] do not overlap.
				mustPut(t, db, "a", "1")
				mustPut(t, db, "b", "1")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
				mustPut(t, db, "y", "1")
				mustPut(t, db, "z", "1")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
				_, nBefore := db.Stats()
				if err := db.Compact(); err != nil {
					t.Fatal(err)
				}
				_, nAfter := db.Stats()
				if nAfter != nBefore {
					t.Fatalf("non-overlapping tables should not merge: before=%d after=%d", nBefore, nAfter)
				}
				mustGet(t, db, "a", "1", true)
				mustGet(t, db, "z", "1", true)
			},
		},
		{
			name: "data correct after compact and reopen",
			run: func(t *testing.T, db *DB) {
				for i := 0; i < 20; i++ {
					mustPut(t, db, fmt.Sprintf("k%02d", i), fmt.Sprintf("v%d", i))
					if i%5 == 4 {
						if err := db.Flush(); err != nil {
							t.Fatal(err)
						}
					}
				}
				mustDelete(t, db, "k05")
				mustPut(t, db, "k10", "updated")
				if err := db.Flush(); err != nil {
					t.Fatal(err)
				}
				if err := db.Compact(); err != nil {
					t.Fatal(err)
				}
				// Reopen via parent test's dir — we need the path.
			},
		},
	}
 
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			db, dir := openTestDB(t, testOpts())
			tc.run(t, db)
			if tc.name == "data correct after compact and reopen" {
				_ = db.Close()
				db2, err := Open(dir, testOpts())
				if err != nil {
					t.Fatal(err)
				}
				defer db2.Close()
				mustGet(t, db2, "k00", "v0", true)
				mustGet(t, db2, "k05", "", false)
				mustGet(t, db2, "k10", "updated", true)
				mustGet(t, db2, "k19", "v19", true)
			}
		})
	}
}
 
func TestRangeIterator(t *testing.T) {
	db, _ := openTestDB(t, testOpts())
	// Populate with flushes so data spans memtable + SSTables.
	for i := 0; i < 30; i++ {
		mustPut(t, db, fmt.Sprintf("k%02d", i), fmt.Sprintf("v%02d", i))
		if i == 10 || i == 20 {
			if err := db.Flush(); err != nil {
				t.Fatal(err)
			}
		}
	}
	mustDelete(t, db, "k05")
	mustPut(t, db, "k15", "new")
 
	it := db.NewIterator([]byte("k03"), []byte("k18"))
	if it.Err() != nil {
		t.Fatal(it.Err())
	}
	var got []string
	for ; it.Valid(); it.Next() {
		got = append(got, fmt.Sprintf("%s=%s", it.Key(), it.Value()))
	}
	if it.Err() != nil {
		t.Fatal(it.Err())
	}
 
	// Expected: k03..k17 except k05 deleted, k15=new
	var want []string
	for i := 3; i < 18; i++ {
		if i == 5 {
			continue
		}
		v := fmt.Sprintf("v%02d", i)
		if i == 15 {
			v = "new"
		}
		want = append(want, fmt.Sprintf("k%02d=%s", i, v))
	}
	if len(got) != len(want) {
		t.Fatalf("range len got %d want %d\ngot=%v\nwant=%v", len(got), len(want), got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("range[%d]=%q want %q", i, got[i], want[i])
		}
	}
}
 
func TestFlushCreatesSSTable(t *testing.T) {
	db, dir := openTestDB(t, testOpts())
	mustPut(t, db, "a", "1")
	if err := db.Flush(); err != nil {
		t.Fatal(err)
	}
	mem, n := db.Stats()
	if mem != 0 {
		t.Fatalf("memtable should be empty after flush, got %d", mem)
	}
	if n != 1 {
		t.Fatalf("expected 1 table, got %d", n)
	}
	// WAL should be empty-ish; SSTable exists
	matches, _ := filepath.Glob(filepath.Join(dir, "*.sst"))
	if len(matches) != 1 {
		t.Fatalf("expected 1 sst file, got %v", matches)
	}
	mustGet(t, db, "a", "1", true)
}
 
func TestEmptyKeyRejected(t *testing.T) {
	db, _ := openTestDB(t, testOpts())
	if err := db.Put(nil, []byte("v")); err == nil {
		t.Fatal("expected error for empty key Put")
	}
	if err := db.Delete([]byte{}); err == nil {
		t.Fatal("expected error for empty key Delete")
	}
}
 
func TestRandomizedAgainstOracle(t *testing.T) {
	const N = 10000
	rng := rand.New(rand.NewSource(42))
 
	db, dir := openTestDB(t, Options{
		MemtableBytes:       2048,
		SparseIndexInterval: 8,
	})
 
	oracle := map[string]string{}
	keys := make([]string, 0, 500)
 
	randKey := func() string {
		if len(keys) > 0 && rng.Intn(3) != 0 {
			return keys[rng.Intn(len(keys))]
		}
		k := fmt.Sprintf("k%05d", rng.Intn(800))
		// track known keys
		found := false
		for _, existing := range keys {
			if existing == k {
				found = true
				break
			}
		}
		if !found {
			keys = append(keys, k)
		}
		return k
	}
 
	for i := 0; i < N; i++ {
		switch rng.Intn(10) {
		case 0, 1, 2, 3, 4: // put 50%
			k := randKey()
			v := fmt.Sprintf("v%d-%d", i, rng.Intn(1000))
			if err := db.Put([]byte(k), []byte(v)); err != nil {
				t.Fatalf("op %d Put: %v", i, err)
			}
			oracle[k] = v
		case 5, 6: // delete 20%
			k := randKey()
			if err := db.Delete([]byte(k)); err != nil {
				t.Fatalf("op %d Delete: %v", i, err)
			}
			delete(oracle, k)
		case 7: // get 10%
			k := randKey()
			want, wantOK := oracle[k]
			got, ok, err := db.Get([]byte(k))
			if err != nil {
				t.Fatalf("op %d Get: %v", i, err)
			}
			if ok != wantOK || (ok && string(got) != want) {
				t.Fatalf("op %d Get(%q): got (%q,%v) want (%q,%v)", i, k, got, ok, want, wantOK)
			}
		case 8: // flush 10%
			if err := db.Flush(); err != nil {
				t.Fatalf("op %d Flush: %v", i, err)
			}
		default: // compact 10%
			if err := db.Compact(); err != nil {
				t.Fatalf("op %d Compact: %v", i, err)
			}
		}
 
		// Periodic full check every 2000 ops
		if (i+1)%2000 == 0 {
			checkOracle(t, db, oracle)
		}
	}
 
	// Final check
	checkOracle(t, db, oracle)
 
	// Crash-reopen and check again
	_ = db.Close()
	db2, err := Open(dir, Options{MemtableBytes: 2048, SparseIndexInterval: 8})
	if err != nil {
		t.Fatal(err)
	}
	defer db2.Close()
	checkOracle(t, db2, oracle)
 
	// Full range scan matches oracle
	it := db2.NewIterator(nil, nil)
	var scanned []string
	for ; it.Valid(); it.Next() {
		scanned = append(scanned, string(it.Key())+"="+string(it.Value()))
	}
	if it.Err() != nil {
		t.Fatal(it.Err())
	}
	var expected []string
	for k, v := range oracle {
		expected = append(expected, k+"="+v)
	}
	sort.Strings(expected)
	sort.Strings(scanned)
	if len(scanned) != len(expected) {
		t.Fatalf("range scan len %d want %d", len(scanned), len(expected))
	}
	for i := range expected {
		if scanned[i] != expected[i] {
			t.Fatalf("range[%d]=%q want %q", i, scanned[i], expected[i])
		}
	}
}
 
func checkOracle(t *testing.T, db *DB, oracle map[string]string) {
	t.Helper()
	for k, want := range oracle {
		got, ok, err := db.Get([]byte(k))
		if err != nil {
			t.Fatalf("Get(%q): %v", k, err)
		}
		if !ok || string(got) != want {
			t.Fatalf("Get(%q): got (%q,%v) want (%q,true)", k, got, ok, want)
		}
	}
	// Sample some missing keys
	for i := 0; i < 50; i++ {
		k := fmt.Sprintf("missing-%d", i)
		if _, exists := oracle[k]; exists {
			continue
		}
		_, ok, err := db.Get([]byte(k))
		if err != nil {
			t.Fatal(err)
		}
		if ok {
			t.Fatalf("Get(%q) should miss", k)
		}
	}
}
 
func TestBloomAndSparseIndex(t *testing.T) {
	// Write many keys, flush, verify gets work (exercises bloom + sparse index).
	db, _ := openTestDB(t, Options{MemtableBytes: 1 << 20, SparseIndexInterval: 8})
	const n = 200
	for i := 0; i < n; i++ {
		mustPut(t, db, fmt.Sprintf("key-%04d", i), fmt.Sprintf("val-%d", i))
	}
	if err := db.Flush(); err != nil {
		t.Fatal(err)
	}
	for i := 0; i < n; i++ {
		mustGet(t, db, fmt.Sprintf("key-%04d", i), fmt.Sprintf("val-%d", i), true)
	}
	// absent keys
	for i := 0; i < 20; i++ {
		mustGet(t, db, fmt.Sprintf("nope-%d", i), "", false)
	}
}
 
func TestMemtableSorted(t *testing.T) {
	m := newMemtable()
	m.put([]byte("c"), []byte("3"))
	m.put([]byte("a"), []byte("1"))
	m.put([]byte("b"), []byte("2"))
	m.delete([]byte("b"))
	snap := m.snapshot()
	if len(snap) != 3 {
		t.Fatalf("len=%d", len(snap))
	}
	if !bytes.Equal(snap[0].key, []byte("a")) || !bytes.Equal(snap[2].key, []byte("c")) {
		t.Fatalf("not sorted: %q %q %q", snap[0].key, snap[1].key, snap[2].key)
	}
	if !snap[1].deleted {
		t.Fatal("b should be tombstone")
	}
}
 
717 linesgo
Implemented package `lsm` — a miniature LSM-tree using only the standard library. `go vet .…
Files at this step
drag the rail to scrub
Discussion

No comments yet. Start the discussion. Recorded by @patrick-toulme.
#go #grok #grok-4.5 #planning #quick #subagents #testing
g6-grok-lsm
← all sessions