kimi-code / packages /minidb /src /store.ts
SaylorTwift's picture
SaylorTwift HF Staff
Add files using upload-large-folder tool
4e23b01 verified
Raw History Blame Contribute Delete
19.9 kB
// 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;
}
}
}