@tungthedev/streams-server 0.2.0

package/src/server.ts

@@ -0,0 +1,264 @@
+import { loadConfig } from "./config";
+import { createApp, createPostgresApp, createPostgresFullApp, type App } from "./app";
+import { StatsCollector, StatsReporter } from "./stats";
+import { LatencyHistogramCollector, HistogramReporter } from "./hist";
+import { MockR2Store } from "./objectstore/mock_r2";
+import type { ObjectStore } from "./objectstore/interface";
+import { R2ObjectStore } from "./objectstore/r2";
+import { bootstrapFromR2 } from "./bootstrap";
+import { bootstrapPostgresFromR2 } from "./postgres/bootstrap";
+import { PostgresDurableStore } from "./postgres/store";
+import { initConsoleLogging } from "./util/log";
+import { applyAutoTune, AutoTuneApplyError, parseAutoTuneArg } from "./server_auto_tune";
+import { parseAuthConfigResult, withAuth } from "./auth";
+import { Result } from "better-result";
+
+initConsoleLogging();
+
+const args = process.argv.slice(2);
+const authConfigResult = parseAuthConfigResult(args);
+if (Result.isError(authConfigResult)) {
+  console.error(authConfigResult.error.message);
+  process.exit(1);
+}
+const authConfig = authConfigResult.value;
+
+const autoTune = parseAutoTuneArg(args);
+if (autoTune.enabled) {
+  try {
+    applyAutoTune(autoTune.valueMb);
+  } catch (error) {
+    if (error instanceof AutoTuneApplyError) {
+      console.error(error.message);
+      process.exit(1);
+    }
+    throw error;
+  }
+}
+
+const cfg = loadConfig();
+
+const statsEnabled = args.includes("--stats");
+const histEnabled = args.includes("--hist");
+const bootstrapEnabled = args.includes("--bootstrap-from-r2");
+const bpBudgetRaw = process.env.DS_BACKPRESSURE_BUDGET_MS;
+const bpBudgetMs = bpBudgetRaw ? Number(bpBudgetRaw) : cfg.ingestFlushIntervalMs + 1;
+if (bpBudgetRaw && !Number.isFinite(bpBudgetMs)) {
+  // eslint-disable-next-line no-console
+  console.error(`invalid DS_BACKPRESSURE_BUDGET_MS: ${bpBudgetRaw}`);
+  process.exit(1);
+}
+const stats = statsEnabled ? new StatsCollector({ backpressureBudgetMs: bpBudgetMs }) : undefined;
+const hist = histEnabled ? new LatencyHistogramCollector() : undefined;
+
+const storeIdx = args.indexOf("--object-store");
+const storeChoice = storeIdx >= 0 ? args[storeIdx + 1] : null;
+
+function requireObjectStoreChoice(): "r2" | "local" {
+  if (!storeChoice || (storeChoice !== "r2" && storeChoice !== "local")) {
+    // eslint-disable-next-line no-console
+    console.error("missing or invalid --object-store (expected: r2 | local)");
+    process.exit(1);
+  }
+  return storeChoice;
+}
+
+function createConfiguredObjectStore(choice: "r2" | "local"): ObjectStore {
+  if (choice === "local") {
+    const memBytesRaw = process.env.DS_MOCK_R2_MAX_INMEM_BYTES;
+    const memMbRaw = process.env.DS_MOCK_R2_MAX_INMEM_MB;
+    const putDelayRaw = process.env.DS_MOCK_R2_PUT_DELAY_MS;
+    const getDelayRaw = process.env.DS_MOCK_R2_GET_DELAY_MS;
+    const headDelayRaw = process.env.DS_MOCK_R2_HEAD_DELAY_MS;
+    const listDelayRaw = process.env.DS_MOCK_R2_LIST_DELAY_MS;
+    const memBytes = memBytesRaw ? Number(memBytesRaw) : memMbRaw ? Number(memMbRaw) * 1024 * 1024 : null;
+    const putDelayMs = putDelayRaw ? Number(putDelayRaw) : 0;
+    const getDelayMs = getDelayRaw ? Number(getDelayRaw) : 0;
+    const headDelayMs = headDelayRaw ? Number(headDelayRaw) : 0;
+    const listDelayMs = listDelayRaw ? Number(listDelayRaw) : 0;
+    if (memBytesRaw && !Number.isFinite(memBytes)) {
+      // eslint-disable-next-line no-console
+      console.error(`invalid DS_MOCK_R2_MAX_INMEM_BYTES: ${memBytesRaw}`);
+      process.exit(1);
+    }
+    if (memMbRaw && !Number.isFinite(Number(memMbRaw))) {
+      // eslint-disable-next-line no-console
+      console.error(`invalid DS_MOCK_R2_MAX_INMEM_MB: ${memMbRaw}`);
+      process.exit(1);
+    }
+    for (const [name, value] of [
+      ["DS_MOCK_R2_PUT_DELAY_MS", putDelayMs],
+      ["DS_MOCK_R2_GET_DELAY_MS", getDelayMs],
+      ["DS_MOCK_R2_HEAD_DELAY_MS", headDelayMs],
+      ["DS_MOCK_R2_LIST_DELAY_MS", listDelayMs],
+    ] as const) {
+      if (!Number.isFinite(value) || value < 0) {
+        // eslint-disable-next-line no-console
+        console.error(`invalid ${name}: ${process.env[name]}`);
+        process.exit(1);
+      }
+    }
+    return new MockR2Store({
+      maxInMemoryBytes: memBytes ?? undefined,
+      spillDir: process.env.DS_MOCK_R2_SPILL_DIR,
+      faults: {
+        putDelayMs,
+        getDelayMs,
+        headDelayMs,
+        listDelayMs,
+      },
+    });
+  }
+
+  const bucket = process.env.DURABLE_STREAMS_R2_BUCKET;
+  const accountId = process.env.DURABLE_STREAMS_R2_ACCOUNT_ID;
+  const accessKeyId = process.env.DURABLE_STREAMS_R2_ACCESS_KEY_ID;
+  const secretAccessKey = process.env.DURABLE_STREAMS_R2_SECRET_ACCESS_KEY;
+  const endpoint = process.env.DURABLE_STREAMS_R2_ENDPOINT;
+  const region = process.env.DURABLE_STREAMS_R2_REGION;
+  if (!bucket || !accountId || !accessKeyId || !secretAccessKey) {
+    // eslint-disable-next-line no-console
+    console.error("missing R2 env vars: DURABLE_STREAMS_R2_BUCKET, DURABLE_STREAMS_R2_ACCOUNT_ID, DURABLE_STREAMS_R2_ACCESS_KEY_ID, DURABLE_STREAMS_R2_SECRET_ACCESS_KEY");
+    process.exit(1);
+  }
+  return new R2ObjectStore({
+    bucket,
+    accountId,
+    accessKeyId,
+    secretAccessKey,
+    endpoint,
+    region,
+  });
+}
+
+let app: App;
+if (cfg.storage === "postgres") {
+  if (cfg.postgresUrl == null) {
+    // loadConfig validates this; keep the local guard for future call-site changes.
+    // eslint-disable-next-line no-console
+    console.error("DS_POSTGRES_URL is required when DS_STORAGE=postgres");
+    process.exit(1);
+  }
+  if (cfg.postgresMode === "wal") {
+    if (storeIdx >= 0) {
+      // eslint-disable-next-line no-console
+      console.error("postgres WAL mode does not support --object-store");
+      process.exit(1);
+    }
+    if (bootstrapEnabled) {
+      // eslint-disable-next-line no-console
+      console.error("postgres WAL mode does not support --bootstrap-from-r2");
+      process.exit(1);
+    }
+    const postgresStore = await PostgresDurableStore.connect(cfg.postgresUrl);
+    app = createPostgresApp(cfg, postgresStore, { stats });
+  } else {
+    const objectStore = createConfiguredObjectStore(requireObjectStoreChoice());
+    if (bootstrapEnabled) {
+      await bootstrapPostgresFromR2(cfg, objectStore, cfg.postgresUrl, { clearLocal: true });
+    }
+    const postgresStore = await PostgresDurableStore.connectFull(cfg.postgresUrl);
+    app = createPostgresFullApp(cfg, postgresStore, objectStore, { stats });
+  }
+} else {
+  const store = createConfiguredObjectStore(requireObjectStoreChoice());
+
+  if (bootstrapEnabled) {
+    await bootstrapFromR2(cfg, store, { clearLocal: true });
+  }
+
+  app = createApp(cfg, store, { stats });
+}
+const statsIntervalMs = process.env.DS_STATS_INTERVAL_MS ? Number(process.env.DS_STATS_INTERVAL_MS) : 60_000;
+if (process.env.DS_STATS_INTERVAL_MS && !Number.isFinite(statsIntervalMs)) {
+  // eslint-disable-next-line no-console
+  console.error(`invalid DS_STATS_INTERVAL_MS: ${process.env.DS_STATS_INTERVAL_MS}`);
+  process.exit(1);
+}
+const statsReporter =
+  statsEnabled && stats && app.deps.storageStats && app.deps.uploader
+    ? new StatsReporter(stats, app.deps.storageStats, app.deps.uploader, app.deps.ingest, app.deps.backpressure, app.deps.memory, statsIntervalMs)
+    : null;
+const histReporter = histEnabled && hist ? new HistogramReporter(hist, statsIntervalMs) : null;
+
+const fetchWithHist = hist
+  ? async (req: Request): Promise<Response> => {
+      const start = Date.now();
+      const resp = await app.fetch(req);
+      const url = req.url;
+      let path: string | null = null;
+      if (url.startsWith("/")) {
+        path = url;
+      } else {
+        const schemeIdx = url.indexOf("://");
+        if (schemeIdx !== -1) {
+          const pathIdx = url.indexOf("/", schemeIdx + 3);
+          path = pathIdx === -1 ? "/" : url.slice(pathIdx);
+        }
+      }
+      if (path) {
+        const isStream = path.startsWith("/v1/stream/") || path.startsWith("/v1/streams");
+        if (isStream) {
+          const ms = Date.now() - start;
+          const method = req.method.toUpperCase();
+          if (method === "GET" || method === "HEAD") hist.recordRead(ms);
+          else if (method === "POST" || method === "PUT" || method === "DELETE") hist.recordWrite(ms);
+        }
+  }
+  return resp;
+    }
+  : app.fetch;
+const fetchWithAuth = withAuth(authConfig, fetchWithHist);
+
+const server = Bun.serve({
+  hostname: cfg.host,
+  port: cfg.port,
+  // Default Bun idleTimeout is 10s, which is too low for long-poll endpoints like /touch/wait.
+  // Bun expects seconds here.
+  idleTimeout: (() => {
+    const raw = process.env.DS_HTTP_IDLE_TIMEOUT_SECONDS;
+    if (raw == null || raw.trim() === "") return 180;
+    const n = Number(raw);
+    if (!Number.isFinite(n) || n <= 0) {
+      // eslint-disable-next-line no-console
+      console.error(`invalid DS_HTTP_IDLE_TIMEOUT_SECONDS: ${raw}`);
+      process.exit(1);
+    }
+    return n;
+  })(),
+  fetch: fetchWithAuth,
+});
+
+statsReporter?.start();
+histReporter?.start();
+
+let shuttingDown = false;
+const shutdown = async (signal: NodeJS.Signals) => {
+  if (shuttingDown) return;
+  shuttingDown = true;
+  // eslint-disable-next-line no-console
+  console.log(`received ${signal}, shutting down prisma-streams server`);
+  statsReporter?.stop();
+  histReporter?.stop();
+  try {
+    server.stop(true);
+  } catch (err) {
+    // eslint-disable-next-line no-console
+    console.error("failed to stop HTTP server cleanly", err);
+  }
+  try {
+    await app.close();
+  } catch (err) {
+    // eslint-disable-next-line no-console
+    console.error("failed to close application cleanly", err);
+    process.exitCode = 1;
+  }
+};
+
+const listenTarget = cfg.host.includes(":") ? `[${cfg.host}]:${server.port}` : `${cfg.host}:${server.port}`;
+
+process.once("SIGINT", () => shutdown("SIGINT"));
+process.once("SIGTERM", () => shutdown("SIGTERM"));
+
+// eslint-disable-next-line no-console
+console.log(`prisma-streams server listening on ${listenTarget}`);

package/src/server_auto_tune.ts

@@ -0,0 +1,158 @@
+import { AUTO_TUNE_PRESETS, memoryLimitForPreset, tuneForPreset, type AutoTuneConfig } from "./auto_tune";
+
+export class AutoTuneApplyError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "AutoTuneApplyError";
+  }
+}
+
+export function parseAutoTuneArg(args: string[]): { enabled: boolean; valueMb: number | null } {
+  let enabled = false;
+  let valueMb: number | null = null;
+  for (let i = 0; i < args.length; i++) {
+    const arg = args[i];
+    if (arg === "--auto-tune") {
+      enabled = true;
+      const next = args[i + 1];
+      if (next && !next.startsWith("--") && /^[0-9]+$/.test(next)) {
+        valueMb = Number(next);
+      }
+    } else if (arg.startsWith("--auto-tune=")) {
+      enabled = true;
+      const raw = arg.split("=", 2)[1] ?? "";
+      if (raw.trim() !== "") valueMb = Number(raw);
+    }
+  }
+  return { enabled, valueMb };
+}
+
+function formatPresetList<T>(presets: number[], selected: number, map: (preset: number) => T, fmt: (val: T) => string): string {
+  return presets
+    .map((preset) => {
+      const value = fmt(map(preset));
+      return preset === selected ? `[${value}]` : value;
+    })
+    .join(", ");
+}
+
+export function applyAutoTune(
+  overrideMb: number | null,
+  opts: {
+    env?: NodeJS.ProcessEnv;
+    log?: (message: string) => void;
+  } = {}
+): void {
+  const env = opts.env ?? process.env;
+  const log = opts.log ?? console.log;
+  const envMemRaw = env.DS_MEMORY_LIMIT_MB;
+  if (overrideMb != null) {
+    if (envMemRaw) {
+      throw new AutoTuneApplyError("--auto-tune with a value cannot be used with DS_MEMORY_LIMIT_MB");
+    }
+  } else if (!envMemRaw) {
+    throw new AutoTuneApplyError("--auto-tune requires DS_MEMORY_LIMIT_MB to be set (or pass a value)");
+  }
+  const memMb = overrideMb != null ? overrideMb : Number(envMemRaw);
+  if (!Number.isFinite(memMb) || memMb <= 0) {
+    const bad = overrideMb != null ? String(overrideMb) : String(envMemRaw);
+    throw new AutoTuneApplyError(`invalid DS_MEMORY_LIMIT_MB: ${bad}`);
+  }
+  if (env.DS_MEMORY_LIMIT_BYTES) {
+    throw new AutoTuneApplyError("--auto-tune does not allow DS_MEMORY_LIMIT_BYTES; use DS_MEMORY_LIMIT_MB");
+  }
+
+  const conflictVars = [
+    "DS_SEGMENT_MAX_BYTES",
+    "DS_SEGMENT_TARGET_ROWS",
+    "DS_SEGMENT_CACHE_MAX_BYTES",
+    "DS_INDEX_CHECK_MS",
+    "DS_SQLITE_CACHE_MB",
+    "DS_SQLITE_CACHE_BYTES",
+    "DS_WORKER_SQLITE_CACHE_MB",
+    "DS_WORKER_SQLITE_CACHE_BYTES",
+    "DS_INDEX_RUN_MEM_CACHE_BYTES",
+    "DS_LEXICON_INDEX_CACHE_MAX_BYTES",
+    "DS_INGEST_MAX_BATCH_BYTES",
+    "DS_INGEST_MAX_QUEUE_BYTES",
+    "DS_INGEST_CONCURRENCY",
+    "DS_READ_CONCURRENCY",
+    "DS_SEARCH_CONCURRENCY",
+    "DS_ASYNC_INDEX_CONCURRENCY",
+    "DS_INDEX_BUILD_CONCURRENCY",
+    "DS_INDEX_COMPACT_CONCURRENCY",
+    "DS_SEGMENTER_WORKERS",
+    "DS_UPLOAD_CONCURRENCY",
+    "DS_SEARCH_COMPANION_TOC_CACHE_BYTES",
+    "DS_SEARCH_COMPANION_SECTION_CACHE_BYTES",
+    "DS_SEARCH_COMPANION_BATCH_SEGMENTS",
+    "DS_SEARCH_COMPANION_YIELD_BLOCKS",
+  ];
+  const conflicts = conflictVars.filter((v) => env[v] != null);
+  if (conflicts.length > 0) {
+    throw new AutoTuneApplyError(`--auto-tune cannot be used with manual memory settings: ${conflicts.join(", ")}`);
+  }
+
+  const presets = [...AUTO_TUNE_PRESETS];
+  const preset = [...presets].reverse().find((v) => v <= memMb);
+  if (!preset) {
+    throw new AutoTuneApplyError(`DS_MEMORY_LIMIT_MB=${memMb} is below the minimum preset (256)`);
+  }
+  const tune: AutoTuneConfig = tuneForPreset(preset);
+
+  const memoryLimitMb = memoryLimitForPreset(preset);
+  env.DS_AUTO_TUNE_REQUESTED_MB = String(memMb);
+  env.DS_AUTO_TUNE_PRESET_MB = String(preset);
+  env.DS_AUTO_TUNE_EFFECTIVE_MEMORY_LIMIT_MB = String(memoryLimitMb);
+  env.DS_MEMORY_LIMIT_MB = String(memoryLimitMb);
+  env.DS_SEGMENT_MAX_BYTES = String(tune.segmentMaxMiB * 1024 * 1024);
+  env.DS_SEGMENT_TARGET_ROWS = String(tune.segmentTargetRows);
+  env.DS_SEGMENT_CACHE_MAX_BYTES = String(tune.segmentCacheMb * 1024 * 1024);
+  env.DS_INDEX_CHECK_MS = String(tune.indexCheckMs);
+  env.DS_SQLITE_CACHE_MB = String(tune.sqliteCacheMb);
+  env.DS_WORKER_SQLITE_CACHE_MB = String(tune.workerSqliteCacheMb);
+  env.DS_INDEX_RUN_MEM_CACHE_BYTES = String(tune.indexMemMb * 1024 * 1024);
+  env.DS_LEXICON_INDEX_CACHE_MAX_BYTES = String(tune.lexiconIndexCacheMb * 1024 * 1024);
+  env.DS_SEARCH_COMPANION_TOC_CACHE_BYTES = String(tune.searchCompanionTocCacheMb * 1024 * 1024);
+  env.DS_SEARCH_COMPANION_SECTION_CACHE_BYTES = String(tune.searchCompanionSectionCacheMb * 1024 * 1024);
+  env.DS_INGEST_MAX_BATCH_BYTES = String(tune.ingestBatchMb * 1024 * 1024);
+  env.DS_INGEST_MAX_QUEUE_BYTES = String(tune.ingestQueueMb * 1024 * 1024);
+  env.DS_INGEST_CONCURRENCY = String(tune.ingestConcurrency);
+  env.DS_READ_CONCURRENCY = String(tune.readConcurrency);
+  env.DS_SEARCH_CONCURRENCY = String(tune.searchConcurrency);
+  env.DS_ASYNC_INDEX_CONCURRENCY = String(tune.asyncIndexConcurrency);
+  env.DS_INDEX_BUILD_CONCURRENCY = String(tune.indexBuildConcurrency);
+  env.DS_INDEX_COMPACT_CONCURRENCY = String(tune.indexCompactConcurrency);
+  env.DS_SEGMENTER_WORKERS = String(tune.segmenterWorkers);
+  env.DS_UPLOAD_CONCURRENCY = String(tune.uploadConcurrency);
+  env.DS_SEARCH_COMPANION_BATCH_SEGMENTS = String(tune.searchCompanionBatchSegments);
+  env.DS_SEARCH_COMPANION_YIELD_BLOCKS = String(tune.searchCompanionYieldBlocks);
+
+  const presetLine = formatPresetList(presets, preset, (v) => v, (v) => String(v));
+  log(`Auto-tuning for memory preset ${presetLine}`);
+  log(
+    `DS_MEMORY_LIMIT_MB presets: ${formatPresetList(presets, preset, (p) => memoryLimitForPreset(p), (v) => String(v))}`
+  );
+  log(`DS_SEGMENT_MAX_MIB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).segmentMaxMiB, (v) => String(v))}`);
+  log(`DS_SEGMENT_TARGET_ROWS presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).segmentTargetRows, (v) => String(v))}`);
+  log(`DS_SEGMENT_CACHE_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).segmentCacheMb, (v) => String(v))}`);
+  log(`DS_INDEX_CHECK_MS presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).indexCheckMs, (v) => String(v))}`);
+  log(`DS_SQLITE_CACHE_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).sqliteCacheMb, (v) => String(v))}`);
+  log(`DS_WORKER_SQLITE_CACHE_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).workerSqliteCacheMb, (v) => String(v))}`);
+  log(`DS_INDEX_RUN_MEM_CACHE_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).indexMemMb, (v) => String(v))}`);
+  log(`DS_LEXICON_INDEX_CACHE_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).lexiconIndexCacheMb, (v) => String(v))}`);
+  log(`DS_SEARCH_COMPANION_TOC_CACHE_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).searchCompanionTocCacheMb, (v) => String(v))}`);
+  log(`DS_SEARCH_COMPANION_SECTION_CACHE_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).searchCompanionSectionCacheMb, (v) => String(v))}`);
+  log(`DS_INGEST_MAX_BATCH_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).ingestBatchMb, (v) => String(v))}`);
+  log(`DS_INGEST_MAX_QUEUE_MB presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).ingestQueueMb, (v) => String(v))}`);
+  log(`DS_INGEST_CONCURRENCY presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).ingestConcurrency, (v) => String(v))}`);
+  log(`DS_READ_CONCURRENCY presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).readConcurrency, (v) => String(v))}`);
+  log(`DS_SEARCH_CONCURRENCY presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).searchConcurrency, (v) => String(v))}`);
+  log(`DS_ASYNC_INDEX_CONCURRENCY presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).asyncIndexConcurrency, (v) => String(v))}`);
+  log(`DS_INDEX_BUILD_CONCURRENCY presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).indexBuildConcurrency, (v) => String(v))}`);
+  log(`DS_INDEX_COMPACT_CONCURRENCY presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).indexCompactConcurrency, (v) => String(v))}`);
+  log(`DS_SEGMENTER_WORKERS presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).segmenterWorkers, (v) => String(v))}`);
+  log(`DS_UPLOAD_CONCURRENCY presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).uploadConcurrency, (v) => String(v))}`);
+  log(`DS_SEARCH_COMPANION_BATCH_SEGMENTS presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).searchCompanionBatchSegments, (v) => String(v))}`);
+  log(`DS_SEARCH_COMPANION_YIELD_BLOCKS presets: ${formatPresetList(presets, preset, (p) => tuneForPreset(p).searchCompanionYieldBlocks, (v) => String(v))}`);
+}