@dotdo/postgres 0.1.0

package/src/streaming/memory-bounded-stream.ts

@@ -0,0 +1,443 @@
+/**
+ * MemoryBoundedStream - Stream with configurable memory limits
+ *
+ * Wraps async iterables with memory tracking and bounded buffering,
+ * preventing OOM by pausing producers when memory limits are reached.
+ *
+ * ## How It Works
+ *
+ * MemoryBoundedStream sits between a producer (e.g., QueryStreamer) and a
+ * consumer, tracking memory usage and applying backpressure when limits
+ * are approached.
+ *
+ * ```
+ * Producer -> MemoryBoundedStream -> Consumer
+ *                |
+ *                +-- Tracks memory
+ *                +-- Pauses when highWaterMark exceeded
+ *                +-- Resumes when lowWaterMark reached
+ *                +-- Throws OOM if memoryLimit exceeded
+ * ```
+ *
+ * ## Water Mark System
+ *
+ * - **highWaterMark**: When memory exceeds this, iteration pauses
+ * - **lowWaterMark**: When memory drops below this, iteration resumes
+ * - **memoryLimit**: Hard limit that throws StreamingOOMError
+ *
+ * The gap between high and low water marks prevents thrashing (rapid
+ * pause/resume cycles).
+ *
+ * @module streaming/memory-bounded-stream
+ */
+
+import { type Row, StreamingOOMError } from './query-streamer'
+
+// ============================================================================
+// Type Definitions
+// ============================================================================
+
+/**
+ * Configuration for MemoryBoundedStream
+ */
+export interface MemoryBoundedStreamConfig {
+  /** Maximum memory limit in bytes (default: 50MB) */
+  memoryLimit?: number
+  /** High water mark - pause when reached (default: 10MB) */
+  highWaterMark?: number
+  /** Low water mark - resume when below (default: 5MB) */
+  lowWaterMark?: number
+  /** Custom size estimator function for items */
+  sizeEstimator?: (item: unknown) => number
+  /** Callback when memory pressure is detected (highWaterMark exceeded) */
+  onMemoryPressure?: () => void | Promise<void>
+  /** Callback when memory is released (dropped below lowWaterMark) */
+  onMemoryRelease?: () => void | Promise<void>
+}
+
+/**
+ * Memory statistics for monitoring stream health
+ */
+export interface MemoryStats {
+  /** Current memory usage in bytes */
+  currentUsage: number
+  /** Peak memory usage seen during stream lifetime */
+  peakUsage: number
+  /** Number of items that have passed through the stream */
+  itemsProcessed: number
+  /** Number of times pause was triggered due to memory pressure */
+  pauseCount: number
+  /** Number of times resume was triggered after memory release */
+  resumeCount: number
+  /** Whether the stream is currently paused */
+  isPaused: boolean
+}
+
+// ============================================================================
+// Imported Configuration
+// ============================================================================
+
+import {
+  DEFAULT_STREAMING_MEMORY_LIMIT,
+  DEFAULT_HIGH_WATER_MARK,
+  DEFAULT_LOW_WATER_MARK,
+} from '../config/streaming'
+
+// ============================================================================
+// Local Aliases (for backward compatibility)
+// ============================================================================
+
+/** Default memory limit (50MB) */
+const DEFAULT_MEMORY_LIMIT = DEFAULT_STREAMING_MEMORY_LIMIT
+
+/**
+ * Default size estimator using JSON serialization.
+ * Falls back to 1KB for items that can't be serialized.
+ */
+function defaultSizeEstimator(item: unknown): number {
+  try {
+    return JSON.stringify(item).length
+  } catch {
+    return 1024
+  }
+}
+
+// ============================================================================
+// Internal Types
+// ============================================================================
+
+/**
+ * Resolved configuration with all defaults applied
+ */
+interface ResolvedConfig {
+  memoryLimit: number
+  highWaterMark: number
+  lowWaterMark: number
+  sizeEstimator: (item: unknown) => number
+  onMemoryPressure: () => void | Promise<void>
+  onMemoryRelease: () => void | Promise<void>
+}
+
+// ============================================================================
+// MemoryBoundedStream Class
+// ============================================================================
+
+/**
+ * MemoryBoundedStream wraps an async iterable with memory tracking.
+ *
+ * Implements backpressure by pausing iteration when memory limits are reached,
+ * ensuring memory usage stays bounded regardless of producer speed.
+ *
+ * ## Memory Tracking
+ *
+ * The stream estimates memory usage using a configurable sizeEstimator function.
+ * By default, this uses JSON.stringify() which provides a reasonable approximation
+ * for most data types.
+ *
+ * ## Automatic vs Manual Release
+ *
+ * The stream automatically releases memory after each yield. For scenarios where
+ * the consumer holds references to items (e.g., accumulating in an array), use
+ * the release() method to manually signal when items are no longer needed.
+ *
+ * @example Basic usage
+ * ```typescript
+ * const source = generateLargeDataset()
+ * const bounded = new MemoryBoundedStream(source, {
+ *   memoryLimit: 10 * 1024 * 1024, // 10MB
+ * })
+ *
+ * for await (const item of bounded) {
+ *   await processItem(item)
+ * }
+ *
+ * console.log('Peak memory:', bounded.getStats().peakUsage)
+ * ```
+ *
+ * @example With callbacks
+ * ```typescript
+ * const bounded = new MemoryBoundedStream(source, {
+ *   highWaterMark: 5 * 1024 * 1024,
+ *   onMemoryPressure: () => console.log('Memory pressure detected'),
+ *   onMemoryRelease: () => console.log('Memory pressure relieved'),
+ * })
+ * ```
+ */
+export class MemoryBoundedStream<T = Row> implements AsyncIterable<T> {
+  private readonly source: AsyncIterable<T>
+  private readonly config: ResolvedConfig
+  private currentUsage = 0
+  private peakUsage = 0
+  private itemsProcessed = 0
+  private pauseCount = 0
+  private resumeCount = 0
+  private isPaused = false
+  private resolveResume: (() => void) | null = null
+
+  /**
+   * Create a new MemoryBoundedStream.
+   *
+   * @param source - Source async iterable to wrap
+   * @param config - Memory bounds configuration
+   */
+  constructor(source: AsyncIterable<T>, config: MemoryBoundedStreamConfig = {}) {
+    this.source = source
+
+    const memoryLimit = config.memoryLimit ?? DEFAULT_MEMORY_LIMIT
+    const highWaterMark = config.highWaterMark ?? DEFAULT_HIGH_WATER_MARK
+    const lowWaterMark = config.lowWaterMark ?? DEFAULT_LOW_WATER_MARK
+
+    // Validate configuration
+    if (lowWaterMark >= highWaterMark) {
+      throw new Error('lowWaterMark must be less than highWaterMark')
+    }
+    if (highWaterMark >= memoryLimit) {
+      throw new Error('highWaterMark must be less than memoryLimit')
+    }
+
+    this.config = {
+      memoryLimit,
+      highWaterMark,
+      lowWaterMark,
+      sizeEstimator: config.sizeEstimator ?? defaultSizeEstimator,
+      onMemoryPressure: config.onMemoryPressure ?? (() => {}),
+      onMemoryRelease: config.onMemoryRelease ?? (() => {}),
+    }
+  }
+
+  // ============================================================================
+  // Internal Memory Management
+  // ============================================================================
+
+  /**
+   * Track memory allocation for an item.
+   * @internal
+   */
+  private async trackAllocation(size: number): Promise<void> {
+    this.currentUsage += size
+    if (this.currentUsage > this.peakUsage) {
+      this.peakUsage = this.currentUsage
+    }
+
+    // Check hard limit first
+    if (this.currentUsage > this.config.memoryLimit) {
+      throw new StreamingOOMError()
+    }
+
+    // Check if we should pause (high water mark)
+    if (!this.isPaused && this.currentUsage > this.config.highWaterMark) {
+      this.isPaused = true
+      this.pauseCount++
+      await this.config.onMemoryPressure()
+    }
+  }
+
+  /**
+   * Wait for memory to be released when paused.
+   * @internal
+   */
+  private async waitForResume(): Promise<void> {
+    if (!this.isPaused) return
+
+    return new Promise<void>((resolve) => {
+      this.resolveResume = resolve
+    })
+  }
+
+  // ============================================================================
+  // Public API
+  // ============================================================================
+
+  /**
+   * Manually release memory for a processed item.
+   *
+   * Call this when the consumer is holding references to items and
+   * automatic memory release after yield is not sufficient.
+   *
+   * @param item - The item to release memory for
+   */
+  release(item: T): void {
+    const size = this.config.sizeEstimator(item)
+    this.currentUsage = Math.max(0, this.currentUsage - size)
+    this.itemsProcessed++
+
+    // Check if we should resume (low water mark)
+    if (this.isPaused && this.currentUsage < this.config.lowWaterMark) {
+      this.isPaused = false
+      this.resumeCount++
+      this.config.onMemoryRelease()
+
+      // Signal to resume if waiting
+      if (this.resolveResume) {
+        this.resolveResume()
+        this.resolveResume = null
+      }
+    }
+  }
+
+  /**
+   * Get current memory statistics for monitoring.
+   */
+  getStats(): MemoryStats {
+    return {
+      currentUsage: this.currentUsage,
+      peakUsage: this.peakUsage,
+      itemsProcessed: this.itemsProcessed,
+      pauseCount: this.pauseCount,
+      resumeCount: this.resumeCount,
+      isPaused: this.isPaused,
+    }
+  }
+
+  /**
+   * Get current estimated memory usage in bytes.
+   */
+  getMemoryUsage(): number {
+    return this.currentUsage
+  }
+
+  /**
+   * Check if currently under memory pressure (above high water mark).
+   */
+  isUnderMemoryPressure(): boolean {
+    return this.currentUsage > this.config.highWaterMark
+  }
+
+  /**
+   * Abort streaming and throw OOM error.
+   * Use this to proactively abort when external memory pressure is detected.
+   */
+  async abortOnMemoryPressure(): Promise<never> {
+    // Reset state before throwing
+    this.currentUsage = 0
+    this.isPaused = false
+    if (this.resolveResume) {
+      this.resolveResume()
+      this.resolveResume = null
+    }
+    throw new StreamingOOMError()
+  }
+
+  /**
+   * Reset statistics while preserving current state.
+   * Useful for monitoring memory usage across multiple operations.
+   */
+  resetStats(): void {
+    this.peakUsage = this.currentUsage
+    this.itemsProcessed = 0
+    this.pauseCount = 0
+    this.resumeCount = 0
+  }
+
+  // ============================================================================
+  // Async Iteration
+  // ============================================================================
+
+  /**
+   * Iterate over the source with memory tracking.
+   *
+   * Memory is automatically released after each yield. For scenarios
+   * where the consumer holds references, call release() manually.
+   */
+  async *[Symbol.asyncIterator](): AsyncGenerator<T> {
+    for await (const item of this.source) {
+      // Wait if paused due to memory pressure
+      await this.waitForResume()
+
+      // Track memory allocation
+      const size = this.config.sizeEstimator(item)
+      await this.trackAllocation(size)
+
+      yield item
+
+      // Auto-release: assume item is no longer referenced after yield
+      this.currentUsage = Math.max(0, this.currentUsage - size)
+      this.itemsProcessed++
+
+      // Check if we should resume (below low water mark)
+      if (this.isPaused && this.currentUsage < this.config.lowWaterMark) {
+        this.isPaused = false
+        this.resumeCount++
+        this.config.onMemoryRelease()
+
+        // Signal to resume if waiting
+        if (this.resolveResume) {
+          this.resolveResume()
+          this.resolveResume = null
+        }
+      }
+    }
+  }
+
+  /**
+   * Create a batched version of this stream.
+   *
+   * Groups items into arrays of the specified size, which can be more
+   * efficient for bulk operations.
+   *
+   * @param size - Maximum number of items per batch
+   * @returns A new MemoryBoundedStream yielding arrays of items
+   */
+  batch(size: number): MemoryBoundedStream<T[]> {
+    if (size <= 0) {
+      throw new Error('Batch size must be positive')
+    }
+
+    const source = this.source
+    const config = this.config
+
+    async function* batchGenerator(): AsyncGenerator<T[]> {
+      let batch: T[] = []
+      for await (const item of source) {
+        batch.push(item)
+        if (batch.length >= size) {
+          yield batch
+          batch = []
+        }
+      }
+      if (batch.length > 0) {
+        yield batch
+      }
+    }
+
+    return new MemoryBoundedStream<T[]>(batchGenerator(), {
+      memoryLimit: config.memoryLimit,
+      highWaterMark: config.highWaterMark,
+      lowWaterMark: config.lowWaterMark,
+    })
+  }
+}
+
+// ============================================================================
+// Factory Functions
+// ============================================================================
+
+/**
+ * Create a memory-bounded wrapper for an async iterable.
+ *
+ * This is a convenience function that creates a new MemoryBoundedStream.
+ *
+ * @param source - Source async iterable to wrap
+ * @param config - Memory bounds configuration (optional)
+ * @returns A new MemoryBoundedStream
+ *
+ * @example
+ * ```typescript
+ * const bounded = createMemoryBoundedStream(generateRows(), {
+ *   memoryLimit: 10 * 1024 * 1024,
+ * })
+ *
+ * for await (const row of bounded) {
+ *   await processRow(row)
+ * }
+ * ```
+ */
+export function createMemoryBoundedStream<T>(
+  source: AsyncIterable<T>,
+  config?: MemoryBoundedStreamConfig
+): MemoryBoundedStream<T> {
+  return new MemoryBoundedStream(source, config)
+}
+
+// Re-export Row type for convenience
+export type { Row }