@dotdo/postgres 0.1.1 → 0.1.3

package/src/worker/background-pglite-manager.ts

@@ -0,0 +1,680 @@
+/**
+ * Background PGLite Manager for PostgresDO
+ *
+ * This module implements a "eager-but-non-blocking" loading pattern:
+ * - WASM loading starts immediately on initialize() but doesn't block
+ * - Non-query endpoints respond instantly while WASM loads in background
+ * - Queries wait for WASM if not ready, but loading already started
+ * - Uses ctx.waitUntil() to keep DO alive during background loading
+ *
+ * ## Why This is Better Than Pure Lazy Loading
+ *
+ * Pure lazy loading just kicks the can down the road:
+ * - First query still pays full ~1200ms WASM load time
+ *
+ * Background loading gives us the best of both worlds:
+ * - Non-query endpoints respond instantly
+ * - WASM starts loading immediately
+ * - First query only waits for remaining load time (often near-zero)
+ * - If query arrives before WASM ready, it waits but loading is already in progress
+ *
+ * ## Usage with Durable Objects
+ *
+ * ```typescript
+ * export class MyDO extends DurableObject {
+ *   private manager: BackgroundPGLiteManager
+ *
+ *   constructor(ctx: DurableObjectState, env: Env) {
+ *     super(ctx, env)
+ *     this.manager = new BackgroundPGLiteManager({
+ *       database: 'mydb',
+ *       waitUntil: (p) => ctx.waitUntil(p), // Keep DO alive during load
+ *     })
+ *   }
+ *
+ *   async init() {
+ *     // Starts WASM loading in background, returns immediately
+ *     await this.manager.initialize()
+ *   }
+ *
+ *   // Health check - responds instantly, triggers background load if needed
+ *   async ping() {
+ *     return { ok: true, wasmLoading: this.manager.isLoading() }
+ *   }
+ *
+ *   // Query - waits for WASM if not ready
+ *   async query(sql: string) {
+ *     return this.manager.query(sql)
+ *   }
+ * }
+ * ```
+ *
+ * @module worker/background-pglite-manager
+ */
+
+import type { PostgresConfig } from './types'
+import { PluginManager, type PluginManagerConfig } from './plugin-manager'
+import type { PGliteLike } from './do-pglite-manager'
+
+// Import WASM module and data bundle for Workers
+import pgliteWasm from '../pglite-assets/pglite.wasm'
+import pgliteData from '../pglite-assets/pglite.data'
+
+// Import Workers-compatible PGLite wrapper
+import { createWorkersPGLite, type WorkersPGLite } from '../pglite/workers-pglite'
+
+// =============================================================================
+// Module-Level WASM State (for hoisting across DO reinstantiation)
+// =============================================================================
+
+/**
+ * Hoisted PGLite instance - survives DO class reinstantiation within same isolate.
+ */
+let bgHoistedPglite: WorkersPGLite | null = null
+
+/**
+ * Promise for in-progress PGLite initialization.
+ * Shared across all manager instances for deduplication.
+ */
+let bgHoistedPglitePromise: Promise<WorkersPGLite> | null = null
+
+/**
+ * Timestamp when WASM loading started
+ */
+let wasmLoadStartedAt: number | null = null
+
+/**
+ * Timestamp when WASM was loaded
+ */
+let wasmLoadedAt: number | null = null
+
+/**
+ * Unique identifier for this module instance
+ */
+const BG_MODULE_INSTANCE_ID = Math.random().toString(36).slice(2, 10)
+
+/**
+ * Check if the hoisted PGLite instance exists.
+ */
+export function hasBgHoistedPglite(): boolean {
+  return bgHoistedPglite !== null
+}
+
+/**
+ * Check if WASM loading is in progress.
+ */
+export function isBgWasmLoading(): boolean {
+  return bgHoistedPglitePromise !== null && bgHoistedPglite === null
+}
+
+/**
+ * Get diagnostics about the background hoisted WASM state.
+ */
+export function getBgHoistedPgliteDiagnostics(): {
+  hasInstance: boolean
+  isLoading: boolean
+  moduleInstanceId: string
+  wasmLoadStartedAt: number | null
+  wasmLoadedAt: number | null
+  loadDurationMs: number | null
+  timeSinceLoadMs: number | null
+} {
+  const now = Date.now()
+  return {
+    hasInstance: bgHoistedPglite !== null,
+    isLoading: bgHoistedPglitePromise !== null && bgHoistedPglite === null,
+    moduleInstanceId: BG_MODULE_INSTANCE_ID,
+    wasmLoadStartedAt,
+    wasmLoadedAt,
+    loadDurationMs: wasmLoadStartedAt && wasmLoadedAt ? wasmLoadedAt - wasmLoadStartedAt : null,
+    timeSinceLoadMs: wasmLoadedAt !== null ? now - wasmLoadedAt : null,
+  }
+}
+
+/**
+ * Reset the hoisted PGLite instance.
+ * WARNING: This should only be used in tests.
+ * @internal
+ */
+export function resetBgHoistedPglite(): void {
+  bgHoistedPglite = null
+  bgHoistedPglitePromise = null
+  wasmLoadStartedAt = null
+  wasmLoadedAt = null
+}
+
+/**
+ * Start WASM loading and return the promise.
+ * This is called by initialize() but not awaited.
+ */
+function startWasmLoading(options: {
+  database?: string
+  debug?: 0 | 1
+}): Promise<WorkersPGLite> {
+  // Fast path: already initialized
+  if (bgHoistedPglite) {
+    return Promise.resolve(bgHoistedPglite)
+  }
+
+  // Deduplication: return existing promise if loading in progress
+  if (bgHoistedPglitePromise) {
+    return bgHoistedPglitePromise
+  }
+
+  wasmLoadStartedAt = Date.now()
+  console.log(`[BackgroundPGLiteManager] Starting WASM load in background - instance: ${BG_MODULE_INSTANCE_ID}`)
+
+  // Start initialization
+  bgHoistedPglitePromise = createWorkersPGLite({
+    wasmModule: pgliteWasm,
+    fsBundle: pgliteData,
+    database: options.database ?? 'postgres',
+    debug: options.debug ?? 0,
+  }).then((pglite) => {
+    bgHoistedPglite = pglite
+    wasmLoadedAt = Date.now()
+    const loadDuration = wasmLoadedAt - (wasmLoadStartedAt ?? wasmLoadedAt)
+    console.log(
+      `[BackgroundPGLiteManager] WASM LOADED - took ${loadDuration}ms, instance: ${BG_MODULE_INSTANCE_ID}`
+    )
+    return pglite
+  }).catch((error) => {
+    console.error(`[BackgroundPGLiteManager] WASM load failed:`, error)
+    // Clear the promise so we can retry
+    bgHoistedPglitePromise = null
+    throw error
+  })
+
+  return bgHoistedPglitePromise
+}
+
+/**
+ * Get the hoisted instance, waiting for loading if in progress.
+ */
+async function getOrAwaitHoistedPglite(options: {
+  database?: string
+  debug?: 0 | 1
+}): Promise<WorkersPGLite> {
+  // Fast path: already loaded
+  if (bgHoistedPglite) {
+    return bgHoistedPglite
+  }
+
+  // Loading in progress - wait for it
+  if (bgHoistedPglitePromise) {
+    return bgHoistedPglitePromise
+  }
+
+  // Not started yet - start now (shouldn't happen if initialize() was called)
+  return startWasmLoading(options)
+}
+
+// =============================================================================
+// Loading State
+// =============================================================================
+
+export type BGWASMLoadingState = 'not_started' | 'loading' | 'loaded' | 'error'
+
+// =============================================================================
+// Configuration
+// =============================================================================
+
+export interface BackgroundPGLiteManagerConfig {
+  /** Database name */
+  database?: string
+  /** Enable debug mode */
+  debug?: boolean
+  /** Plugin configuration */
+  plugins?: PostgresConfig['plugins']
+  /** Custom PGLite factory for testing */
+  createPGLite?: () => Promise<PGliteLike>
+  /**
+   * Disable WASM hoisting optimization.
+   * @default false
+   */
+  disableHoisting?: boolean
+  /**
+   * waitUntil function from Durable Object context.
+   * Used to keep the DO alive while WASM loads in background.
+   *
+   * @example
+   * ```typescript
+   * new BackgroundPGLiteManager({
+   *   waitUntil: (p) => ctx.waitUntil(p)
+   * })
+   * ```
+   */
+  waitUntil?: (promise: Promise<unknown>) => void
+}
+
+// =============================================================================
+// BackgroundPGLiteManager Class
+// =============================================================================
+
+/**
+ * BackgroundPGLiteManager starts WASM loading immediately but doesn't block.
+ *
+ * This gives the best of both worlds:
+ * - Non-query endpoints respond instantly
+ * - WASM starts loading on first init
+ * - Queries wait only for remaining load time (often zero if loaded)
+ * - Uses ctx.waitUntil() to keep DO alive during background loading
+ */
+export class BackgroundPGLiteManager {
+  private pglite: PGliteLike | null = null
+  private initialized = false
+  private loadingState: BGWASMLoadingState = 'not_started'
+  private loadError: Error | null = null
+  private config: BackgroundPGLiteManagerConfig
+  private pluginManager: PluginManager
+  private usingHoistedInstance = false
+  private loadPromise: Promise<PGliteLike> | null = null
+
+  /** Timing metrics for diagnostics */
+  private metrics = {
+    initializeCalledAt: null as number | null,
+    wasmLoadStartAt: null as number | null,
+    wasmLoadEndAt: null as number | null,
+    firstQueryAt: null as number | null,
+    firstQueryWaitedMs: null as number | null,
+  }
+
+  constructor(config: BackgroundPGLiteManagerConfig = {}) {
+    this.config = {
+      database: config.database ?? 'postgres',
+      debug: config.debug ?? false,
+      ...config,
+    }
+
+    const pluginManagerConfig: PluginManagerConfig = {}
+    if (config.plugins !== undefined) {
+      pluginManagerConfig.plugins = config.plugins
+    }
+    if (config.debug !== undefined) {
+      pluginManagerConfig.debug = config.debug
+    }
+    this.pluginManager = new PluginManager(pluginManagerConfig)
+  }
+
+  /**
+   * Initialize the manager.
+   *
+   * This starts WASM loading in the background but returns immediately.
+   * Non-query operations can proceed while WASM loads.
+   *
+   * If waitUntil is provided, it's used to keep the DO alive during loading.
+   */
+  async initialize(): Promise<void> {
+    if (this.initialized) {
+      return
+    }
+
+    this.metrics.initializeCalledAt = Date.now()
+    this.initialized = true
+
+    // If custom factory provided, use it (for testing)
+    if (this.config.createPGLite) {
+      this.loadingState = 'loading'
+      this.loadPromise = this.config.createPGLite()
+
+      this.loadPromise.then((pg) => {
+        this.pglite = pg
+        this.loadingState = 'loaded'
+        this.metrics.wasmLoadEndAt = Date.now()
+      }).catch((err) => {
+        this.loadingState = 'error'
+        this.loadError = err
+      })
+
+      if (this.config.waitUntil) {
+        this.config.waitUntil(this.loadPromise.catch(() => {}))
+      }
+
+      console.log(`[BackgroundPGLiteManager] Initialized (custom factory, loading in background)`)
+      return
+    }
+
+    // If hoisting is disabled, start loading but don't hoist
+    if (this.config.disableHoisting) {
+      this.loadingState = 'loading'
+      this.metrics.wasmLoadStartAt = Date.now()
+
+      this.loadPromise = createWorkersPGLite({
+        wasmModule: pgliteWasm,
+        fsBundle: pgliteData,
+        database: this.config.database ?? 'postgres',
+        debug: this.config.debug ? 1 : 0,
+      })
+
+      this.loadPromise.then((pg) => {
+        this.pglite = pg
+        this.loadingState = 'loaded'
+        this.metrics.wasmLoadEndAt = Date.now()
+        console.log(`[BackgroundPGLiteManager] WASM LOADED (non-hoisted) - took ${this.metrics.wasmLoadEndAt - (this.metrics.wasmLoadStartAt ?? 0)}ms`)
+      }).catch((err) => {
+        this.loadingState = 'error'
+        this.loadError = err
+        console.error(`[BackgroundPGLiteManager] WASM load failed:`, err)
+      })
+
+      if (this.config.waitUntil) {
+        this.config.waitUntil(this.loadPromise.catch(() => {}))
+      }
+
+      console.log(`[BackgroundPGLiteManager] Initialized (hoisting disabled, loading in background)`)
+      return
+    }
+
+    // Use hoisted instance
+    this.usingHoistedInstance = true
+
+    // Check if already loaded
+    if (bgHoistedPglite) {
+      this.pglite = bgHoistedPglite
+      this.loadingState = 'loaded'
+      console.log(`[BackgroundPGLiteManager] Initialized (WASM already loaded, reusing hoisted instance)`)
+      return
+    }
+
+    // Start loading in background
+    this.loadingState = 'loading'
+    this.metrics.wasmLoadStartAt = Date.now()
+
+    const loadPromise = startWasmLoading({
+      database: this.config.database,
+      debug: this.config.debug ? 1 : 0,
+    })
+
+    // Track when loading completes for this manager
+    loadPromise.then((pg) => {
+      this.pglite = pg
+      this.loadingState = 'loaded'
+      this.metrics.wasmLoadEndAt = Date.now()
+    }).catch((err) => {
+      this.loadingState = 'error'
+      this.loadError = err
+    })
+
+    // Use waitUntil to keep DO alive during loading
+    if (this.config.waitUntil) {
+      this.config.waitUntil(loadPromise.catch(() => {}))
+    }
+
+    console.log(`[BackgroundPGLiteManager] Initialized (starting WASM load in background)`)
+  }
+
+  /**
+   * Check if WASM is fully loaded and ready for queries.
+   */
+  isWASMLoaded(): boolean {
+    return this.loadingState === 'loaded' && this.pglite !== null
+  }
+
+  /**
+   * Check if WASM is currently loading.
+   */
+  isLoading(): boolean {
+    return this.loadingState === 'loading'
+  }
+
+  /**
+   * Get the current loading state.
+   */
+  getLoadingState(): BGWASMLoadingState {
+    return this.loadingState
+  }
+
+  /**
+   * Get the PGLite instance, waiting for loading to complete if needed.
+   * This is called internally by query methods.
+   */
+  private async getPglite(): Promise<PGliteLike> {
+    // Already loaded
+    if (this.pglite) {
+      return this.pglite
+    }
+
+    // Using hoisted instance
+    if (this.usingHoistedInstance && !this.config.createPGLite) {
+      const waitStart = performance.now()
+      const pg = await getOrAwaitHoistedPglite({
+        database: this.config.database,
+        debug: this.config.debug ? 1 : 0,
+      })
+      this.pglite = pg
+      this.loadingState = 'loaded'
+      const waitMs = performance.now() - waitStart
+      if (waitMs > 1) {
+        console.log(`[BackgroundPGLiteManager] Query waited ${waitMs.toFixed(2)}ms for WASM to finish loading`)
+      }
+      return pg
+    }
+
+    // Using custom factory or non-hoisted
+    if (this.loadPromise) {
+      const waitStart = performance.now()
+      const pg = await this.loadPromise
+      const waitMs = performance.now() - waitStart
+      if (waitMs > 1) {
+        console.log(`[BackgroundPGLiteManager] Query waited ${waitMs.toFixed(2)}ms for WASM to finish loading`)
+      }
+      return pg
+    }
+
+    throw new Error('PGLite not initialized - call initialize() first')
+  }
+
+  /**
+   * Execute a query.
+   *
+   * If WASM is not loaded yet, this will wait for loading to complete.
+   * If WASM is already loaded, this returns immediately after query execution.
+   */
+  async query<T = unknown>(
+    sql: string,
+    params?: unknown[]
+  ): Promise<{
+    rows: T[]
+    fields: { name: string; dataTypeID: number }[]
+    affectedRows?: number
+  }> {
+    if (!this.initialized) {
+      throw new Error('Manager not initialized - call initialize() first')
+    }
+
+    if (this.loadingState === 'error') {
+      throw this.loadError ?? new Error('WASM loading failed')
+    }
+
+    // Track first query timing
+    if (this.metrics.firstQueryAt === null) {
+      this.metrics.firstQueryAt = Date.now()
+    }
+
+    const pg = await this.getPglite()
+    return pg.query<T>(sql, params)
+  }
+
+  /**
+   * Execute a SQL statement without returning rows.
+   */
+  async exec(sql: string): Promise<void> {
+    const pg = await this.getPglite()
+    if (pg.exec) {
+      await pg.exec(sql)
+    } else {
+      await pg.query(sql)
+    }
+  }
+
+  /**
+   * Check if the manager is initialized (not the same as WASM loaded).
+   */
+  isInitialized(): boolean {
+    return this.initialized
+  }
+
+  /**
+   * Get the PGLite instance if loaded, null otherwise.
+   * Does NOT wait for loading - use for checking readiness.
+   */
+  getInstanceOrNull(): PGliteLike | null {
+    return this.pglite
+  }
+
+  /**
+   * Get the PGLite instance (throws if not loaded).
+   * Use query() instead to automatically wait for loading.
+   */
+  getInstance(): PGliteLike {
+    if (!this.pglite) {
+      throw new Error('PGLite not loaded yet')
+    }
+    return this.pglite
+  }
+
+  /**
+   * Wait for WASM to be loaded.
+   * Call this if you want to ensure WASM is ready before proceeding.
+   */
+  async ensureWASMLoaded(): Promise<void> {
+    await this.getPglite()
+  }
+
+  /**
+   * Check if PGLite is responsive.
+   */
+  async checkLiveness(): Promise<boolean> {
+    if (!this.isWASMLoaded()) {
+      return false
+    }
+
+    try {
+      const result = await this.pglite!.query('SELECT 1 as alive')
+      return result.rows.length > 0
+    } catch {
+      return false
+    }
+  }
+
+  /**
+   * Get timing metrics for diagnostics.
+   */
+  getMetrics(): {
+    initializeCalledAt: number | null
+    wasmLoadStartAt: number | null
+    wasmLoadEndAt: number | null
+    wasmLoadDurationMs: number | null
+    firstQueryAt: number | null
+    timeFromInitToFirstQueryMs: number | null
+    isUsingHoistedInstance: boolean
+  } {
+    const loadDuration = this.metrics.wasmLoadStartAt && this.metrics.wasmLoadEndAt
+      ? this.metrics.wasmLoadEndAt - this.metrics.wasmLoadStartAt
+      : null
+    const initToFirstQuery = this.metrics.initializeCalledAt && this.metrics.firstQueryAt
+      ? this.metrics.firstQueryAt - this.metrics.initializeCalledAt
+      : null
+
+    return {
+      ...this.metrics,
+      wasmLoadDurationMs: loadDuration,
+      timeFromInitToFirstQueryMs: initToFirstQuery,
+      isUsingHoistedInstance: this.usingHoistedInstance,
+    }
+  }
+
+  /**
+   * Get comprehensive diagnostics.
+   */
+  getDiagnostics(): {
+    manager: {
+      initialized: boolean
+      loadingState: BGWASMLoadingState
+      hasInstance: boolean
+      usingHoistedInstance: boolean
+      error: string | null
+    }
+    module: ReturnType<typeof getBgHoistedPgliteDiagnostics>
+    metrics: ReturnType<BackgroundPGLiteManager['getMetrics']>
+  } {
+    return {
+      manager: {
+        initialized: this.initialized,
+        loadingState: this.loadingState,
+        hasInstance: this.pglite !== null,
+        usingHoistedInstance: this.usingHoistedInstance,
+        error: this.loadError?.message ?? null,
+      },
+      module: getBgHoistedPgliteDiagnostics(),
+      metrics: this.getMetrics(),
+    }
+  }
+
+  /**
+   * Close the manager.
+   * Note: Does NOT close hoisted instances (they're reused across DO reinstantiations).
+   */
+  async close(timeoutMs: number = 5000): Promise<void> {
+    // Don't close hoisted instances
+    if (this.usingHoistedInstance) {
+      console.log(`[BackgroundPGLiteManager] Skipping close of hoisted WASM instance (will be reused)`)
+      this.pglite = null
+      this.initialized = false
+      return
+    }
+
+    // Close non-hoisted instance
+    if (this.pglite?.close) {
+      try {
+        await Promise.race([
+          this.pglite.close(),
+          new Promise<void>((_, reject) =>
+            setTimeout(() => reject(new Error('Close timeout')), timeoutMs)
+          ),
+        ])
+      } catch (error) {
+        console.error('[BackgroundPGLiteManager] Error closing PGLite:', error)
+      }
+    }
+
+    this.pglite = null
+    this.initialized = false
+    this.loadingState = 'not_started'
+  }
+
+  /**
+   * Reset the manager state.
+   */
+  reset(): void {
+    this.pglite = null
+    this.initialized = false
+    this.loadingState = 'not_started'
+    this.loadError = null
+    this.loadPromise = null
+    this.usingHoistedInstance = false
+    this.metrics = {
+      initializeCalledAt: null,
+      wasmLoadStartAt: null,
+      wasmLoadEndAt: null,
+      firstQueryAt: null,
+      firstQueryWaitedMs: null,
+    }
+  }
+
+  /**
+   * Get the plugin manager.
+   */
+  getPluginManager(): PluginManager {
+    return this.pluginManager
+  }
+}
+
+/**
+ * Factory function to create a BackgroundPGLiteManager.
+ */
+export function createBackgroundPGLiteManager(
+  config: BackgroundPGLiteManagerConfig = {}
+): BackgroundPGLiteManager {
+  return new BackgroundPGLiteManager(config)
+}