@dotdo/postgres 0.1.0

package/src/routing/websocket-pool.ts

@@ -0,0 +1,577 @@
+/**
+ * WebSocket Connection Pool
+ *
+ * Task: postgres-d6g.8
+ *
+ * Manages WebSocket connections to Durable Objects for cost optimization.
+ * WebSocket hibernation provides ~95% cost reduction for DO connections.
+ */
+
+export type ConnectionState = 'connecting' | 'open' | 'idle' | 'hibernated' | 'reconnecting' | 'closing' | 'closed'
+
+export interface PoolConfig {
+  /** Default WebSocket endpoint for DO connections */
+  doEndpoint: string
+
+  /** Maximum number of connections in the pool */
+  maxConnections?: number
+
+  /** Idle timeout before closing connection (ms) */
+  idleTimeoutMs?: number
+
+  /** Enable WebSocket hibernation for idle connections */
+  enableHibernation?: boolean
+
+  /** Timeout before hibernating idle connections (ms) */
+  hibernationTimeoutMs?: number
+
+  /** Validate connection health before returning from pool */
+  validateOnAcquire?: boolean
+
+  /** Interval for periodic health checks (ms) */
+  healthCheckIntervalMs?: number
+
+  /** Enable automatic reconnection on connection errors */
+  autoReconnect?: boolean
+
+  /** Maximum number of reconnection attempts */
+  maxReconnectAttempts?: number
+
+  /** Base backoff time for reconnection (ms) */
+  reconnectBackoffMs?: number
+
+  /** Maximum backoff time for reconnection (ms) */
+  maxReconnectBackoffMs?: number
+
+  /** Wait for connection when pool is exhausted */
+  waitForConnection?: boolean
+}
+
+export interface AcquireOptions {
+  /** Timeout for acquiring a connection (ms) */
+  timeout?: number
+
+  /** Custom endpoint override */
+  endpoint?: string
+
+  /** Tags for connection routing */
+  tags?: string[]
+}
+
+export interface PoolMetrics {
+  activeConnections: number
+  pooledConnections: number
+  hibernatedConnections: number
+  totalConnectionsCreated: number
+  messagesSent: number
+  messagesReceived: number
+  errors: number
+  healthChecksRun: number
+}
+
+type MessageHandler = (data: string) => void
+type StateChangeHandler = (state: ConnectionState) => void
+
+let connectionIdCounter = 0
+
+export class PooledConnection {
+  readonly id: string
+  readonly endpoint: string
+
+  private ws: WebSocket | null = null
+  private _state: ConnectionState = 'connecting'
+  private messageHandlers: MessageHandler[] = []
+  private stateHandlers: StateChangeHandler[] = []
+  private pendingRequests: Map<string, { resolve: (data: string) => void; reject: (err: Error) => void }> = new Map()
+  private requestIdCounter = 0
+  private _healthy = true
+  private reconnectAttempts = 0
+  private tags: string[]
+  private pool: WebSocketPool
+
+  constructor(
+    endpoint: string,
+    tags: string[],
+    pool: WebSocketPool,
+    private config: Required<PoolConfig>
+  ) {
+    this.id = `conn-${++connectionIdCounter}`
+    this.endpoint = endpoint
+    this.tags = tags
+    this.pool = pool
+    this.connect()
+  }
+
+  get state(): ConnectionState {
+    return this._state
+  }
+
+  private set state(newState: ConnectionState) {
+    if (this._state !== newState) {
+      this._state = newState
+      this.stateHandlers.forEach(handler => handler(newState))
+    }
+  }
+
+  private connect(): Promise<void> {
+    return new Promise((resolve, reject) => {
+      try {
+        this.ws = new WebSocket(this.endpoint)
+
+        this.ws.addEventListener('open', () => {
+          this.state = 'open'
+          this._healthy = true
+          this.reconnectAttempts = 0
+          resolve()
+        })
+
+        this.ws.addEventListener('close', () => {
+          this.state = 'closed'
+        })
+
+        this.ws.addEventListener('error', (event: Event) => {
+          this.pool.recordError()
+          if (this._state === 'connecting') {
+            reject(event)
+          } else if (this.config.autoReconnect && this.reconnectAttempts < this.config.maxReconnectAttempts) {
+            this.attemptReconnect()
+          } else {
+            this.state = 'closed'
+          }
+        })
+
+        this.ws.addEventListener('message', (event: MessageEvent) => {
+          this.pool.recordMessageReceived()
+          const data = typeof event.data === 'string' ? event.data : String(event.data)
+          this.handleMessage(data)
+        })
+      } catch (err) {
+        reject(err)
+      }
+    })
+  }
+
+  private attemptReconnect(): void {
+    this.reconnectAttempts++
+    this.state = 'reconnecting'
+
+    const backoff = Math.min(
+      this.config.reconnectBackoffMs * Math.pow(2, this.reconnectAttempts - 1),
+      this.config.maxReconnectBackoffMs
+    )
+
+    setTimeout(() => {
+      if (this._state === 'reconnecting') {
+        void this.connect().catch(() => {
+          if (this.reconnectAttempts >= this.config.maxReconnectAttempts) {
+            this.state = 'closed'
+          }
+        })
+      }
+    }, backoff)
+  }
+
+  private handleMessage(data: string): void {
+    // Check if this is a response to a pending request
+    try {
+      const parsed = JSON.parse(data)
+      if (parsed.id && this.pendingRequests.has(parsed.id)) {
+        const { resolve } = this.pendingRequests.get(parsed.id)!
+        this.pendingRequests.delete(parsed.id)
+        resolve(data)
+        return
+      }
+    } catch {
+      // Not JSON or no ID, pass to message handlers
+    }
+
+    this.messageHandlers.forEach(handler => handler(data))
+  }
+
+  async send(data: string): Promise<void> {
+    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
+      throw new Error('WebSocket is not open')
+    }
+    this.ws.send(data)
+    this.pool.recordMessageSent()
+  }
+
+  async request(data: string, timeout = 30000): Promise<string> {
+    const requestId = `req-${++this.requestIdCounter}`
+
+    // Inject request ID into the message
+    let messageWithId: string
+    try {
+      const parsed = JSON.parse(data)
+      parsed.id = requestId
+      messageWithId = JSON.stringify(parsed)
+    } catch {
+      messageWithId = data
+    }
+
+    return new Promise((resolve, reject) => {
+      const timeoutId = setTimeout(() => {
+        this.pendingRequests.delete(requestId)
+        reject(new Error('Request timeout'))
+      }, timeout)
+
+      this.pendingRequests.set(requestId, {
+        resolve: (response) => {
+          clearTimeout(timeoutId)
+          resolve(response)
+        },
+        reject: (error) => {
+          clearTimeout(timeoutId)
+          reject(error)
+        },
+      })
+
+      this.send(messageWithId).catch(reject)
+    })
+  }
+
+  onMessage(handler: MessageHandler): void {
+    this.messageHandlers.push(handler)
+  }
+
+  onStateChange(handler: StateChangeHandler): void {
+    this.stateHandlers.push(handler)
+  }
+
+  isHealthy(): boolean {
+    return this._healthy && this.ws?.readyState === WebSocket.OPEN
+  }
+
+  markUnhealthy(): void {
+    this._healthy = false
+  }
+
+  markIdle(): void {
+    this.state = 'idle'
+  }
+
+  markHibernated(): void {
+    this.state = 'hibernated'
+  }
+
+  wake(): void {
+    if (this._state === 'hibernated' || this._state === 'idle') {
+      this.state = 'open'
+    }
+  }
+
+  simulateError(_error: Error): void {
+    this.pool.recordError()
+    if (this.config.autoReconnect && this.reconnectAttempts < this.config.maxReconnectAttempts) {
+      this.attemptReconnect()
+    } else {
+      this.state = 'closed'
+    }
+  }
+
+  simulateMessage(data: string): void {
+    this.pool.recordMessageReceived()
+    this.handleMessage(data)
+  }
+
+  close(): void {
+    this.state = 'closing'
+    this.ws?.close()
+    this.state = 'closed'
+  }
+
+  hasTags(tags: string[]): boolean {
+    return tags.every(tag => this.tags.includes(tag))
+  }
+}
+
+interface WaitingRequest {
+  resolve: (conn: PooledConnection) => void
+  reject: (error: Error) => void
+  options: AcquireOptions
+}
+
+export class WebSocketPool {
+  private config: Required<PoolConfig>
+  private connections: Map<string, PooledConnection> = new Map()
+  private idleConnections: Set<string> = new Set()
+  private hibernatedConnections: Set<string> = new Set()
+  private waitingRequests: WaitingRequest[] = []
+  private metrics: PoolMetrics
+  private healthCheckInterval: ReturnType<typeof setInterval> | null = null
+  private idleCheckInterval: ReturnType<typeof setInterval> | null = null
+  private connectionTimers: Map<string, ReturnType<typeof setTimeout>> = new Map()
+
+  constructor(config: PoolConfig) {
+    this.config = {
+      doEndpoint: config.doEndpoint,
+      maxConnections: config.maxConnections ?? 10,
+      idleTimeoutMs: config.idleTimeoutMs ?? 60000,
+      enableHibernation: config.enableHibernation ?? false,
+      hibernationTimeoutMs: config.hibernationTimeoutMs ?? 30000,
+      validateOnAcquire: config.validateOnAcquire ?? false,
+      healthCheckIntervalMs: config.healthCheckIntervalMs ?? 30000,
+      autoReconnect: config.autoReconnect ?? true,
+      maxReconnectAttempts: config.maxReconnectAttempts ?? 3,
+      reconnectBackoffMs: config.reconnectBackoffMs ?? 1000,
+      maxReconnectBackoffMs: config.maxReconnectBackoffMs ?? 30000,
+      waitForConnection: config.waitForConnection ?? false,
+    }
+
+    this.metrics = this.createEmptyMetrics()
+
+    // Start health check interval if configured
+    if (config.healthCheckIntervalMs) {
+      this.healthCheckInterval = setInterval(() => {
+        this.runHealthCheck()
+      }, this.config.healthCheckIntervalMs)
+    }
+
+    // Start idle check interval
+    this.idleCheckInterval = setInterval(() => {
+      this.checkIdleConnections()
+    }, Math.min(this.config.idleTimeoutMs, this.config.hibernationTimeoutMs) / 2)
+  }
+
+  async acquire(options: AcquireOptions = {}): Promise<PooledConnection> {
+    const endpoint = options.endpoint ?? this.config.doEndpoint
+    const tags = options.tags ?? []
+
+    // Try to find an existing idle connection with matching tags
+    for (const connId of this.idleConnections) {
+      const conn = this.connections.get(connId)
+      if (conn && conn.endpoint === endpoint && conn.hasTags(tags)) {
+        // Always check health - unhealthy connections should be removed
+        if (!conn.isHealthy()) {
+          this.removeConnection(connId)
+          continue
+        }
+        this.idleConnections.delete(connId)
+        this.hibernatedConnections.delete(connId)
+        this.clearConnectionTimer(connId)
+        conn.wake()
+        return conn
+      }
+    }
+
+    // Check if we can create a new connection
+    const activeCount = this.connections.size - this.idleConnections.size
+    if (activeCount >= this.config.maxConnections) {
+      if (this.config.waitForConnection) {
+        return this.waitForConnection(options)
+      }
+
+      if (options.timeout !== undefined) {
+        throw new Error('timeout: max connections reached')
+      }
+      throw new Error('max connections reached')
+    }
+
+    // Create a new connection
+    const conn = new PooledConnection(endpoint, tags, this, this.config)
+    this.connections.set(conn.id, conn)
+    this.metrics.totalConnectionsCreated++
+
+    // Wait for connection to open
+    await new Promise<void>((resolve, reject) => {
+      const timeout = options.timeout ?? 30000
+      const timeoutId = setTimeout(() => {
+        reject(new Error('Connection timeout'))
+      }, timeout)
+
+      const checkState = () => {
+        if (conn.state === 'open') {
+          clearTimeout(timeoutId)
+          resolve()
+        } else if (conn.state === 'closed') {
+          clearTimeout(timeoutId)
+          reject(new Error('Connection failed'))
+        } else {
+          setTimeout(checkState, 10)
+        }
+      }
+      checkState()
+    })
+
+    return conn
+  }
+
+  private waitForConnection(options: AcquireOptions): Promise<PooledConnection> {
+    return new Promise((resolve, reject) => {
+      const timeout = options.timeout ?? 30000
+      const timeoutId = setTimeout(() => {
+        const index = this.waitingRequests.findIndex(req => req.resolve === resolve)
+        if (index !== -1) {
+          this.waitingRequests.splice(index, 1)
+        }
+        reject(new Error('timeout waiting for connection'))
+      }, timeout)
+
+      this.waitingRequests.push({
+        resolve: (conn) => {
+          clearTimeout(timeoutId)
+          resolve(conn)
+        },
+        reject: (error) => {
+          clearTimeout(timeoutId)
+          reject(error)
+        },
+        options,
+      })
+    })
+  }
+
+  release(conn: PooledConnection): void {
+    if (!this.connections.has(conn.id)) {
+      return
+    }
+
+    conn.markIdle()
+    this.idleConnections.add(conn.id)
+
+    // Check if there are waiting requests
+    if (this.waitingRequests.length > 0) {
+      const waitingReq = this.waitingRequests.shift()!
+      this.idleConnections.delete(conn.id)
+      conn.wake()
+      waitingReq.resolve(conn)
+      return
+    }
+
+    // Start idle/hibernation timer
+    this.startConnectionTimer(conn)
+  }
+
+  private startConnectionTimer(conn: PooledConnection): void {
+    this.clearConnectionTimer(conn.id)
+
+    if (this.config.enableHibernation) {
+      // First hibernate, then close after idle timeout
+      const hibernationTimer = setTimeout(() => {
+        if (this.idleConnections.has(conn.id)) {
+          conn.markHibernated()
+          this.hibernatedConnections.add(conn.id)
+
+          const idleTimer = setTimeout(() => {
+            this.removeConnection(conn.id)
+          }, this.config.idleTimeoutMs - this.config.hibernationTimeoutMs)
+
+          this.connectionTimers.set(conn.id, idleTimer)
+        }
+      }, this.config.hibernationTimeoutMs)
+
+      this.connectionTimers.set(conn.id, hibernationTimer)
+    } else {
+      const idleTimer = setTimeout(() => {
+        this.removeConnection(conn.id)
+      }, this.config.idleTimeoutMs)
+
+      this.connectionTimers.set(conn.id, idleTimer)
+    }
+  }
+
+  private clearConnectionTimer(connId: string): void {
+    const timer = this.connectionTimers.get(connId)
+    if (timer) {
+      clearTimeout(timer)
+      this.connectionTimers.delete(connId)
+    }
+  }
+
+  private removeConnection(connId: string): void {
+    const conn = this.connections.get(connId)
+    if (conn) {
+      conn.close()
+      this.connections.delete(connId)
+      this.idleConnections.delete(connId)
+      this.hibernatedConnections.delete(connId)
+      this.clearConnectionTimer(connId)
+    }
+  }
+
+  private checkIdleConnections(): void {
+    // This is handled by individual timers now
+  }
+
+  private runHealthCheck(): void {
+    this.metrics.healthChecksRun++
+
+    for (const [connId, conn] of this.connections) {
+      if (this.idleConnections.has(connId) && !conn.isHealthy()) {
+        this.removeConnection(connId)
+      }
+    }
+  }
+
+  async close(): Promise<void> {
+    // Clear intervals
+    if (this.healthCheckInterval) {
+      clearInterval(this.healthCheckInterval)
+    }
+    if (this.idleCheckInterval) {
+      clearInterval(this.idleCheckInterval)
+    }
+
+    // Clear all connection timers
+    for (const timer of this.connectionTimers.values()) {
+      clearTimeout(timer)
+    }
+    this.connectionTimers.clear()
+
+    // Close all connections
+    for (const conn of this.connections.values()) {
+      conn.close()
+    }
+    this.connections.clear()
+    this.idleConnections.clear()
+    this.hibernatedConnections.clear()
+
+    // Reject waiting requests
+    for (const req of this.waitingRequests) {
+      req.reject(new Error('Pool closed'))
+    }
+    this.waitingRequests = []
+  }
+
+  getMetrics(): PoolMetrics {
+    return {
+      ...this.metrics,
+      activeConnections: this.connections.size - this.idleConnections.size,
+      pooledConnections: this.idleConnections.size,
+      hibernatedConnections: this.hibernatedConnections.size,
+    }
+  }
+
+  resetMetrics(): void {
+    this.metrics = this.createEmptyMetrics()
+  }
+
+  getConfig(): Required<PoolConfig> {
+    return { ...this.config }
+  }
+
+  recordMessageSent(): void {
+    this.metrics.messagesSent++
+  }
+
+  recordMessageReceived(): void {
+    this.metrics.messagesReceived++
+  }
+
+  recordError(): void {
+    this.metrics.errors++
+  }
+
+  private createEmptyMetrics(): PoolMetrics {
+    return {
+      activeConnections: 0,
+      pooledConnections: 0,
+      hibernatedConnections: 0,
+      totalConnectionsCreated: 0,
+      messagesSent: 0,
+      messagesReceived: 0,
+      errors: 0,
+      healthChecksRun: 0,
+    }
+  }
+}