go-kv/memtable/memtable.go

package memtable

import (
	"errors"
	"sync"
	"sync/atomic"
)

// ErrMemTableFull is returned when the memtable cannot reserve enough space.
var ErrMemTableFull = errors.New("memtable: insufficient capacity")

// ReserveEntry describes a single key-value pair to be reserved.
type ReserveEntry struct {
	Key      []byte
	Value    []byte
	IsDelete bool
}

// GetResult is returned by MemTable.Get.
type GetResult struct {
	Found    bool
	Value    []byte
	Sequence uint64
}

// MemTable wraps an Arena and SkipList with publish/abort semantics.
//
// Writers first reserve capacity, then insert pending entries, then either
// publish (making them visible) or abort (leaving them invisible forever).
// Readers only see published entries that have not been aborted.
type MemTable struct {
	arena     *Arena
	skiplist  *SkipList
	published atomic.Uint64 // highest published sequence
	aborted   sync.Map      // map[uint64]struct{} — set of aborted sequences
}

// NewMemTable creates a MemTable backed by an arena of the given capacity.
func NewMemTable(capacity uint32) *MemTable {
	arena := NewArena(capacity)
	return &MemTable{
		arena:    arena,
		skiplist: NewSkipList(arena),
	}
}

// Reserve estimates whether the arena has enough space for all entries.
// It returns the total estimated size for caller tracking.
//
// The estimate is conservative: key + value bytes plus a fixed metadata
// overhead per entry. Since Phase 1 stores nodes on the Go heap (not in
// the arena), this reservation is primarily a capacity gate for future
// arena-backed phases.
func (mt *MemTable) Reserve(entries []ReserveEntry) (uint32, error) {
	const metadataOverhead uint32 = 32 // per-entry overhead estimate

	var totalSize uint32
	for _, e := range entries {
		sz := metadataOverhead + uint32(len(e.Key)) + uint32(len(e.Value))
		// Align to 8 bytes
		sz = (sz + 7) &^ uint32(7)
		totalSize += sz
	}

	if err := mt.arena.Reserve(totalSize); err != nil {
		return 0, ErrMemTableFull
	}
	return totalSize, nil
}

// PutPending inserts a key-value pair as pending (invisible to Get/Iterator).
func (mt *MemTable) PutPending(key []byte, value []byte, sequence uint64) error {
	return mt.skiplist.Put(key, value, sequence, true)
}

// DeletePending inserts a tombstone entry (nil value) as pending.
func (mt *MemTable) DeletePending(key []byte, sequence uint64) error {
	return mt.skiplist.Put(key, nil, sequence, true)
}

// Publish makes all pending entries with sequence <= upToSequence visible.
//
// It iterates the skip list and re-inserts each pending entry with
// pending=false, which the lock-free reader path will then observe.
// Aborted entries are skipped and remain invisible forever.
func (mt *MemTable) Publish(upToSequence uint64) {
	// Collect entries to publish under the iterator (which skips pending).
	// We need a raw walk, so we access the skiplist directly.
	type entry struct {
		key      []byte
		value    []byte
		sequence uint64
	}

	// Walk the raw skip list level-0 chain to find pending entries.
	// We cannot use NewIterator because it skips pending entries.
	var toPublish []entry

	mt.skiplist.mu.Lock()
	node := mt.skiplist.head.next[0].Load()
	for node != nil {
		if node.pending && node.sequence <= upToSequence {
			if _, aborted := mt.aborted.Load(node.sequence); !aborted {
				toPublish = append(toPublish, entry{
					key:      node.key,
					value:    node.value,
					sequence: node.sequence,
				})
			}
		}
		node = node.next[0].Load()
	}
	mt.skiplist.mu.Unlock()

	// Re-put each entry as published. Each call acquires the skiplist mutex.
	for _, e := range toPublish {
		_ = mt.skiplist.Put(e.key, e.value, e.sequence, false)
	}

	// Update high-water mark after all entries are visible.
	mt.published.Store(upToSequence)
}

// Abort marks a sequence as aborted. Aborted entries are never visible to readers.
func (mt *MemTable) Abort(sequence uint64) {
	mt.aborted.Store(sequence, struct{}{})
}

// Get retrieves the latest visible value for key.
//
// A value is visible if it is published (pending=false in the skip list)
// and not in the aborted set.
func (mt *MemTable) Get(key []byte) *GetResult {
	found, value, sequence := mt.skiplist.Get(key)
	if !found {
		return &GetResult{Found: false}
	}
	return &GetResult{
		Found:    true,
		Value:    value,
		Sequence: sequence,
	}
}

// NewIterator returns a forward iterator over all published (visible) entries.
// Pending and aborted entries are automatically skipped.
func (mt *MemTable) NewIterator() *Iterator {
	return mt.skiplist.NewIterator()
}

// ApproximateSize returns the number of bytes used in the arena.
func (mt *MemTable) ApproximateSize() uint64 {
	return uint64(mt.arena.Capacity() - mt.arena.Remaining())
}

// UsableCapacity returns the remaining bytes available in the arena.
func (mt *MemTable) UsableCapacity() uint32 {
	return mt.arena.Remaining()
}