go-kv/db.go

package go_kv

import (
	"fmt"
	"os"
	"sync"
	"sync/atomic"

	"github.com/dailz/go-kv/config"
	"github.com/dailz/go-kv/memtable"
	"github.com/dailz/go-kv/wal"
)

// DB is the public database handle for the storage engine.
type DB struct {
	cfg         *config.WalConfig
	dir         string
	writer      *wal.WalWriter
	recoveryMem *memtable.MemTable
	closed      atomic.Bool
}

// Open opens a database directory, recovers durable WAL entries, and starts a
// fresh WAL writer at the next recovered segment and sequence.
func Open(dir string, cfg *config.WalConfig) (*DB, error) {
	if cfg == nil {
		defaults := config.Defaults()
		cfg = &defaults
	}
	if err := cfg.Validate(); err != nil {
		return nil, fmt.Errorf("db: validate config: %w", err)
	}
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, fmt.Errorf("db: create directory: %w", err)
	}

	recoveryMem := memtable.NewMemTable(cfg.MemTableSize)
	replayer := &memTableReplayer{mem: recoveryMem}
	result, err := wal.Recover(dir, replayer)
	if err != nil {
		return nil, fmt.Errorf("db: recover WAL: %w", err)
	}
	if err := replayer.Err(); err != nil {
		return nil, fmt.Errorf("db: replay WAL: %w", err)
	}
	recoveryMem.Publish(^uint64(0))

	writer, err := wal.NewWalWriter(cfg, dir, result.NextSegmentID, result.NextSequence)
	if err != nil {
		return nil, fmt.Errorf("db: create WAL writer: %w", err)
	}

	return &DB{
		cfg:         cfg,
		dir:         dir,
		writer:      writer,
		recoveryMem: recoveryMem,
	}, nil
}

// Close drains outstanding writes and releases database resources.
func (db *DB) Close() error {
	if db.closed.Swap(true) {
		return nil
	}
	if err := db.writer.Close(); err != nil {
		return fmt.Errorf("db: close WAL writer: %w", err)
	}
	return nil
}

// Put stores key with value.
func (db *DB) Put(key, value []byte) error {
	if db.closed.Load() {
		return ErrWriteStopped
	}
	if err := db.writer.Put(key, value); err != nil {
		return fmt.Errorf("db: put: %w", err)
	}
	return nil
}

// Delete removes key.
func (db *DB) Delete(key []byte) error {
	if db.closed.Load() {
		return ErrWriteStopped
	}
	if err := db.writer.Delete(key); err != nil {
		return fmt.Errorf("db: delete: %w", err)
	}
	return nil
}

// Get returns the newest visible value for key across recovered and new writes.
func (db *DB) Get(key []byte) *memtable.GetResult {
	recovered := db.recoveryMem.Get(key)
	current := db.writer.Get(key)

	if current.Found || current.Sequence > recovered.Sequence {
		return current
	}
	if !recovered.Found {
		return &memtable.GetResult{Found: false}
	}
	if recovered.Value == nil {
		return &memtable.GetResult{Found: false, Sequence: recovered.Sequence}
	}
	return recovered
}

// GetDurableSequence returns the highest sequence durably fsynced to WAL.
func (db *DB) GetDurableSequence() uint64 {
	return db.writer.GetDurableSequence()
}

// IsWriteStopped reports whether the DB rejects new writes.
func (db *DB) IsWriteStopped() bool {
	return db.closed.Load() || db.writer.IsWriteStopped()
}

type memTableReplayer struct {
	mem *memtable.MemTable
	mu  sync.Mutex
	err error
}

func (r *memTableReplayer) ReplayPut(key, value []byte, sequence uint64) {
	if err := r.mem.PutPending(cloneBytes(key), cloneBytes(value), sequence); err != nil {
		r.setErr(fmt.Errorf("put pending: %w", err))
	}
}

func (r *memTableReplayer) ReplayDelete(key []byte, sequence uint64) {
	if err := r.mem.DeletePending(cloneBytes(key), sequence); err != nil {
		r.setErr(fmt.Errorf("delete pending: %w", err))
	}
}

func (r *memTableReplayer) Err() error {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.err
}

func (r *memTableReplayer) setErr(err error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.err == nil {
		r.err = err
	}
}

func cloneBytes(src []byte) []byte {
	if src == nil {
		return nil
	}
	dst := make([]byte, len(src))
	copy(dst, src)
	return dst
}