@dotdo/postgres 0.1.0

package/src/pglite/parallel-loader.ts

@@ -0,0 +1,481 @@
+/**
+ * Parallel WASM + VFS Loader
+ *
+ * Optimizes PGLite cold start latency by loading WASM and cold storage data
+ * in parallel instead of sequentially.
+ *
+ * Current Flow (Sequential):
+ * 1. Load WASM module (blocking)
+ * 2. Initialize PGLite
+ * 3. Load data from R2 (if cold)
+ * 4. Hydrate VFS
+ * 5. Ready for queries
+ *
+ * Optimized Flow (Parallel):
+ * 1. Start WASM loading ─────────────────┐
+ * 2. Start R2/Cache data fetch ────────┐ │
+ * 3. Wait for both ←───────────────────┴─┘
+ * 4. Initialize PGLite with pre-fetched data
+ * 5. Ready for queries
+ *
+ * @module pglite/parallel-loader
+ */
+
+import type { R2StorageLayer } from './r2-layer'
+import type { CacheLayer } from './cache-layer'
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/**
+ * Timing metrics for parallel loading
+ */
+export interface ParallelLoadMetrics {
+  /** Total time for parallel load in milliseconds */
+  totalMs: number
+  /** Time to prepare WASM (if any preparation needed) */
+  wasmPrepMs: number
+  /** Time to fetch cold storage data */
+  coldFetchMs: number
+  /** Time to fetch warm cache data */
+  warmFetchMs: number
+  /** Which source provided the data */
+  dataSource: 'hot' | 'warm' | 'cold' | 'empty'
+  /** Size of loaded data in bytes */
+  dataSize: number
+  /** Whether loading was done in parallel */
+  parallelized: boolean
+}
+
+/**
+ * Result of parallel loading
+ */
+export interface ParallelLoadResult {
+  /** Pre-compiled WASM module (passed through) */
+  wasmModule: WebAssembly.Module
+  /** Filesystem bundle as Blob for PGLite */
+  fsBundle: Blob
+  /** Loading metrics */
+  metrics: ParallelLoadMetrics
+}
+
+/**
+ * Options for parallel loading
+ */
+export interface ParallelLoadOptions {
+  /**
+   * Pre-compiled WASM module
+   * This should be statically imported and pre-compiled by Wrangler
+   */
+  wasmModule: WebAssembly.Module
+
+  /**
+   * Key for cold storage data in R2
+   * If provided, will attempt to fetch from R2
+   */
+  coldStorageKey?: string
+
+  /**
+   * R2 bucket for cold storage
+   */
+  r2Bucket?: R2Bucket
+
+  /**
+   * Cache key for warm storage
+   * If provided, will check cache first
+   */
+  warmCacheKey?: string
+
+  /**
+   * Cache layer for warm tier
+   */
+  cacheLayer?: CacheLayer
+
+  /**
+   * R2 storage layer (alternative to raw bucket)
+   */
+  r2Layer?: R2StorageLayer
+
+  /**
+   * Hot data from DO SQLite (synchronous, already loaded)
+   * If provided, skips cold/warm fetches
+   */
+  hotData?: ArrayBuffer | Uint8Array
+
+  /**
+   * Timeout for cold fetch in milliseconds (default: 10000)
+   */
+  coldFetchTimeoutMs?: number
+
+  /**
+   * Enable debug logging
+   */
+  debug?: boolean
+}
+
+// ============================================================================
+// Implementation
+// ============================================================================
+
+/**
+ * Load WASM and VFS data in parallel for optimized cold starts
+ *
+ * This function coordinates parallel fetching of:
+ * 1. Any WASM preparation (usually a no-op since WASM is pre-compiled)
+ * 2. VFS data from tiered storage (hot -> warm -> cold)
+ *
+ * @example
+ * ```typescript
+ * import pgliteWasm from './pglite.wasm'
+ *
+ * const result = await parallelLoad({
+ *   wasmModule: pgliteWasm,
+ *   coldStorageKey: 'databases/my-db/data.tar',
+ *   r2Bucket: env.R2_BUCKET,
+ * })
+ *
+ * const pg = await PGlite.create({
+ *   wasmModule: result.wasmModule,
+ *   fsBundle: result.fsBundle,
+ * })
+ * ```
+ */
+export async function parallelLoad(
+  options: ParallelLoadOptions
+): Promise<ParallelLoadResult> {
+  const {
+    wasmModule,
+    coldStorageKey,
+    r2Bucket,
+    warmCacheKey,
+    cacheLayer,
+    r2Layer,
+    hotData,
+    coldFetchTimeoutMs = 10000,
+    debug = false,
+  } = options
+
+  const startTime = performance.now()
+  const log = debug ? console.log.bind(console, '[ParallelLoader]') : () => {}
+
+  const metrics: ParallelLoadMetrics = {
+    totalMs: 0,
+    wasmPrepMs: 0,
+    coldFetchMs: 0,
+    warmFetchMs: 0,
+    dataSource: 'empty',
+    dataSize: 0,
+    parallelized: false,
+  }
+
+  // If hot data is provided, use it immediately (synchronous path)
+  if (hotData) {
+    log('Using hot data from DO SQLite')
+    const data = hotData instanceof Uint8Array ? hotData.buffer : hotData
+    metrics.dataSource = 'hot'
+    metrics.dataSize = data.byteLength
+    metrics.totalMs = performance.now() - startTime
+    metrics.parallelized = false
+
+    return {
+      wasmModule,
+      fsBundle: new Blob([data]),
+      metrics,
+    }
+  }
+
+  // Set up parallel fetches
+  const fetchPromises: Promise<{ source: 'warm' | 'cold'; data: Uint8Array | null; timeMs: number }>[] = []
+
+  // Warm fetch (cache layer)
+  if (warmCacheKey && cacheLayer) {
+    const warmStartTime = performance.now()
+    const warmPromise = cacheLayer
+      .get(warmCacheKey)
+      .then((data) => ({
+        source: 'warm' as const,
+        data,
+        timeMs: performance.now() - warmStartTime,
+      }))
+      .catch((err) => {
+        log('Warm fetch error:', err)
+        return { source: 'warm' as const, data: null, timeMs: performance.now() - warmStartTime }
+      })
+    fetchPromises.push(warmPromise)
+  }
+
+  // Cold fetch (R2 layer or bucket)
+  if (coldStorageKey && (r2Bucket || r2Layer)) {
+    const coldStartTime = performance.now()
+    const coldPromise = (async () => {
+      try {
+        let data: Uint8Array | null = null
+
+        if (r2Layer) {
+          // Use R2StorageLayer if provided
+          data = await Promise.race([
+            r2Layer.get(coldStorageKey),
+            new Promise<null>((_, reject) =>
+              setTimeout(() => reject(new Error('Cold fetch timeout')), coldFetchTimeoutMs)
+            ),
+          ])
+        } else if (r2Bucket) {
+          // Use raw R2Bucket
+          const obj = await Promise.race([
+            r2Bucket.get(coldStorageKey),
+            new Promise<null>((_, reject) =>
+              setTimeout(() => reject(new Error('Cold fetch timeout')), coldFetchTimeoutMs)
+            ),
+          ])
+
+          if (obj) {
+            data = new Uint8Array(await obj.arrayBuffer())
+          }
+        }
+
+        return {
+          source: 'cold' as const,
+          data,
+          timeMs: performance.now() - coldStartTime,
+        }
+      } catch (err) {
+        log('Cold fetch error:', err)
+        return {
+          source: 'cold' as const,
+          data: null,
+          timeMs: performance.now() - coldStartTime,
+        }
+      }
+    })()
+    fetchPromises.push(coldPromise)
+  }
+
+  // Execute all fetches in parallel
+  if (fetchPromises.length > 0) {
+    metrics.parallelized = fetchPromises.length > 1
+    log(`Starting ${fetchPromises.length} parallel fetches`)
+
+    const results = await Promise.all(fetchPromises)
+
+    // Process results in priority order: warm first, then cold
+    for (const result of results) {
+      if (result.source === 'warm') {
+        metrics.warmFetchMs = result.timeMs
+        if (result.data) {
+          log(`Warm cache hit: ${result.data.length} bytes in ${result.timeMs.toFixed(2)}ms`)
+          metrics.dataSource = 'warm'
+          metrics.dataSize = result.data.length
+          metrics.totalMs = performance.now() - startTime
+
+          return {
+            wasmModule,
+            fsBundle: new Blob([result.data]),
+            metrics,
+          }
+        }
+      }
+    }
+
+    // No warm hit, check cold results
+    for (const result of results) {
+      if (result.source === 'cold') {
+        metrics.coldFetchMs = result.timeMs
+        if (result.data) {
+          log(`Cold storage hit: ${result.data.length} bytes in ${result.timeMs.toFixed(2)}ms`)
+          metrics.dataSource = 'cold'
+          metrics.dataSize = result.data.length
+          metrics.totalMs = performance.now() - startTime
+
+          return {
+            wasmModule,
+            fsBundle: new Blob([result.data]),
+            metrics,
+          }
+        }
+      }
+    }
+  }
+
+  // No data found anywhere, return empty bundle
+  log('No data found, returning empty bundle')
+  metrics.totalMs = performance.now() - startTime
+
+  return {
+    wasmModule,
+    fsBundle: new Blob([]),
+    metrics,
+  }
+}
+
+/**
+ * Pre-fetch cold storage data without blocking
+ *
+ * Use this to start fetching cold data early (e.g., in the Worker)
+ * before the DO is even instantiated.
+ *
+ * @example
+ * ```typescript
+ * // In Worker entry point
+ * export default {
+ *   async fetch(request, env) {
+ *     // Start cold fetch immediately
+ *     const coldFetch = prefetchColdData({
+ *       key: 'databases/my-db/data.tar',
+ *       r2Bucket: env.R2_BUCKET,
+ *     })
+ *
+ *     // Route to DO
+ *     const id = env.POSTGRES_DO.idFromName('my-db')
+ *     const stub = env.POSTGRES_DO.get(id)
+ *
+ *     // Pass pre-fetched data to DO
+ *     const response = await stub.fetch(request, {
+ *       cf: { coldData: await coldFetch }
+ *     })
+ *
+ *     return response
+ *   }
+ * }
+ * ```
+ */
+export async function prefetchColdData(options: {
+  key: string
+  r2Bucket?: R2Bucket
+  r2Layer?: R2StorageLayer
+  timeoutMs?: number
+}): Promise<Uint8Array | null> {
+  const { key, r2Bucket, r2Layer, timeoutMs = 10000 } = options
+
+  try {
+    if (r2Layer) {
+      return await Promise.race([
+        r2Layer.get(key),
+        new Promise<null>((_, reject) =>
+          setTimeout(() => reject(new Error('Prefetch timeout')), timeoutMs)
+        ),
+      ])
+    }
+
+    if (r2Bucket) {
+      const obj = await Promise.race([
+        r2Bucket.get(key),
+        new Promise<null>((_, reject) =>
+          setTimeout(() => reject(new Error('Prefetch timeout')), timeoutMs)
+        ),
+      ])
+
+      if (obj) {
+        return new Uint8Array(await obj.arrayBuffer())
+      }
+    }
+
+    return null
+  } catch {
+    return null
+  }
+}
+
+/**
+ * Create a warm-up function that pre-loads data into cache
+ *
+ * @example
+ * ```typescript
+ * const warmUp = createWarmUpFunction({
+ *   cacheLayer,
+ *   r2Layer,
+ * })
+ *
+ * // Warm up specific database data
+ * await warmUp('databases/my-db/data.tar', 'my-db-cache-key')
+ * ```
+ */
+export function createWarmUpFunction(options: {
+  cacheLayer: CacheLayer
+  r2Layer: R2StorageLayer
+}): (coldKey: string, warmKey: string) => Promise<boolean> {
+  const { cacheLayer, r2Layer } = options
+
+  return async (coldKey: string, warmKey: string): Promise<boolean> => {
+    try {
+      // Check if already in cache
+      const cached = await cacheLayer.get(warmKey)
+      if (cached) {
+        return true // Already warm
+      }
+
+      // Fetch from cold storage
+      const coldData = await r2Layer.get(coldKey)
+      if (!coldData) {
+        return false // No data to warm
+      }
+
+      // Promote to warm tier
+      await cacheLayer.put(warmKey, coldData, {
+        lastAccessed: Date.now(),
+        accessCount: 1,
+      })
+
+      return true
+    } catch {
+      return false
+    }
+  }
+}
+
+/**
+ * Batch parallel load for multiple databases
+ *
+ * Useful when a Worker handles multiple DO instances and can
+ * pre-fetch data for all of them in parallel.
+ *
+ * @example
+ * ```typescript
+ * const results = await batchParallelLoad({
+ *   wasmModule: pgliteWasm,
+ *   databases: [
+ *     { id: 'db1', coldKey: 'databases/db1/data.tar' },
+ *     { id: 'db2', coldKey: 'databases/db2/data.tar' },
+ *   ],
+ *   r2Bucket: env.R2_BUCKET,
+ * })
+ *
+ * // results is a Map<string, ParallelLoadResult>
+ * const db1Data = results.get('db1')
+ * ```
+ */
+export async function batchParallelLoad(options: {
+  wasmModule: WebAssembly.Module
+  databases: Array<{ id: string; coldKey: string; warmKey?: string }>
+  r2Bucket?: R2Bucket
+  r2Layer?: R2StorageLayer
+  cacheLayer?: CacheLayer
+  debug?: boolean
+}): Promise<Map<string, ParallelLoadResult>> {
+  const { wasmModule, databases, r2Bucket, r2Layer, cacheLayer, debug } = options
+
+  const results = new Map<string, ParallelLoadResult>()
+
+  // Load all databases in parallel
+  const loadPromises = databases.map(async (db) => {
+    const result = await parallelLoad({
+      wasmModule,
+      coldStorageKey: db.coldKey,
+      warmCacheKey: db.warmKey,
+      r2Bucket,
+      r2Layer,
+      cacheLayer,
+      debug,
+    })
+
+    return { id: db.id, result }
+  })
+
+  const loadResults = await Promise.all(loadPromises)
+
+  for (const { id, result } of loadResults) {
+    results.set(id, result)
+  }
+
+  return results
+}