kv.go

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
	"path"
	"sync"
	"syscall"
)

// The signature identifies database file format
const DB_SIG = "GopherSQL"

const (
	FLAG_DELETED = byte(1)  // key was deleted in this transaction
	FLAG_UPDATED = byte(2)  // key was inserted or updated
)

// KeyRange represents the start and stop boundaries of a read or write dependency
type KeyRange struct {
	start []byte
	stop  []byte 
}

// CommittedTX represents a successfully completed transaction in the database's history
type CommittedTX struct {
	version uint64
	writes  []KeyRange
}

// KV represents the database instance
type KV struct {
	Path string
	fd   int               
	tree BTree
	free FreeList

	mmap struct {
		total  int        
		chunks [][]byte   
	}
	
	page struct {
		flushed uint64     
		nappend int
		updates map[uint64][]byte
	}
	
	failed bool 

	// --- Concurrency Fields ---
	mutex   sync.Mutex    // Serializes Begin, Commit, Abort
	version uint64        // Monotonic database version (persisted)
	ongoing []uint64      // Version numbers of active transactions
	history []CommittedTX // Log of recent commits for conflict detection
}

// KVTX represents a single database transaction.
type KVTX struct {
	db       *KV
	version  uint64
	snapshot BTree    // Read-only view of the database
	pending  BTree    // Private, in-memory tree for this transaction's writes
	reads    []KeyRange // Dependency tracking for Conflict Detection
	writes   []KeyRange // Log of keys changed by this transaction
	pages    [][]byte   // Temporary storage for pending pages
}

// --- CONCURRENCY CONTROL ---

func rangesOverlap(reads []KeyRange, writes []KeyRange) bool {
	for _, r := range reads {
		for _, w := range writes {
			if bytes.Compare(r.start, w.stop) <= 0 && bytes.Compare(r.stop, w.start) >= 0 {
				return true
			}
		}
	}
	return false
}

func detectConflicts(kv *KV, tx *KVTX) bool {
	for i := len(kv.history) - 1; i >= 0; i-- {
		if kv.history[i].version <= tx.version {
			break 
		}
		if rangesOverlap(tx.reads, kv.history[i].writes) {
			return true
		}
	}
	return false
}

func (kv *KV) removeOngoing(version uint64) {
	for i, v := range kv.ongoing {
		if v == version {
			kv.ongoing = append(kv.ongoing[:i], kv.ongoing[i+1:]...)
			break
		}
	}
}

// --- TRANSACTION LIFECYCLE ---

func (kv *KV) Begin(tx *KVTX) {
	kv.mutex.Lock()
	defer kv.mutex.Unlock()

	tx.db = kv
	tx.version = kv.version
	kv.ongoing = append(kv.ongoing, tx.version)

	// 1. Setup the Read-Only Snapshot
	tx.snapshot.root = kv.tree.root
	chunks := kv.mmap.chunks 
	tx.snapshot.get = func(ptr uint64) []byte {
		start := uint64(0)
		for _, chunk := range chunks {
			end := start + uint64(len(chunk))/BTREE_PAGE_SIZE
			if ptr < end {
				offset := BTREE_PAGE_SIZE * (ptr - start)
				return chunk[offset : offset+BTREE_PAGE_SIZE]
			}
			start = end
		}
		panic("bad page pointer")
	}

	// 2. Setup the Private In-Memory Pending Tree
	tx.pages = nil
	tx.pending.get = func(ptr uint64) []byte { return tx.pages[ptr-1] }
	tx.pending.new = func(node []byte) uint64 {
		tx.pages = append(tx.pages, node)
		return uint64(len(tx.pages))
	}
	tx.pending.del = func(uint64) {} 
}

func (kv *KV) Commit(tx *KVTX) error {
	kv.mutex.Lock()
	defer kv.mutex.Unlock()

	// 1. Conflict Detection
	if detectConflicts(kv, tx) {
		kv.removeOngoing(tx.version)
		return fmt.Errorf("transaction aborted due to conflict")
	}

	// 2. Apply pending changes to the real tree
	meta := saveMeta(kv)
	for _, w := range tx.writes {
		key := w.start
		val, ok := tx.pending.Get(key)
		if ok {
			if val[0] == FLAG_UPDATED {
				kv.tree.Insert(key, val[1:])
			} else if val[0] == FLAG_DELETED {
				kv.tree.Delete(key)
			}
		}
	}

	// 3. Flush to disk
	err := updateOrRevert(kv, meta)
	if err != nil {
		kv.removeOngoing(tx.version)
		return err
	}

	// 4. Update History and Versions
	kv.version++
	if len(tx.writes) > 0 {
		kv.history = append(kv.history, CommittedTX{version: kv.version, writes: tx.writes})
	}
	
	kv.removeOngoing(tx.version)
	
	kv.free.curVer = kv.version
	if len(kv.ongoing) > 0 {
		kv.free.maxVer = kv.ongoing[0] 
	} else {
		kv.free.maxVer = kv.version
	}

	return nil
}

func (kv *KV) Abort(tx *KVTX) {
	kv.mutex.Lock()
	defer kv.mutex.Unlock()
	kv.removeOngoing(tx.version)
}

// --- TRANSACTION OPERATIONS ---

func (tx *KVTX) Set(key []byte, val []byte) {
	buf := make([]byte, len(val)+1)
	buf[0] = FLAG_UPDATED
	copy(buf[1:], val)
	
	tx.pending.Insert(key, buf)
	tx.writes = append(tx.writes, KeyRange{start: key, stop: key})
}

func (tx *KVTX) Del(key []byte) bool {
	tx.pending.Insert(key, []byte{FLAG_DELETED})
	tx.writes = append(tx.writes, KeyRange{start: key, stop: key})
	return true
}

func (tx *KVTX) Get(key []byte) ([]byte, bool) {
	tx.reads = append(tx.reads, KeyRange{start: key, stop: key})

	val, ok := tx.pending.Get(key)
	if ok {
		if val[0] == FLAG_UPDATED {
			return val[1:], true
		} else if val[0] == FLAG_DELETED {
			return nil, false
		}
	}

	return tx.snapshot.Get(key)
}

func (tx *KVTX) Seek(key []byte, cmp int) *BIter {
	tx.reads = append(tx.reads, KeyRange{start: key, stop: key})
	return tx.snapshot.Seek(key, cmp)
}

// --- DISK MANAGEMENT ---

func (db *KV) pageReadFile(ptr uint64) []byte {
	start := uint64(0)
	for _, chunk := range db.mmap.chunks {
		end := start + uint64(len(chunk))/BTREE_PAGE_SIZE
		if ptr < end {
			offset := BTREE_PAGE_SIZE * (ptr - start)
			return chunk[offset : offset+BTREE_PAGE_SIZE]
		}
		start = end
	}
	panic("bad page pointer")
}

func (db *KV) pageRead(ptr uint64) []byte {
	if node, ok := db.page.updates[ptr]; ok {
		return node
	}
	return db.pageReadFile(ptr) 
}

func (db *KV) pageWrite(ptr uint64) []byte {
	if node, ok := db.page.updates[ptr]; ok {
		return node
	}
	node := make([]byte, BTREE_PAGE_SIZE)
	copy(node, db.pageReadFile(ptr)) 
	db.page.updates[ptr] = node
	return node
}

func (db *KV) pageAlloc(node []byte) uint64 {
	if ptr := db.free.PopHead(); ptr != 0 {
		db.page.updates[ptr] = node
		return ptr
	}
	return db.pageAppend(node) 
}

func (db *KV) pageAppend(node []byte) uint64 {
	ptr := db.page.flushed + uint64(db.page.nappend)
	db.page.nappend++
	db.page.updates[ptr] = node
	return ptr
}

func extendMmap(db *KV, size int) error {
	if size <= db.mmap.total {
		return nil  
	}

	alloc := max(db.mmap.total, 64<<20)
	for db.mmap.total+alloc < size {
		alloc *= 2           
	}
	
	chunk, err := syscall.Mmap(
		db.fd, int64(db.mmap.total), alloc,
		syscall.PROT_READ, syscall.MAP_SHARED,
	)

	if err != nil {
		return fmt.Errorf("mmap: %w", err)
	}

	db.mmap.total += alloc
	db.mmap.chunks = append(db.mmap.chunks, chunk)
	return nil
}

func writePages(db *KV) error {
	size := (int(db.page.flushed) + db.page.nappend) * BTREE_PAGE_SIZE
	if err := extendMmap(db, size); err != nil {
		return err
	}

	for ptr, node := range db.page.updates {
		offset := int64(ptr * BTREE_PAGE_SIZE)
		if _, err := syscall.Pwrite(db.fd, node, offset); err != nil {
			return err
		}
	}

	db.page.flushed += uint64(db.page.nappend)
	db.page.nappend = 0
	db.page.updates = make(map[uint64][]byte) 
	return nil
}

func saveMeta(db *KV) []byte {
	var data [128]byte // Increased to store version
	copy(data[:16], []byte(DB_SIG))
	binary.LittleEndian.PutUint64(data[16:24], db.tree.root)
	binary.LittleEndian.PutUint64(data[24:32], db.page.flushed)
	binary.LittleEndian.PutUint64(data[32:40], db.free.headPage)
	binary.LittleEndian.PutUint64(data[40:48], db.free.headSeq)
	binary.LittleEndian.PutUint64(data[48:56], db.free.tailPage)
	binary.LittleEndian.PutUint64(data[56:64], db.free.tailSeq)
	binary.LittleEndian.PutUint64(data[64:72], db.version) // Persist Version!
	return data[:]
}

func loadMeta(db *KV, data []byte) {
	db.tree.root = binary.LittleEndian.Uint64(data[16:24])
	db.page.flushed = binary.LittleEndian.Uint64(data[24:32])
	db.free.headPage = binary.LittleEndian.Uint64(data[32:40])
	db.free.headSeq = binary.LittleEndian.Uint64(data[40:48])
	db.free.tailPage = binary.LittleEndian.Uint64(data[48:56])
	db.free.tailSeq = binary.LittleEndian.Uint64(data[56:64])
	if len(data) >= 72 {
		db.version = binary.LittleEndian.Uint64(data[64:72])
	}
}

func updateRoot(db *KV) error {
	if _, err := syscall.Pwrite(db.fd, saveMeta(db), 0); err != nil {
		return fmt.Errorf("write meta page: %w", err)
	}
	return nil
}

func updateFile(db *KV) error {
	if err := writePages(db); err != nil {
		return err
	}
	if err := syscall.Fsync(db.fd); err != nil {
		return err
	}
	if err := updateRoot(db); err != nil {
		return err
    }
	if err := syscall.Fsync(db.fd); err != nil {
		return err
	}
	
	db.free.SetMaxSeq()
	return nil
}

func updateOrRevert(db *KV, meta []byte) error {
	if db.failed {
		if _, err := syscall.Pwrite(db.fd, meta, 0); err != nil {
			return err
		}
		if err := syscall.Fsync(db.fd); err != nil {
			return err
		}
		db.failed = false
	}

	err := updateFile(db)
	if err != nil {
		db.failed = true
		loadMeta(db, meta)
		db.page.updates = make(map[uint64][]byte)
		db.page.nappend = 0
	}
	return err
}

func createFileSync(file string) (int, error) {
	flags := os.O_RDONLY | syscall.O_DIRECTORY
	dirfd, err := syscall.Open(path.Dir(file), flags, 0o644)
	if err != nil {
		return -1, fmt.Errorf("open directory: %w", err)
	}
	defer syscall.Close(dirfd)

	flags = os.O_RDWR | os.O_CREATE
	fd, err := syscall.Openat(dirfd, path.Base(file), flags, 0o644)
	if err != nil {
		return -1, fmt.Errorf("open file: %w", err)
	}

	if err = syscall.Fsync(dirfd); err != nil {
		_ = syscall.Close(fd) 
		return -1, fmt.Errorf("fsync directory: %w", err)
	}

	return fd, nil
}

func readRoot(db *KV, fileSize int64) error {
	if fileSize == 0 {
		db.page.flushed = 2
		db.free.headPage = 1 
		db.free.tailPage = 1
		db.version = 1 // Start versioning at 1
		return nil 
	}
	data := db.mmap.chunks[0]
	loadMeta(db, data)
	return nil
}

func (db *KV) Open() error {
	fd, err := createFileSync(db.Path)
	if err != nil {
		return err
	}
	db.fd = fd

	fi, err := os.Stat(db.Path)
	if err != nil {
		return err
	}
	fileSize := fi.Size()

	if err := extendMmap(db, int(fileSize)); err != nil {
		return err
	}

	if err := readRoot(db, fileSize); err != nil {
		return err
	}

	db.page.updates = make(map[uint64][]byte)
	if fileSize == 0 {
		db.page.updates[1] = make([]byte, BTREE_PAGE_SIZE)
	}
	db.tree.get = db.pageRead
	db.tree.new = db.pageAlloc     
	db.tree.del = db.free.PushTail 

	db.free.get = db.pageRead
	db.free.new = db.pageAppend    
	db.free.set = db.pageWrite     
	
	// Initialize FreeList tracking
	db.free.maxVer = db.version
	db.free.curVer = db.version

	return nil
}

func (db *KV) Close() {
	for _, chunk := range db.mmap.chunks {
		syscall.Munmap(chunk)
	}
	syscall.Close(db.fd)
}