@dotdo/postgres 0.1.1 → 0.1.3

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

@@ -0,0 +1,595 @@
+/**
+ * Lazy PGLite Manager for PostgresDO
+ *
+ * SPIKE: Lazy WASM Loading Experiment
+ *
+ * This module implements a lazy loading variant of PGLiteManager that defers
+ * WASM loading until the first query is executed, rather than loading on
+ * DO instantiation.
+ *
+ * ## Hypothesis
+ *
+ * Loading WASM only on first query (not on DO instantiation) may improve:
+ * 1. Cold start perception (faster initial response for non-query endpoints)
+ * 2. Memory usage (don't load if no queries made)
+ *
+ * ## Trade-offs
+ *
+ * - Eager loading: ~1200ms cold start, all subsequent requests fast
+ * - Lazy loading: ~0ms until first query, first query pays ~1200ms, subsequent fast
+ *
+ * For workloads where:
+ * - Many requests hit non-query endpoints (health checks, metadata) -> lazy wins
+ * - All requests are queries -> eager wins (same total latency, better predictability)
+ * - Mixed workload -> depends on query ratio
+ *
+ * ## Implementation
+ *
+ * The manager provides:
+ * - `isWASMLoaded()` - Check if WASM is loaded without triggering load
+ * - `ensureWASMLoaded()` - Explicitly trigger WASM loading
+ * - `query()` - Execute query (lazy loads WASM if needed)
+ * - Promise deduplication for concurrent first queries
+ *
+ * @module worker/lazy-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
+// These are bundled at build time as CompiledWasm and Data modules
+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 Lazy WASM State (for hoisting across DO reinstantiation)
+// =============================================================================
+
+/**
+ * Hoisted PGLite instance - survives DO class reinstantiation within same isolate.
+ * Same as eager loading, but only populated on first query.
+ */
+let lazyHoistedPglite: WorkersPGLite | null = null
+
+/**
+ * Promise for in-progress PGLite initialization.
+ * Prevents duplicate WASM loading when multiple queries execute concurrently.
+ */
+let lazyHoistedPglitePromise: Promise<WorkersPGLite> | null = null
+
+/**
+ * Timestamp when WASM was first loaded (for diagnostics)
+ */
+let wasmLoadedAt: number | null = null
+
+/**
+ * Unique identifier for this module instance (for debugging/diagnostics).
+ */
+const LAZY_MODULE_INSTANCE_ID = Math.random().toString(36).slice(2, 10)
+
+/**
+ * Check if the lazy hoisted PGLite instance exists.
+ */
+export function hasLazyHoistedPglite(): boolean {
+  return lazyHoistedPglite !== null
+}
+
+/**
+ * Get diagnostics about the lazy hoisted WASM state.
+ */
+export function getLazyHoistedPgliteDiagnostics(): {
+  hasInstance: boolean
+  hasPendingPromise: boolean
+  moduleInstanceId: string
+  wasmLoadedAt: number | null
+  timeSinceLoad: number | null
+} {
+  const now = Date.now()
+  return {
+    hasInstance: lazyHoistedPglite !== null,
+    hasPendingPromise: lazyHoistedPglitePromise !== null,
+    moduleInstanceId: LAZY_MODULE_INSTANCE_ID,
+    wasmLoadedAt,
+    timeSinceLoad: wasmLoadedAt !== null ? now - wasmLoadedAt : null,
+  }
+}
+
+/**
+ * Reset the lazy hoisted PGLite instance.
+ * WARNING: This should only be used in tests.
+ * @internal
+ */
+export function resetLazyHoistedPglite(): void {
+  lazyHoistedPglite = null
+  lazyHoistedPglitePromise = null
+  wasmLoadedAt = null
+}
+
+/**
+ * Lazily get or create the hoisted PGLite instance.
+ *
+ * This function is called ONLY when a query is executed, not during
+ * manager initialization.
+ */
+async function lazyGetOrCreateHoistedPglite(options: {
+  database?: string
+  debug?: 0 | 1
+}): Promise<WorkersPGLite> {
+  // Fast path: already initialized
+  if (lazyHoistedPglite) {
+    return lazyHoistedPglite
+  }
+
+  // Deduplication: return existing promise if initialization in progress
+  if (lazyHoistedPglitePromise) {
+    return lazyHoistedPglitePromise
+  }
+
+  const loadStartTime = performance.now()
+
+  // Start initialization
+  lazyHoistedPglitePromise = createWorkersPGLite({
+    wasmModule: pgliteWasm,
+    fsBundle: pgliteData,
+    database: options.database ?? 'postgres',
+    debug: options.debug ?? 0,
+  })
+
+  try {
+    lazyHoistedPglite = await lazyHoistedPglitePromise
+    wasmLoadedAt = Date.now()
+
+    const loadDuration = performance.now() - loadStartTime
+    console.log(
+      `[LazyPGLiteManager] WASM LOADED (lazy, on first query) - took ${loadDuration.toFixed(2)}ms, instance: ${LAZY_MODULE_INSTANCE_ID}`
+    )
+
+    return lazyHoistedPglite
+  } finally {
+    // Clear the promise after resolution (success or failure)
+    lazyHoistedPglitePromise = null
+  }
+}
+
+// =============================================================================
+// Loading State Enum
+// =============================================================================
+
+/**
+ * WASM loading state for diagnostics
+ */
+export type WASMLoadingState = 'not_loaded' | 'loading' | 'loaded' | 'error'
+
+// =============================================================================
+// Configuration
+// =============================================================================
+
+/**
+ * Configuration for Lazy PGLite manager
+ */
+export interface LazyPGLiteManagerConfig {
+  /** 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.
+   * When false (default), the PGLite WASM instance is cached at module level.
+   * @default false
+   */
+  disableHoisting?: boolean
+}
+
+// =============================================================================
+// LazyPGLiteManager Class
+// =============================================================================
+
+/**
+ * LazyPGLiteManager defers WASM loading until the first query.
+ *
+ * Unlike the eager PGLiteManager, this manager does NOT load WASM during
+ * initialize(). WASM is only loaded when query() or ensureWASMLoaded() is called.
+ *
+ * Benefits:
+ * - Non-query endpoints (health, metadata, stats) respond immediately
+ * - Memory is not consumed until queries are actually executed
+ * - Better cold start perception for mixed workloads
+ *
+ * Trade-offs:
+ * - First query incurs the full WASM loading time (~1200ms)
+ * - Less predictable query latency (first vs subsequent)
+ *
+ * @example
+ * ```typescript
+ * const manager = new LazyPGLiteManager({ database: 'mydb' })
+ *
+ * // Initialize (fast - no WASM loading)
+ * await manager.initialize()
+ *
+ * // Check if WASM is loaded (without triggering load)
+ * console.log(manager.isWASMLoaded()) // false
+ *
+ * // First query triggers WASM load (~1200ms)
+ * const result1 = await manager.query('SELECT 1')
+ *
+ * // Now WASM is loaded
+ * console.log(manager.isWASMLoaded()) // true
+ *
+ * // Subsequent queries are fast (~ms)
+ * const result2 = await manager.query('SELECT 2')
+ * ```
+ */
+export class LazyPGLiteManager {
+  private pglite: PGliteLike | null = null
+  private initialized = false
+  private loadingState: WASMLoadingState = 'not_loaded'
+  private loadError: Error | null = null
+  private config: LazyPGLiteManagerConfig
+  private pluginManager: PluginManager
+
+  /** Track whether this manager is using the hoisted instance */
+  private usingHoistedInstance = false
+
+  /** Timing metrics for diagnostics */
+  private metrics = {
+    initializeCalledAt: null as number | null,
+    firstQueryAt: null as number | null,
+    wasmLoadStartAt: null as number | null,
+    wasmLoadEndAt: null as number | null,
+  }
+
+  constructor(config: LazyPGLiteManagerConfig = {}) {
+    this.config = {
+      database: config.database ?? 'postgres',
+      debug: config.debug ?? false,
+    }
+    if (config.plugins !== undefined) {
+      this.config.plugins = config.plugins
+    }
+    if (config.createPGLite !== undefined) {
+      this.config.createPGLite = config.createPGLite
+    }
+    if (config.disableHoisting !== undefined) {
+      this.config.disableHoisting = config.disableHoisting
+    }
+
+    // Build plugin manager 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)
+  }
+
+  /**
+   * Check if the manager is initialized.
+   * Note: This returns true after initialize() is called, even if WASM is not loaded.
+   */
+  isInitialized(): boolean {
+    return this.initialized
+  }
+
+  /**
+   * Check if WASM is loaded (without triggering load).
+   * This is the key difference from eager loading.
+   */
+  isWASMLoaded(): boolean {
+    return this.pglite !== null
+  }
+
+  /**
+   * Get the current WASM loading state.
+   */
+  getLoadingState(): WASMLoadingState {
+    return this.loadingState
+  }
+
+  /**
+   * Get any error that occurred during WASM loading.
+   */
+  getLoadError(): Error | null {
+    return this.loadError
+  }
+
+  /**
+   * Get the PGLite instance (throws if not loaded).
+   * For lazy loading, prefer using query() which handles loading automatically.
+   */
+  getInstance(): PGliteLike {
+    if (!this.pglite) {
+      throw new Error('PGLite WASM not loaded. Call ensureWASMLoaded() or query() first.')
+    }
+    return this.pglite
+  }
+
+  /**
+   * Get the PGLite instance or null if not loaded.
+   */
+  getInstanceOrNull(): PGliteLike | null {
+    return this.pglite
+  }
+
+  /**
+   * Get the plugin manager.
+   */
+  getPluginManager(): PluginManager {
+    return this.pluginManager
+  }
+
+  /**
+   * Get timing metrics for diagnostics.
+   */
+  getMetrics(): {
+    initializeCalledAt: number | null
+    firstQueryAt: number | null
+    wasmLoadStartAt: number | null
+    wasmLoadEndAt: number | null
+    wasmLoadDurationMs: number | null
+    timeToFirstQuery: number | null
+  } {
+    const wasmLoadDurationMs =
+      this.metrics.wasmLoadStartAt !== null && this.metrics.wasmLoadEndAt !== null
+        ? this.metrics.wasmLoadEndAt - this.metrics.wasmLoadStartAt
+        : null
+
+    const timeToFirstQuery =
+      this.metrics.initializeCalledAt !== null && this.metrics.firstQueryAt !== null
+        ? this.metrics.firstQueryAt - this.metrics.initializeCalledAt
+        : null
+
+    return {
+      ...this.metrics,
+      wasmLoadDurationMs,
+      timeToFirstQuery,
+    }
+  }
+
+  /**
+   * Initialize the manager (does NOT load WASM).
+   *
+   * This is intentionally fast - WASM loading is deferred to first query.
+   * Safe to call multiple times.
+   */
+  async initialize(): Promise<void> {
+    if (this.initialized) {
+      return
+    }
+
+    this.metrics.initializeCalledAt = performance.now()
+    this.initialized = true
+
+    console.log(`[LazyPGLiteManager] Initialized (WASM not loaded yet)`)
+  }
+
+  /**
+   * Explicitly ensure WASM is loaded.
+   *
+   * Use this if you want to pre-warm WASM loading before the first query,
+   * but still want the option to skip loading entirely for non-query requests.
+   */
+  async ensureWASMLoaded(): Promise<void> {
+    if (this.pglite) {
+      return
+    }
+
+    await this.loadWASM()
+  }
+
+  /**
+   * Internal method to load WASM.
+   */
+  private async loadWASM(): Promise<PGliteLike> {
+    if (this.pglite) {
+      return this.pglite
+    }
+
+    if (this.loadingState === 'loading') {
+      // Wait for existing load to complete
+      if (lazyHoistedPglitePromise) {
+        const result = await lazyHoistedPglitePromise
+        this.pglite = result
+        this.loadingState = 'loaded'
+        return result
+      }
+    }
+
+    this.loadingState = 'loading'
+    this.metrics.wasmLoadStartAt = performance.now()
+
+    try {
+      // Use custom factory if provided (for testing)
+      if (this.config.createPGLite) {
+        this.usingHoistedInstance = false
+        this.pglite = await this.config.createPGLite()
+      } else if (this.config.disableHoisting) {
+        // Create new instance without hoisting
+        this.usingHoistedInstance = false
+        console.log(`[LazyPGLiteManager] WASM hoisting disabled - creating new instance`)
+        this.pglite = await createWorkersPGLite({
+          wasmModule: pgliteWasm,
+          fsBundle: pgliteData,
+          database: this.config.database ?? 'postgres',
+          debug: this.config.debug ? 1 : 0,
+        })
+      } else {
+        // Use the lazy hoisted instance
+        const wasmWasAlreadyLoaded = lazyHoistedPglite !== null
+        this.pglite = await lazyGetOrCreateHoistedPglite({
+          database: this.config.database,
+          debug: this.config.debug ? 1 : 0,
+        })
+        this.usingHoistedInstance = true
+
+        if (wasmWasAlreadyLoaded) {
+          const diagnostics = getLazyHoistedPgliteDiagnostics()
+          console.log(
+            `[LazyPGLiteManager] WASM REUSED (lazy hoisted) - loaded ${diagnostics.timeSinceLoad}ms ago, instance: ${diagnostics.moduleInstanceId}`
+          )
+        }
+      }
+
+      this.metrics.wasmLoadEndAt = performance.now()
+      this.loadingState = 'loaded'
+
+      // Wait for ready if needed
+      if (this.pglite.waitReady) {
+        await this.pglite.waitReady
+      }
+
+      // Run auto-create for plugins
+      await this.pluginManager.runAutoCreate(this.pglite)
+
+      return this.pglite
+    } catch (error) {
+      this.loadingState = 'error'
+      this.loadError = error instanceof Error ? error : new Error(String(error))
+      throw error
+    }
+  }
+
+  /**
+   * Execute a query (lazy loads WASM if needed).
+   *
+   * This is the primary method for interacting with the database.
+   * On first call, it will trigger WASM loading.
+   */
+  async query<T = unknown>(
+    sql: string,
+    params?: unknown[]
+  ): Promise<{
+    rows: T[]
+    fields: { name: string; dataTypeID: number }[]
+    affectedRows?: number
+  }> {
+    // Track first query time
+    if (this.metrics.firstQueryAt === null) {
+      this.metrics.firstQueryAt = performance.now()
+    }
+
+    // Lazy load WASM if not loaded
+    const pglite = await this.loadWASM()
+
+    // Execute query
+    return pglite.query<T>(sql, params)
+  }
+
+  /**
+   * Check if PGLite is responsive by running a simple query.
+   * This will trigger WASM loading if not already loaded.
+   */
+  async checkLiveness(): Promise<boolean> {
+    try {
+      const result = await this.query('SELECT 1 as alive')
+      return result.rows.length > 0
+    } catch {
+      return false
+    }
+  }
+
+  /**
+   * Close PGLite instance.
+   */
+  async close(timeoutMs: number = 5000): Promise<void> {
+    // If using hoisted instance, don't close it
+    if (this.usingHoistedInstance) {
+      console.log(`[LazyPGLiteManager] Skipping close of hoisted WASM instance (will be reused)`)
+      this.pglite = null
+      this.initialized = false
+      this.usingHoistedInstance = false
+      this.loadingState = 'not_loaded'
+      return
+    }
+
+    if (!this.pglite?.close) {
+      this.pglite = null
+      this.initialized = false
+      this.loadingState = 'not_loaded'
+      return
+    }
+
+    try {
+      await Promise.race([
+        this.pglite.close(),
+        new Promise<void>((_, reject) =>
+          setTimeout(() => reject(new Error('PGlite close timeout')), timeoutMs)
+        ),
+      ])
+    } catch (error) {
+      console.error('Error closing PGlite:', error)
+    }
+
+    this.pglite = null
+    this.initialized = false
+    this.loadingState = 'not_loaded'
+  }
+
+  /**
+   * Reset manager state (for testing or hibernation recovery).
+   */
+  reset(): void {
+    this.pglite = null
+    this.initialized = false
+    this.usingHoistedInstance = false
+    this.loadingState = 'not_loaded'
+    this.loadError = null
+    this.metrics = {
+      initializeCalledAt: null,
+      firstQueryAt: null,
+      wasmLoadStartAt: null,
+      wasmLoadEndAt: null,
+    }
+  }
+
+  /**
+   * Check if this manager is using the hoisted WASM instance.
+   */
+  isUsingHoistedInstance(): boolean {
+    return this.usingHoistedInstance
+  }
+
+  /**
+   * Get comprehensive diagnostics about lazy loading state.
+   */
+  getDiagnostics(): {
+    initialized: boolean
+    wasmLoaded: boolean
+    loadingState: WASMLoadingState
+    usingHoistedInstance: boolean
+    metrics: ReturnType<LazyPGLiteManager['getMetrics']>
+    hoisted: {
+      hasInstance: boolean
+      hasPendingPromise: boolean
+      moduleInstanceId: string
+      wasmLoadedAt: number | null
+      timeSinceLoad: number | null
+    }
+  } {
+    return {
+      initialized: this.initialized,
+      wasmLoaded: this.pglite !== null,
+      loadingState: this.loadingState,
+      usingHoistedInstance: this.usingHoistedInstance,
+      metrics: this.getMetrics(),
+      hoisted: getLazyHoistedPgliteDiagnostics(),
+    }
+  }
+}
+
+/**
+ * Factory function to create a LazyPGLiteManager.
+ */
+export function createLazyPGLiteManager(config?: LazyPGLiteManagerConfig): LazyPGLiteManager {
+  return new LazyPGLiteManager(config)
+}