From 546a7added01421c7dd2040a03c0a9f7aa661abc Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar Date: Wed, 19 Aug 2026 20:45:20 +0530 Subject: [PATCH 1/9] Update README.md --- README.md | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index bb66cf4..dc9f5d8 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,12 @@ BenchmarkParseHeader-16 206667542 5.822 ns/op 0 B/op 0 allocs/op ``` *Result: `5.8ns` parsing latency with **zero heap allocations** per request, allowing the engine to sustain massive throughput without triggering GC pauses.* +**CPU Profiling (pprof):** + +![CPU Flame Graph proving zero GC overhead and async WAL syncing](./assets/cpu-flamegraph.png) + +*The flame graph above demonstrates the engine under heavy concurrent write load. Notice the complete absence of `mallocgc` (Garbage Collection) blocks on the TCP read/write paths, and the heavily isolated `engine.(*WAL).syncLoop` handling disk persistence entirely asynchronously.* + --- ## System Architecture @@ -82,8 +88,8 @@ Forest is built to survive catastrophic failure. You can run the Chaos Engineeri This script compiles the server, blasts it with concurrent writes, abruptly kills the process, and validates that the WAL recovers 100% of the acknowledged data on reboot. ```bash -git clone [https://github.com/Awesome06/forest.git](https://github.com/Awesome06/forest.git) -cd forest +git clone [https://github.com/Awesome06/Forest_DatabaseEngine.git](https://github.com/Awesome06/Forest_DatabaseEngine.git) +cd Forest_DatabaseEngine chmod +x ./scripts/chaos_test.sh ./scripts/chaos_test.sh ``` @@ -99,4 +105,4 @@ go build -o bin/forest ./cmd/server ## Author **Mrigank Bhatnagar** * GitHub: [@Awesome06](https://github.com/Awesome06) -* Connect regarding distributed systems, data engineering, and backend infrastructure. +* Connect regarding distributed systems, data engineering, and backend infrastructure. \ No newline at end of file From 62d80e7422027c256c56bd7c2305a28212d2e4cd Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar Date: Wed, 19 Aug 2026 21:02:24 +0530 Subject: [PATCH 2/9] Harden engine concurrency and TCP handling Refactors core engine paths to use atomic MemTable pointers for lock-free reads, improves WAL/SSTable durability and error propagation, and adds stricter SSTable/Bloom validation constants. Updates network protocol/server behavior with clearer packet validation (new magic byte), safer request/response writes, configurable server startup flags, and graceful shutdown that stops the listener before closing the DB. --- cmd/server/main.go | 53 ++++++++++--- internal/engine/bloom.go | 52 ++++++++----- internal/engine/compactor.go | 127 ++++++++++++++++++-------------- internal/engine/db.go | 139 +++++++++++++++++++---------------- internal/engine/iterator.go | 66 +++++++++++------ internal/engine/memtable.go | 89 ++++++++++++++-------- internal/engine/sstable.go | 121 +++++++++++++++++------------- internal/engine/version.go | 19 +++-- internal/engine/wal.go | 119 +++++++++++++++++------------- internal/network/handler.go | 62 +++++++++++----- internal/network/protocol.go | 22 ++++-- internal/network/server.go | 78 +++++++++++++------- 12 files changed, 587 insertions(+), 360 deletions(-) diff --git a/cmd/server/main.go b/cmd/server/main.go index 2626349..a2cf017 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -1,6 +1,11 @@ +// Package main is the entry point for the Forest Database Engine server. +// It initializes the storage engine, sets up the TCP network listener, +// and manages graceful shutdown procedures. package main import ( + "flag" + "fmt" "log" "os" "os/signal" @@ -11,25 +16,55 @@ import ( ) func main() { - db, err := engine.NewDB("forest.wal") + // Parse command line flags for server configuration. + portFlag := flag.Int("port", 9000, "TCP port to listen on for incoming client connections") + walDirFlag := flag.String("wal-dir", "data/wal", "Directory path for Write-Ahead Log (WAL) files") + // sstDirFlag is provided to match the CLI usage from the README. + // Currently unused by engine.NewDB, but available for future use. + _ = flag.String("sst-dir", "data/sst", "Directory path for Sorted String Table (SSTable) files") + flag.Parse() + + // Ensure the WAL directory exists before initializing the engine. + if err := os.MkdirAll(*walDirFlag, 0755); err != nil { + log.Fatalf("failed to create WAL directory: %v", err) + } + + walPath := fmt.Sprintf("%s/forest.wal", *walDirFlag) + + // Initialize the storage engine. + dbEngine, err := engine.NewDB(walPath) if err != nil { - log.Fatalf("Failed to initialize engine: %v", err) + log.Fatalf("failed to initialize storage engine at %s: %v", walPath, err) } - server := network.NewServer("127.0.0.1:9000", db) + // Initialize the TCP network server. + addr := fmt.Sprintf("127.0.0.1:%d", *portFlag) + tcpServer := network.NewServer(addr, dbEngine) - // Handle graceful shutdown + // Handle graceful shutdown via OS signals. go func() { sigChan := make(chan os.Signal, 1) signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) <-sigChan - - log.Println("Shutting down database...") - db.Close() + + log.Println("Received termination signal. Shutting down database...") + + // Stop accepting new TCP connections. + if err := tcpServer.Stop(); err != nil { + log.Printf("error stopping TCP server: %v", err) + } + + // Flush remaining MemTables and close WAL. + if err := dbEngine.Close(); err != nil { + log.Printf("error closing storage engine: %v", err) + } + + log.Println("Database shut down successfully.") os.Exit(0) }() - if err := server.Start(); err != nil { - log.Fatalf("Server failed: %v", err) + // Start accepting connections. This call blocks indefinitely. + if err := tcpServer.Start(); err != nil { + log.Fatalf("TCP server encountered a fatal error: %v", err) } } diff --git a/internal/engine/bloom.go b/internal/engine/bloom.go index cfae93e..16e2621 100644 --- a/internal/engine/bloom.go +++ b/internal/engine/bloom.go @@ -2,18 +2,22 @@ package engine import ( "encoding/binary" + "fmt" "hash/fnv" "math" ) -// BloomFilter is a custom bit-array bloom filter. +// BloomFilter is a highly optimized bit-array used to probabilistically test whether a key +// is present in an SSTable. It mathematically guarantees no false negatives, allowing us to +// bypass expensive disk seeks for keys that definitively do not exist. type BloomFilter struct { - bits []uint64 // The underlying bit array - k uint8 // Number of hash functions - m uint64 // Total number of bits + bits []uint64 // The underlying bit array packed into 64-bit words for fast bitwise math. + k uint8 // The optimal number of hash functions applied per key. + m uint64 // The total number of bits in the filter. } -// NewBloomFilter calculates the optimal number of bits and hash functions. +// NewBloomFilter calculates the optimal bit array size (m) and hash functions (k) +// for an expected number of elements (n) and a target false positive probability (p). func NewBloomFilter(n int, p float64) *BloomFilter { if n <= 0 { n = 1 @@ -22,20 +26,23 @@ func NewBloomFilter(n int, p float64) *BloomFilter { p = 0.01 // default to 1% false positive rate } - // m = ceil((n * log(p)) / log(1 / pow(2, log(2)))) + // m = - (n * ln(p)) / (ln(2)^2) + // We allocate the exact mathematical minimum number of bits required. mFloat := math.Ceil((float64(n) * math.Log(p)) / math.Log(1.0/math.Pow(2.0, math.Log(2.0)))) m := uint64(mFloat) if m == 0 { m = 64 } - // k = round((m / n) * log(2)) + // k = (m / n) * ln(2) + // Balances the bit array saturation so we don't overwrite too many bits per insert. kFloat := math.Round((float64(m) / float64(n)) * math.Log(2.0)) k := uint8(kFloat) if k == 0 { k = 1 } + // Pack bits into 64-bit words. The +63 ensures we round up to the nearest word boundary. numUint64s := (m + 63) / 64 return &BloomFilter{ bits: make([]uint64, numUint64s), @@ -44,32 +51,43 @@ func NewBloomFilter(n int, p float64) *BloomFilter { } } -// hash generates two 32-bit hashes from a single 64-bit FNV hash. +// hash generates two distinct 32-bit hashes from a single 64-bit FNV hash. +// This implements the Kirsch-Mitzenmacher optimization, which proves mathematically +// that we can simulate 'k' independent hash functions using just two base hashes: +// h_i = h1 + i * h2. This drastically cuts down CPU hashing overhead. func hash(key []byte) (uint32, uint32) { h := fnv.New64a() h.Write(key) sum := h.Sum64() + // Split the 64-bit hash into two 32-bit halves. return uint32(sum), uint32(sum >> 32) } -// Add inserts a key into the Bloom Filter using Kirsch-Mitzenmacher double-hashing. +// Add sets 'k' bits in the array for the given key. func (bf *BloomFilter) Add(key []byte) { h1, h2 := hash(key) for i := uint8(0); i < bf.k; i++ { + // Calculate the compound hash for this iteration. idx := (uint64(h1) + uint64(i)*uint64(h2)) % bf.m + + // Map the bit index to a specific uint64 word and flip the specific bit. wordIdx := idx / 64 bitIdx := idx % 64 bf.bits[wordIdx] |= (1 << bitIdx) } } -// Contains checks if a key might be in the set. +// Contains checks if all 'k' bits for the key are set. +// If any bit is 0, the key absolutely does NOT exist (saving a disk read). +// If all bits are 1, the key MIGHT exist (subject to false positive rate 'p'). func (bf *BloomFilter) Contains(key []byte) bool { h1, h2 := hash(key) for i := uint8(0); i < bf.k; i++ { idx := (uint64(h1) + uint64(i)*uint64(h2)) % bf.m wordIdx := idx / 64 bitIdx := idx % 64 + + // If the specific bit is zero via bitwise AND, this key was never added. if (bf.bits[wordIdx] & (1 << bitIdx)) == 0 { return false } @@ -77,8 +95,8 @@ func (bf *BloomFilter) Contains(key []byte) bool { return true } -// Encode serializes the Bloom Filter to a byte slice for SSTable storage. -// Format: [k (1B)] | [m (8B)] | [bit array...] +// Encode serializes the Bloom Filter state to a raw byte slice for SSTable persistence. +// Layout: k (1B) | m (8B) | bit array (m/8 B) func (bf *BloomFilter) Encode() []byte { buf := make([]byte, 1+8+(len(bf.bits)*8)) buf[0] = bf.k @@ -92,10 +110,10 @@ func (bf *BloomFilter) Encode() []byte { return buf } -// DecodeBloomFilter reconstructs a BloomFilter from a byte slice. -func DecodeBloomFilter(data []byte) *BloomFilter { +// DecodeBloomFilter reconstructs a BloomFilter from a byte slice loaded from disk. +func DecodeBloomFilter(data []byte) (*BloomFilter, error) { if len(data) < 9 { - return nil + return nil, fmt.Errorf("bloom filter data too short: expected at least 9 bytes, got %d", len(data)) } k := data[0] @@ -107,7 +125,7 @@ func DecodeBloomFilter(data []byte) *BloomFilter { offset := 9 for i := 0; i < int(numUint64s); i++ { if offset+8 > len(data) { - break + break // Allow partial reads in case of corruption padding. } bits[i] = binary.BigEndian.Uint64(data[offset : offset+8]) offset += 8 @@ -117,5 +135,5 @@ func DecodeBloomFilter(data []byte) *BloomFilter { bits: bits, k: k, m: m, - } + }, nil } diff --git a/internal/engine/compactor.go b/internal/engine/compactor.go index 2071d43..ac96b91 100644 --- a/internal/engine/compactor.go +++ b/internal/engine/compactor.go @@ -5,6 +5,7 @@ import ( "bytes" "container/heap" "encoding/binary" + "fmt" "log" "os" "time" @@ -12,7 +13,9 @@ import ( "github.com/Forest_DatabaseEngine/internal/network" ) -// HeapItem represents the current head of an SSTIterator in the Min-Heap. +// HeapItem represents the current head of an SSTIterator inside the Min-Heap. +// We track the FileSeq to ensure that when keys collide, the mutation from the +// newer file (higher sequence) shadows the older one. type HeapItem struct { Key []byte Value []byte @@ -21,22 +24,27 @@ type HeapItem struct { Iterator *SSTIterator } -// MinHeap implements heap.Interface +// MinHeap implements heap.Interface for K-Way Merging of SSTables. type MinHeap []HeapItem func (h MinHeap) Len() int { return len(h) } func (h MinHeap) Less(i, j int) bool { cmp := bytes.Compare(h[i].Key, h[j].Key) if cmp == 0 { - // Higher sequence (newer file) comes first + // If keys are identical, prioritize the newer file (higher sequence number) + // so it pops first and shadows the older entry. return h[i].FileSeq > h[j].FileSeq } return cmp < 0 } func (h MinHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +// Push adds an item to the heap. func (h *MinHeap) Push(x interface{}) { *h = append(*h, x.(HeapItem)) } + +// Pop removes and returns the smallest item from the heap. func (h *MinHeap) Pop() interface{} { old := *h n := len(old) @@ -45,23 +53,23 @@ func (h *MinHeap) Pop() interface{} { return item } -// CompactL0toL1 performs a K-Way Merge on the provided L0 files and produces an L1 file. +// CompactL0toL1 performs a K-Way Merge on multiple overlapping L0 files to produce a single sorted L1 file. +// This prevents read-amplification by ensuring that older values and duplicates are aggressively purged. func CompactL0toL1(l0Files []string, l1Filepath string) error { var iterators []*SSTIterator h := &MinHeap{} heap.Init(h) - // Open all iterators and push their first element + // 1. Initialize iterators for all L0 files and push their first elements to seed the heap. for i, fpath := range l0Files { - // Use the index as a proxy for file sequence (assuming l0Files is ordered oldest to newest) - // For robustness, file names could contain true timestamps, but slice order suffices for V1. + // File sequence is proxied by the slice index (ordered oldest to newest). it, err := NewSSTIterator(fpath, uint64(i)) if err != nil { - // Cleanup previously opened iterators + // Clean up safely if we fail mid-initialization. for _, prev := range iterators { prev.Close() } - return err + return fmt.Errorf("failed to open iterator for %s: %w", fpath, err) } iterators = append(iterators, it) @@ -76,25 +84,30 @@ func CompactL0toL1(l0Files []string, l1Filepath string) error { } } - // We will manually build the L1 SSTable using an SSTableWriter abstraction. - // (For simplicity here, we create a mini inline writer mimicking FlushMemTable) - f, err := os.OpenFile(l1Filepath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644) + // Ensure all iterators are closed when compaction completes. + defer func() { + for _, it := range iterators { + it.Close() + } + }() + + file, err := os.OpenFile(l1Filepath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644) if err != nil { - return err + return fmt.Errorf("failed to create L1 SSTable file %s: %w", l1Filepath, err) } - defer f.Close() - - // 1. Setup Writer - bw, bf, index, currentOffset, lastIndexOffset := initSSTableWriter(f) + defer file.Close() + // 2. Setup the target SSTable writer. + bufWriter, bf, index, currentOffset, lastIndexOffset := initSSTableWriter(file) var lastKey []byte - // 2. K-Way Merge Loop + // 3. K-Way Merge Loop for h.Len() > 0 { item := heap.Pop(h).(HeapItem) - // Purge Duplicates: If the next item in the heap has the exact same key, - // it is older (due to Less() logic), so we advance its iterator and discard it. + // Purge Duplicates: If the next items in the heap have the exact same key, + // they are older (enforced by MinHeap.Less), so we aggressively advance their + // iterators and discard them, keeping only the freshest mutation. for h.Len() > 0 { nextItem := (*h)[0] if bytes.Equal(nextItem.Key, item.Key) { @@ -113,14 +126,15 @@ func CompactL0toL1(l0Files []string, l1Filepath string) error { } } - // Write item to L1 - // We preserve tombstones by writing OpCode exactly as it is. + // Write the freshest item to the L1 file. + // Note: We intentionally write tombstones (OpDelete) into L1 so they can continue + // to shadow values in L2+ if multi-level compaction is implemented later. if !bytes.Equal(lastKey, item.Key) { - currentOffset, lastIndexOffset = writeSSTableEntry(bw, bf, &index, currentOffset, lastIndexOffset, item) - lastKey = append([]byte(nil), item.Key...) // copy + currentOffset, lastIndexOffset = writeSSTableEntry(bufWriter, bf, &index, currentOffset, lastIndexOffset, item) + lastKey = append([]byte(nil), item.Key...) // copy to avoid memory aliasing } - // Advance the iterator of the chosen item + // Advance the iterator of the chosen item. if item.Iterator.Next() { heap.Push(h, HeapItem{ Key: item.Iterator.Key(), @@ -132,41 +146,40 @@ func CompactL0toL1(l0Files []string, l1Filepath string) error { } } - // Close all iterators - for _, it := range iterators { - it.Close() + // 4. Finalize the L1 SSTable (Indexes, Bloom Filter, Footer). + if err := finalizeSSTable(file, bufWriter, bf, index, currentOffset); err != nil { + return fmt.Errorf("failed to finalize L1 SSTable: %w", err) } - - // 3. Finalize L1 SSTable - return finalizeSSTable(f, bw, bf, index, currentOffset) + return nil } -// Background deletion worker to delete files once refs hit 0 +// PurgeObsoleteFiles executes the physical deletion of old SSTable files in the background. +// It uses an RCU (Read-Copy-Update) spinlock to guarantee that no active network readers +// are referencing the file before the OS unlinks it. func PurgeObsoleteFiles(v *Version, files []string) { - // Spinlock wait with sleep to prevent busy waiting. - // In production, sync.Cond or channel-based notification is better, - // but atomic polling is robust for V1 file deletion. + // Wait until all active readers that acquired this version have released it. for v.Refs.Load() > 0 { time.Sleep(10 * time.Millisecond) } for _, f := range files { if err := os.Remove(f); err != nil { - log.Printf("Failed to purge obsolete file %s: %v", f, err) - } else { - log.Printf("Purged obsolete file %s", f) + log.Printf("critical: failed to purge obsolete file %s (possible disk leak): %v", f, err) } } } -func initSSTableWriter(f *os.File) (*bufio.Writer, *BloomFilter, []IndexEntry, uint64, uint64) { - bw := bufio.NewWriterSize(f, 4*1024*1024) +// initSSTableWriter primes the internal buffers and structures for writing an SSTable. +func initSSTableWriter(file *os.File) (*bufio.Writer, *BloomFilter, []IndexEntry, uint64, uint64) { + bufWriter := bufio.NewWriterSize(file, 4*1024*1024) bf := NewBloomFilter(4000, 0.01) - return bw, bf, nil, 0, 0 + return bufWriter, bf, nil, 0, 0 } -func writeSSTableEntry(bw *bufio.Writer, bf *BloomFilter, index *[]IndexEntry, currentOffset, lastIndexOffset uint64, item HeapItem) (uint64, uint64) { +// writeSSTableEntry serializes a single merged record into the data block phase of the SSTable. +func writeSSTableEntry(bufWriter *bufio.Writer, bf *BloomFilter, index *[]IndexEntry, currentOffset, lastIndexOffset uint64, item HeapItem) (uint64, uint64) { bf.Add(item.Key) + if currentOffset-lastIndexOffset >= 4096 || currentOffset == 0 { *index = append(*index, IndexEntry{Key: item.Key, Offset: currentOffset}) lastIndexOffset = currentOffset @@ -177,47 +190,47 @@ func writeSSTableEntry(bw *bufio.Writer, bf *BloomFilter, index *[]IndexEntry, c binary.BigEndian.PutUint16(header[1:3], uint16(len(item.Key))) binary.BigEndian.PutUint32(header[3:7], uint32(len(item.Value))) - n, _ := bw.Write(header[:]) + n, _ := bufWriter.Write(header[:]) currentOffset += uint64(n) - n, _ = bw.Write(item.Key) + n, _ = bufWriter.Write(item.Key) currentOffset += uint64(n) - n, _ = bw.Write(item.Value) + n, _ = bufWriter.Write(item.Value) currentOffset += uint64(n) return currentOffset, lastIndexOffset } -func finalizeSSTable(f *os.File, bw *bufio.Writer, bf *BloomFilter, index []IndexEntry, currentOffset uint64) error { +// finalizeSSTable appends the trailing metadata layers to seal the SSTable. +func finalizeSSTable(file *os.File, bufWriter *bufio.Writer, bf *BloomFilter, index []IndexEntry, currentOffset uint64) error { indexBlockOffset := currentOffset var countBuf [4]byte binary.BigEndian.PutUint32(countBuf[:], uint32(len(index))) - n, _ := bw.Write(countBuf[:]) + n, _ := bufWriter.Write(countBuf[:]) currentOffset += uint64(n) for _, entry := range index { var idxHeader [10]byte binary.BigEndian.PutUint16(idxHeader[0:2], uint16(len(entry.Key))) binary.BigEndian.PutUint64(idxHeader[2:10], entry.Offset) - n1, _ := bw.Write(idxHeader[0:2]) - n2, _ := bw.Write(entry.Key) - n3, _ := bw.Write(idxHeader[2:10]) + n1, _ := bufWriter.Write(idxHeader[0:2]) + n2, _ := bufWriter.Write(entry.Key) + n3, _ := bufWriter.Write(idxHeader[2:10]) currentOffset += uint64(n1 + n2 + n3) } bloomFilterOffset := currentOffset - bloomData := bf.Encode() - bw.Write(bloomData) + bufWriter.Write(bloomData) - var footer [20]byte + var footer [SSTableFooterSize]byte binary.BigEndian.PutUint64(footer[0:8], indexBlockOffset) binary.BigEndian.PutUint64(footer[8:16], bloomFilterOffset) - binary.BigEndian.PutUint32(footer[16:20], 0xF0E57DBB) - bw.Write(footer[:]) + binary.BigEndian.PutUint32(footer[16:20], SSTableMagic) + bufWriter.Write(footer[:]) - if err := bw.Flush(); err != nil { + if err := bufWriter.Flush(); err != nil { return err } - return f.Sync() + return file.Sync() } diff --git a/internal/engine/db.go b/internal/engine/db.go index 8ee6210..ecbb038 100644 --- a/internal/engine/db.go +++ b/internal/engine/db.go @@ -1,3 +1,4 @@ +// Package engine contains the core data structures and logic for the LSM-Tree database. package engine import ( @@ -11,37 +12,48 @@ import ( ) // DB coordinates the active MemTable, WAL, and background SSTable flushes. +// It uses atomic pointers and RCU (Read-Copy-Update) semantics to ensure +// that network readers are never blocked by background compactions or flushes. type DB struct { - activeMT *MemTable - immutableMT *MemTable // Holds the frozen MemTable during background flush - wal *WAL - mu sync.RWMutex + // Active and immutable MemTables are stored as atomic pointers. + // This eliminates the need for reader mutexes during the hot Get() path. + activeMT atomic.Pointer[MemTable] + immutableMT atomic.Pointer[MemTable] + + wal *WAL + + // mu is strictly used to serialize flush requests and background coordination. + // It is purposefully bypassed by the hot read/write paths. + mu sync.Mutex + flushQ chan *MemTable wg sync.WaitGroup stopCh chan struct{} flushThreshold int64 + + // currentVersion uses RCU to swap the active SSTable manifest atomically. currentVersion atomic.Pointer[Version] } -// NewDB initializes the engine state and starts background workers. +// NewDB initializes the engine state, recovers from the WAL, and starts background workers. func NewDB(walPath string) (*DB, error) { wal, err := NewWAL(walPath) if err != nil { - return nil, err + return nil, fmt.Errorf("engine failed to initialize WAL: %w", err) } mt := NewMemTable() if err := Recover(walPath, mt); err != nil { - return nil, err + return nil, fmt.Errorf("engine failed to recover from WAL: %w", err) } db := &DB{ - activeMT: mt, wal: wal, flushQ: make(chan *MemTable, 1), stopCh: make(chan struct{}), flushThreshold: 4 * 1024 * 1024, // 4MB } + db.activeMT.Store(mt) initialVersion := &Version{ Manifest: Manifest{ @@ -57,22 +69,21 @@ func NewDB(walPath string) (*DB, error) { return db, nil } -// Get executes the complete LSM read path. +// Get executes the complete LSM read path (MemTable -> Immutable -> SSTables). +// It is strictly lock-free and utilizes atomic pointers and RCU to prevent read stalls. func (db *DB) Get(key []byte) ([]byte, bool, error) { - // 1. Active MemTable (and Immutable MemTable) - db.mu.RLock() - active := db.activeMT - immutable := db.immutableMT - db.mu.RUnlock() - - if val, ok := active.Get(key); ok { - if len(val) == 0 { // Tombstone check - return nil, false, nil + // 1. Search Active MemTable (Lock-Free) + if active := db.activeMT.Load(); active != nil { + if val, ok := active.Get(key); ok { + if len(val) == 0 { // Tombstone check + return nil, false, nil + } + return val, true, nil } - return val, true, nil } - if immutable != nil { + // 2. Search Immutable MemTable (Lock-Free) + if immutable := db.immutableMT.Load(); immutable != nil { if val, ok := immutable.Get(key); ok { if len(val) == 0 { return nil, false, nil @@ -81,7 +92,8 @@ func (db *DB) Get(key []byte) ([]byte, bool, error) { } } - // 2. RCU Version SSTables + // 3. Search RCU Version SSTables + // Acquire a read reference to prevent files from being deleted during read. version := db.currentVersion.Load() version.Acquire() defer version.Release() @@ -90,7 +102,7 @@ func (db *DB) Get(key []byte) ([]byte, bool, error) { for i := len(version.Manifest.L0Files) - 1; i >= 0; i-- { val, found, err := ReadSSTable(version.Manifest.L0Files[i], key) if err != nil { - return nil, false, err + return nil, false, fmt.Errorf("failed to read L0 SSTable: %w", err) } if found { if len(val) == 0 { @@ -104,7 +116,7 @@ func (db *DB) Get(key []byte) ([]byte, bool, error) { for i := len(version.Manifest.L1Files) - 1; i >= 0; i-- { val, found, err := ReadSSTable(version.Manifest.L1Files[i], key) if err != nil { - return nil, false, err + return nil, false, fmt.Errorf("failed to read L1 SSTable: %w", err) } if found { if len(val) == 0 { @@ -117,62 +129,62 @@ func (db *DB) Get(key []byte) ([]byte, bool, error) { return nil, false, nil } -// Put writes the key-value pair to the WAL and active MemTable. +// Put appends the mutation to the WAL and inserts it into the active MemTable lock-free. func (db *DB) Put(key, val []byte) error { if err := db.wal.Append(network.OpPut, key, val); err != nil { - return err + return fmt.Errorf("engine failed to append Put to WAL: %w", err) } - db.mu.RLock() - mt := db.activeMT - db.mu.RUnlock() - - mt.Put(key, val) + active := db.activeMT.Load() + active.Put(key, val) - if mt.Size() >= db.flushThreshold { + if active.Size() >= db.flushThreshold { db.CheckFlush() } return nil } -// Delete appends a tombstone for the key. +// Delete appends a tombstone to the WAL and inserts it into the active MemTable lock-free. func (db *DB) Delete(key []byte) error { if err := db.wal.Append(network.OpDelete, key, nil); err != nil { - return err + return fmt.Errorf("engine failed to append Delete to WAL: %w", err) } - db.mu.RLock() - mt := db.activeMT - db.mu.RUnlock() + active := db.activeMT.Load() + active.Put(key, nil) - mt.Put(key, nil) - - if mt.Size() >= db.flushThreshold { + if active.Size() >= db.flushThreshold { db.CheckFlush() } return nil } -// CheckFlush safely swaps the active MemTable if it exceeds the threshold. +// CheckFlush safely rotates the active MemTable into the flush queue if it exceeds the capacity threshold. func (db *DB) CheckFlush() { db.mu.Lock() defer db.mu.Unlock() - if db.activeMT.Size() < db.flushThreshold { + active := db.activeMT.Load() + if active.Size() < db.flushThreshold { return } - if db.immutableMT != nil { - log.Println("WARNING: Flush triggered while immutableMT is still flushing. Stalling.") + if db.immutableMT.Load() != nil { + // This acts as a backpressure signal. If disk I/O cannot keep up with network inserts, + // the active table will temporarily exceed 4MB until the current flush completes. + log.Println("warn: MemTable flush triggered while background flush is still active. Write delayed.") return } - db.immutableMT = db.activeMT - db.activeMT = NewMemTable() - db.flushQ <- db.immutableMT + // Atomically rotate the MemTables. + db.immutableMT.Store(active) + db.activeMT.Store(NewMemTable()) + + // Queue the frozen table for disk persistence. + db.flushQ <- active } -// flushWorker runs in the background and writes frozen MemTables to SSTable files. +// flushWorker acts as a background sink, persisting frozen MemTables to L0 SSTables. func (db *DB) flushWorker() { defer db.wg.Done() flushCount := time.Now().UnixNano() @@ -181,13 +193,14 @@ func (db *DB) flushWorker() { select { case mt := <-db.flushQ: filepath := fmt.Sprintf("data-L0-%d.sst", flushCount) - log.Printf("Flushing MemTable to %s...", filepath) if err := FlushMemTable(mt, filepath); err != nil { - log.Printf("FATAL: failed to flush SSTable: %v", err) + log.Printf("critical: failed to flush MemTable to %s: %v", filepath, err) } flushCount++ - // RCU Manifest Update + // Perform an RCU (Read-Copy-Update) replacement of the manifest. + // This allows lock-free readers to safely continue using the old manifest + // while we introduce the newly created SSTable. oldVersion := db.currentVersion.Load() newManifest := Manifest{ L0Files: append([]string(nil), oldVersion.Manifest.L0Files...), @@ -200,13 +213,13 @@ func (db *DB) flushWorker() { } db.currentVersion.Store(newVersion) - // Cleanly remove immutableMT so next flush can happen + // Cleanly remove the immutable table so the next flush can trigger. db.mu.Lock() - db.immutableMT = nil + db.immutableMT.Store(nil) db.mu.Unlock() case <-db.stopCh: - // Drain remaining flush queue on shutdown + // Drain remaining flush queue on shutdown to guarantee crash durability. close(db.flushQ) for mt := range db.flushQ { filepath := fmt.Sprintf("data-L0-%d.sst", flushCount) @@ -218,27 +231,29 @@ func (db *DB) flushWorker() { } } -// Close cleanly shuts down the WAL and flushes remaining MemTables. +// Close gracefully stops the database. It halts the WAL and forces remaining MemTables to disk. func (db *DB) Close() error { - // First close WAL so no more Puts can happen + // First close WAL so no more mutations can be received. if err := db.wal.Close(); err != nil { - return err + return fmt.Errorf("failed to cleanly close WAL: %w", err) } - // Force flush whatever is in activeMT if it's not empty + // Force a final flush of whatever remains in the active MemTable. db.mu.Lock() - if db.activeMT.Size() > 0 { - // Wait if immutable is currently flushing - for db.immutableMT != nil { + active := db.activeMT.Load() + if active.Size() > 0 { + // Stably wait if the background worker is currently busy flushing. + for db.immutableMT.Load() != nil { db.mu.Unlock() time.Sleep(10 * time.Millisecond) db.mu.Lock() } - db.immutableMT = db.activeMT - db.flushQ <- db.immutableMT + db.immutableMT.Store(active) + db.flushQ <- active } db.mu.Unlock() + // Signal the worker to stop and wait for it to drain the queue. close(db.stopCh) db.wg.Wait() return nil diff --git a/internal/engine/iterator.go b/internal/engine/iterator.go index 9496a89..715cebd 100644 --- a/internal/engine/iterator.go +++ b/internal/engine/iterator.go @@ -2,15 +2,17 @@ package engine import ( "encoding/binary" + "fmt" "io" "os" "github.com/Forest_DatabaseEngine/internal/network" ) -// SSTIterator streams key-value pairs from a single SSTable sequentially. +// SSTIterator streams key-value pairs sequentially from the data block phase of a single SSTable. +// It is heavily used during compaction (K-Way merge) to read large contiguous chunks of data without loading the entire file into memory. type SSTIterator struct { - f *os.File + file *os.File dataEnd int64 current int64 key []byte @@ -20,34 +22,42 @@ type SSTIterator struct { err error } -// NewSSTIterator opens an SSTable and parses the footer to find the end of the Data Blocks. +// NewSSTIterator opens an SSTable and parses the footer to determine the exact byte boundary +// where the Data Blocks end and the Sparse Index begins. func NewSSTIterator(filepath string, fileSeq uint64) (*SSTIterator, error) { - f, err := os.Open(filepath) + file, err := os.Open(filepath) if err != nil { - return nil, err + return nil, fmt.Errorf("failed to open SSTable iterator for %s: %w", filepath, err) } - stat, err := f.Stat() + stat, err := file.Stat() if err != nil { - f.Close() - return nil, err + file.Close() + return nil, fmt.Errorf("failed to stat SSTable for iterator: %w", err) } - if stat.Size() < 20 { - f.Close() - return nil, io.ErrUnexpectedEOF + if stat.Size() < SSTableFooterSize { + file.Close() + return nil, fmt.Errorf("SSTable %s is too small (corrupt)", filepath) } - // Read Footer to find IndexBlockOffset (which marks the end of Data Blocks) - var footer [20]byte - if _, err := f.ReadAt(footer[:], stat.Size()-20); err != nil { - f.Close() - return nil, err + // 1. Read Footer to find IndexBlockOffset (which marks the exact end of Data Blocks) + var footer [SSTableFooterSize]byte + if _, err := file.ReadAt(footer[:], stat.Size()-SSTableFooterSize); err != nil { + file.Close() + return nil, fmt.Errorf("failed to read SSTable footer in iterator: %w", err) } + + magic := binary.BigEndian.Uint32(footer[16:20]) + if magic != SSTableMagic { + file.Close() + return nil, fmt.Errorf("invalid SSTable magic number in iterator: expected %X, got %X", SSTableMagic, magic) + } + dataEnd := int64(binary.BigEndian.Uint64(footer[0:8])) return &SSTIterator{ - f: f, + file: file, dataEnd: dataEnd, current: 0, fileSeq: fileSeq, @@ -55,15 +65,16 @@ func NewSSTIterator(filepath string, fileSeq uint64) (*SSTIterator, error) { } // Next reads the next key-value pair from the data blocks. Returns true if successful. +// It strictly halts when the read head reaches the Sparse Index phase (it.dataEnd). func (it *SSTIterator) Next() bool { if it.current >= it.dataEnd { return false } var header [7]byte - n, err := it.f.ReadAt(header[:], it.current) + n, err := it.file.ReadAt(header[:], it.current) if err != nil { - it.err = err + it.err = fmt.Errorf("iterator failed to read block header at offset %d: %w", it.current, err) return false } if n < 7 { @@ -76,9 +87,9 @@ func (it *SSTIterator) Next() bool { valLen := int(binary.BigEndian.Uint32(header[3:7])) payload := make([]byte, keyLen+valLen) - n, err = it.f.ReadAt(payload, it.current+7) + n, err = it.file.ReadAt(payload, it.current+7) if err != nil { - it.err = err + it.err = fmt.Errorf("iterator failed to read block payload at offset %d: %w", it.current+7, err) return false } if n < keyLen+valLen { @@ -92,9 +103,20 @@ func (it *SSTIterator) Next() bool { return true } +// Key returns the current key. func (it *SSTIterator) Key() []byte { return it.key } + +// Value returns the current value. func (it *SSTIterator) Value() []byte { return it.val } + +// OpCode returns the current operation code (e.g., OpPut or OpDelete). func (it *SSTIterator) OpCode() network.OpCode { return it.op } + +// FileSeq returns the sequence identifier for tie-breaking during compaction. func (it *SSTIterator) FileSeq() uint64 { return it.fileSeq } + +// Error returns the last encountered read error, if any. func (it *SSTIterator) Error() error { return it.err } -func (it *SSTIterator) Close() error { return it.f.Close() } + +// Close cleanly releases the underlying file handle. +func (it *SSTIterator) Close() error { return it.file.Close() } diff --git a/internal/engine/memtable.go b/internal/engine/memtable.go index f627efd..74f534d 100644 --- a/internal/engine/memtable.go +++ b/internal/engine/memtable.go @@ -7,31 +7,44 @@ import ( "sync/atomic" ) +// maxLevel defines the maximum height of the SkipList. +// 16 is mathematically sufficient for up to 2^16 elements with a branching factor of 0.5. +// For p=0.25, it comfortably supports standard 4MB MemTable capacities. const maxLevel = 16 -// Node represents a node in the lock-free SkipList MemTable. +// Node represents a single key-value entry in the lock-free SkipList. +// Readers traverse this structure without holding any locks. type Node struct { - key []byte - // val is an atomic pointer to a byte slice to allow lock-free, race-free value updates. + key []byte + // val is stored as an atomic pointer to a byte slice. + // This ensures that readers never observe partially updated values + // when a key is overwritten by a concurrent writer. val atomic.Pointer[[]byte] + // nexts holds the forward pointers for each level of the SkipList. + // They are updated atomically from bottom-to-top during insertion + // to guarantee lock-free readers always see a structurally valid list. nexts [maxLevel]atomic.Pointer[Node] } -// MemTable is a concurrent, lock-free SkipList (readers are lock-free). +// MemTable is a thread-safe, lock-free SkipList serving as the active write buffer. +// It allows single-writer serialization (via mutex) while sustaining infinite +// concurrent readers (via atomic pointer chasing) with zero contention. type MemTable struct { head *Node + // mu serializes writers. Readers completely bypass this lock. mu sync.RWMutex size atomic.Int64 } -// NewMemTable initializes an empty MemTable. +// NewMemTable initializes and returns an empty MemTable with a dummy head node. func NewMemTable() *MemTable { return &MemTable{ head: &Node{}, } } -// randomLevel generates a random height for a new node. +// randomLevel generates a randomized height for a newly inserted node. +// It uses a geometric distribution with p=0.25 to balance search speed and memory footprint. func randomLevel() int { level := 1 for level < maxLevel && rand.Float32() < 0.25 { @@ -41,15 +54,17 @@ func randomLevel() int { } // Put inserts or updates a key-value pair in the MemTable. -// It uses a mutex to serialize writers but does not block readers. +// It locks the structure to prevent write-write conflicts but relies on atomic +// operations so read-write conflicts are handled lock-free. func (m *MemTable) Put(key, val []byte) { + // Serialize writers to prevent structural corruption of the SkipList. m.mu.Lock() defer m.mu.Unlock() var preds [maxLevel]*Node curr := m.head - // Traverse to find the predecessors at each level + // Traverse from the top level down, recording the predecessor node at each level. for i := maxLevel - 1; i >= 0; i-- { for { next := curr.nexts[i].Load() @@ -62,72 +77,87 @@ func (m *MemTable) Put(key, val []byte) { preds[i] = curr } - // Check if the key already exists next := preds[0].nexts[0].Load() if next != nil && bytes.Equal(next.key, key) { - // Key exists: atomically update the value to prevent races with lock-free readers. - // Note: This does not update the approximate size for the new/old value difference - // for simplicity in this phase. - newVal := make([]byte, len(val)) - copy(newVal, val) + // Key already exists: we only need to update the value. + // By doing an atomic pointer swap, active readers will either see the old + // value or the new one, but never corrupted memory. + var newVal []byte + if val != nil { + newVal = make([]byte, len(val)) + copy(newVal, val) + } next.val.Store(&newVal) + + // Note: Approximate size is not adjusted on updates to avoid heavy calculation + // on the hot path. The 4MB threshold is a soft limit. return } - // Create new node + // Key does not exist: create a new node and determine its height. level := randomLevel() newNode := &Node{ key: key, } - newVal := make([]byte, len(val)) - copy(newVal, val) + + var newVal []byte + if val != nil { + newVal = make([]byte, len(val)) + copy(newVal, val) + } newNode.val.Store(&newVal) - // Increment size (key + val) + // Increment approximate memory size (useful for flush triggers). m.size.Add(int64(len(key) + len(val))) - // Wire the new node to the next nodes + // Wire the new node's forward pointers to the existing nodes. + // This step is completely invisible to readers since newNode is not yet linked. for i := 0; i < level; i++ { newNode.nexts[i].Store(preds[i].nexts[i].Load()) } - // Atomically swap the predecessor's next pointers from bottom to top. - // This ensures lock-free readers traversing from top-to-bottom or left-to-right - // will see a fully linked node if they encounter it. + // Publish the new node to readers by atomically swapping the predecessors' pointers. + // We MUST link from bottom (Level 0) to top. If a reader finds the node at a higher level, + // it will inevitably drop down and expect the lower levels to be linked as well. for i := 0; i < level; i++ { preds[i].nexts[i].Store(newNode) } } -// Get performs a lock-free search for a key in the MemTable. +// Get performs a lock-free search across the SkipList. +// It returns the value and a boolean indicating if the key was found. +// A nil/empty value with true indicates a tombstone (deleted key). func (m *MemTable) Get(key []byte) ([]byte, bool) { curr := m.head - // Traverse lock-free from top to bottom + // Traverse lock-free from top to bottom. for i := maxLevel - 1; i >= 0; i-- { for { next := curr.nexts[i].Load() if next == nil { break } + cmp := bytes.Compare(next.key, key) if cmp < 0 { + // Move right if the next key is smaller than our target. curr = next } else if cmp == 0 { - // Exact match found + // Exact match found. Load the atomic value pointer. valPtr := next.val.Load() if valPtr != nil { return *valPtr, true } + // Failsafe (should never happen by design) return nil, false } else { - // next.key > key, drop down a level + // next.key > target key, drop down one level. break } } } - // Check the level 0 next node just in case + // Check the base level (Level 0) next node just in case we overshot. next := curr.nexts[0].Load() if next != nil && bytes.Equal(next.key, key) { valPtr := next.val.Load() @@ -136,10 +166,11 @@ func (m *MemTable) Get(key []byte) ([]byte, bool) { } } - return nil, false + return nil, false // Key truly not found } -// Size returns the approximate byte size of the MemTable. +// Size returns the approximate memory footprint of the MemTable in bytes. +// It is used by the Engine to trigger background SSTable flushes. func (m *MemTable) Size() int64 { return m.size.Load() } diff --git a/internal/engine/sstable.go b/internal/engine/sstable.go index 82ce8db..bfa0014 100644 --- a/internal/engine/sstable.go +++ b/internal/engine/sstable.go @@ -10,25 +10,35 @@ import ( "github.com/Forest_DatabaseEngine/internal/network" ) -// IndexEntry represents a pointer to a data block inside the SSTable. +// SSTableMagic is the validation signature for SSTable files. +const SSTableMagic uint32 = 0xF0E57DBB + +// SSTableFooterSize is the fixed byte length of the trailing metadata. +// Layout: IndexOffset(8) + BloomOffset(8) + Magic(4) = 20 bytes +const SSTableFooterSize = 20 + +// IndexEntry represents a sparse index pointer mapping a key to a 4KB data block offset. type IndexEntry struct { Key []byte Offset uint64 } // FlushMemTable writes a frozen MemTable to an immutable, append-only SSTable file. +// The file layout is heavily optimized for minimal disk I/O: +// [Data Blocks] -> [Sparse Index] -> [Bloom Filter] -> [Footer] func FlushMemTable(mt *MemTable, filepath string) error { - f, err := os.OpenFile(filepath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644) + file, err := os.OpenFile(filepath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644) if err != nil { - return err + return fmt.Errorf("failed to create SSTable file %s: %w", filepath, err) } - defer f.Close() + defer file.Close() - // Use a 4MB buffer to prevent excessive syscalls - bw := bufio.NewWriterSize(f, 4*1024*1024) + // Buffer syscalls to aggressively minimize write overhead. + bufWriter := bufio.NewWriterSize(file, 4*1024*1024) var index []IndexEntry // We size the Bloom Filter for ~4000 items (4MB / ~1KB per item) with a 1% false positive rate. + // This mathematically eliminates >99% of disk reads for non-existent keys. bf := NewBloomFilter(4000, 0.01) var currentOffset uint64 = 0 @@ -39,32 +49,32 @@ func FlushMemTable(mt *MemTable, filepath string) error { for curr != nil { key := curr.key valPtr := curr.val.Load() - if valPtr == nil { + if valPtr == nil { // Nil pointer shouldn't happen by design, but we guard. curr = curr.nexts[0].Load() continue } val := *valPtr - // Add to Bloom Filter + // Note: We write tombstones (len(val)==0) to disk so they can shadow older values in L1+. bf.Add(key) - // Create Sparse Index Entry every 4KB of data + // Create Sparse Index Entry roughly every 4KB of data to bound linear scan time. if currentOffset-lastIndexOffset >= 4096 || currentOffset == 0 { index = append(index, IndexEntry{Key: key, Offset: currentOffset}) lastIndexOffset = currentOffset } - // Write Data Block Entry: [OpCode (1B) | KeyLen (2B) | ValLen (4B) | Key | Val] + // Write Data Block Entry: [OpCode(1B) | KeyLen(2B) | ValLen(4B) | Key | Val] var header [7]byte - header[0] = byte(network.OpEcho) // We use OpEcho to signify a standard Put for this phase + header[0] = byte(network.OpPut) // Standardize on OpPut for disk entries. binary.BigEndian.PutUint16(header[1:3], uint16(len(key))) binary.BigEndian.PutUint32(header[3:7], uint32(len(val))) - n, _ := bw.Write(header[:]) + n, _ := bufWriter.Write(header[:]) currentOffset += uint64(n) - n, _ = bw.Write(key) + n, _ = bufWriter.Write(key) currentOffset += uint64(n) - n, _ = bw.Write(val) + n, _ = bufWriter.Write(val) currentOffset += uint64(n) curr = curr.nexts[0].Load() @@ -76,16 +86,16 @@ func FlushMemTable(mt *MemTable, filepath string) error { // NumIndexEntries (4B) var countBuf [4]byte binary.BigEndian.PutUint32(countBuf[:], uint32(len(index))) - n, _ := bw.Write(countBuf[:]) + n, _ := bufWriter.Write(countBuf[:]) currentOffset += uint64(n) for _, entry := range index { var idxHeader [10]byte // KeyLen(2) + Offset(8) binary.BigEndian.PutUint16(idxHeader[0:2], uint16(len(entry.Key))) binary.BigEndian.PutUint64(idxHeader[2:10], entry.Offset) - n1, _ := bw.Write(idxHeader[0:2]) - n2, _ := bw.Write(entry.Key) - n3, _ := bw.Write(idxHeader[2:10]) + n1, _ := bufWriter.Write(idxHeader[0:2]) + n2, _ := bufWriter.Write(entry.Key) + n3, _ := bufWriter.Write(idxHeader[2:10]) currentOffset += uint64(n1 + n2 + n3) } @@ -93,64 +103,71 @@ func FlushMemTable(mt *MemTable, filepath string) error { // 3. Write Bloom Filter bloomData := bf.Encode() - bw.Write(bloomData) + bufWriter.Write(bloomData) - // 4. Write Footer (Fixed 20 bytes) - var footer [20]byte + // 4. Write Footer (Fixed Size) + var footer [SSTableFooterSize]byte binary.BigEndian.PutUint64(footer[0:8], indexBlockOffset) binary.BigEndian.PutUint64(footer[8:16], bloomFilterOffset) - binary.BigEndian.PutUint32(footer[16:20], 0xF0E57DBB) // Magic Number - bw.Write(footer[:]) + binary.BigEndian.PutUint32(footer[16:20], SSTableMagic) + bufWriter.Write(footer[:]) // Flush bufio to disk - if err := bw.Flush(); err != nil { - return err + if err := bufWriter.Flush(); err != nil { + return fmt.Errorf("failed to flush SSTable buffers: %w", err) } - // Final fsync - return f.Sync() + // Final fsync guarantees durability before updating the manifest. + if err := file.Sync(); err != nil { + return fmt.Errorf("failed to fsync SSTable: %w", err) + } + return nil } // ReadSSTable performs a point lookup for a key in the SSTable file. +// It aggressively avoids disk I/O by checking the Bloom Filter first. func ReadSSTable(filepath string, searchKey []byte) ([]byte, bool, error) { - f, err := os.Open(filepath) + file, err := os.Open(filepath) if err != nil { - return nil, false, err + return nil, false, fmt.Errorf("failed to open SSTable for read: %w", err) } - defer f.Close() + defer file.Close() - stat, err := f.Stat() + stat, err := file.Stat() if err != nil { - return nil, false, err + return nil, false, fmt.Errorf("failed to stat SSTable: %w", err) } - if stat.Size() < 20 { - return nil, false, fmt.Errorf("file too small") + if stat.Size() < SSTableFooterSize { + return nil, false, fmt.Errorf("SSTable %s is too small (corrupt)", filepath) } // 1. Read Footer - var footer [20]byte - if _, err := f.ReadAt(footer[:], stat.Size()-20); err != nil { - return nil, false, err + var footer [SSTableFooterSize]byte + if _, err := file.ReadAt(footer[:], stat.Size()-SSTableFooterSize); err != nil { + return nil, false, fmt.Errorf("failed to read SSTable footer: %w", err) } magic := binary.BigEndian.Uint32(footer[16:20]) - if magic != 0xF0E57DBB { - return nil, false, fmt.Errorf("invalid SSTable magic number") + if magic != SSTableMagic { + return nil, false, fmt.Errorf("invalid SSTable magic number: expected %X, got %X", SSTableMagic, magic) } indexOffset := binary.BigEndian.Uint64(footer[0:8]) bloomOffset := binary.BigEndian.Uint64(footer[8:16]) // 2. Read Bloom Filter - bloomLen := stat.Size() - 20 - int64(bloomOffset) + bloomLen := stat.Size() - SSTableFooterSize - int64(bloomOffset) bloomData := make([]byte, bloomLen) - if _, err := f.ReadAt(bloomData, int64(bloomOffset)); err != nil { - return nil, false, err + if _, err := file.ReadAt(bloomData, int64(bloomOffset)); err != nil { + return nil, false, fmt.Errorf("failed to read Bloom filter: %w", err) } - bf := DecodeBloomFilter(bloomData) - if bf == nil || !bf.Contains(searchKey) { + bf, err := DecodeBloomFilter(bloomData) + if err != nil { + return nil, false, fmt.Errorf("failed to decode Bloom filter: %w", err) + } + if !bf.Contains(searchKey) { // Bloom filter says NO - skip reading the data blocks entirely! return nil, false, nil } @@ -158,12 +175,12 @@ func ReadSSTable(filepath string, searchKey []byte) ([]byte, bool, error) { // 3. Read Sparse Index indexLen := bloomOffset - indexOffset indexData := make([]byte, indexLen) - if _, err := f.ReadAt(indexData, int64(indexOffset)); err != nil { - return nil, false, err + if _, err := file.ReadAt(indexData, int64(indexOffset)); err != nil { + return nil, false, fmt.Errorf("failed to read Sparse Index: %w", err) } if len(indexData) < 4 { - return nil, false, fmt.Errorf("invalid index block") + return nil, false, fmt.Errorf("invalid sparse index block size") } numEntries := binary.BigEndian.Uint32(indexData[0:4]) @@ -197,8 +214,8 @@ func ReadSSTable(filepath string, searchKey []byte) ([]byte, bool, error) { // 4. Read Data Block if blockOffset == 0 && numEntries > 0 { - // We are looking for a key smaller than the very first key in the index - // This means it doesn't exist. + // We are looking for a key smaller than the very first key in the index. + // Because the SSTable is sorted, this means the key definitely does not exist. firstKeyLen := binary.BigEndian.Uint16(indexData[4:6]) firstKey := indexData[6 : 6+firstKeyLen] if bytes.Compare(firstKey, searchKey) > 0 { @@ -208,8 +225,8 @@ func ReadSSTable(filepath string, searchKey []byte) ([]byte, bool, error) { blockLen := nextBlockOffset - blockOffset blockData := make([]byte, blockLen) - if _, err := f.ReadAt(blockData, int64(blockOffset)); err != nil { - return nil, false, err + if _, err := file.ReadAt(blockData, int64(blockOffset)); err != nil { + return nil, false, fmt.Errorf("failed to read Data Block: %w", err) } // 5. Linear scan within the 4KB block @@ -229,7 +246,7 @@ func ReadSSTable(filepath string, searchKey []byte) ([]byte, bool, error) { if cmp == 0 { return currVal, true, nil } else if cmp > 0 { - // Because it's sorted, if we pass it, it's not here. + // Because the block is sorted, if we pass the target key, it's not here. break } diff --git a/internal/engine/version.go b/internal/engine/version.go index 307626e..026786e 100644 --- a/internal/engine/version.go +++ b/internal/engine/version.go @@ -4,29 +4,32 @@ import ( "sync/atomic" ) -// Manifest tracks active files for an Engine version. -// Using this structure allows readers to lock-free access files, while -// writers (compactors) create a new Version and atomically swap the pointer. +// Manifest tracks all active SSTable files comprising a specific point-in-time state of the database. +// By isolating the file lists inside an immutable Manifest, background compaction workers can freely +// generate a new Manifest and atomically swap the engine pointer, leaving concurrent readers completely unaffected. type Manifest struct { L0Files []string L1Files []string } -// Version represents an immutable point-in-time view of the Manifest. +// Version represents an immutable, reference-counted snapshot of the Manifest. +// This is the core of our Multi-Version Concurrency Control (MVCC) and RCU design. type Version struct { Manifest Manifest Refs atomic.Int32 } -// Acquire increments the reference count. -// Called by TCP readers before reading from SSTables. +// Acquire increments the reference count for this specific version. +// It must be called by readers (e.g., Get operations) before they begin traversing the SSTables +// to ensure the underlying files are not physically unlinked by a background compaction worker. func (v *Version) Acquire() { v.Refs.Add(1) } // Release decrements the reference count. -// Called by TCP readers when they are done. When Refs drops to 0, -// obsolete files from this version can be safely purged. +// It must be called via defer by readers when their operation completes. +// Once a Version is detached from the active database pointer and its Refs drop to 0, +// the compaction workers are safely permitted to unlink its obsolete files from disk. func (v *Version) Release() { v.Refs.Add(-1) } diff --git a/internal/engine/wal.go b/internal/engine/wal.go index 67c8e87..c71e8fc 100644 --- a/internal/engine/wal.go +++ b/internal/engine/wal.go @@ -1,9 +1,12 @@ +// Package engine contains the core data structures and logic for the LSM-Tree database. +// It includes the MemTable, Write-Ahead Log (WAL), SSTables, and background compaction. package engine import ( "bufio" "encoding/binary" "errors" + "fmt" "hash/crc32" "io" "os" @@ -13,28 +16,34 @@ import ( "github.com/Forest_DatabaseEngine/internal/network" ) -const WALHeaderSize = 12 // Magic(1) + CRC32(4) + OpCode(1) + KeyLen(2) + ValLen(4) +// WALHeaderSize defines the fixed byte length of a WAL entry header. +// Layout: Magic(1) + CRC32(4) + OpCode(1) + KeyLen(2) + ValLen(4) = 12 bytes. +const WALHeaderSize = 12 -// WAL represents an append-only Write-Ahead Log. +// WAL represents a crash-safe, append-only Write-Ahead Log. +// It utilizes double-buffering (via bufio) and asynchronous fsyncs +// to prevent disk I/O from blocking the fast network parser. type WAL struct { - file *os.File - bw *bufio.Writer - mu sync.Mutex - stopCh chan struct{} - wg sync.WaitGroup + file *os.File + bufWriter *bufio.Writer + mu sync.Mutex + stopCh chan struct{} + wg sync.WaitGroup } -// NewWAL creates or opens a WAL file and starts the background syncer. +// NewWAL creates or opens a WAL file and starts the background fsync loop. func NewWAL(filepath string) (*WAL, error) { - f, err := os.OpenFile(filepath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644) + file, err := os.OpenFile(filepath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644) if err != nil { - return nil, err + return nil, fmt.Errorf("failed to open WAL file %s: %w", filepath, err) } wal := &WAL{ - file: f, - bw: bufio.NewWriterSize(f, 4*1024*1024), // 4MB buffer - stopCh: make(chan struct{}), + file: file, + // 4MB buffer prevents frequent small disk writes, relying on the background + // syncer to persist chunks asynchronously. + bufWriter: bufio.NewWriterSize(file, 4*1024*1024), + stopCh: make(chan struct{}), } wal.wg.Add(1) @@ -43,42 +52,41 @@ func NewWAL(filepath string) (*WAL, error) { return wal, nil } -// Append writes a new entry to the WAL buffer. +// Append serializes a new operation into the WAL buffer. +// The lock ensures thread-safe sequential appending across concurrent client requests. func (w *WAL) Append(op network.OpCode, key, val []byte) error { w.mu.Lock() defer w.mu.Unlock() var header [WALHeaderSize]byte header[0] = network.MagicByte - - // OpCode, KeyLen, ValLen header[5] = byte(op) binary.BigEndian.PutUint16(header[6:8], uint16(len(key))) binary.BigEndian.PutUint32(header[8:12], uint32(len(val))) - // Compute CRC32 of OpCode + KeyLen + ValLen + Key + Val + // Compute CRC32 over the mutable fields and payload to detect disk corruption or torn writes. h := crc32.NewIEEE() h.Write(header[5:12]) h.Write(key) h.Write(val) - crc := h.Sum32() - binary.BigEndian.PutUint32(header[1:5], crc) + binary.BigEndian.PutUint32(header[1:5], h.Sum32()) - // Write header and payload - if _, err := w.bw.Write(header[:]); err != nil { - return err + // Write header and payload sequentially to the memory buffer. + if _, err := w.bufWriter.Write(header[:]); err != nil { + return fmt.Errorf("failed to write WAL header: %w", err) } - if _, err := w.bw.Write(key); err != nil { - return err + if _, err := w.bufWriter.Write(key); err != nil { + return fmt.Errorf("failed to write WAL key: %w", err) } - if _, err := w.bw.Write(val); err != nil { - return err + if _, err := w.bufWriter.Write(val); err != nil { + return fmt.Errorf("failed to write WAL value: %w", err) } return nil } -// syncLoop flushes the buffer to disk every 10ms. +// syncLoop is a dedicated goroutine that flushes the memory buffer to the OS, +// and forces an fsync to persistent storage every 10ms. func (w *WAL) syncLoop() { defer w.wg.Done() ticker := time.NewTicker(10 * time.Millisecond) @@ -89,82 +97,89 @@ func (w *WAL) syncLoop() { case <-ticker.C: w.flush() case <-w.stopCh: - // Final flush on shutdown + // Ensure final writes are persisted during graceful shutdown. w.flush() return } } } +// flush safely pushes buffered data to the underlying file and triggers an fsync. func (w *WAL) flush() { w.mu.Lock() - err := w.bw.Flush() + err := w.bufWriter.Flush() w.mu.Unlock() if err == nil { - // Sync to disk outside the lock to avoid blocking incoming appends. + // Sync is called outside the mutex so disk I/O latency does not block active Appends. _ = w.file.Sync() } } -// Close gracefully stops the background syncer and closes the file. +// Close gracefully stops the background syncer, forcing a final flush, and closes the file descriptor. func (w *WAL) Close() error { close(w.stopCh) w.wg.Wait() - return w.file.Close() + if err := w.file.Close(); err != nil { + return fmt.Errorf("failed to close WAL file: %w", err) + } + return nil } -// Recover reads a WAL file and repopulates a MemTable. +// Recover replays a WAL file sequentially to repopulate the MemTable upon startup. +// It gracefully halts on torn writes (EOF or CRC mismatch) indicating the end of safe data. func Recover(filepath string, mt *MemTable) error { - f, err := os.Open(filepath) + file, err := os.Open(filepath) if err != nil { if os.IsNotExist(err) { - return nil // Nothing to recover + return nil // A new database instance, nothing to recover. } - return err + return fmt.Errorf("failed to open WAL for recovery: %w", err) } - defer f.Close() + defer file.Close() var header [WALHeaderSize]byte for { - _, err := io.ReadFull(f, header[:]) - if err != nil { + if _, err := io.ReadFull(file, header[:]); err != nil { if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { - return nil // End of file or torn write at header + return nil // Reached the end of the fully synced data. } - return err + return fmt.Errorf("failed reading WAL header: %w", err) } if header[0] != network.MagicByte { - return errors.New("corrupted WAL: invalid magic byte") + return fmt.Errorf("corrupted WAL: expected magic byte 0x%X, got 0x%X", network.MagicByte, header[0]) } expectedCRC := binary.BigEndian.Uint32(header[1:5]) - _ = network.OpCode(header[5]) + op := network.OpCode(header[5]) keyLen := binary.BigEndian.Uint16(header[6:8]) valLen := binary.BigEndian.Uint32(header[8:12]) key := make([]byte, keyLen) - if _, err := io.ReadFull(f, key); err != nil { - return nil // Torn write at key + if _, err := io.ReadFull(file, key); err != nil { + return nil // Torn write at key payload, stop recovery. } val := make([]byte, valLen) - if _, err := io.ReadFull(f, val); err != nil { - return nil // Torn write at val + if _, err := io.ReadFull(file, val); err != nil { + return nil // Torn write at value payload, stop recovery. } + // Verify data integrity before applying state. h := crc32.NewIEEE() h.Write(header[5:12]) h.Write(key) h.Write(val) if h.Sum32() != expectedCRC { - return nil // Torn write / corruption detected, stop reading + return nil // CRC mismatch implies corrupted tail data, stop recovery cleanly. } - // Apply to MemTable - // For a real DB, we'd handle deletes as well. - // Assuming OpEcho or OpPut acts as a Put for this phase. - mt.Put(key, val) + // Apply the recovered state to the MemTable. + if op == network.OpDelete { + mt.Put(key, nil) // Tombstone for delete + } else { + mt.Put(key, val) // Echo or Put + } } } diff --git a/internal/network/handler.go b/internal/network/handler.go index dce6231..b38dfd6 100644 --- a/internal/network/handler.go +++ b/internal/network/handler.go @@ -2,18 +2,20 @@ package network import ( "encoding/binary" + "fmt" "io" "net" ) -// HandleRequest processes the TCP packet using zero-allocation strategies where possible. +// HandleRequest processes the parsed TCP packet header and payload using zero-allocation strategies. // payloadBuf is pre-allocated by the caller to avoid heap allocations on the hot path. func HandleRequest(conn net.Conn, header RequestHeader, payloadBuf []byte, db Storage) error { totalLen := int(header.KeyLen) + int(header.ValueLen) + + // Read the variable length payload directly into the pre-allocated buffer if totalLen > 0 { - _, err := io.ReadFull(conn, payloadBuf[:totalLen]) - if err != nil { - return err + if _, err := io.ReadFull(conn, payloadBuf[:totalLen]); err != nil { + return fmt.Errorf("failed to read payload from connection: %w", err) } } @@ -22,69 +24,89 @@ func HandleRequest(conn net.Conn, header RequestHeader, payloadBuf []byte, db St switch header.Op { case OpEcho: - // Echo response: [KeyLen(2B)] [ValLen(4B)] [Key] [Val] - // For Echo, we write directly from the payload buffer + // Echo response: [MagicByte(1B)] [Op(1B)] [KeyLen(2B)] [ValLen(4B)] [Key] [Val] + // We construct the response header directly on the stack to prevent allocations. var headerBuf [8]byte headerBuf[0] = MagicByte headerBuf[1] = byte(OpEcho) binary.BigEndian.PutUint16(headerBuf[2:4], header.KeyLen) binary.BigEndian.PutUint32(headerBuf[4:8], header.ValueLen) - conn.Write(headerBuf[:]) + if _, err := conn.Write(headerBuf[:]); err != nil { + return fmt.Errorf("failed to write echo header: %w", err) + } if totalLen > 0 { - conn.Write(payloadBuf[:totalLen]) + if _, err := conn.Write(payloadBuf[:totalLen]); err != nil { + return fmt.Errorf("failed to write echo payload: %w", err) + } } return nil case OpPut: if err := db.Put(key, val); err != nil { - return err + return fmt.Errorf("storage engine failed to process Put: %w", err) } return sendAck(conn) case OpGet: foundVal, found, err := db.Get(key) if err != nil { - return err + return fmt.Errorf("storage engine failed to process Get: %w", err) } if !found { return sendNotFound(conn) } - // Send found value + // Construct the response header with the actual value length retrieved from storage var headerBuf [8]byte headerBuf[0] = MagicByte headerBuf[1] = byte(OpGet) binary.BigEndian.PutUint16(headerBuf[2:4], header.KeyLen) binary.BigEndian.PutUint32(headerBuf[4:8], uint32(len(foundVal))) - conn.Write(headerBuf[:]) - conn.Write(key) - conn.Write(foundVal) + + if _, err := conn.Write(headerBuf[:]); err != nil { + return fmt.Errorf("failed to write get header: %w", err) + } + if _, err := conn.Write(key); err != nil { + return fmt.Errorf("failed to write get key: %w", err) + } + if _, err := conn.Write(foundVal); err != nil { + return fmt.Errorf("failed to write get value: %w", err) + } return nil case OpDelete: if err := db.Delete(key); err != nil { - return err + return fmt.Errorf("storage engine failed to process Delete: %w", err) } return sendAck(conn) default: - // Unknown op code return sendNotFound(conn) } } +// sendAck sends a lightweight acknowledgment packet back to the client. +// It is used for successful Put and Delete operations. func sendAck(conn net.Conn) error { var resp [8]byte resp[0] = MagicByte - _, err := conn.Write(resp[:]) - return err + // The rest of the header is 0x00, indicating an empty payload response + if _, err := conn.Write(resp[:]); err != nil { + return fmt.Errorf("failed to send ACK: %w", err) + } + return nil } +// sendNotFound sends a standard "key not found" response. +// The KeyLen field is set to 0xFFFF as an indicator to the client. func sendNotFound(conn net.Conn) error { var resp [8]byte resp[0] = MagicByte binary.BigEndian.PutUint16(resp[2:4], 0xFFFF) - _, err := conn.Write(resp[:]) - return err + + if _, err := conn.Write(resp[:]); err != nil { + return fmt.Errorf("failed to send NotFound response: %w", err) + } + return nil } diff --git a/internal/network/protocol.go b/internal/network/protocol.go index c365c9a..8eceefb 100644 --- a/internal/network/protocol.go +++ b/internal/network/protocol.go @@ -3,13 +3,18 @@ package network import ( "encoding/binary" "errors" + "fmt" "io" ) -// MagicByte is used to identify our protocol packets. -const MagicByte byte = 0xFE +// MagicByte is used to validate our protocol packets, immediately rejecting malformed data. +// It is set to 0xA1 as per the binary protocol specification. +const MagicByte byte = 0xA1 + +// HeaderSize defines the fixed size of the protocol header in bytes. const HeaderSize = 8 +// OpCode represents the type of operation in the binary protocol. type OpCode byte const ( @@ -19,6 +24,7 @@ const ( OpDelete OpCode = 0x04 ) +// RequestHeader represents the parsed 8-byte TCP packet header. type RequestHeader struct { Op OpCode KeyLen uint16 @@ -26,30 +32,34 @@ type RequestHeader struct { } var ( + // ErrInvalidMagic is returned when a packet does not start with the expected MagicByte. ErrInvalidMagic = errors.New("invalid magic byte") + // ErrShortBuffer is returned if the provided parsing buffer is smaller than HeaderSize. ErrShortBuffer = errors.New("provided buffer is too short for header") ) // Storage defines the interface for the database engine to prevent cyclic dependencies. +// It mandates the core LSM-Tree operations required by the network handlers. type Storage interface { Put(key, val []byte) error Get(key []byte) ([]byte, bool, error) Delete(key []byte) error } +// ParseHeader reads exactly HeaderSize bytes from the io.Reader into the provided buffer +// and decodes them into a RequestHeader. The buffer must be pre-allocated to avoid heap allocations. func ParseHeader(r io.Reader, buf []byte) (RequestHeader, error) { if len(buf) < HeaderSize { return RequestHeader{}, ErrShortBuffer } headerBuf := buf[:HeaderSize] - _, err := io.ReadFull(r, headerBuf) - if err != nil { - return RequestHeader{}, err + if _, err := io.ReadFull(r, headerBuf); err != nil { + return RequestHeader{}, err // Do not wrap EOF so the caller can handle client disconnects } if headerBuf[0] != MagicByte { - return RequestHeader{}, ErrInvalidMagic + return RequestHeader{}, fmt.Errorf("%w: expected 0x%X, got 0x%X", ErrInvalidMagic, MagicByte, headerBuf[0]) } return RequestHeader{ diff --git a/internal/network/server.go b/internal/network/server.go index f5c94d0..7b6ada6 100644 --- a/internal/network/server.go +++ b/internal/network/server.go @@ -1,72 +1,98 @@ +// Package network implements the zero-copy, binary TCP protocol for the Forest engine. +// It is designed to read packets directly from socket buffers to eliminate heap allocations +// and Garbage Collection (GC) pauses during high-throughput workloads. package network import ( + "errors" + "fmt" "log" "net" ) +// Server handles incoming TCP connections and routes them to the storage engine. type Server struct { - addr string - db Storage - l net.Listener + addr string + db Storage + listener net.Listener } +// NewServer initializes a new TCP server bound to the given address. func NewServer(addr string, db Storage) *Server { - return &Server{addr: addr, db: db} + return &Server{ + addr: addr, + db: db, + } } +// Start opens the TCP listener and begins accepting client connections. +// This method blocks until the server is stopped or encounters a fatal error. func (s *Server) Start() error { - l, err := net.Listen("tcp", s.addr) + listener, err := net.Listen("tcp", s.addr) if err != nil { - return err + return fmt.Errorf("failed to bind TCP listener on %s: %w", s.addr, err) } - s.l = l + + s.listener = listener log.Printf("Forest Database Engine TCP server listening on %s\n", s.addr) for { - conn, err := s.l.Accept() + conn, err := s.listener.Accept() if err != nil { - // If the server was stopped, the listener is closed and we should gracefully exit - if err == net.ErrClosed || err.Error() == "use of closed network connection" { - return nil - } - // In Go 1.16+, net.ErrClosed is the standard, but we fallback to string matching just in case - if err != nil && (err.Error() == "use of closed network connection" || err.Error() == "accept tcp 127.0.0.1:9002: use of closed network connection") { + // If the server was stopped, the listener is closed and we should gracefully exit. + if errors.Is(err, net.ErrClosed) || isClosedNetworkError(err) { return nil } - // For all other errors, we can log and return (or continue if transient) - return err + // For all other errors, we return which acts as a fatal error for the server loop. + return fmt.Errorf("failed to accept TCP connection: %w", err) } + + // Handle each client concurrently. go s.handleConnection(conn) } } +// isClosedNetworkError is a fallback helper to detect closed listeners. +func isClosedNetworkError(err error) bool { + if err == nil { + return false + } + return err.Error() == "use of closed network connection" || err.Error() == "accept tcp 127.0.0.1:9002: use of closed network connection" +} + +// Stop closes the active TCP listener, signaling the Start loop to exit gracefully. func (s *Server) Stop() error { - if s.l != nil { - return s.l.Close() + if s.listener != nil { + if err := s.listener.Close(); err != nil { + return fmt.Errorf("failed to close TCP listener: %w", err) + } } return nil } +// handleConnection manages the lifecycle of a single client TCP socket. +// It utilizes pre-allocated buffers to maintain a zero-allocation hot path. func (s *Server) handleConnection(conn net.Conn) { defer conn.Close() - // Pre-allocate buffer to strictly avoid allocations in the hot loop + // Pre-allocate buffers for the connection lifecycle to strictly avoid + // heap allocations (mallocgc) in the hot request/response loop. headerBuf := make([]byte, HeaderSize) - payloadBuf := make([]byte, 4*1024*1024) // 4MB max payload for simplicity + // 4MB max payload for simplicity, matching the MemTable flush threshold. + payloadBuf := make([]byte, 4*1024*1024) for { header, err := ParseHeader(conn, headerBuf) if err != nil { - if err.Error() == "EOF" { - return // Client disconnected normally - } - log.Printf("Error parsing header: %v", err) - return + // Expected client disconnects (EOF) or short reads terminate the session cleanly. + // Malformed packets also drop the connection. We avoid logging in the hot path + // to prevent IO bottlenecks during heavy loads or chaos testing. + return } if err := HandleRequest(conn, header, payloadBuf, s.db); err != nil { - log.Printf("Error handling request: %v", err) + // Log critical storage/engine errors, but avoid logging routine connection drops. + log.Printf("critical: failed to handle request for operation %v: %v", header.Op, err) return } } From 8c525b395be1011d163bc80f7651b3f61aca2ce0 Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar Date: Wed, 19 Aug 2026 21:10:26 +0530 Subject: [PATCH 3/9] Harden CI and WAL/Bloom test correctness Adds `go vet ./...` to the build workflow to catch idiom and bug-prone issues earlier, and introduces a Go-focused `.gitignore` for binaries, coverage outputs, workspace files, and local WAL/SST artifacts. Updates engine tests to match current APIs/semantics: Bloom decode now checks returned errors, WAL tests use `OpPut` for realistic recovery paths, and corruption recovery assertions are tightened around torn-write/CRC behavior. The chaos script is also aligned with the current wire protocol by switching the request magic byte to `0xA1`. --- .github/workflows/build.yml | 5 +++++ .gitignore | 24 ++++++++++++++++++++++++ internal/engine/bloom_test.go | 7 +++++-- internal/engine/wal_test.go | 31 +++++++++++++++---------------- scripts/chaos_test.sh | 4 ++-- 5 files changed, 51 insertions(+), 20 deletions(-) create mode 100644 .gitignore diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 72d720f..455a57a 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -27,6 +27,11 @@ jobs: go-version: '1.22' cache: false + - name: Run Go Vet + # Enforce standard Go idioms and detect subtle concurrency/shadowing bugs + run: go vet ./... + shell: bash + - name: Run Unit & Race Tests # The -race flag ensures your lock-free atomic pointers are safe # Note: We omit -race and set CGO_ENABLED=0 on macOS due to a known Apple dyld LC_UUID linker bug diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..72f70b1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +# Binaries for programs and plugins +*.exe +*.exe~ +*.dll +*.so +*.dylib + +# Test binary, built with `go test -c` +*.test +network.test.exe + +# Output of the go coverage tool, specifically when used with LiteIDE +*.out +cpu.out + +# Dependency directories (remove the comment below to include it) +# vendor/ + +# Go workspace file +go.work + +# Local DB files +*.wal +*.sst diff --git a/internal/engine/bloom_test.go b/internal/engine/bloom_test.go index a675bab..d6d0b83 100644 --- a/internal/engine/bloom_test.go +++ b/internal/engine/bloom_test.go @@ -63,10 +63,13 @@ func TestBloomFilter_EncodeDecode(t *testing.T) { bf.Add([]byte("test_key")) encoded := bf.Encode() - decoded := DecodeBloomFilter(encoded) + decoded, err := DecodeBloomFilter(encoded) + if err != nil { + t.Fatalf("failed to decode bloom filter: %v", err) + } if decoded == nil { - t.Fatalf("failed to decode bloom filter") + t.Fatalf("failed to decode bloom filter: returned nil without error") } if decoded.k != bf.k { t.Errorf("k mismatch: got %d, want %d", decoded.k, bf.k) diff --git a/internal/engine/wal_test.go b/internal/engine/wal_test.go index 74f680a..e817f20 100644 --- a/internal/engine/wal_test.go +++ b/internal/engine/wal_test.go @@ -10,6 +10,8 @@ import ( "github.com/Forest_DatabaseEngine/internal/network" ) +// TestWAL_AppendAndRecover verifies that standard Put operations survive a crash +// by persisting to the WAL and correctly reconstructing the MemTable. func TestWAL_AppendAndRecover(t *testing.T) { tempDir := t.TempDir() walPath := filepath.Join(tempDir, "test.wal") @@ -22,10 +24,10 @@ func TestWAL_AppendAndRecover(t *testing.T) { key1, val1 := []byte("key1"), []byte("val1") key2, val2 := []byte("key2"), []byte("val2") - _ = wal.Append(network.OpEcho, key1, val1) - _ = wal.Append(network.OpEcho, key2, val2) + _ = wal.Append(network.OpPut, key1, val1) + _ = wal.Append(network.OpPut, key2, val2) - // Wait for background syncer to flush + // Wait for background syncer to fsync (mechanical sympathy test) time.Sleep(50 * time.Millisecond) err = wal.Close() @@ -39,7 +41,7 @@ func TestWAL_AppendAndRecover(t *testing.T) { t.Fatalf("failed to recover WAL: %v", err) } - // Verify MemTable + // Verify MemTable lock-free reconstruction gotVal1, ok := mt.Get(key1) if !ok || !bytes.Equal(gotVal1, val1) { t.Errorf("expected val1, got %s", string(gotVal1)) @@ -50,6 +52,8 @@ func TestWAL_AppendAndRecover(t *testing.T) { } } +// TestWAL_CorruptionRecovery proves that torn writes (simulated by byte flipping) +// are safely detected via CRC checksums, halting recovery cleanly without panicking. func TestWAL_CorruptionRecovery(t *testing.T) { tempDir := t.TempDir() walPath := filepath.Join(tempDir, "corrupt.wal") @@ -60,37 +64,32 @@ func TestWAL_CorruptionRecovery(t *testing.T) { } key, val := []byte("corrupt_key"), []byte("corrupt_val") - _ = wal.Append(network.OpEcho, key, val) + _ = wal.Append(network.OpPut, key, val) - // Wait for background syncer to flush time.Sleep(50 * time.Millisecond) _ = wal.Close() - // Corrupt the file manually by flipping a byte at the end of the file data, err := os.ReadFile(walPath) if err != nil { t.Fatalf("failed to read WAL: %v", err) } - // Flip the last byte + // Flip the last byte to invalidate the CRC checksum of the final record data[len(data)-1] ^= 0xFF - err = os.WriteFile(walPath, data, 0644) - if err != nil { + if err := os.WriteFile(walPath, data, 0644); err != nil { t.Fatalf("failed to write corrupted WAL: %v", err) } mt := NewMemTable() // Recovery should succeed but silently stop before the corrupted record - // (as requested: "treat it as a torn write (EOF) and stop reading"). - err = Recover(walPath, mt) - if err != nil { - t.Fatalf("recovery returned an error instead of stopping: %v", err) + // (as it is treated as a torn write / EOF). + if err := Recover(walPath, mt); err != nil { + t.Fatalf("recovery returned an error instead of stopping at EOF/Torn Write: %v", err) } // The record should NOT be in the MemTable - _, ok := mt.Get(key) - if ok { + if _, ok := mt.Get(key); ok { t.Errorf("record should not have been recovered due to CRC mismatch") } } diff --git a/scripts/chaos_test.sh b/scripts/chaos_test.sh index 08917e7..e565d3c 100644 --- a/scripts/chaos_test.sh +++ b/scripts/chaos_test.sh @@ -39,7 +39,7 @@ func main() { val := []byte(fmt.Sprintf("val-%d", i)) var header [8]byte - header[0] = 0xFE + header[0] = 0xA1 // MagicByte header[1] = 0x02 // OpPut binary.BigEndian.PutUint16(header[2:4], uint16(len(key))) binary.BigEndian.PutUint32(header[4:8], uint32(len(val))) @@ -92,7 +92,7 @@ func main() { key := []byte("key-0") // Assuming key-0 was written before kill var header [8]byte - header[0] = 0xFE + header[0] = 0xA1 // MagicByte header[1] = 0x03 // OpGet binary.BigEndian.PutUint16(header[2:4], uint16(len(key))) binary.BigEndian.PutUint32(header[4:8], 0) From 7c89a1c3e85d714eb3a818c65b13de9709132c89 Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar <47363219+Awesome06@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:11:36 +0530 Subject: [PATCH 4/9] Add MIT License to the project --- LICENSE | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 LICENSE diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..f06f10b --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Mrigank Bhatnagar + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. From 0d53d018214677df534deaca5cf4dc922c0b1fb4 Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar Date: Wed, 19 Aug 2026 21:12:17 +0530 Subject: [PATCH 5/9] Create Makefile --- Makefile | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 Makefile diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..a3f36af --- /dev/null +++ b/Makefile @@ -0,0 +1,24 @@ +.PHONY: build test bench chaos clean + +# Build the Forest Database Engine server +build: + go build -o forest_server ./cmd/server + +# Run the standard unit test suite with race detection +test: + go vet ./... + go test -v -race ./... + +# Run the performance benchmarks to profile zero-allocation read/write paths +bench: + go test -bench=. -benchmem ./internal/network/... + +# Execute the Chaos Engineering test script +chaos: build + bash ./scripts/chaos_test.sh + +# Clean up build artifacts and database files +clean: + rm -f forest_server + rm -f *.wal *.sst + rm -f cpu.out network.test.exe From 1c0a49bb370f52a4c6d4b2256acd7103531f9823 Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar Date: Wed, 19 Aug 2026 21:16:24 +0530 Subject: [PATCH 6/9] Use errors.Is in protocol header test Update `TestParseHeader` to check expected error types with `errors.Is` instead of direct equality. This makes the test resilient to wrapped errors while preserving the same failure expectations. --- internal/network/protocol_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/internal/network/protocol_test.go b/internal/network/protocol_test.go index ad7f1ff..93b45ea 100644 --- a/internal/network/protocol_test.go +++ b/internal/network/protocol_test.go @@ -3,6 +3,7 @@ package network import ( "bytes" "encoding/binary" + "errors" "io" "net" "reflect" @@ -69,7 +70,7 @@ func TestParseHeader(t *testing.T) { t.Errorf("ParseHeader() error = %v, wantErr %v", err, tt.wantErr) return } - if tt.wantErr && tt.errType != nil && err != tt.errType { + if tt.wantErr && tt.errType != nil && !errors.Is(err, tt.errType) { t.Errorf("ParseHeader() error = %v, want errType %v", err, tt.errType) } if !reflect.DeepEqual(got, tt.want) { From fb41357fdf0184619dd0eadec505303cfb72da9d Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar Date: Wed, 19 Aug 2026 21:30:49 +0530 Subject: [PATCH 7/9] Harden writes and memtable flush concurrency Adds stronger safety checks across the engine and network path. MemTables are now marked frozen during flush rotation, and DB Put/Delete use a lock-free retry loop so concurrent writes move to the new active table instead of writing into a table being flushed. SSTable/compaction write paths now propagate header/key/value write errors instead of ignoring them. The request handler also validates payload length against the buffer to prevent out-of-bounds reads on malformed input. --- internal/engine/compactor.go | 27 +++++++++++++++++++++------ internal/engine/db.go | 36 ++++++++++++++++++++++++++---------- internal/engine/memtable.go | 11 +++++++++-- internal/engine/sstable.go | 17 ++++++++++++++--- internal/network/handler.go | 5 +++++ 5 files changed, 75 insertions(+), 21 deletions(-) diff --git a/internal/engine/compactor.go b/internal/engine/compactor.go index ac96b91..2473f38 100644 --- a/internal/engine/compactor.go +++ b/internal/engine/compactor.go @@ -130,7 +130,11 @@ func CompactL0toL1(l0Files []string, l1Filepath string) error { // Note: We intentionally write tombstones (OpDelete) into L1 so they can continue // to shadow values in L2+ if multi-level compaction is implemented later. if !bytes.Equal(lastKey, item.Key) { - currentOffset, lastIndexOffset = writeSSTableEntry(bufWriter, bf, &index, currentOffset, lastIndexOffset, item) + var err error + currentOffset, lastIndexOffset, err = writeSSTableEntry(bufWriter, bf, &index, currentOffset, lastIndexOffset, item) + if err != nil { + return fmt.Errorf("compaction failed to write entry: %w", err) + } lastKey = append([]byte(nil), item.Key...) // copy to avoid memory aliasing } @@ -177,7 +181,7 @@ func initSSTableWriter(file *os.File) (*bufio.Writer, *BloomFilter, []IndexEntry } // writeSSTableEntry serializes a single merged record into the data block phase of the SSTable. -func writeSSTableEntry(bufWriter *bufio.Writer, bf *BloomFilter, index *[]IndexEntry, currentOffset, lastIndexOffset uint64, item HeapItem) (uint64, uint64) { +func writeSSTableEntry(bufWriter *bufio.Writer, bf *BloomFilter, index *[]IndexEntry, currentOffset, lastIndexOffset uint64, item HeapItem) (uint64, uint64, error) { bf.Add(item.Key) if currentOffset-lastIndexOffset >= 4096 || currentOffset == 0 { @@ -190,14 +194,25 @@ func writeSSTableEntry(bufWriter *bufio.Writer, bf *BloomFilter, index *[]IndexE binary.BigEndian.PutUint16(header[1:3], uint16(len(item.Key))) binary.BigEndian.PutUint32(header[3:7], uint32(len(item.Value))) - n, _ := bufWriter.Write(header[:]) + n, err := bufWriter.Write(header[:]) + if err != nil { + return 0, 0, fmt.Errorf("failed to write header: %w", err) + } currentOffset += uint64(n) - n, _ = bufWriter.Write(item.Key) + + n, err = bufWriter.Write(item.Key) + if err != nil { + return 0, 0, fmt.Errorf("failed to write key: %w", err) + } currentOffset += uint64(n) - n, _ = bufWriter.Write(item.Value) + + n, err = bufWriter.Write(item.Value) + if err != nil { + return 0, 0, fmt.Errorf("failed to write value: %w", err) + } currentOffset += uint64(n) - return currentOffset, lastIndexOffset + return currentOffset, lastIndexOffset, nil } // finalizeSSTable appends the trailing metadata layers to seal the SSTable. diff --git a/internal/engine/db.go b/internal/engine/db.go index ecbb038..2b2fb13 100644 --- a/internal/engine/db.go +++ b/internal/engine/db.go @@ -135,13 +135,20 @@ func (db *DB) Put(key, val []byte) error { return fmt.Errorf("engine failed to append Put to WAL: %w", err) } - active := db.activeMT.Load() - active.Put(key, val) + // Lock-free retry loop + for { + active := db.activeMT.Load() + err := active.Put(key, val) + if err != nil { + // Table was frozen by a concurrent CheckFlush, retry to get the new table + continue + } - if active.Size() >= db.flushThreshold { - db.CheckFlush() + if active.Size() >= db.flushThreshold { + db.CheckFlush() + } + return nil } - return nil } // Delete appends a tombstone to the WAL and inserts it into the active MemTable lock-free. @@ -150,13 +157,20 @@ func (db *DB) Delete(key []byte) error { return fmt.Errorf("engine failed to append Delete to WAL: %w", err) } - active := db.activeMT.Load() - active.Put(key, nil) + // Lock-free retry loop + for { + active := db.activeMT.Load() + err := active.Put(key, nil) + if err != nil { + // Table was frozen by a concurrent CheckFlush, retry + continue + } - if active.Size() >= db.flushThreshold { - db.CheckFlush() + if active.Size() >= db.flushThreshold { + db.CheckFlush() + } + return nil } - return nil } // CheckFlush safely rotates the active MemTable into the flush queue if it exceeds the capacity threshold. @@ -176,6 +190,8 @@ func (db *DB) CheckFlush() { return } + active.frozen.Store(true) + // Atomically rotate the MemTables. db.immutableMT.Store(active) db.activeMT.Store(NewMemTable()) diff --git a/internal/engine/memtable.go b/internal/engine/memtable.go index 74f534d..3d6ca4b 100644 --- a/internal/engine/memtable.go +++ b/internal/engine/memtable.go @@ -2,6 +2,7 @@ package engine import ( "bytes" + "errors" "math/rand" "sync" "sync/atomic" @@ -34,6 +35,7 @@ type MemTable struct { // mu serializes writers. Readers completely bypass this lock. mu sync.RWMutex size atomic.Int64 + frozen atomic.Bool } // NewMemTable initializes and returns an empty MemTable with a dummy head node. @@ -56,7 +58,11 @@ func randomLevel() int { // Put inserts or updates a key-value pair in the MemTable. // It locks the structure to prevent write-write conflicts but relies on atomic // operations so read-write conflicts are handled lock-free. -func (m *MemTable) Put(key, val []byte) { +func (m *MemTable) Put(key, val []byte) error { + if m.frozen.Load() { + return errors.New("memtable is frozen") + } + // Serialize writers to prevent structural corruption of the SkipList. m.mu.Lock() defer m.mu.Unlock() @@ -91,7 +97,7 @@ func (m *MemTable) Put(key, val []byte) { // Note: Approximate size is not adjusted on updates to avoid heavy calculation // on the hot path. The 4MB threshold is a soft limit. - return + return nil } // Key does not exist: create a new node and determine its height. @@ -122,6 +128,7 @@ func (m *MemTable) Put(key, val []byte) { for i := 0; i < level; i++ { preds[i].nexts[i].Store(newNode) } + return nil } // Get performs a lock-free search across the SkipList. diff --git a/internal/engine/sstable.go b/internal/engine/sstable.go index bfa0014..845052a 100644 --- a/internal/engine/sstable.go +++ b/internal/engine/sstable.go @@ -70,11 +70,22 @@ func FlushMemTable(mt *MemTable, filepath string) error { binary.BigEndian.PutUint16(header[1:3], uint16(len(key))) binary.BigEndian.PutUint32(header[3:7], uint32(len(val))) - n, _ := bufWriter.Write(header[:]) + n, err := bufWriter.Write(header[:]) + if err != nil { + return fmt.Errorf("failed to write header: %w", err) + } currentOffset += uint64(n) - n, _ = bufWriter.Write(key) + + n, err = bufWriter.Write(key) + if err != nil { + return fmt.Errorf("failed to write key: %w", err) + } currentOffset += uint64(n) - n, _ = bufWriter.Write(val) + + n, err = bufWriter.Write(val) + if err != nil { + return fmt.Errorf("failed to write value: %w", err) + } currentOffset += uint64(n) curr = curr.nexts[0].Load() diff --git a/internal/network/handler.go b/internal/network/handler.go index b38dfd6..7b43709 100644 --- a/internal/network/handler.go +++ b/internal/network/handler.go @@ -12,6 +12,11 @@ import ( func HandleRequest(conn net.Conn, header RequestHeader, payloadBuf []byte, db Storage) error { totalLen := int(header.KeyLen) + int(header.ValueLen) + // Add this explicit bounds check + if totalLen > len(payloadBuf) { + return fmt.Errorf("protocol error: payload length %d exceeds buffer capacity %d", totalLen, len(payloadBuf)) + } + // Read the variable length payload directly into the pre-allocated buffer if totalLen > 0 { if _, err := io.ReadFull(conn, payloadBuf[:totalLen]); err != nil { From 7330a263ac0829db81e9c07e2c56f4dee1454280 Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar Date: Wed, 19 Aug 2026 21:55:21 +0530 Subject: [PATCH 8/9] Harden memtable freeze and SSTable writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tightens engine durability and concurrency behavior by surfacing write failures during SSTable/compaction finalization and by failing on truncated bloom-filter payloads instead of silently accepting partial data. It also fixes memtable flush rotation safety by introducing a lock-based `Freeze()` path, re-checking frozen state under lock in `Put`, and copying inserted keys to avoid aliasing issues. Tests and WAL recovery call sites were updated to handle `Put`’s error return explicitly. --- internal/engine/bloom.go | 2 +- internal/engine/compactor.go | 22 ++++++++++++++++------ internal/engine/db.go | 3 ++- internal/engine/memtable.go | 20 +++++++++++++++++++- internal/engine/memtable_test.go | 20 ++++++++++---------- internal/engine/sstable.go | 22 ++++++++++++++++------ internal/engine/wal.go | 4 ++-- 7 files changed, 66 insertions(+), 27 deletions(-) diff --git a/internal/engine/bloom.go b/internal/engine/bloom.go index 16e2621..8b433b3 100644 --- a/internal/engine/bloom.go +++ b/internal/engine/bloom.go @@ -125,7 +125,7 @@ func DecodeBloomFilter(data []byte) (*BloomFilter, error) { offset := 9 for i := 0; i < int(numUint64s); i++ { if offset+8 > len(data) { - break // Allow partial reads in case of corruption padding. + return nil, fmt.Errorf("bloom filter data truncated: missing bytes") } bits[i] = binary.BigEndian.Uint64(data[offset : offset+8]) offset += 8 diff --git a/internal/engine/compactor.go b/internal/engine/compactor.go index 2473f38..1ae33df 100644 --- a/internal/engine/compactor.go +++ b/internal/engine/compactor.go @@ -221,28 +221,38 @@ func finalizeSSTable(file *os.File, bufWriter *bufio.Writer, bf *BloomFilter, in var countBuf [4]byte binary.BigEndian.PutUint32(countBuf[:], uint32(len(index))) - n, _ := bufWriter.Write(countBuf[:]) + n, err := bufWriter.Write(countBuf[:]) + if err != nil { + return fmt.Errorf("failed to write index count: %w", err) + } currentOffset += uint64(n) for _, entry := range index { var idxHeader [10]byte binary.BigEndian.PutUint16(idxHeader[0:2], uint16(len(entry.Key))) binary.BigEndian.PutUint64(idxHeader[2:10], entry.Offset) - n1, _ := bufWriter.Write(idxHeader[0:2]) - n2, _ := bufWriter.Write(entry.Key) - n3, _ := bufWriter.Write(idxHeader[2:10]) + n1, err := bufWriter.Write(idxHeader[0:2]) + if err != nil { return fmt.Errorf("failed to write index header: %w", err) } + n2, err := bufWriter.Write(entry.Key) + if err != nil { return fmt.Errorf("failed to write index key: %w", err) } + n3, err := bufWriter.Write(idxHeader[2:10]) + if err != nil { return fmt.Errorf("failed to write index offset: %w", err) } currentOffset += uint64(n1 + n2 + n3) } bloomFilterOffset := currentOffset bloomData := bf.Encode() - bufWriter.Write(bloomData) + if _, err := bufWriter.Write(bloomData); err != nil { + return fmt.Errorf("failed to write bloom filter: %w", err) + } var footer [SSTableFooterSize]byte binary.BigEndian.PutUint64(footer[0:8], indexBlockOffset) binary.BigEndian.PutUint64(footer[8:16], bloomFilterOffset) binary.BigEndian.PutUint32(footer[16:20], SSTableMagic) - bufWriter.Write(footer[:]) + if _, err := bufWriter.Write(footer[:]); err != nil { + return fmt.Errorf("failed to write footer: %w", err) + } if err := bufWriter.Flush(); err != nil { return err diff --git a/internal/engine/db.go b/internal/engine/db.go index 2b2fb13..4b3143c 100644 --- a/internal/engine/db.go +++ b/internal/engine/db.go @@ -190,7 +190,8 @@ func (db *DB) CheckFlush() { return } - active.frozen.Store(true) + // Drain in-flight writers and mark the active table as frozen. + active.Freeze() // Atomically rotate the MemTables. db.immutableMT.Store(active) diff --git a/internal/engine/memtable.go b/internal/engine/memtable.go index 3d6ca4b..630f54d 100644 --- a/internal/engine/memtable.go +++ b/internal/engine/memtable.go @@ -67,6 +67,12 @@ func (m *MemTable) Put(key, val []byte) error { m.mu.Lock() defer m.mu.Unlock() + // Re-check frozen status under lock in case a concurrent CheckFlush froze it + // while we were waiting for m.mu.Lock(). + if m.frozen.Load() { + return errors.New("memtable is frozen") + } + var preds [maxLevel]*Node curr := m.head @@ -102,8 +108,12 @@ func (m *MemTable) Put(key, val []byte) error { // Key does not exist: create a new node and determine its height. level := randomLevel() + + keyCopy := make([]byte, len(key)) + copy(keyCopy, key) + newNode := &Node{ - key: key, + key: keyCopy, } var newVal []byte @@ -181,3 +191,11 @@ func (m *MemTable) Get(key []byte) ([]byte, bool) { func (m *MemTable) Size() int64 { return m.size.Load() } + +// Freeze marks the MemTable as immutable. It acquires the write lock to ensure +// all in-flight writers have completed their inserts before it returns. +func (m *MemTable) Freeze() { + m.mu.Lock() + m.frozen.Store(true) + m.mu.Unlock() +} diff --git a/internal/engine/memtable_test.go b/internal/engine/memtable_test.go index 3fc41a2..291b03a 100644 --- a/internal/engine/memtable_test.go +++ b/internal/engine/memtable_test.go @@ -12,9 +12,9 @@ func TestMemTable_BasicOperations(t *testing.T) { mt := NewMemTable() // Put some keys - mt.Put([]byte("apple"), []byte("red")) - mt.Put([]byte("banana"), []byte("yellow")) - mt.Put([]byte("cherry"), []byte("red")) + _ = mt.Put([]byte("apple"), []byte("red")) + _ = mt.Put([]byte("banana"), []byte("yellow")) + _ = mt.Put([]byte("cherry"), []byte("red")) // Get existing keys val, ok := mt.Get([]byte("apple")) @@ -34,7 +34,7 @@ func TestMemTable_BasicOperations(t *testing.T) { } // Update existing key - mt.Put([]byte("apple"), []byte("green")) + _ = mt.Put([]byte("apple"), []byte("green")) val, ok = mt.Get([]byte("apple")) if !ok || string(val) != "green" { t.Errorf("expected apple=green, got %s", string(val)) @@ -45,11 +45,11 @@ func TestMemTable_SortedOrder(t *testing.T) { mt := NewMemTable() // Insert in reverse order - mt.Put([]byte("e"), []byte("5")) - mt.Put([]byte("d"), []byte("4")) - mt.Put([]byte("c"), []byte("3")) - mt.Put([]byte("b"), []byte("2")) - mt.Put([]byte("a"), []byte("1")) + _ = mt.Put([]byte("e"), []byte("5")) + _ = mt.Put([]byte("d"), []byte("4")) + _ = mt.Put([]byte("c"), []byte("3")) + _ = mt.Put([]byte("b"), []byte("2")) + _ = mt.Put([]byte("a"), []byte("1")) // Traverse level 0 and ensure sorted order curr := mt.head.nexts[0].Load() @@ -91,7 +91,7 @@ func TestMemTable_Concurrency(t *testing.T) { // Randomly Put or Get if rand.Float32() < 0.5 { - mt.Put(key, val) + _ = mt.Put(key, val) } else { mt.Get(key) } diff --git a/internal/engine/sstable.go b/internal/engine/sstable.go index 845052a..6a1447b 100644 --- a/internal/engine/sstable.go +++ b/internal/engine/sstable.go @@ -97,16 +97,22 @@ func FlushMemTable(mt *MemTable, filepath string) error { // NumIndexEntries (4B) var countBuf [4]byte binary.BigEndian.PutUint32(countBuf[:], uint32(len(index))) - n, _ := bufWriter.Write(countBuf[:]) + n, err := bufWriter.Write(countBuf[:]) + if err != nil { + return fmt.Errorf("failed to write index count: %w", err) + } currentOffset += uint64(n) for _, entry := range index { var idxHeader [10]byte // KeyLen(2) + Offset(8) binary.BigEndian.PutUint16(idxHeader[0:2], uint16(len(entry.Key))) binary.BigEndian.PutUint64(idxHeader[2:10], entry.Offset) - n1, _ := bufWriter.Write(idxHeader[0:2]) - n2, _ := bufWriter.Write(entry.Key) - n3, _ := bufWriter.Write(idxHeader[2:10]) + n1, err := bufWriter.Write(idxHeader[0:2]) + if err != nil { return fmt.Errorf("failed to write index header: %w", err) } + n2, err := bufWriter.Write(entry.Key) + if err != nil { return fmt.Errorf("failed to write index key: %w", err) } + n3, err := bufWriter.Write(idxHeader[2:10]) + if err != nil { return fmt.Errorf("failed to write index offset: %w", err) } currentOffset += uint64(n1 + n2 + n3) } @@ -114,14 +120,18 @@ func FlushMemTable(mt *MemTable, filepath string) error { // 3. Write Bloom Filter bloomData := bf.Encode() - bufWriter.Write(bloomData) + if _, err := bufWriter.Write(bloomData); err != nil { + return fmt.Errorf("failed to write bloom filter: %w", err) + } // 4. Write Footer (Fixed Size) var footer [SSTableFooterSize]byte binary.BigEndian.PutUint64(footer[0:8], indexBlockOffset) binary.BigEndian.PutUint64(footer[8:16], bloomFilterOffset) binary.BigEndian.PutUint32(footer[16:20], SSTableMagic) - bufWriter.Write(footer[:]) + if _, err := bufWriter.Write(footer[:]); err != nil { + return fmt.Errorf("failed to write footer: %w", err) + } // Flush bufio to disk if err := bufWriter.Flush(); err != nil { diff --git a/internal/engine/wal.go b/internal/engine/wal.go index c71e8fc..80c79e4 100644 --- a/internal/engine/wal.go +++ b/internal/engine/wal.go @@ -177,9 +177,9 @@ func Recover(filepath string, mt *MemTable) error { // Apply the recovered state to the MemTable. if op == network.OpDelete { - mt.Put(key, nil) // Tombstone for delete + _ = mt.Put(key, nil) // Tombstone for delete (error is impossible, table not frozen) } else { - mt.Put(key, val) // Echo or Put + _ = mt.Put(key, val) // Echo or Put (error is impossible, table not frozen) } } } From 2592dbfe4bbf161261e269cf1ca989d18238b1d9 Mon Sep 17 00:00:00 2001 From: Mrigank Bhatnagar Date: Wed, 19 Aug 2026 22:04:54 +0530 Subject: [PATCH 9/9] Harden concurrency and error handling paths This change fixes race-prone lifecycle handling in both storage and networking layers. Version refs now use a CAS-based acquire/purge protocol (including a purge sentinel) so readers cannot attach to a version being deleted, and DB reads retry until a safe acquire succeeds. The TCP server now guards listener access with a mutex and uses a stable local listener in the accept loop to avoid Stop/Start races. It also improves reliability by surfacing or logging previously ignored I/O errors: memtable flush on shutdown, WAL fsync, WAL recovery writes, SSTable flush/fsync, and protocol header read failures (while preserving raw EOF behavior for disconnect handling). --- internal/engine/compactor.go | 15 ++++++++++++--- internal/engine/db.go | 13 ++++++++++--- internal/engine/version.go | 15 +++++++++++---- internal/engine/wal.go | 13 ++++++++++--- internal/network/protocol.go | 5 ++++- internal/network/server.go | 11 ++++++++++- 6 files changed, 57 insertions(+), 15 deletions(-) diff --git a/internal/engine/compactor.go b/internal/engine/compactor.go index 1ae33df..33b0523 100644 --- a/internal/engine/compactor.go +++ b/internal/engine/compactor.go @@ -162,7 +162,13 @@ func CompactL0toL1(l0Files []string, l1Filepath string) error { // are referencing the file before the OS unlinks it. func PurgeObsoleteFiles(v *Version, files []string) { // Wait until all active readers that acquired this version have released it. - for v.Refs.Load() > 0 { + for { + refs := v.Refs.Load() + if refs == 0 { + if v.Refs.CompareAndSwap(0, -1) { + break + } + } time.Sleep(10 * time.Millisecond) } @@ -255,7 +261,10 @@ func finalizeSSTable(file *os.File, bufWriter *bufio.Writer, bf *BloomFilter, in } if err := bufWriter.Flush(); err != nil { - return err + return fmt.Errorf("failed to flush buffer: %w", err) } - return file.Sync() + if err := file.Sync(); err != nil { + return fmt.Errorf("failed to fsync file: %w", err) + } + return nil } diff --git a/internal/engine/db.go b/internal/engine/db.go index 4b3143c..4a525f0 100644 --- a/internal/engine/db.go +++ b/internal/engine/db.go @@ -94,8 +94,13 @@ func (db *DB) Get(key []byte) ([]byte, bool, error) { // 3. Search RCU Version SSTables // Acquire a read reference to prevent files from being deleted during read. - version := db.currentVersion.Load() - version.Acquire() + var version *Version + for { + version = db.currentVersion.Load() + if version.Acquire() { + break + } + } defer version.Release() // Search L0 files (newer files are appended to the end, so we search backwards) @@ -240,7 +245,9 @@ func (db *DB) flushWorker() { close(db.flushQ) for mt := range db.flushQ { filepath := fmt.Sprintf("data-L0-%d.sst", flushCount) - FlushMemTable(mt, filepath) + if err := FlushMemTable(mt, filepath); err != nil { + log.Printf("critical: failed to flush MemTable to %s during shutdown: %v", filepath, err) + } flushCount++ } return diff --git a/internal/engine/version.go b/internal/engine/version.go index 026786e..0a38510 100644 --- a/internal/engine/version.go +++ b/internal/engine/version.go @@ -20,10 +20,17 @@ type Version struct { } // Acquire increments the reference count for this specific version. -// It must be called by readers (e.g., Get operations) before they begin traversing the SSTables -// to ensure the underlying files are not physically unlinked by a background compaction worker. -func (v *Version) Acquire() { - v.Refs.Add(1) +// It returns true if successful, or false if the version has already been marked for purge. +func (v *Version) Acquire() bool { + for { + refs := v.Refs.Load() + if refs < 0 { + return false + } + if v.Refs.CompareAndSwap(refs, refs+1) { + return true + } + } } // Release decrements the reference count. diff --git a/internal/engine/wal.go b/internal/engine/wal.go index 80c79e4..e53b30c 100644 --- a/internal/engine/wal.go +++ b/internal/engine/wal.go @@ -9,6 +9,7 @@ import ( "fmt" "hash/crc32" "io" + "log" "os" "sync" "time" @@ -112,7 +113,9 @@ func (w *WAL) flush() { if err == nil { // Sync is called outside the mutex so disk I/O latency does not block active Appends. - _ = w.file.Sync() + if syncErr := w.file.Sync(); syncErr != nil { + log.Printf("critical: failed to fsync WAL: %v", syncErr) + } } } @@ -177,9 +180,13 @@ func Recover(filepath string, mt *MemTable) error { // Apply the recovered state to the MemTable. if op == network.OpDelete { - _ = mt.Put(key, nil) // Tombstone for delete (error is impossible, table not frozen) + if err := mt.Put(key, nil); err != nil { + return fmt.Errorf("failed to put tombstone during recovery: %w", err) + } } else { - _ = mt.Put(key, val) // Echo or Put (error is impossible, table not frozen) + if err := mt.Put(key, val); err != nil { + return fmt.Errorf("failed to put value during recovery: %w", err) + } } } } diff --git a/internal/network/protocol.go b/internal/network/protocol.go index 8eceefb..d774e59 100644 --- a/internal/network/protocol.go +++ b/internal/network/protocol.go @@ -55,7 +55,10 @@ func ParseHeader(r io.Reader, buf []byte) (RequestHeader, error) { headerBuf := buf[:HeaderSize] if _, err := io.ReadFull(r, headerBuf); err != nil { - return RequestHeader{}, err // Do not wrap EOF so the caller can handle client disconnects + if errors.Is(err, io.EOF) { + return RequestHeader{}, err // Do not wrap EOF so the caller can handle client disconnects + } + return RequestHeader{}, fmt.Errorf("failed to read protocol header: %w", err) } if headerBuf[0] != MagicByte { diff --git a/internal/network/server.go b/internal/network/server.go index 7b6ada6..e850bf7 100644 --- a/internal/network/server.go +++ b/internal/network/server.go @@ -8,6 +8,7 @@ import ( "fmt" "log" "net" + "sync" ) // Server handles incoming TCP connections and routes them to the storage engine. @@ -15,6 +16,7 @@ type Server struct { addr string db Storage listener net.Listener + mu sync.Mutex } // NewServer initializes a new TCP server bound to the given address. @@ -33,11 +35,14 @@ func (s *Server) Start() error { return fmt.Errorf("failed to bind TCP listener on %s: %w", s.addr, err) } + s.mu.Lock() s.listener = listener + s.mu.Unlock() + log.Printf("Forest Database Engine TCP server listening on %s\n", s.addr) for { - conn, err := s.listener.Accept() + conn, err := listener.Accept() if err != nil { // If the server was stopped, the listener is closed and we should gracefully exit. if errors.Is(err, net.ErrClosed) || isClosedNetworkError(err) { @@ -62,10 +67,14 @@ func isClosedNetworkError(err error) bool { // Stop closes the active TCP listener, signaling the Start loop to exit gracefully. func (s *Server) Stop() error { + s.mu.Lock() + defer s.mu.Unlock() + if s.listener != nil { if err := s.listener.Close(); err != nil { return fmt.Errorf("failed to close TCP listener: %w", err) } + s.listener = nil } return nil }