packages/minidb/src/store.ts

// src/store.ts
//
// In-memory primary index with TTL and an ordered key index. Values may be held
// either inline in memory or as a pointer into the on-disk snapshot/WAL; in the
// latter case the caller injects a synchronous positioned reader.

import { SkipList, cmpString } from './skiplist.js';
import type { RangeEntry, RangeOptions } from './skiplist.js';

export interface ValueLoc {
  file: 'snapshot' | 'wal';
  off: number;
  len: number;
}

export type ValueRef = { kind: 'memory'; value: Buffer } | { kind: 'disk'; loc: ValueLoc };

export interface StoreRecord {
  ref: ValueRef;
  expireAt: number;
  seq: number;
  dt: Record<string, number> | null;
}

export interface StoreEntry {
  key: Buffer;
  value: Buffer;
  expireAt: number;
  dt: Record<string, number> | null;
}

export type ValueReader = (loc: ValueLoc) => Buffer;

/** Internal raw-record view used by derived-index rebuilds: the canonical
 *  (byte-string) key, the dt metadata, and a lazy value reader. Nothing is
 *  materialized unless readValue() is called, so a rebuild that only needs
 *  metadata never copies a buffer (memory mode) or issues a positioned read
 *  (disk mode). */
export interface RawRecord {
  kstr: string;
  dt: Record<string, number> | null;
  readValue: () => Buffer;
}

const toKStr = (key: string | Buffer): string =>
  typeof key === 'string' ? key : Buffer.from(key).toString('binary');
const fromKStr = (kstr: string): Buffer => Buffer.from(kstr, 'binary');
const DISK_REF_BYTES = 64;

interface HeapEntry {
  t: number;
  k: string;
  seq: number;
}

class MinHeap {
  private a: HeapEntry[] = [];
  get size(): number {
    return this.a.length;
  }
  clear(): void {
    this.a = [];
  }
  peek(): HeapEntry | undefined {
    return this.a[0];
  }
  push(item: HeapEntry): void {
    const a = this.a;
    a.push(item);
    let i = a.length - 1;
    while (i > 0) {
      const p = (i - 1) >> 1;
      if (a[p]!.t <= a[i]!.t) break;
      [a[p], a[i]] = [a[i]!, a[p]!];
      i = p;
    }
  }
  pop(): HeapEntry | undefined {
    const a = this.a;
    const top = a[0];
    const last = a.pop();
    if (a.length && last !== undefined) {
      a[0] = last;
      let i = 0;
      while (true) {
        let s = i;
        const l = 2 * i + 1;
        const r = l + 1;
        if (l < a.length && a[l]!.t < a[s]!.t) s = l;
        if (r < a.length && a[r]!.t < a[s]!.t) s = r;
        if (s === i) break;
        [a[s], a[i]] = [a[i]!, a[s]!];
        i = s;
      }
    }
    return top;
  }
}

export interface StoreOptions {
  activeExpireIntervalMs?: number;
  activeExpireMaxPerTick?: number;
  /** Time budget (ms) per active-expiration tick, used once
   *  `activeExpireMaxPerTick` entries have been reaped: a simultaneously-expired
   *  burst keeps being reaped within this budget, so TTL storms drain far
   *  faster than the fixed per-tick count alone, while the event-loop pause per
   *  tick stays bounded. */
  activeExpireTimeBudgetMs?: number;
  /** Read a value back from disk for disk-backed records. Required whenever a
   *  StoreRecord may hold a disk ref. */
  readValue?: ValueReader;
  /** Called after a key is removed due to TTL expiration (lazy or active),
   *  so the owner can drop derived state (secondary/dt/text indexes). Not
   *  called for explicit del(), which the owner handles itself. */
  onExpire?: (key: string, record: StoreRecord) => void;
}

export class Store {
  readonly map = new Map<string, StoreRecord>(); // kstr -> record
  private order = new SkipList<string, string>({ compareKey: cmpString }); // kstr ordered
  private readonly heap = new MinHeap();
  private seq = 0;
  /** Approximate bytes held by live + expired-not-yet-reaped records. In
   *  valueMode:'disk' this counts keys/metadata/refs, not the value bulk. */
  bytes = 0;
  /** Number of records with an expiry set. Enables an O(1) size fast path when
   *  TTL is not in use. */
  private expiring = 0;
  private readonly maxPerTick: number;
  private readonly expireTimeBudgetMs: number;
  /** Sticky flag: set when the last tick reaped ≥ maxPerTick (i.e. there is an
   *  expiry backlog), making the next ticks aggressive until the storm drains
   *  — like Redis's aggressive expire cycle. */
  private expireAggressive = false;
  /** In-flight async bulk load (bulkLoadRefsAsync only — the sync bulk load
   *  yields nothing, so no timer can fire mid-load there). Active expiry is
   *  paused for the load's duration: a mid-load reap would delete the key
   *  from the map (and from the about-to-be-replaced OLD ordered index, a
   *  no-op) while the load still rebuilds `order` from its accumulated
   *  entries — resurrecting the reaped key in the ordered index, where a
   *  later re-set would duplicate it in ordered scans. Whatever elapsed
   *  during the load is reaped by the first tick after it, which is exactly
   *  the sync bulkLoadRefs behavior. */
  private bulkLoading = false;
  private timer: ReturnType<typeof setInterval> | null = null;
  private readonly onExpire?: (key: string, record: StoreRecord) => void;
  private readonly readValue?: ValueReader;

  constructor(opts: StoreOptions = {}) {
    this.maxPerTick = opts.activeExpireMaxPerTick ?? 100;
    this.expireTimeBudgetMs = opts.activeExpireTimeBudgetMs ?? 2;
    this.onExpire = opts.onExpire;
    this.readValue = opts.readValue;
    const interval = opts.activeExpireIntervalMs ?? 100;
    this.timer = interval > 0 ? setInterval(() => this.activeExpire(), interval) : null;
    this.timer?.unref?.();
  }

  get size(): number {
    // O(1) fast path: with no TTL ever set, every map entry is live.
    if (this.expiring === 0) return this.map.size;
    // Count only logically-live keys: expired-but-not-yet-reaped entries stay
    // in the map until lazy/active expiration removes them, so map.size would
    // otherwise over-report.
    let n = 0;
    const now = Date.now();
    for (const [, r] of this.map) if (!r.expireAt || r.expireAt > now) n++;
    return n;
  }

  private metaBytes(dt: Record<string, number> | null): number {
    return dt ? Buffer.byteLength(JSON.stringify({ dt }), 'utf8') : 0;
  }

  private refBytes(ref: ValueRef): number {
    return ref.kind === 'memory' ? ref.value.length : DISK_REF_BYTES;
  }

  private cloneRef(ref: ValueRef): ValueRef {
    return ref.kind === 'memory' ? { kind: 'memory', value: Buffer.from(ref.value) } : { kind: 'disk', loc: { ...ref.loc } };
  }

  private materialize(ref: ValueRef): Buffer {
    if (ref.kind === 'memory') return ref.value;
    if (!this.readValue) throw new Error('Store cannot read disk-backed value without a ValueReader');
    return this.readValue(ref.loc);
  }

  /** Approximate bytes used by one record, matching the bytes tracked on set(). */
  recordBytes(k: string): number {
    const r = this.map.get(k);
    return r ? Buffer.byteLength(k, 'binary') + this.refBytes(r.ref) + this.metaBytes(r.dt) : 0;
  }

  /** Approximate bytes a SET would store for this key/value/dt. Pass
   *  countValue:false for valueMode:'disk', where only key/metadata/ref bytes
   *  stay in RAM and the value bulk lives in the snapshot/WAL. */
  estimateSetBytes(
    key: string | Buffer,
    value: Buffer,
    dt: Record<string, number> | null,
    opts: { countValue?: boolean } = {},
  ): number {
    const k = toKStr(key);
    const countValue = opts.countValue ?? true;
    return Buffer.byteLength(k, 'binary') + (countValue ? value.length : DISK_REF_BYTES) + this.metaBytes(dt);
  }

  private remove(k: string): boolean {
    const r = this.map.get(k);
    const ok = this.map.delete(k);
    if (ok) {
      if (r) {
        this.bytes -= Buffer.byteLength(k, 'binary') + this.refBytes(r.ref) + this.metaBytes(r.dt);
        if (r.expireAt) this.expiring--;
      }
      this.order.delete(k, k);
    }
    return ok;
  }

  /** Remove a key that has expired and notify the owner so derived indexes
   *  stay in sync. */
  private expireKey(k: string, rec: StoreRecord): void {
    if (this.remove(k)) this.onExpire?.(k, rec);
  }

  set(key: string | Buffer, value: Buffer, expireAt = 0, dt: Record<string, number> | null = null): void {
    this.setRef(key, { kind: 'memory', value: Buffer.from(value) }, expireAt, dt);
  }

  setRef(key: string | Buffer, ref: ValueRef, expireAt = 0, dt: Record<string, number> | null = null): void {
    const k = toKStr(key);
    const prev = this.map.get(k);
    if (prev) {
      if (prev.expireAt) this.expiring--;
      this.bytes -= this.recordBytes(k);
    }
    const seq = ++this.seq;
    const stored = this.cloneRef(ref);
    this.map.set(k, { ref: stored, expireAt: expireAt || 0, seq, dt });
    this.bytes += Buffer.byteLength(k, 'binary') + this.refBytes(stored) + this.metaBytes(dt);
    if (!prev) this.order.insert(k, k);
    if (expireAt) {
      this.expiring++;
      this.heap.push({ t: expireAt, k, seq });
      // Overwriting a TTL key leaves a stale heap entry that is only reaped
      // when its (possibly far-future) timestamp passes, so the heap can grow
      // without bound under frequent TTL updates. Rebuild it once stale entries
      // clearly dominate the live set, bounding memory at ~2x live TTL keys.
      if (this.heap.size > this.map.size * 2 + 64) this.rebuildHeap();
    }
  }

  private rebuildHeap(): void {
    this.heap.clear();
    for (const [k, r] of this.map) {
      if (r.expireAt) this.heap.push({ t: r.expireAt, k, seq: r.seq });
    }
  }

  get(key: string | Buffer): Buffer | undefined {
    const k = toKStr(key);
    const r = this.map.get(k);
    if (!r) return undefined;
    if (r.expireAt && r.expireAt <= Date.now()) {
      this.expireKey(k, r); // lazy expiration
      return undefined;
    }
    return this.materialize(r.ref);
  }

  /** Read the full raw record. Like get(), an expired record is lazily reaped
   *  here (notifying the owner so derived indexes stay in sync) rather than
   *  being left behind as a ghost that read paths such as scan/query/dtRange
   *  would otherwise skip-but-not-clean. The returned record is raw: its ref may
   *  point at disk rather than holding a Buffer. */
  getRecord(key: string | Buffer): StoreRecord | undefined {
    const k = toKStr(key);
    const r = this.map.get(k);
    if (!r) return undefined;
    if (r.expireAt && r.expireAt <= Date.now()) {
      this.expireKey(k, r); // lazy expiration (same as get)
      return undefined;
    }
    return r;
  }

  del(key: string | Buffer): boolean {
    return this.remove(toKStr(key));
  }

  /** Metadata-only existence check: no value materialization (no buffer copy in
   *  memory mode, no positioned disk read for disk-backed records). Applies the
   *  same lazy-expiration semantics as get(). */
  has(key: string | Buffer): boolean {
    const k = toKStr(key);
    const r = this.map.get(k);
    if (!r) return false;
    if (r.expireAt && r.expireAt <= Date.now()) {
      this.expireKey(k, r); // lazy expiration, same as get()
      return false;
    }
    return true;
  }

  *entries(): Generator<StoreEntry> {
    const now = Date.now();
    for (const [k, r] of this.map) {
      if (r.expireAt && r.expireAt <= now) continue;
      yield { key: fromKStr(k), value: this.materialize(r.ref), expireAt: r.expireAt, dt: r.dt };
    }
  }

  /** Walk live records without materializing values: yields the canonical key,
   *  dt metadata, and a lazy value reader. Expired records are skipped (not
   *  reaped) exactly as in entries(). Internal to the package — derived-index
   *  rebuilds use it to share a single walk and a single decode per record. */
  *rawRecords(): Generator<RawRecord> {
    const now = Date.now();
    for (const [k, r] of this.map) {
      if (r.expireAt && r.expireAt <= now) continue;
      yield { kstr: k, dt: r.dt, readValue: () => this.materialize(r.ref) };
    }
  }

  /** Walk live records with their RAW value ref (never materialized): the
   *  stage-6 snapshot writer reads disk-backed values through its own async
   *  path (grouped, bounded-concurrency) instead of one synchronous
   *  positioned read per record. Expired records are skipped exactly as in
   *  entries(). Internal to the package. */
  *rawRefRecords(): Generator<{ kstr: string; ref: ValueRef; expireAt: number; dt: Record<string, number> | null }> {
    const now = Date.now();
    for (const [k, r] of this.map) {
      if (r.expireAt && r.expireAt <= now) continue;
      yield { kstr: k, ref: r.ref, expireAt: r.expireAt, dt: r.dt };
    }
  }

  /** Ordered scan over keys. */
  *scan(opts: RangeOptions<string> = {}): Generator<StoreEntry> {
    for (const n of this.order.range(opts) as Iterable<RangeEntry<string, string>>) {
      const r = this.getRecord(n.key);
      if (!r) continue;
      yield { key: fromKStr(n.key), value: this.materialize(r.ref), expireAt: r.expireAt, dt: r.dt };
    }
  }

  /** Prefix scan over keys. */
  *prefix(p: string, limit = Infinity): Generator<StoreEntry> {
    const pk = toKStr(p);
    yield* this.scan({ gte: pk, lt: pk + '\uffff', count: limit });
  }

  /** Ordered scan yielding canonical keys only, without materializing values
   *  (no buffer copies in memory mode, no positioned reads in disk mode).
   *  Expired records are lazily reaped, exactly as in scan(). */
  *rawKeys(opts: RangeOptions<string> = {}): Generator<string> {
    for (const n of this.order.range(opts) as Iterable<RangeEntry<string, string>>) {
      if (!this.getRecord(n.key)) continue;
      yield n.key;
    }
  }

  /** Rewrite disk-backed value locations after compaction rotates the
   *  snapshot/WAL files. Memory refs are left untouched. */
  remapLocs(remap: (k: string, loc: ValueLoc, rec: StoreRecord) => ValueLoc | undefined): void {
    for (const [k, r] of this.map) {
      if (r.ref.kind !== 'disk') continue;
      const next = remap(k, r.ref.loc, r);
      if (next) r.ref = { kind: 'disk', loc: { ...next } };
    }
  }

  /** Stage-5 generation load: populate the store wholesale from a recovered
   *  generation store image. `records` must be expiry-filtered by the caller
   *  (expired-past records dropped) and sorted by canonical key ascending (the
   *  image's write order), so the ordered index is bulk-built in O(N) instead
   *  of per-record inserts.
   *
   *  OWNERSHIP: the records' refs are adopted as-is (no defensive clone) —
   *  the image parser produced fresh buffers for exactly this purpose.
   *  `metaBytes` is the precomputed dt accounting value (0 = none), so the
   *  load never re-stringifies per record. */
  bulkLoadRefs(
    records: Iterable<{ kstr: string; ref: ValueRef; expireAt: number; dt: Record<string, number> | null; metaBytes?: number }>,
  ): void {
    const orderEntries: RangeEntry<string, string>[] = [];
    for (const { kstr, ref, expireAt, dt, metaBytes } of records) {
      const seq = ++this.seq;
      this.map.set(kstr, { ref, expireAt: expireAt || 0, seq, dt });
      this.bytes += Buffer.byteLength(kstr, 'binary') + this.refBytes(ref) + (metaBytes ?? 0);
      if (expireAt) {
        this.expiring++;
        this.heap.push({ t: expireAt, k: kstr, seq });
      }
      orderEntries.push({ key: kstr, val: kstr });
    }
    this.order = SkipList.bulkLoad(orderEntries, { compareKey: cmpString });
  }

  /** Sliced variant of bulkLoadRefs (the open-time main-thread path):
   *  identical resulting state, but the per-record map insertions yield to
   *  the event loop every `sliceEvery` records, so a large store image never
   *  loads in one synchronous run. The final ordered-index bulk build is a
   *  single O(N) pass over the sorted entries and is likewise sliced via
   *  SkipList.bulkLoadAsync. Safe to yield mid-load: the store is not
   *  published until open() returns, and active expiry is paused for the
   *  load's duration (see `bulkLoading`) so a mid-load reap cannot diverge
   *  the map from the not-yet-rebuilt ordered index. */
  async bulkLoadRefsAsync(
    records: Iterable<{ kstr: string; ref: ValueRef; expireAt: number; dt: Record<string, number> | null; metaBytes?: number }>,
    opts: { sliceEvery?: number } = {},
  ): Promise<void> {
    const sliceEvery = opts.sliceEvery ?? 8192;
    const orderEntries: RangeEntry<string, string>[] = [];
    this.bulkLoading = true;
    try {
      let n = 0;
      for (const { kstr, ref, expireAt, dt, metaBytes } of records) {
        const seq = ++this.seq;
        this.map.set(kstr, { ref, expireAt: expireAt || 0, seq, dt });
        this.bytes += Buffer.byteLength(kstr, 'binary') + this.refBytes(ref) + (metaBytes ?? 0);
        if (expireAt) {
          this.expiring++;
          this.heap.push({ t: expireAt, k: kstr, seq });
        }
        orderEntries.push({ key: kstr, val: kstr });
        if (++n % sliceEvery === 0) await new Promise((r) => setImmediate(r));
      }
      this.order = await SkipList.bulkLoadAsync(orderEntries, { compareKey: cmpString }, { sliceEvery });
    } finally {
      this.bulkLoading = false;
    }
  }

  private activeExpire(): void {
    // A sliced bulk load holds the map/order consistency contract: reaping
    // mid-load would drop the key from the map while the load still rebuilds
    // `order` from its accumulated entries. Skip this tick; the next one
    // reaps whatever elapsed meanwhile (identical to the sync bulk load,
    // during which no timer can fire at all).
    if (this.bulkLoading) return;
    const now = Date.now();
    // Normal ticks stay within the small budget; a tick that still finds a
    // full quota of expired keys flips to aggressive mode (larger budget, like
    // Redis's fast expire cycle) until the storm drains.
    const deadline = now + (this.expireAggressive ? Math.max(this.expireTimeBudgetMs, 10) : this.expireTimeBudgetMs);
    let n = 0;
    let reaped = 0;
    while (this.heap.size && this.heap.peek()!.t <= now) {
      // Once the guaranteed per-tick quota is reaped, keep draining within the
      // time budget: a fixed ~1000/s rate let a large simultaneous-expiry storm
      // (e.g. 100k keys) linger in memory for minutes. The budget bounds the
      // pause each tick instead of bounding the work.
      if (n >= this.maxPerTick && (n & 15) === 0 && Date.now() >= deadline) break;
      const e = this.heap.pop()!;
      const r = this.map.get(e.k);
      if (r && r.seq === e.seq && r.expireAt && r.expireAt <= now) {
        this.expireKey(e.k, r);
        reaped++;
      }
      n++;
    }
    this.expireAggressive = reaped >= this.maxPerTick;
  }

  /** Synchronously reap every expired record. Returns the number removed. */
  reapExpired(): number {
    const now = Date.now();
    let n = 0;
    for (const [k, r] of this.map) {
      if (r.expireAt && r.expireAt <= now) {
        this.expireKey(k, r);
        n++;
      }
    }
    return n;
  }

  /** Reap already-expired records via the TTL min-heap: O(due + stale heap
   *  entries) instead of the O(store) full scan of reapExpired(), so callers
   *  on the write hot path do not pay a full-store sweep per call. Falls back
   *  to the full scan only when the heap provably diverged from the map (live
   *  TTL records remain but the heap is empty — an invariant no code path may
   *  produce), resyncing instead of leaking expired bytes. */
  reapExpiredDue(): number {
    const now = Date.now();
    let n = 0;
    while (this.heap.size && this.heap.peek()!.t <= now) {
      const e = this.heap.pop()!;
      const r = this.map.get(e.k);
      if (r && r.seq === e.seq && r.expireAt && r.expireAt <= now) {
        this.expireKey(e.k, r);
        n++;
      }
    }
    if (this.expiring > 0 && this.heap.size === 0) return n + this.reapExpired();
    return n;
  }

  /** Stop the active-expiration timer. */
  close(): void {
    if (this.timer) {
      clearInterval(this.timer);
      this.timer = null;
    }
  }
}