@prisma/streams-server 0.1.1 → 0.1.3

package/src/index/lexicon_file_cache.ts

@@ -0,0 +1,261 @@
+import {
+  existsSync,
+  mkdirSync,
+  readdirSync,
+  renameSync,
+  statSync,
+  unlinkSync,
+  writeFileSync,
+} from "node:fs";
+import { dirname, join, relative } from "node:path";
+import { Result } from "better-result";
+import { LruCache } from "../util/lru";
+
+export type LexiconFileCacheError = { kind: "invalid_lexicon_cache"; message: string };
+
+export type MappedLexiconFile = {
+  objectKey: string;
+  path: string;
+  bytes: Uint8Array;
+  sizeBytes: number;
+};
+
+type FileEntry = {
+  path: string;
+  size: number;
+  mtimeMs: number;
+};
+
+export type LexiconFileCacheStats = {
+  usedBytes: number;
+  entryCount: number;
+  mappedBytes: number;
+  mappedEntryCount: number;
+  pinnedEntryCount: number;
+};
+
+function invalidLexiconCache<T = never>(message: string): Result<T, LexiconFileCacheError> {
+  return Result.err({ kind: "invalid_lexicon_cache", message });
+}
+
+export class LexiconFileCache {
+  private readonly entries = new Map<string, FileEntry>();
+  private readonly pinnedKeys = new Set<string>();
+  private readonly mappedFiles: LruCache<string, MappedLexiconFile>;
+  private totalBytes = 0;
+
+  constructor(
+    private readonly rootDir: string,
+    private readonly maxBytes: number,
+    mappedEntries: number
+  ) {
+    if (this.maxBytes > 0) mkdirSync(this.rootDir, { recursive: true });
+    this.mappedFiles = new LruCache(Math.max(1, mappedEntries));
+    this.loadIndex();
+    this.pruneForBudget(0);
+  }
+
+  clearMapped(): void {
+    this.mappedFiles.clear();
+  }
+
+  bytesForObjectKeyPrefix(prefix: string): number {
+    let total = 0;
+    for (const [objectKey, entry] of this.entries.entries()) {
+      if (objectKey.startsWith(prefix)) total += entry.size;
+    }
+    return total;
+  }
+
+  storeBytesResult(objectKey: string, bytes: Uint8Array): Result<string, LexiconFileCacheError> {
+    if (this.maxBytes <= 0) return Result.ok("");
+    if (bytes.byteLength > this.maxBytes) {
+      return invalidLexiconCache(`lexicon cache object exceeds budget: ${objectKey}`);
+    }
+    this.pruneForBudget(bytes.byteLength);
+    const path = this.pathFor(objectKey);
+    mkdirSync(dirname(path), { recursive: true });
+    const tmpPath = `${path}.tmp-${process.pid}-${Date.now()}`;
+    try {
+      writeFileSync(tmpPath, bytes);
+      renameSync(tmpPath, path);
+    } catch (e: unknown) {
+      try {
+        unlinkSync(tmpPath);
+      } catch {
+        // ignore temp cleanup failures
+      }
+      return invalidLexiconCache(String((e as any)?.message ?? e));
+    }
+    const existing = this.entries.get(objectKey);
+    if (existing) this.totalBytes = Math.max(0, this.totalBytes - existing.size);
+    this.entries.delete(objectKey);
+    this.entries.set(objectKey, { path, size: bytes.byteLength, mtimeMs: Date.now() });
+    this.totalBytes += bytes.byteLength;
+    return Result.ok(path);
+  }
+
+  async loadMappedFileResult(args: {
+    objectKey: string;
+    expectedSize: number;
+    loadBytes: () => Promise<Uint8Array>;
+  }): Promise<Result<MappedLexiconFile, LexiconFileCacheError>> {
+    const cached = this.mappedFiles.get(args.objectKey);
+    if (cached) {
+      this.pinnedKeys.add(args.objectKey);
+      this.touch(args.objectKey);
+      return Result.ok(cached);
+    }
+
+    const localPathRes = await this.ensureLocalFileResult(args.objectKey, args.expectedSize, args.loadBytes);
+    if (Result.isError(localPathRes)) return localPathRes;
+
+    const mappedRes = this.mapFileResult(args.objectKey, localPathRes.value, args.expectedSize);
+    if (Result.isError(mappedRes)) return mappedRes;
+    this.pinnedKeys.add(args.objectKey);
+    this.mappedFiles.set(args.objectKey, mappedRes.value);
+    this.touch(args.objectKey);
+    return mappedRes;
+  }
+
+  stats(): LexiconFileCacheStats {
+    let mappedBytes = 0;
+    let mappedEntryCount = 0;
+    for (const mapped of this.mappedFiles.values()) {
+      mappedBytes += mapped.sizeBytes;
+      mappedEntryCount += 1;
+    }
+    return {
+      usedBytes: this.totalBytes,
+      entryCount: this.entries.size,
+      mappedBytes,
+      mappedEntryCount,
+      pinnedEntryCount: this.pinnedKeys.size,
+    };
+  }
+
+  private async ensureLocalFileResult(
+    objectKey: string,
+    expectedSize: number,
+    loadBytes: () => Promise<Uint8Array>
+  ): Promise<Result<string, LexiconFileCacheError>> {
+    if (this.maxBytes <= 0) return invalidLexiconCache("lexicon cache disabled");
+    const existing = this.entries.get(objectKey);
+    if (existing) {
+      const stat = this.safeStat(existing.path);
+      if (stat && stat.size === expectedSize) {
+        this.touch(objectKey);
+        return Result.ok(existing.path);
+      }
+      this.removeEntry(objectKey);
+    }
+    let bytes: Uint8Array;
+    try {
+      bytes = await loadBytes();
+    } catch (e: unknown) {
+      return invalidLexiconCache(String((e as any)?.message ?? e));
+    }
+    if (bytes.byteLength !== expectedSize) {
+      return invalidLexiconCache(`unexpected lexicon object size for ${objectKey}`);
+    }
+    return this.storeBytesResult(objectKey, bytes);
+  }
+
+  private mapFileResult(objectKey: string, path: string, expectedSize: number): Result<MappedLexiconFile, LexiconFileCacheError> {
+    const stat = this.safeStat(path);
+    if (!stat) return invalidLexiconCache(`missing cached lexicon file ${path}`);
+    if (stat.size !== expectedSize) return invalidLexiconCache(`unexpected cached lexicon size for ${objectKey}`);
+    let bytes: Uint8Array;
+    try {
+      bytes = (Bun as any).mmap(path, { shared: true }) as Uint8Array;
+    } catch (e: unknown) {
+      return invalidLexiconCache(String((e as any)?.message ?? e));
+    }
+    if (bytes.byteLength !== expectedSize) {
+      return invalidLexiconCache(`unexpected mapped lexicon size for ${objectKey}`);
+    }
+    return Result.ok({ objectKey, path, bytes, sizeBytes: expectedSize });
+  }
+
+  private pathFor(objectKey: string): string {
+    return join(this.rootDir, objectKey);
+  }
+
+  private loadIndex(): void {
+    if (!existsSync(this.rootDir)) return;
+    const files: Array<{ key: string; path: string; size: number; mtimeMs: number }> = [];
+    const walk = (dir: string) => {
+      for (const entry of readdirSync(dir, { withFileTypes: true })) {
+        const full = join(dir, entry.name);
+        if (entry.isDirectory()) {
+          walk(full);
+          continue;
+        }
+        if (!entry.isFile()) continue;
+        const stat = statSync(full);
+        files.push({
+          key: relative(this.rootDir, full),
+          path: full,
+          size: stat.size,
+          mtimeMs: stat.mtimeMs,
+        });
+      }
+    };
+    walk(this.rootDir);
+    files.sort((left, right) => left.mtimeMs - right.mtimeMs);
+    for (const file of files) {
+      this.entries.set(file.key, { path: file.path, size: file.size, mtimeMs: file.mtimeMs });
+      this.totalBytes += file.size;
+    }
+  }
+
+  private pruneForBudget(incomingBytes: number): void {
+    if (this.maxBytes <= 0) return;
+    while (this.totalBytes + incomingBytes > this.maxBytes) {
+      const next = this.entries.keys().next();
+      if (next.done) break;
+      const objectKey = next.value as string;
+      if (this.pinnedKeys.has(objectKey)) {
+        let removed = false;
+        for (const candidateKey of this.entries.keys()) {
+          if (this.pinnedKeys.has(candidateKey)) continue;
+          this.removeEntry(candidateKey);
+          removed = true;
+          break;
+        }
+        if (!removed) break;
+        continue;
+      }
+      this.removeEntry(objectKey);
+    }
+  }
+
+  private touch(objectKey: string): void {
+    const entry = this.entries.get(objectKey);
+    if (!entry) return;
+    this.entries.delete(objectKey);
+    this.entries.set(objectKey, { ...entry });
+  }
+
+  private removeEntry(objectKey: string): void {
+    if (this.pinnedKeys.has(objectKey)) return;
+    const entry = this.entries.get(objectKey);
+    if (!entry) return;
+    try {
+      unlinkSync(entry.path);
+    } catch {
+      // ignore remove failures
+    }
+    this.totalBytes = Math.max(0, this.totalBytes - entry.size);
+    this.entries.delete(objectKey);
+    this.mappedFiles.delete(objectKey);
+  }
+
+  private safeStat(path: string): { size: number } | null {
+    try {
+      return existsSync(path) ? { size: statSync(path).size } : null;
+    } catch {
+      return null;
+    }
+  }
+}

package/src/index/lexicon_format.ts

@@ -0,0 +1,93 @@
+import { Result } from "better-result";
+import { RestartStringTableView, encodeRestartStringTable } from "../search/binary/restart_strings";
+import { concatBytes, readU32BE, readU64BE, writeU32BE, writeU64BE } from "../util/endian";
+import { dsError } from "../util/ds_error.ts";
+
+export const LEXICON_RUN_MAGIC = "LRN1";
+export const LEXICON_RUN_VERSION = 1;
+
+export type LexiconRunMeta = {
+  runId: string;
+  level: number;
+  startSegment: number;
+  endSegment: number;
+  objectKey: string;
+  recordCount: number;
+};
+
+export type LexiconRun = {
+  meta: LexiconRunMeta;
+  payloadBytes: Uint8Array;
+  terms: RestartStringTableView;
+};
+
+export type LexiconRunFormatError = {
+  kind: "invalid_lexicon_run";
+  message: string;
+};
+
+function invalidLexiconRun<T = never>(message: string): Result<T, LexiconRunFormatError> {
+  return Result.err({ kind: "invalid_lexicon_run", message });
+}
+
+export function encodeLexiconRun(run: LexiconRun): Uint8Array {
+  const res = encodeLexiconRunResult(run);
+  if (Result.isError(res)) throw dsError(res.error.message);
+  return res.value;
+}
+
+export function encodeLexiconRunResult(run: LexiconRun): Result<Uint8Array, LexiconRunFormatError> {
+  const payload = run.payloadBytes;
+  const header = new Uint8Array(32);
+  header[0] = LEXICON_RUN_MAGIC.charCodeAt(0);
+  header[1] = LEXICON_RUN_MAGIC.charCodeAt(1);
+  header[2] = LEXICON_RUN_MAGIC.charCodeAt(2);
+  header[3] = LEXICON_RUN_MAGIC.charCodeAt(3);
+  header[4] = LEXICON_RUN_VERSION;
+  header[5] = run.meta.level & 0xff;
+  header[6] = 0;
+  header[7] = 0;
+  writeU64BE(header, 8, BigInt(run.meta.startSegment));
+  writeU64BE(header, 16, BigInt(run.meta.endSegment));
+  writeU32BE(header, 24, run.meta.recordCount);
+  writeU32BE(header, 28, payload.byteLength);
+  return Result.ok(concatBytes([header, payload]));
+}
+
+export function decodeLexiconRun(data: Uint8Array): LexiconRun {
+  const res = decodeLexiconRunResult(data);
+  if (Result.isError(res)) throw dsError(res.error.message);
+  return res.value;
+}
+
+export function decodeLexiconRunResult(data: Uint8Array): Result<LexiconRun, LexiconRunFormatError> {
+  if (data.byteLength < 32) return invalidLexiconRun("run too short");
+  const magic = String.fromCharCode(data[0], data[1], data[2], data[3]);
+  if (magic !== LEXICON_RUN_MAGIC) return invalidLexiconRun("invalid run magic");
+  if (data[4] !== LEXICON_RUN_VERSION) return invalidLexiconRun("unsupported run version");
+  const level = data[5] ?? 0;
+  const startSegment = Number(readU64BE(data, 8));
+  const endSegment = Number(readU64BE(data, 16));
+  const recordCount = readU32BE(data, 24);
+  const payloadLen = readU32BE(data, 28);
+  if (32 + payloadLen > data.byteLength) return invalidLexiconRun("run payload truncated");
+  const payloadBytes = data.subarray(32, 32 + payloadLen);
+  const terms = new RestartStringTableView(payloadBytes);
+  if (terms.count() !== recordCount) return invalidLexiconRun("record count mismatch");
+  return Result.ok({
+    meta: {
+      runId: "",
+      level,
+      startSegment,
+      endSegment,
+      objectKey: "",
+      recordCount,
+    },
+    payloadBytes,
+    terms,
+  });
+}
+
+export function buildLexiconRunPayload(keys: string[]): Uint8Array {
+  return encodeRestartStringTable(keys, 16);
+}