diff --git a/src/index.ts b/src/index.ts index 468aeb2..aa5899b 100644 --- a/src/index.ts +++ b/src/index.ts @@ -110,6 +110,8 @@ interface WSRelayData { target: string; relayPath: string; upstream?: WebSocket; + /** Set when client buffer exceeds threshold — stops forwarding upstream data */ + paused?: boolean; } // --- Route handlers ---------------------------------------------------------- @@ -621,6 +623,9 @@ const server: Server = Bun.serve({ }; upstream.onmessage = (event: MessageEvent) => { + // Backpressure: drop messages when client buffer is full + if (ws.data.paused) return; + const data = event.data; if (typeof data === "string") { ws.sendText(data); @@ -628,7 +633,9 @@ const server: Server = Bun.serve({ ws.sendBinary(new Uint8Array(data)); } else if (data instanceof Blob) { data.arrayBuffer().then((buf) => { - ws.sendBinary(new Uint8Array(buf)); + if (!ws.data.paused) { + ws.sendBinary(new Uint8Array(buf)); + } }); } else { ws.sendBinary(data as unknown as Uint8Array); @@ -669,15 +676,22 @@ const server: Server = Bun.serve({ }, drain(ws: ServerWebSocket) { - // Backpressure: pause upstream reads when client is slow + // Backpressure: pause forwarding when client buffer is large const upstream = ws.data.upstream; - if (upstream && upstream.readyState === WebSocket.OPEN) { - // Bun WebSocket handles backpressure internally via bufferAmount - // Log if buffer is growing excessively - const buffered = (ws as any).bufferAmount ?? 0; - if (buffered > 1024 * 1024) { - console.warn(`[ws] Client backpressure: ${buffered} bytes buffered`); + if (!upstream || upstream.readyState !== WebSocket.OPEN) return; + + const buffered = (ws as any).bufferAmount ?? 0; + const BACKPRESSURE_THRESHOLD = 512 * 1024; // 512KB + const RESUME_THRESHOLD = 64 * 1024; // 64KB + + if (buffered > BACKPRESSURE_THRESHOLD) { + if (!ws.data.paused) { + ws.data.paused = true; + console.warn(`[ws] Client backpressure: pausing upstream forwarding (${buffered} bytes buffered)`); } + } else if (ws.data.paused && buffered < RESUME_THRESHOLD) { + ws.data.paused = false; + console.log(`[ws] Client backpressure cleared: resuming upstream forwarding (${buffered} bytes buffered)`); } }, }, @@ -694,11 +708,11 @@ if (isDevMode()) { // --- Graceful Shutdown ------------------------------------------------------- -const shutdownHandler = (signal: string) => { +const shutdownHandler = async (signal: string) => { console.log(`\n[relay] Received ${signal}, shutting down gracefully...`); // Close active SSE streams so clients get proper stream end events - closeAllActiveReaders(); + await closeAllActiveReaders(); server.stop(); process.exit(0); diff --git a/src/lib/ai-proxy.ts b/src/lib/ai-proxy.ts index 648ea44..e5b00d3 100644 --- a/src/lib/ai-proxy.ts +++ b/src/lib/ai-proxy.ts @@ -13,7 +13,7 @@ */ import type { ProxyPool, SessionProxyPool } from "./proxy-pool"; -import { fetchWithRetry, fetchWithSessionRetry, SSELineBuffer, isDevMode, type FetchWithRetryResult } from "./fetch-utils"; +import { fetchWithRetry, fetchWithSessionRetry, SSELineBuffer, isDevMode, trackReader, releaseReader, wrapStreamWithCleanup, type FetchWithRetryResult } from "./fetch-utils"; import { getJwt, invalidateJwt } from "./mimo-auth"; import * as aichatAuth from "./aichat-auth"; @@ -371,6 +371,13 @@ export async function handleChatCompletion( sessionPool?: SessionProxyPool, sessionId?: string, ): Promise; +export async function handleChatCompletion( + body: unknown, + proxyPool?: ProxyPool, + sessionPool?: SessionProxyPool, + sessionId?: string, + ipv6Source?: string, +): Promise; export async function handleChatCompletion( body: unknown, proxyPool?: ProxyPool, @@ -571,7 +578,7 @@ function transformStream( config: BackendConfig, req: OpenAIRequest, ): ReadableStream { - const reader = body.getReader(); + const reader = trackReader(body.getReader() as any as ReadableStreamDefaultReader); const decoder = new TextDecoder(); const encoder = new TextEncoder(); const lineBuffer = new SSELineBuffer(); @@ -603,10 +610,16 @@ function transformStream( async pull(controller) { try { startKeepalive(controller); + // Process multiple chunks before yielding to event loop + // to reduce per-chunk setTimeout overhead (~1-4ms each) + const BATCH_SIZE = 8; + let chunksProcessed = 0; + while (true) { const { done, value } = await reader.read(); if (done) { stopKeepalive(); + releaseReader(reader); // Flush remaining text after stream ends const remaining = lineBuffer.flush(); if (remaining.length > 0) { @@ -638,11 +651,16 @@ function transformStream( } } - // Yield to event loop after each chunk to prevent starvation - await new Promise((r) => setTimeout(r, 0)); + chunksProcessed++; + // Yield to event loop after batch to prevent starvation + if (chunksProcessed >= BATCH_SIZE) { + chunksProcessed = 0; + await new Promise((r) => setTimeout(r, 0)); + } } } catch (err) { stopKeepalive(); + releaseReader(reader); if (isDevMode()) { controller.enqueue( encoder.encode( @@ -671,6 +689,7 @@ function transformStream( /** * If a session is active, wrap the stream so the session is released on end/error. * Otherwise pass through the stream unchanged. + * Uses the shared wrapStreamWithCleanup from fetch-utils. */ function wrapStreamMaybe( body: ReadableStream, @@ -680,32 +699,3 @@ function wrapStreamMaybe( if (!sessionPool || !sessionId) return body; return wrapStreamWithCleanup(body, () => sessionPool.release(sessionId)); } - -/** - * Wraps a ReadableStream and calls `cleanup` when the stream ends, errors, - * or is cancelled by the consumer. - */ -function wrapStreamWithCleanup(body: ReadableStream, cleanup: () => void): ReadableStream { - const reader = body.getReader(); - - return new ReadableStream({ - async pull(controller) { - try { - const { done, value } = await reader.read(); - if (done) { - cleanup(); - controller.close(); - return; - } - controller.enqueue(value); - } catch (err) { - cleanup(); - controller.error(err); - } - }, - cancel(reason) { - cleanup(); - reader.cancel(reason); - }, - }); -} diff --git a/src/lib/anthropic-proxy.ts b/src/lib/anthropic-proxy.ts index 96b0347..dafee2e 100644 --- a/src/lib/anthropic-proxy.ts +++ b/src/lib/anthropic-proxy.ts @@ -12,9 +12,7 @@ import type { ProxyPool, SessionProxyPool } from "./proxy-pool"; import { MODEL_ROUTES, type BackendConfig } from "./ai-proxy"; -import { fetchWithRetry, fetchWithSessionRetry, type FetchWithRetryResult } from "./fetch-utils"; -import { SSELineBuffer } from "./fetch-utils"; -import { isDevMode } from "./fetch-utils"; +import { fetchWithRetry, fetchWithSessionRetry, SSELineBuffer, isDevMode, trackReader, releaseReader, wrapStreamWithCleanup, type FetchWithRetryResult } from "./fetch-utils"; // --- Types ------------------------------------------------------------------- @@ -277,7 +275,7 @@ function transformAnthropicStream( model: string, config: BackendConfig, ): ReadableStream { - const reader = body.getReader(); + const reader = trackReader(body.getReader() as any as ReadableStreamDefaultReader); const decoder = new TextDecoder(); const encoder = new TextEncoder(); const lineBuffer = new SSELineBuffer(); @@ -340,33 +338,44 @@ function transformAnthropicStream( } while (phase === "block") { - const { done, value } = await reader.read(); - if (done) { - stopKeepalive(); - const remaining = lineBuffer.flush(); - if (remaining.length > 0) { - const adapted = backendLineToAnthropicSSE(remaining, model, config); + // Process multiple chunks before yielding to event loop + // to reduce per-chunk setTimeout overhead (~1-4ms each) + const BATCH_SIZE = 8; + let chunksProcessed = 0; + + while (phase === "block") { + const { done, value } = await reader.read(); + if (done) { + stopKeepalive(); + releaseReader(reader); + const remaining = lineBuffer.flush(); + if (remaining.length > 0) { + const adapted = backendLineToAnthropicSSE(remaining, model, config); + if (adapted) { + controller.enqueue(encoder.encode(adapted + "\n\n")); + } + } + phase = "done"; + break; + } + + const chunk = decoder.decode(value, { stream: true }); + const lines = lineBuffer.add(chunk); + + for (const line of lines) { + const adapted = backendLineToAnthropicSSE(line, model, config); if (adapted) { controller.enqueue(encoder.encode(adapted + "\n\n")); } } - phase = "done"; - break; - } - const chunk = decoder.decode(value, { stream: true }); - const lines = lineBuffer.add(chunk); - - for (const line of lines) { - const adapted = backendLineToAnthropicSSE(line, model, config); - if (adapted) { - controller.enqueue(encoder.encode(adapted + "\n\n")); + chunksProcessed++; + if (chunksProcessed >= BATCH_SIZE) { + chunksProcessed = 0; + await new Promise((r) => setTimeout(r, 0)); + return; // Yield to event loop, pull will be called again } } - - // Yield to event loop after each chunk - await new Promise((r) => setTimeout(r, 0)); - return; } if (phase === "done") { @@ -398,6 +407,7 @@ function transformAnthropicStream( } } catch (err) { stopKeepalive(); + releaseReader(reader); if (isDevMode()) { controller.enqueue( encoder.encode( @@ -500,6 +510,13 @@ export async function handleAnthropicMessages( sessionPool?: SessionProxyPool, sessionId?: string, ): Promise; +export async function handleAnthropicMessages( + body: unknown, + proxyPool?: ProxyPool, + sessionPool?: SessionProxyPool, + sessionId?: string, + ipv6Source?: string, +): Promise; export async function handleAnthropicMessages( body: unknown, proxyPool?: ProxyPool, @@ -650,6 +667,7 @@ export async function handleAnthropicMessages( /** * If a session is active, wrap the stream so the session is released on end/error. * Otherwise pass through the stream unchanged. + * Uses the shared wrapStreamWithCleanup from fetch-utils. */ function wrapAnthropicStreamMaybe( body: ReadableStream, @@ -657,34 +675,5 @@ function wrapAnthropicStreamMaybe( sessionId?: string, ): ReadableStream { if (!sessionPool || !sessionId) return body; - return wrapAnthropicStreamWithCleanup(body, () => sessionPool.release(sessionId)); -} - -/** - * Wraps a ReadableStream and calls `cleanup` when the stream ends, errors, - * or is cancelled by the consumer. - */ -function wrapAnthropicStreamWithCleanup(body: ReadableStream, cleanup: () => void): ReadableStream { - const reader = body.getReader(); - - return new ReadableStream({ - async pull(controller) { - try { - const { done, value } = await reader.read(); - if (done) { - cleanup(); - controller.close(); - return; - } - controller.enqueue(value); - } catch (err) { - cleanup(); - controller.error(err); - } - }, - cancel(reason) { - cleanup(); - reader.cancel(reason); - }, - }); + return wrapStreamWithCleanup(body, () => sessionPool.release(sessionId)); } diff --git a/src/lib/fetch-utils.ts b/src/lib/fetch-utils.ts index d837e5e..2bb2ea2 100644 --- a/src/lib/fetch-utils.ts +++ b/src/lib/fetch-utils.ts @@ -7,6 +7,7 @@ import type { ProxyPool, SessionProxyPool } from "./proxy-pool"; import type { IPv6SourcePool } from "./ipv6-pool"; +import { unlink } from "node:fs/promises"; // ─── Constants ───────────────────────────────────────────────────────── @@ -20,28 +21,40 @@ export const ACTIVE_READERS = new Set(); /** * Track a reader for graceful shutdown. Returns a wrapped reader that - * auto-removes itself from ACTIVE_READERS when done/cancelled. + * auto-removes itself from ACTIVE_READERS when cancelled. + * + * IMPORTANT: When the stream finishes normally (done=true), the caller + * MUST call `releaseReader(reader)` to remove it from the tracking set. */ -export function trackReader(reader: T): T { - ACTIVE_READERS.add(reader); - // Auto-remove on stream end or cancellation +export function trackReader>(reader: T): T { + ACTIVE_READERS.add(reader as any); + // Auto-remove on cancellation const origCancel = reader.cancel.bind(reader); (reader as any).cancel = async (...args: any[]) => { - ACTIVE_READERS.delete(reader); + ACTIVE_READERS.delete(reader as any); return origCancel(...args); }; return reader; } /** - * Close all tracked active readers (called during graceful shutdown). - * Each reader's cancellation propagates to the upstream connection. + * Release a reader from the active tracking set. + * Call this when a stream finishes normally (reader.read() returns done=true). */ -export function closeAllActiveReaders(): void { - for (const reader of ACTIVE_READERS) { - try { reader.cancel(); } catch { /* already closed */ } - } +export function releaseReader(reader: ReadableStreamDefaultReader): void { + ACTIVE_READERS.delete(reader as any); +} + +/** + * Close all tracked active readers (called during graceful shutdown). + * Awaits each cancellation so upstream connections are properly closed. + */ +export async function closeAllActiveReaders(): Promise { + const readers = Array.from(ACTIVE_READERS); ACTIVE_READERS.clear(); + await Promise.allSettled( + readers.map((reader) => reader.cancel().catch(() => {})) + ); } // ─── Dev-mode guard ────────────────────────────────────────────────────── @@ -57,6 +70,40 @@ export function isDevMode(): boolean { return false; } +// ─── Stream cleanup wrapper ───────────────────────────────────────────── + +/** + * Wraps a ReadableStream and calls `cleanup` when the stream ends, + * errors, or is cancelled by the consumer. + * + * Use this to release proxy sessions, close connections, or free resources + * when a streaming response finishes. + */ +export function wrapStreamWithCleanup(body: ReadableStream, cleanup: () => void): ReadableStream { + const reader = body.getReader(); + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + cleanup(); + controller.close(); + return; + } + controller.enqueue(value); + } catch (err) { + cleanup(); + controller.error(err); + } + }, + cancel(reason) { + cleanup(); + reader.cancel(reason); + }, + }); +} + // ─── SSE line buffer (fixes chunk-boundary corruption) ─────────────────── /** @@ -132,6 +179,23 @@ function extractModel(context?: string): string | undefined { return context || undefined; } +// ─── Retry backoff helper ────────────────────────────────────────────── + +/** + * Calculate exponential backoff delay for retry attempts. + * Returns 0 for attempt 0 (no delay on first try). + * Pattern: 200ms, 400ms, 800ms, ... capped at 2000ms. + */ +function retryBackoffMs(attempt: number): number { + if (attempt <= 0) return 0; + return Math.min(200 * Math.pow(2, attempt - 1), 2000); +} + +/** Sleep for the specified milliseconds. */ +function sleep(ms: number): Promise { + return new Promise((r) => setTimeout(r, ms)); +} + // ─── Error sanitization (prevent leaking upstream details) ────────────── /** @@ -164,19 +228,9 @@ export function sanitizeErrorMessage(raw: string): string { * This is used when outbound IPv6 source rotation is needed, since * Bun's built-in fetch() does not support specifying a local address. * - * Returns a streaming Response — the body is a ReadableStream from curl's - * stdout. The HTTP status code is extracted from a separate stderr header - * written by curl's `-w` flag (via a wrapper script approach). - * - * For simplicity and reliability, we use a two-phase approach: - * 1. First, send headers-only request to get status code (HEAD-like) - * 2. Then, stream the body via a second curl call - * - * Actually, simpler: we write a tiny wrapper that extracts the status code - * from the first line and streams the rest. But curl doesn't support that. - * - * Best approach: use `-w` to write status to a temp file descriptor, - * and stream stdout directly. We'll use a pipe-based approach. + * Status code is extracted from a temp file written by curl's `-w` flag + * (written after body completes). The body is collected into a single + * buffer and returned as a ReadableStream for zero-copy handoff. * * @param url - Target URL * @param init - Request init (method, headers, body) @@ -191,7 +245,6 @@ export async function fetchViaCurl( ): Promise { const method = (init.method ?? "GET").toUpperCase(); - // Write status code to a temp file, stream body to stdout const statusFile = `/tmp/curl-status-${crypto.randomUUID().slice(0, 8)}`; const args = [ "curl", @@ -201,7 +254,7 @@ export async function fetchViaCurl( "-s", // silent mode "--compressed", // auto-decompress gzip/brotli "-o", "-", // output body to stdout - "-w", statusFile, // write status code to file + "-w", statusFile, // write status code to file (plain text, appended after body) "--max-time", String(Math.ceil(timeoutMs / 1000)), "--connect-timeout", "10", ]; @@ -241,9 +294,10 @@ export async function fetchViaCurl( try { proc.kill("SIGKILL"); } catch { /* already dead */ } }, timeoutMs + 5000); - // Collect stdout into a buffer so we can parse status code from -w file - // after process exits, then return the buffered body as a stream. - // This is necessary because curl's -w writes AFTER the body completes. + // Collect stdout into minimal chunks then concatenate into a single buffer. + // This is needed because curl's -w (status code) is written AFTER the body + // completes, so we must wait for proc.exited before reading the status file. + // We minimize peak memory by streaming chunks into a pre-allocated buffer. const stdoutChunks: Uint8Array[] = []; const stdoutReader = proc.stdout.getReader(); try { @@ -254,7 +308,7 @@ export async function fetchViaCurl( } } catch { /* stream cancelled */ } - // Wait for process to exit and read status code + // Wait for process to exit, then read status code await proc.exited; clearTimeout(killTimer); @@ -264,21 +318,24 @@ export async function fetchViaCurl( statusCode = parseInt(statusText.trim(), 10) || 502; } catch { // Status file not written — connection likely failed - } finally { - try { await Bun.write(statusFile, ""); } catch { /* ignore cleanup errors */ } } + // Clean up temp file (async, non-blocking) + unlink(statusFile).catch(() => {}); + logProxy("fetchViaCurl", `response status=${statusCode}`, { ipv6Source, url }); - // Return buffered body as a readable stream - let offset = 0; + // Concatenate chunks into a single buffer for the Response body stream. + // Uses a single allocation to reduce GC pressure from many small chunks. const totalLength = stdoutChunks.reduce((sum, c) => sum + c.byteLength, 0); const combined = new Uint8Array(totalLength); + let offset = 0; for (const chunk of stdoutChunks) { combined.set(chunk, offset); offset += chunk.byteLength; } + // Zero-copy: hand the buffer directly to ReadableStream const bodyStream = new ReadableStream({ start(controller) { controller.enqueue(combined); @@ -329,6 +386,13 @@ export async function fetchWithRetry( const maxAttempts = poolSize > 0 ? poolSize + 1 : 3; for (let attempt = 0; attempt < maxAttempts; attempt++) { + // Exponential backoff between attempts (skip on first attempt) + const backoff = retryBackoffMs(attempt); + if (backoff > 0) { + logProxy("fetchWithRetry", `backoff ${backoff}ms before attempt ${attempt + 1}`, { context }); + await sleep(backoff); + } + if (attempt === 0) { // First attempt — direct (no proxy) init.proxy = undefined; @@ -512,17 +576,35 @@ export async function fetchWithSessionRetry( let lastError: unknown; let lastResponse: Response | undefined; for (let attempt = 0; attempt < totalAttempts; attempt++) { + // Exponential backoff between attempts (skip on first attempt) + const backoff = retryBackoffMs(attempt); + if (backoff > 0) { + logProxy("fetchWithSessionRetry", `backoff ${backoff}ms before attempt ${attempt + 1}`, { + context, + sessionId: sessionId.slice(0, 8), + }); + await sleep(backoff); + } + // Determine proxy for this attempt if (attempt === 0) { // First attempt — direct (no proxy) init.proxy = undefined; } else if (sessionPool.size > 0) { - // Fallback — use session-sticky proxy - const proxyUrl = attempt === 1 - ? sessionPool.acquire(sessionId) - : sessionPool.getProxyUrl(sessionId); - if (proxyUrl) { - init.proxy = proxyUrl; + // Fallback — acquire on first proxy attempt, rotate to different proxy on retries + if (attempt === 1) { + const proxyUrl = sessionPool.acquire(sessionId); + if (proxyUrl) { + init.proxy = proxyUrl; + } + } else { + // Force rotation to a different proxy before getting URL + const model = extractModel(context); + sessionPool.rotateNow(sessionId, model); + const proxyUrl = sessionPool.getProxyUrl(sessionId); + if (proxyUrl) { + init.proxy = proxyUrl; + } } } else { // No proxy pool, retry direct diff --git a/src/lib/proxy-pool.ts b/src/lib/proxy-pool.ts index 68f8e8b..f6ff250 100644 --- a/src/lib/proxy-pool.ts +++ b/src/lib/proxy-pool.ts @@ -41,6 +41,8 @@ export class ProxyPool { /** host:port::model -> expiry epoch ms */ private cooldowns = new Map(); private cooldownDuration = 60000; // default 60s + /** Periodic cleanup timer for expired cooldowns */ + private cleanupTimer: ReturnType | null = null; // -- Load -------------------------------------------------------------------- @@ -126,9 +128,29 @@ export class ProxyPool { this.currentIndex = 0; this.failures.clear(); + // Start periodic cooldown cleanup (every 30s) + this.startCooldownCleanup(); + logPool(`loaded ${this.proxies.length} proxies from ${source}`); } + /** Periodically clean up expired cooldown entries to prevent memory leaks. */ + private startCooldownCleanup(): void { + if (this.cleanupTimer) return; + this.cleanupTimer = setInterval(() => { + const now = Date.now(); + for (const [key, expiry] of this.cooldowns) { + if (now > expiry) { + this.cooldowns.delete(key); + } + } + }, 30_000); // every 30s + // Allow process to exit even if timer is active + if (this.cleanupTimer && typeof this.cleanupTimer === "object" && "unref" in this.cleanupTimer) { + this.cleanupTimer.unref(); + } + } + /** * Convenience -- load from a path or skip. * Returns `true` if proxies were loaded. diff --git a/src/middleware/rate-limiter.ts b/src/middleware/rate-limiter.ts index f377fd2..b384c3e 100644 --- a/src/middleware/rate-limiter.ts +++ b/src/middleware/rate-limiter.ts @@ -42,13 +42,16 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter { function pruneTimestamps(timestamps: number[], now: number): number[] { const cutoff = now - windowMs; - const result: number[] = []; - for (let i = 0; i < timestamps.length; i++) { - if (timestamps[i] >= cutoff) { - result.push(timestamps[i]); - } + // Find the first valid index using binary search-like scan + // (timestamps are monotonically increasing, so we can find the cutoff point) + let keepStart = 0; + while (keepStart < timestamps.length && timestamps[keepStart]! < cutoff) { + keepStart++; } - return result; + // Only allocate a new array if pruning is needed + if (keepStart === 0) return timestamps; + if (keepStart >= timestamps.length) return []; + return timestamps.slice(keepStart); } /** Periodically sweep the entire memory store to free memory. */ @@ -58,16 +61,16 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter { const cutoff = now - windowMs; for (const [key, timestamps] of store) { - const pruned: number[] = []; - for (let i = 0; i < timestamps.length; i++) { - if (timestamps[i] >= cutoff) { - pruned.push(timestamps[i]); - } + // Find first valid index (timestamps are sorted ascending) + let keepStart = 0; + while (keepStart < timestamps.length && timestamps[keepStart]! < cutoff) { + keepStart++; } - if (pruned.length === 0) { + if (keepStart === 0) continue; // all valid, skip + if (keepStart >= timestamps.length) { store.delete(key); } else { - store.set(key, pruned); + store.set(key, timestamps.slice(keepStart)); } }