@prisma/streams-server 0.1.1 → 0.1.3

package/src/sqlite/adapter.ts

@@ -1,3 +1,6 @@
+import { createRequire } from "node:module";
+import type { HostRuntime } from "../runtime/host_runtime.ts";
+import { detectHostRuntime } from "../runtime/host_runtime.ts";
 import { dsError } from "../util/ds_error.ts";
 export interface SqliteStatement {
   get(...params: any[]): any;
@@ -10,15 +13,49 @@ export interface SqliteStatement {
 export interface SqliteDatabase {
   exec(sql: string): void;
   query(sql: string): SqliteStatement;
+  prepare(sql: string): SqliteStatement;
   transaction<T>(fn: () => T): () => T;
   close(): void;
 }
 
+type SqliteAdapterRuntimeCounts = {
+  open_connections: number;
+  prepared_statements: number;
+};
+
+const sqliteAdapterRuntimeCounts: SqliteAdapterRuntimeCounts = {
+  open_connections: 0,
+  prepared_statements: 0,
+};
+
+function incrementSqliteConnection(): void {
+  sqliteAdapterRuntimeCounts.open_connections += 1;
+}
+
+function decrementSqliteConnection(): void {
+  sqliteAdapterRuntimeCounts.open_connections = Math.max(0, sqliteAdapterRuntimeCounts.open_connections - 1);
+}
+
+function incrementPreparedStatement(): void {
+  sqliteAdapterRuntimeCounts.prepared_statements += 1;
+}
+
+function decrementPreparedStatement(): void {
+  sqliteAdapterRuntimeCounts.prepared_statements = Math.max(0, sqliteAdapterRuntimeCounts.prepared_statements - 1);
+}
+
+export function getSqliteAdapterRuntimeCounts(): SqliteAdapterRuntimeCounts {
+  return { ...sqliteAdapterRuntimeCounts };
+}
+
 class BunStatementAdapter implements SqliteStatement {
   private readonly stmt: any;
+  private readonly onFinalize?: () => void;
+  private finalized = false;
 
-  constructor(stmt: any) {
+  constructor(stmt: any, onFinalize?: () => void) {
     this.stmt = stmt;
+    this.onFinalize = onFinalize;
   }
 
   get(...params: any[]): any {
@@ -38,23 +75,58 @@ class BunStatementAdapter implements SqliteStatement {
   }
 
   finalize(): void {
-    if (typeof this.stmt.finalize === "function") this.stmt.finalize();
+    if (this.finalized) return;
+    this.finalized = true;
+    try {
+      if (typeof this.stmt.finalize === "function") this.stmt.finalize();
+    } finally {
+      this.onFinalize?.();
+    }
   }
 }
 
 class BunDatabaseAdapter implements SqliteDatabase {
   private readonly db: any;
+  private preparedStatementCount = 0;
+  private closed = false;
+  private readonly statementCache = new Map<string, BunStatementAdapter>();
 
   constructor(db: any) {
     this.db = db;
+    incrementSqliteConnection();
+  }
+
+  private trackStatement(): () => void {
+    let released = false;
+    this.preparedStatementCount += 1;
+    incrementPreparedStatement();
+    return () => {
+      if (released) return;
+      released = true;
+      this.preparedStatementCount = Math.max(0, this.preparedStatementCount - 1);
+      decrementPreparedStatement();
+    };
   }
 
   exec(sql: string): void {
     this.db.exec(sql);
   }
 
+  prepare(sql: string): SqliteStatement {
+    return new BunStatementAdapter(this.db.query(sql), this.trackStatement());
+  }
+
   query(sql: string): SqliteStatement {
-    return new BunStatementAdapter(this.db.query(sql));
+    const cached = this.statementCache.get(sql);
+    if (cached) return cached;
+    let adapter: BunStatementAdapter;
+    const release = this.trackStatement();
+    adapter = new BunStatementAdapter(this.db.query(sql), () => {
+      this.statementCache.delete(sql);
+      release();
+    });
+    this.statementCache.set(sql, adapter);
+    return adapter;
   }
 
   transaction<T>(fn: () => T): () => T {
@@ -62,15 +134,37 @@ class BunDatabaseAdapter implements SqliteDatabase {
   }
 
   close(): void {
-    this.db.close();
+    if (this.closed) return;
+    this.closed = true;
+    const cachedStatements = Array.from(this.statementCache.values());
+    this.statementCache.clear();
+    for (const stmt of cachedStatements) {
+      try {
+        stmt.finalize?.();
+      } catch {
+        // Ignore finalizer failures during shutdown.
+      }
+    }
+    while (this.preparedStatementCount > 0) {
+      this.preparedStatementCount -= 1;
+      decrementPreparedStatement();
+    }
+    try {
+      this.db.close();
+    } finally {
+      decrementSqliteConnection();
+    }
   }
 }
 
 class NodeStatementAdapter implements SqliteStatement {
   private readonly stmt: any;
+  private readonly onFinalize?: () => void;
+  private finalized = false;
 
-  constructor(stmt: any) {
+  constructor(stmt: any, onFinalize?: () => void) {
     this.stmt = stmt;
+    this.onFinalize = onFinalize;
   }
 
   get(...params: any[]): any {
@@ -90,7 +184,13 @@ class NodeStatementAdapter implements SqliteStatement {
   }
 
   finalize(): void {
-    if (typeof this.stmt.finalize === "function") this.stmt.finalize();
+    if (this.finalized) return;
+    this.finalized = true;
+    try {
+      if (typeof this.stmt.finalize === "function") this.stmt.finalize();
+    } finally {
+      this.onFinalize?.();
+    }
   }
 }
 
@@ -98,19 +198,50 @@ class NodeDatabaseAdapter implements SqliteDatabase {
   private txDepth = 0;
   private txCounter = 0;
   private readonly db: any;
+  private preparedStatementCount = 0;
+  private closed = false;
+  private readonly statementCache = new Map<string, NodeStatementAdapter>();
 
   constructor(db: any) {
     this.db = db;
+    incrementSqliteConnection();
+  }
+
+  private trackStatement(): () => void {
+    let released = false;
+    this.preparedStatementCount += 1;
+    incrementPreparedStatement();
+    return () => {
+      if (released) return;
+      released = true;
+      this.preparedStatementCount = Math.max(0, this.preparedStatementCount - 1);
+      decrementPreparedStatement();
+    };
   }
 
   exec(sql: string): void {
     this.db.exec(sql);
   }
 
+  prepare(sql: string): SqliteStatement {
+    const stmt = this.db.prepare(sql);
+    if (typeof stmt?.setReadBigInts === "function") stmt.setReadBigInts(true);
+    return new NodeStatementAdapter(stmt, this.trackStatement());
+  }
+
   query(sql: string): SqliteStatement {
+    const cached = this.statementCache.get(sql);
+    if (cached) return cached;
     const stmt = this.db.prepare(sql);
     if (typeof stmt?.setReadBigInts === "function") stmt.setReadBigInts(true);
-    return new NodeStatementAdapter(stmt);
+    let adapter: NodeStatementAdapter;
+    const release = this.trackStatement();
+    adapter = new NodeStatementAdapter(stmt, () => {
+      this.statementCache.delete(sql);
+      release();
+    });
+    this.statementCache.set(sql, adapter);
+    return adapter;
   }
 
   transaction<T>(fn: () => T): () => T {
@@ -144,21 +275,61 @@ class NodeDatabaseAdapter implements SqliteDatabase {
   }
 
   close(): void {
-    this.db.close();
+    if (this.closed) return;
+    this.closed = true;
+    const cachedStatements = Array.from(this.statementCache.values());
+    this.statementCache.clear();
+    for (const stmt of cachedStatements) {
+      try {
+        stmt.finalize?.();
+      } catch {
+        // Ignore finalizer failures during shutdown.
+      }
+    }
+    while (this.preparedStatementCount > 0) {
+      this.preparedStatementCount -= 1;
+      decrementPreparedStatement();
+    }
+    try {
+      this.db.close();
+    } finally {
+      decrementSqliteConnection();
+    }
   }
 }
 
 let openImpl: ((path: string) => SqliteDatabase) | null = null;
+let openImplRuntime: HostRuntime | null = null;
+let runtimeOverride: HostRuntime | null = null;
+const require = createRequire(import.meta.url);
 
-if (typeof (globalThis as any).Bun !== "undefined") {
-  const { Database } = await import("bun:sqlite");
-  openImpl = (path: string) => new BunDatabaseAdapter(new Database(path));
-} else {
-  const { DatabaseSync } = await import("node:sqlite");
-  openImpl = (path: string) => new NodeDatabaseAdapter(new DatabaseSync(path));
+function selectedRuntime(): HostRuntime {
+  return runtimeOverride ?? detectHostRuntime();
+}
+
+function buildOpenImpl(runtime: HostRuntime): (path: string) => SqliteDatabase {
+  if (runtime === "bun") {
+    const { Database } = require("bun:sqlite") as { Database: new (path: string) => any };
+    return (path: string) => new BunDatabaseAdapter(new Database(path));
+  }
+  const { DatabaseSync } = require("node:sqlite") as { DatabaseSync: new (path: string) => any };
+  return (path: string) => new NodeDatabaseAdapter(new DatabaseSync(path));
+}
+
+export function setSqliteRuntimeOverride(runtime: HostRuntime | null): void {
+  runtimeOverride = runtime;
+  if (runtimeOverride && openImplRuntime && runtimeOverride !== openImplRuntime) {
+    openImpl = null;
+    openImplRuntime = null;
+  }
 }
 
 export function openSqliteDatabase(path: string): SqliteDatabase {
+  const runtime = selectedRuntime();
+  if (!openImpl || openImplRuntime !== runtime) {
+    openImpl = buildOpenImpl(runtime);
+    openImplRuntime = runtime;
+  }
   if (!openImpl) throw dsError("sqlite adapter not initialized");
   return openImpl(path);
 }

package/src/sqlite/runtime_stats.ts

@@ -0,0 +1,163 @@
+import { createRequire } from "node:module";
+import { detectHostRuntime } from "../runtime/host_runtime.ts";
+import { getSqliteAdapterRuntimeCounts } from "./adapter.ts";
+import type { SqliteRuntimeMemoryStats } from "../runtime_memory.ts";
+
+const SQLITE_STATUS_MEMORY_USED = 0;
+const SQLITE_STATUS_PAGECACHE_USED = 1;
+const SQLITE_STATUS_PAGECACHE_OVERFLOW = 2;
+const SQLITE_STATUS_MALLOC_COUNT = 9;
+
+type SqliteStatus64Fn = (op: number, currentPtr: number, highwaterPtr: number, resetFlag: number) => number;
+type BunFfiModule = {
+  dlopen: (
+    path: string,
+    symbols: {
+      sqlite3_status64: {
+        args: ["i32", "ptr", "ptr", "i32"];
+        returns: "i32";
+      };
+    }
+  ) => { symbols: { sqlite3_status64: SqliteStatus64Fn }; close: () => void };
+  ptr: (value: TypedArray) => number;
+};
+
+type TypedArray = BigInt64Array | BigUint64Array | Int32Array | Uint32Array;
+
+type SqliteLibraryHandle = {
+  sqlite3_status64: SqliteStatus64Fn;
+  close: () => void;
+};
+
+type CachedSqliteRuntimeStats = {
+  at_ms: number;
+  value: SqliteRuntimeMemoryStats;
+};
+
+let ffiModule: BunFfiModule | null | undefined;
+let sqliteLibrary: SqliteLibraryHandle | null | undefined;
+let cachedRuntimeStats: CachedSqliteRuntimeStats | null = null;
+const require = createRequire(import.meta.url);
+
+function sqliteLibraryNames(): string[] {
+  if (process.platform === "darwin") {
+    return ["libsqlite3.dylib", "/usr/lib/libsqlite3.dylib"];
+  }
+  if (process.platform === "linux") {
+    return ["libsqlite3.so.0", "libsqlite3.so", "/usr/lib/x86_64-linux-gnu/libsqlite3.so.0", "/lib/x86_64-linux-gnu/libsqlite3.so.0"];
+  }
+  return ["libsqlite3.so", "libsqlite3.dylib"];
+}
+
+function loadBunFfi(): BunFfiModule | null {
+  if (ffiModule !== undefined) return ffiModule;
+  if (detectHostRuntime() !== "bun") {
+    ffiModule = null;
+    return ffiModule;
+  }
+  try {
+    ffiModule = require("bun:ffi") as BunFfiModule;
+  } catch {
+    ffiModule = null;
+  }
+  return ffiModule;
+}
+
+function loadSqliteLibrary(): SqliteLibraryHandle | null {
+  if (sqliteLibrary !== undefined) return sqliteLibrary;
+  const ffi = loadBunFfi();
+  if (!ffi) {
+    sqliteLibrary = null;
+    return sqliteLibrary;
+  }
+  for (const name of sqliteLibraryNames()) {
+    try {
+      const lib = ffi.dlopen(name, {
+        sqlite3_status64: {
+          args: ["i32", "ptr", "ptr", "i32"],
+          returns: "i32",
+        },
+      });
+      sqliteLibrary = {
+        sqlite3_status64: lib.symbols.sqlite3_status64,
+        close: () => lib.close(),
+      };
+      return sqliteLibrary;
+    } catch {
+      // Try the next library name.
+    }
+  }
+  sqliteLibrary = null;
+  return sqliteLibrary;
+}
+
+function readStatus64(
+  lib: SqliteLibraryHandle,
+  ffi: BunFfiModule,
+  op: number
+): { current: number; highwater: number } | null {
+  const current = new BigInt64Array(1);
+  const highwater = new BigInt64Array(1);
+  try {
+    const rc = lib.sqlite3_status64(op, ffi.ptr(current), ffi.ptr(highwater), 0);
+    if (rc !== 0) return null;
+    return {
+      current: Number(current[0]),
+      highwater: Number(highwater[0]),
+    };
+  } catch {
+    return null;
+  }
+}
+
+export function readSqliteRuntimeMemoryStats(ttlMs = 1_000): SqliteRuntimeMemoryStats {
+  const now = Date.now();
+  if (cachedRuntimeStats && now - cachedRuntimeStats.at_ms < Math.max(0, ttlMs)) {
+    return cachedRuntimeStats.value;
+  }
+
+  const adapterCounts = getSqliteAdapterRuntimeCounts();
+  const unavailable: SqliteRuntimeMemoryStats = {
+    available: false,
+    source: "unavailable",
+    memory_used_bytes: 0,
+    memory_highwater_bytes: 0,
+    pagecache_used_slots: 0,
+    pagecache_used_slots_highwater: 0,
+    pagecache_overflow_bytes: 0,
+    pagecache_overflow_highwater_bytes: 0,
+    malloc_count: 0,
+    malloc_count_highwater: 0,
+    open_connections: adapterCounts.open_connections,
+    prepared_statements: adapterCounts.prepared_statements,
+  };
+
+  const ffi = loadBunFfi();
+  const lib = loadSqliteLibrary();
+  if (!ffi || !lib) {
+    cachedRuntimeStats = { at_ms: now, value: unavailable };
+    return unavailable;
+  }
+
+  const memoryUsed = readStatus64(lib, ffi, SQLITE_STATUS_MEMORY_USED);
+  const pagecacheUsed = readStatus64(lib, ffi, SQLITE_STATUS_PAGECACHE_USED);
+  const pagecacheOverflow = readStatus64(lib, ffi, SQLITE_STATUS_PAGECACHE_OVERFLOW);
+  const mallocCount = readStatus64(lib, ffi, SQLITE_STATUS_MALLOC_COUNT);
+
+  const stats: SqliteRuntimeMemoryStats = {
+    available: memoryUsed != null,
+    source: memoryUsed != null ? "sqlite3_status64" : "unavailable",
+    memory_used_bytes: Math.max(0, Math.floor(memoryUsed?.current ?? 0)),
+    memory_highwater_bytes: Math.max(0, Math.floor(memoryUsed?.highwater ?? 0)),
+    pagecache_used_slots: Math.max(0, Math.floor(pagecacheUsed?.current ?? 0)),
+    pagecache_used_slots_highwater: Math.max(0, Math.floor(pagecacheUsed?.highwater ?? 0)),
+    pagecache_overflow_bytes: Math.max(0, Math.floor(pagecacheOverflow?.current ?? 0)),
+    pagecache_overflow_highwater_bytes: Math.max(0, Math.floor(pagecacheOverflow?.highwater ?? 0)),
+    malloc_count: Math.max(0, Math.floor(mallocCount?.current ?? 0)),
+    malloc_count_highwater: Math.max(0, Math.floor(mallocCount?.highwater ?? 0)),
+    open_connections: adapterCounts.open_connections,
+    prepared_statements: adapterCounts.prepared_statements,
+  };
+  cachedRuntimeStats = { at_ms: now, value: stats };
+  return stats;
+}

package/src/stats.ts

@@ -1,6 +1,6 @@
 import type { SqliteDurableStore } from "./db/db";
 import type { UploaderController } from "./uploader";
-import type { MemoryGuard } from "./memory";
+import type { MemoryPressureMonitor } from "./memory";
 import type { BackpressureGate } from "./backpressure";
 import type { IngestQueue } from "./ingest";
 
@@ -110,7 +110,7 @@ export class StatsReporter {
   private readonly uploader: UploaderController;
   private readonly ingest?: IngestQueue;
   private readonly backpressure?: BackpressureGate;
-  private readonly memory?: MemoryGuard;
+  private readonly memory?: MemoryPressureMonitor;
 
   constructor(
     stats: StatsCollector,
@@ -118,7 +118,7 @@ export class StatsReporter {
     uploader: UploaderController,
     ingest?: IngestQueue,
     backpressure?: BackpressureGate,
-    memory?: MemoryGuard,
+    memory?: MemoryPressureMonitor,
     intervalMs = 60_000
   ) {
     this.stats = stats;

package/src/stream_size_reconciler.ts

@@ -0,0 +1,100 @@
+import { existsSync } from "node:fs";
+import { readFile } from "node:fs/promises";
+import { Result } from "better-result";
+import type { SqliteDurableStore, SegmentRow } from "./db/db";
+import type { ObjectStore } from "./objectstore/interface";
+import type { SegmentDiskCache } from "./segment/cache";
+import { loadSegmentBytesCached } from "./segment/cached_segment";
+import { iterateBlocksResult } from "./segment/format";
+import { dsError } from "./util/ds_error";
+import { yieldToEventLoop } from "./util/yield";
+
+export class StreamSizeReconciler {
+  private stopped = false;
+  private running: Promise<void> | null = null;
+
+  constructor(
+    private readonly db: SqliteDurableStore,
+    private readonly os: ObjectStore,
+    private readonly segmentCache?: SegmentDiskCache,
+    private readonly onMetadataChanged?: (stream: string) => void
+  ) {}
+
+  start(): void {
+    if (this.running) return;
+    this.running = this.run().finally(() => {
+      this.running = null;
+    });
+  }
+
+  stop(): void {
+    this.stopped = true;
+  }
+
+  private async run(): Promise<void> {
+    while (!this.stopped) {
+      const streams = this.db.listStreamsMissingLogicalSize(8);
+      if (streams.length === 0) return;
+      for (const stream of streams) {
+        if (this.stopped) return;
+        try {
+          await this.reconcileStream(stream);
+        } catch (e) {
+          const msg = String((e as any)?.message ?? e);
+          if (!this.stopped && !msg.includes("Statement has finalized")) {
+            // eslint-disable-next-line no-console
+            console.error("stream size reconcile failed", stream, e);
+          }
+        }
+      }
+    }
+  }
+
+  private async reconcileStream(stream: string): Promise<void> {
+    for (let attempt = 0; attempt < 3 && !this.stopped; attempt++) {
+      const before = this.db.getStream(stream);
+      if (!before || this.db.isDeleted(before) || before.next_offset <= 0n) return;
+      if (before.logical_size_bytes > 0n) return;
+
+      const segments = this.db.listSegmentsForStream(stream);
+      let total = 0n;
+      for (const segment of segments) {
+        if (this.stopped) return;
+        total += await this.sumSegmentPayloadBytes(segment);
+        await yieldToEventLoop();
+      }
+
+      const after = this.db.getStream(stream);
+      if (!after || this.db.isDeleted(after)) return;
+      if (after.logical_size_bytes > 0n) return;
+
+      if (segments.length !== this.db.countSegmentsForStream(stream)) continue;
+
+      const finalTotal = total + after.wal_bytes;
+      if (finalTotal > after.logical_size_bytes) {
+        this.db.setStreamLogicalSizeBytes(stream, finalTotal);
+        this.onMetadataChanged?.(stream);
+      }
+      return;
+    }
+  }
+
+  private async sumSegmentPayloadBytes(segment: SegmentRow): Promise<bigint> {
+    const bytes = await this.loadSegmentBytes(segment);
+    let total = 0n;
+    for (const blockRes of iterateBlocksResult(bytes)) {
+      if (Result.isError(blockRes)) throw dsError(blockRes.error.message);
+      for (const record of blockRes.value.decoded.records) {
+        total += BigInt(record.payload.byteLength);
+      }
+    }
+    return total;
+  }
+
+  private async loadSegmentBytes(segment: SegmentRow): Promise<Uint8Array> {
+    if (existsSync(segment.local_path)) {
+      return new Uint8Array(await readFile(segment.local_path));
+    }
+    return loadSegmentBytesCached(this.os, segment, this.segmentCache);
+  }
+}

package/src/touch/canonical_change.ts

@@ -0,0 +1,7 @@
+export type CanonicalChange = {
+  entity: string;
+  key?: string;
+  op: "insert" | "update" | "delete";
+  before?: unknown;
+  after?: unknown;
+};