perf: stability and performance improvements across proxy
- Add exponential backoff (200ms-2s) between retry attempts - Fix session retry bug: rotate proxy before retry to avoid same proxy - Fix ACTIVE_READERS memory leak: release reader on stream done/error - Make closeAllActiveReaders async with proper await - WebSocket backpressure: pause upstream forwarding at 512KB buffer - Optimize rate-limiter pruneTimestamps: avoid array reallocation - Batch stream chunks (8 per yield) to reduce event loop overhead - Add periodic cooldown cleanup in ProxyPool (30s interval) - Extract shared wrapStreamWithCleanup utility to reduce duplication - Fix fetchViaCurl temp file cleanup with proper unlink() - Add missing ipv6Source overload to AI proxy handlers
This commit is contained in:
+24
-10
@@ -110,6 +110,8 @@ interface WSRelayData {
|
|||||||
target: string;
|
target: string;
|
||||||
relayPath: string;
|
relayPath: string;
|
||||||
upstream?: WebSocket;
|
upstream?: WebSocket;
|
||||||
|
/** Set when client buffer exceeds threshold — stops forwarding upstream data */
|
||||||
|
paused?: boolean;
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- Route handlers ----------------------------------------------------------
|
// --- Route handlers ----------------------------------------------------------
|
||||||
@@ -621,6 +623,9 @@ const server: Server<WSRelayData> = Bun.serve<WSRelayData>({
|
|||||||
};
|
};
|
||||||
|
|
||||||
upstream.onmessage = (event: MessageEvent) => {
|
upstream.onmessage = (event: MessageEvent) => {
|
||||||
|
// Backpressure: drop messages when client buffer is full
|
||||||
|
if (ws.data.paused) return;
|
||||||
|
|
||||||
const data = event.data;
|
const data = event.data;
|
||||||
if (typeof data === "string") {
|
if (typeof data === "string") {
|
||||||
ws.sendText(data);
|
ws.sendText(data);
|
||||||
@@ -628,7 +633,9 @@ const server: Server<WSRelayData> = Bun.serve<WSRelayData>({
|
|||||||
ws.sendBinary(new Uint8Array(data));
|
ws.sendBinary(new Uint8Array(data));
|
||||||
} else if (data instanceof Blob) {
|
} else if (data instanceof Blob) {
|
||||||
data.arrayBuffer().then((buf) => {
|
data.arrayBuffer().then((buf) => {
|
||||||
ws.sendBinary(new Uint8Array(buf));
|
if (!ws.data.paused) {
|
||||||
|
ws.sendBinary(new Uint8Array(buf));
|
||||||
|
}
|
||||||
});
|
});
|
||||||
} else {
|
} else {
|
||||||
ws.sendBinary(data as unknown as Uint8Array);
|
ws.sendBinary(data as unknown as Uint8Array);
|
||||||
@@ -669,15 +676,22 @@ const server: Server<WSRelayData> = Bun.serve<WSRelayData>({
|
|||||||
},
|
},
|
||||||
|
|
||||||
drain(ws: ServerWebSocket<WSRelayData>) {
|
drain(ws: ServerWebSocket<WSRelayData>) {
|
||||||
// Backpressure: pause upstream reads when client is slow
|
// Backpressure: pause forwarding when client buffer is large
|
||||||
const upstream = ws.data.upstream;
|
const upstream = ws.data.upstream;
|
||||||
if (upstream && upstream.readyState === WebSocket.OPEN) {
|
if (!upstream || upstream.readyState !== WebSocket.OPEN) return;
|
||||||
// Bun WebSocket handles backpressure internally via bufferAmount
|
|
||||||
// Log if buffer is growing excessively
|
const buffered = (ws as any).bufferAmount ?? 0;
|
||||||
const buffered = (ws as any).bufferAmount ?? 0;
|
const BACKPRESSURE_THRESHOLD = 512 * 1024; // 512KB
|
||||||
if (buffered > 1024 * 1024) {
|
const RESUME_THRESHOLD = 64 * 1024; // 64KB
|
||||||
console.warn(`[ws] Client backpressure: ${buffered} bytes buffered`);
|
|
||||||
|
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 -------------------------------------------------------
|
// --- Graceful Shutdown -------------------------------------------------------
|
||||||
|
|
||||||
const shutdownHandler = (signal: string) => {
|
const shutdownHandler = async (signal: string) => {
|
||||||
console.log(`\n[relay] Received ${signal}, shutting down gracefully...`);
|
console.log(`\n[relay] Received ${signal}, shutting down gracefully...`);
|
||||||
|
|
||||||
// Close active SSE streams so clients get proper stream end events
|
// Close active SSE streams so clients get proper stream end events
|
||||||
closeAllActiveReaders();
|
await closeAllActiveReaders();
|
||||||
|
|
||||||
server.stop();
|
server.stop();
|
||||||
process.exit(0);
|
process.exit(0);
|
||||||
|
|||||||
+23
-33
@@ -13,7 +13,7 @@
|
|||||||
*/
|
*/
|
||||||
|
|
||||||
import type { ProxyPool, SessionProxyPool } from "./proxy-pool";
|
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 { getJwt, invalidateJwt } from "./mimo-auth";
|
||||||
import * as aichatAuth from "./aichat-auth";
|
import * as aichatAuth from "./aichat-auth";
|
||||||
|
|
||||||
@@ -371,6 +371,13 @@ export async function handleChatCompletion(
|
|||||||
sessionPool?: SessionProxyPool,
|
sessionPool?: SessionProxyPool,
|
||||||
sessionId?: string,
|
sessionId?: string,
|
||||||
): Promise<Response>;
|
): Promise<Response>;
|
||||||
|
export async function handleChatCompletion(
|
||||||
|
body: unknown,
|
||||||
|
proxyPool?: ProxyPool,
|
||||||
|
sessionPool?: SessionProxyPool,
|
||||||
|
sessionId?: string,
|
||||||
|
ipv6Source?: string,
|
||||||
|
): Promise<Response>;
|
||||||
export async function handleChatCompletion(
|
export async function handleChatCompletion(
|
||||||
body: unknown,
|
body: unknown,
|
||||||
proxyPool?: ProxyPool,
|
proxyPool?: ProxyPool,
|
||||||
@@ -571,7 +578,7 @@ function transformStream(
|
|||||||
config: BackendConfig,
|
config: BackendConfig,
|
||||||
req: OpenAIRequest,
|
req: OpenAIRequest,
|
||||||
): ReadableStream {
|
): ReadableStream {
|
||||||
const reader = body.getReader();
|
const reader = trackReader(body.getReader() as any as ReadableStreamDefaultReader);
|
||||||
const decoder = new TextDecoder();
|
const decoder = new TextDecoder();
|
||||||
const encoder = new TextEncoder();
|
const encoder = new TextEncoder();
|
||||||
const lineBuffer = new SSELineBuffer();
|
const lineBuffer = new SSELineBuffer();
|
||||||
@@ -603,10 +610,16 @@ function transformStream(
|
|||||||
async pull(controller) {
|
async pull(controller) {
|
||||||
try {
|
try {
|
||||||
startKeepalive(controller);
|
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) {
|
while (true) {
|
||||||
const { done, value } = await reader.read();
|
const { done, value } = await reader.read();
|
||||||
if (done) {
|
if (done) {
|
||||||
stopKeepalive();
|
stopKeepalive();
|
||||||
|
releaseReader(reader);
|
||||||
// Flush remaining text after stream ends
|
// Flush remaining text after stream ends
|
||||||
const remaining = lineBuffer.flush();
|
const remaining = lineBuffer.flush();
|
||||||
if (remaining.length > 0) {
|
if (remaining.length > 0) {
|
||||||
@@ -638,11 +651,16 @@ function transformStream(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Yield to event loop after each chunk to prevent starvation
|
chunksProcessed++;
|
||||||
await new Promise((r) => setTimeout(r, 0));
|
// Yield to event loop after batch to prevent starvation
|
||||||
|
if (chunksProcessed >= BATCH_SIZE) {
|
||||||
|
chunksProcessed = 0;
|
||||||
|
await new Promise((r) => setTimeout(r, 0));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
stopKeepalive();
|
stopKeepalive();
|
||||||
|
releaseReader(reader);
|
||||||
if (isDevMode()) {
|
if (isDevMode()) {
|
||||||
controller.enqueue(
|
controller.enqueue(
|
||||||
encoder.encode(
|
encoder.encode(
|
||||||
@@ -671,6 +689,7 @@ function transformStream(
|
|||||||
/**
|
/**
|
||||||
* If a session is active, wrap the stream so the session is released on end/error.
|
* If a session is active, wrap the stream so the session is released on end/error.
|
||||||
* Otherwise pass through the stream unchanged.
|
* Otherwise pass through the stream unchanged.
|
||||||
|
* Uses the shared wrapStreamWithCleanup from fetch-utils.
|
||||||
*/
|
*/
|
||||||
function wrapStreamMaybe(
|
function wrapStreamMaybe(
|
||||||
body: ReadableStream,
|
body: ReadableStream,
|
||||||
@@ -680,32 +699,3 @@ function wrapStreamMaybe(
|
|||||||
if (!sessionPool || !sessionId) return body;
|
if (!sessionPool || !sessionId) return body;
|
||||||
return wrapStreamWithCleanup(body, () => sessionPool.release(sessionId));
|
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);
|
|
||||||
},
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|||||||
+43
-54
@@ -12,9 +12,7 @@
|
|||||||
|
|
||||||
import type { ProxyPool, SessionProxyPool } from "./proxy-pool";
|
import type { ProxyPool, SessionProxyPool } from "./proxy-pool";
|
||||||
import { MODEL_ROUTES, type BackendConfig } from "./ai-proxy";
|
import { MODEL_ROUTES, type BackendConfig } from "./ai-proxy";
|
||||||
import { fetchWithRetry, fetchWithSessionRetry, type FetchWithRetryResult } from "./fetch-utils";
|
import { fetchWithRetry, fetchWithSessionRetry, SSELineBuffer, isDevMode, trackReader, releaseReader, wrapStreamWithCleanup, type FetchWithRetryResult } from "./fetch-utils";
|
||||||
import { SSELineBuffer } from "./fetch-utils";
|
|
||||||
import { isDevMode } from "./fetch-utils";
|
|
||||||
|
|
||||||
// --- Types -------------------------------------------------------------------
|
// --- Types -------------------------------------------------------------------
|
||||||
|
|
||||||
@@ -277,7 +275,7 @@ function transformAnthropicStream(
|
|||||||
model: string,
|
model: string,
|
||||||
config: BackendConfig,
|
config: BackendConfig,
|
||||||
): ReadableStream {
|
): ReadableStream {
|
||||||
const reader = body.getReader();
|
const reader = trackReader(body.getReader() as any as ReadableStreamDefaultReader);
|
||||||
const decoder = new TextDecoder();
|
const decoder = new TextDecoder();
|
||||||
const encoder = new TextEncoder();
|
const encoder = new TextEncoder();
|
||||||
const lineBuffer = new SSELineBuffer();
|
const lineBuffer = new SSELineBuffer();
|
||||||
@@ -340,33 +338,44 @@ function transformAnthropicStream(
|
|||||||
}
|
}
|
||||||
|
|
||||||
while (phase === "block") {
|
while (phase === "block") {
|
||||||
const { done, value } = await reader.read();
|
// Process multiple chunks before yielding to event loop
|
||||||
if (done) {
|
// to reduce per-chunk setTimeout overhead (~1-4ms each)
|
||||||
stopKeepalive();
|
const BATCH_SIZE = 8;
|
||||||
const remaining = lineBuffer.flush();
|
let chunksProcessed = 0;
|
||||||
if (remaining.length > 0) {
|
|
||||||
const adapted = backendLineToAnthropicSSE(remaining, model, config);
|
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) {
|
if (adapted) {
|
||||||
controller.enqueue(encoder.encode(adapted + "\n\n"));
|
controller.enqueue(encoder.encode(adapted + "\n\n"));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
phase = "done";
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
const chunk = decoder.decode(value, { stream: true });
|
chunksProcessed++;
|
||||||
const lines = lineBuffer.add(chunk);
|
if (chunksProcessed >= BATCH_SIZE) {
|
||||||
|
chunksProcessed = 0;
|
||||||
for (const line of lines) {
|
await new Promise((r) => setTimeout(r, 0));
|
||||||
const adapted = backendLineToAnthropicSSE(line, model, config);
|
return; // Yield to event loop, pull will be called again
|
||||||
if (adapted) {
|
|
||||||
controller.enqueue(encoder.encode(adapted + "\n\n"));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Yield to event loop after each chunk
|
|
||||||
await new Promise((r) => setTimeout(r, 0));
|
|
||||||
return;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if (phase === "done") {
|
if (phase === "done") {
|
||||||
@@ -398,6 +407,7 @@ function transformAnthropicStream(
|
|||||||
}
|
}
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
stopKeepalive();
|
stopKeepalive();
|
||||||
|
releaseReader(reader);
|
||||||
if (isDevMode()) {
|
if (isDevMode()) {
|
||||||
controller.enqueue(
|
controller.enqueue(
|
||||||
encoder.encode(
|
encoder.encode(
|
||||||
@@ -500,6 +510,13 @@ export async function handleAnthropicMessages(
|
|||||||
sessionPool?: SessionProxyPool,
|
sessionPool?: SessionProxyPool,
|
||||||
sessionId?: string,
|
sessionId?: string,
|
||||||
): Promise<Response>;
|
): Promise<Response>;
|
||||||
|
export async function handleAnthropicMessages(
|
||||||
|
body: unknown,
|
||||||
|
proxyPool?: ProxyPool,
|
||||||
|
sessionPool?: SessionProxyPool,
|
||||||
|
sessionId?: string,
|
||||||
|
ipv6Source?: string,
|
||||||
|
): Promise<Response>;
|
||||||
export async function handleAnthropicMessages(
|
export async function handleAnthropicMessages(
|
||||||
body: unknown,
|
body: unknown,
|
||||||
proxyPool?: ProxyPool,
|
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.
|
* If a session is active, wrap the stream so the session is released on end/error.
|
||||||
* Otherwise pass through the stream unchanged.
|
* Otherwise pass through the stream unchanged.
|
||||||
|
* Uses the shared wrapStreamWithCleanup from fetch-utils.
|
||||||
*/
|
*/
|
||||||
function wrapAnthropicStreamMaybe(
|
function wrapAnthropicStreamMaybe(
|
||||||
body: ReadableStream,
|
body: ReadableStream,
|
||||||
@@ -657,34 +675,5 @@ function wrapAnthropicStreamMaybe(
|
|||||||
sessionId?: string,
|
sessionId?: string,
|
||||||
): ReadableStream {
|
): ReadableStream {
|
||||||
if (!sessionPool || !sessionId) return body;
|
if (!sessionPool || !sessionId) return body;
|
||||||
return wrapAnthropicStreamWithCleanup(body, () => sessionPool.release(sessionId));
|
return wrapStreamWithCleanup(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);
|
|
||||||
},
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
|
|||||||
+122
-40
@@ -7,6 +7,7 @@
|
|||||||
|
|
||||||
import type { ProxyPool, SessionProxyPool } from "./proxy-pool";
|
import type { ProxyPool, SessionProxyPool } from "./proxy-pool";
|
||||||
import type { IPv6SourcePool } from "./ipv6-pool";
|
import type { IPv6SourcePool } from "./ipv6-pool";
|
||||||
|
import { unlink } from "node:fs/promises";
|
||||||
|
|
||||||
// ─── Constants ─────────────────────────────────────────────────────────
|
// ─── Constants ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
@@ -20,28 +21,40 @@ export const ACTIVE_READERS = new Set<ReadableStreamDefaultReader>();
|
|||||||
|
|
||||||
/**
|
/**
|
||||||
* Track a reader for graceful shutdown. Returns a wrapped reader that
|
* 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<T extends ReadableStreamDefaultReader>(reader: T): T {
|
export function trackReader<T extends ReadableStreamDefaultReader<any>>(reader: T): T {
|
||||||
ACTIVE_READERS.add(reader);
|
ACTIVE_READERS.add(reader as any);
|
||||||
// Auto-remove on stream end or cancellation
|
// Auto-remove on cancellation
|
||||||
const origCancel = reader.cancel.bind(reader);
|
const origCancel = reader.cancel.bind(reader);
|
||||||
(reader as any).cancel = async (...args: any[]) => {
|
(reader as any).cancel = async (...args: any[]) => {
|
||||||
ACTIVE_READERS.delete(reader);
|
ACTIVE_READERS.delete(reader as any);
|
||||||
return origCancel(...args);
|
return origCancel(...args);
|
||||||
};
|
};
|
||||||
return reader;
|
return reader;
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Close all tracked active readers (called during graceful shutdown).
|
* Release a reader from the active tracking set.
|
||||||
* Each reader's cancellation propagates to the upstream connection.
|
* Call this when a stream finishes normally (reader.read() returns done=true).
|
||||||
*/
|
*/
|
||||||
export function closeAllActiveReaders(): void {
|
export function releaseReader(reader: ReadableStreamDefaultReader<any>): void {
|
||||||
for (const reader of ACTIVE_READERS) {
|
ACTIVE_READERS.delete(reader as any);
|
||||||
try { reader.cancel(); } catch { /* already closed */ }
|
}
|
||||||
}
|
|
||||||
|
/**
|
||||||
|
* Close all tracked active readers (called during graceful shutdown).
|
||||||
|
* Awaits each cancellation so upstream connections are properly closed.
|
||||||
|
*/
|
||||||
|
export async function closeAllActiveReaders(): Promise<void> {
|
||||||
|
const readers = Array.from(ACTIVE_READERS);
|
||||||
ACTIVE_READERS.clear();
|
ACTIVE_READERS.clear();
|
||||||
|
await Promise.allSettled(
|
||||||
|
readers.map((reader) => reader.cancel().catch(() => {}))
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// ─── Dev-mode guard ──────────────────────────────────────────────────────
|
// ─── Dev-mode guard ──────────────────────────────────────────────────────
|
||||||
@@ -57,6 +70,40 @@ export function isDevMode(): boolean {
|
|||||||
return false;
|
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) ───────────────────
|
// ─── SSE line buffer (fixes chunk-boundary corruption) ───────────────────
|
||||||
|
|
||||||
/**
|
/**
|
||||||
@@ -132,6 +179,23 @@ function extractModel(context?: string): string | undefined {
|
|||||||
return context || 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<void> {
|
||||||
|
return new Promise((r) => setTimeout(r, ms));
|
||||||
|
}
|
||||||
|
|
||||||
// ─── Error sanitization (prevent leaking upstream details) ──────────────
|
// ─── 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
|
* This is used when outbound IPv6 source rotation is needed, since
|
||||||
* Bun's built-in fetch() does not support specifying a local address.
|
* Bun's built-in fetch() does not support specifying a local address.
|
||||||
*
|
*
|
||||||
* Returns a streaming Response — the body is a ReadableStream from curl's
|
* Status code is extracted from a temp file written by curl's `-w` flag
|
||||||
* stdout. The HTTP status code is extracted from a separate stderr header
|
* (written after body completes). The body is collected into a single
|
||||||
* written by curl's `-w` flag (via a wrapper script approach).
|
* buffer and returned as a ReadableStream for zero-copy handoff.
|
||||||
*
|
|
||||||
* 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.
|
|
||||||
*
|
*
|
||||||
* @param url - Target URL
|
* @param url - Target URL
|
||||||
* @param init - Request init (method, headers, body)
|
* @param init - Request init (method, headers, body)
|
||||||
@@ -191,7 +245,6 @@ export async function fetchViaCurl(
|
|||||||
): Promise<Response> {
|
): Promise<Response> {
|
||||||
const method = (init.method ?? "GET").toUpperCase();
|
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 statusFile = `/tmp/curl-status-${crypto.randomUUID().slice(0, 8)}`;
|
||||||
const args = [
|
const args = [
|
||||||
"curl",
|
"curl",
|
||||||
@@ -201,7 +254,7 @@ export async function fetchViaCurl(
|
|||||||
"-s", // silent mode
|
"-s", // silent mode
|
||||||
"--compressed", // auto-decompress gzip/brotli
|
"--compressed", // auto-decompress gzip/brotli
|
||||||
"-o", "-", // output body to stdout
|
"-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)),
|
"--max-time", String(Math.ceil(timeoutMs / 1000)),
|
||||||
"--connect-timeout", "10",
|
"--connect-timeout", "10",
|
||||||
];
|
];
|
||||||
@@ -241,9 +294,10 @@ export async function fetchViaCurl(
|
|||||||
try { proc.kill("SIGKILL"); } catch { /* already dead */ }
|
try { proc.kill("SIGKILL"); } catch { /* already dead */ }
|
||||||
}, timeoutMs + 5000);
|
}, timeoutMs + 5000);
|
||||||
|
|
||||||
// Collect stdout into a buffer so we can parse status code from -w file
|
// Collect stdout into minimal chunks then concatenate into a single buffer.
|
||||||
// after process exits, then return the buffered body as a stream.
|
// This is needed because curl's -w (status code) is written AFTER the body
|
||||||
// This is necessary because curl's -w writes AFTER the body completes.
|
// 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 stdoutChunks: Uint8Array[] = [];
|
||||||
const stdoutReader = proc.stdout.getReader();
|
const stdoutReader = proc.stdout.getReader();
|
||||||
try {
|
try {
|
||||||
@@ -254,7 +308,7 @@ export async function fetchViaCurl(
|
|||||||
}
|
}
|
||||||
} catch { /* stream cancelled */ }
|
} catch { /* stream cancelled */ }
|
||||||
|
|
||||||
// Wait for process to exit and read status code
|
// Wait for process to exit, then read status code
|
||||||
await proc.exited;
|
await proc.exited;
|
||||||
clearTimeout(killTimer);
|
clearTimeout(killTimer);
|
||||||
|
|
||||||
@@ -264,21 +318,24 @@ export async function fetchViaCurl(
|
|||||||
statusCode = parseInt(statusText.trim(), 10) || 502;
|
statusCode = parseInt(statusText.trim(), 10) || 502;
|
||||||
} catch {
|
} catch {
|
||||||
// Status file not written — connection likely failed
|
// 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 });
|
logProxy("fetchViaCurl", `response status=${statusCode}`, { ipv6Source, url });
|
||||||
|
|
||||||
// Return buffered body as a readable stream
|
// Concatenate chunks into a single buffer for the Response body stream.
|
||||||
let offset = 0;
|
// Uses a single allocation to reduce GC pressure from many small chunks.
|
||||||
const totalLength = stdoutChunks.reduce((sum, c) => sum + c.byteLength, 0);
|
const totalLength = stdoutChunks.reduce((sum, c) => sum + c.byteLength, 0);
|
||||||
const combined = new Uint8Array(totalLength);
|
const combined = new Uint8Array(totalLength);
|
||||||
|
let offset = 0;
|
||||||
for (const chunk of stdoutChunks) {
|
for (const chunk of stdoutChunks) {
|
||||||
combined.set(chunk, offset);
|
combined.set(chunk, offset);
|
||||||
offset += chunk.byteLength;
|
offset += chunk.byteLength;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Zero-copy: hand the buffer directly to ReadableStream
|
||||||
const bodyStream = new ReadableStream({
|
const bodyStream = new ReadableStream({
|
||||||
start(controller) {
|
start(controller) {
|
||||||
controller.enqueue(combined);
|
controller.enqueue(combined);
|
||||||
@@ -329,6 +386,13 @@ export async function fetchWithRetry(
|
|||||||
const maxAttempts = poolSize > 0 ? poolSize + 1 : 3;
|
const maxAttempts = poolSize > 0 ? poolSize + 1 : 3;
|
||||||
|
|
||||||
for (let attempt = 0; attempt < maxAttempts; attempt++) {
|
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) {
|
if (attempt === 0) {
|
||||||
// First attempt — direct (no proxy)
|
// First attempt — direct (no proxy)
|
||||||
init.proxy = undefined;
|
init.proxy = undefined;
|
||||||
@@ -512,17 +576,35 @@ export async function fetchWithSessionRetry(
|
|||||||
let lastError: unknown;
|
let lastError: unknown;
|
||||||
let lastResponse: Response | undefined;
|
let lastResponse: Response | undefined;
|
||||||
for (let attempt = 0; attempt < totalAttempts; attempt++) {
|
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
|
// Determine proxy for this attempt
|
||||||
if (attempt === 0) {
|
if (attempt === 0) {
|
||||||
// First attempt — direct (no proxy)
|
// First attempt — direct (no proxy)
|
||||||
init.proxy = undefined;
|
init.proxy = undefined;
|
||||||
} else if (sessionPool.size > 0) {
|
} else if (sessionPool.size > 0) {
|
||||||
// Fallback — use session-sticky proxy
|
// Fallback — acquire on first proxy attempt, rotate to different proxy on retries
|
||||||
const proxyUrl = attempt === 1
|
if (attempt === 1) {
|
||||||
? sessionPool.acquire(sessionId)
|
const proxyUrl = sessionPool.acquire(sessionId);
|
||||||
: sessionPool.getProxyUrl(sessionId);
|
if (proxyUrl) {
|
||||||
if (proxyUrl) {
|
init.proxy = 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 {
|
} else {
|
||||||
// No proxy pool, retry direct
|
// No proxy pool, retry direct
|
||||||
|
|||||||
@@ -41,6 +41,8 @@ export class ProxyPool {
|
|||||||
/** host:port::model -> expiry epoch ms */
|
/** host:port::model -> expiry epoch ms */
|
||||||
private cooldowns = new Map<string, number>();
|
private cooldowns = new Map<string, number>();
|
||||||
private cooldownDuration = 60000; // default 60s
|
private cooldownDuration = 60000; // default 60s
|
||||||
|
/** Periodic cleanup timer for expired cooldowns */
|
||||||
|
private cleanupTimer: ReturnType<typeof setInterval> | null = null;
|
||||||
|
|
||||||
// -- Load --------------------------------------------------------------------
|
// -- Load --------------------------------------------------------------------
|
||||||
|
|
||||||
@@ -126,9 +128,29 @@ export class ProxyPool {
|
|||||||
this.currentIndex = 0;
|
this.currentIndex = 0;
|
||||||
this.failures.clear();
|
this.failures.clear();
|
||||||
|
|
||||||
|
// Start periodic cooldown cleanup (every 30s)
|
||||||
|
this.startCooldownCleanup();
|
||||||
|
|
||||||
logPool(`loaded ${this.proxies.length} proxies from ${source}`);
|
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.
|
* Convenience -- load from a path or skip.
|
||||||
* Returns `true` if proxies were loaded.
|
* Returns `true` if proxies were loaded.
|
||||||
|
|||||||
@@ -42,13 +42,16 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter {
|
|||||||
|
|
||||||
function pruneTimestamps(timestamps: number[], now: number): number[] {
|
function pruneTimestamps(timestamps: number[], now: number): number[] {
|
||||||
const cutoff = now - windowMs;
|
const cutoff = now - windowMs;
|
||||||
const result: number[] = [];
|
// Find the first valid index using binary search-like scan
|
||||||
for (let i = 0; i < timestamps.length; i++) {
|
// (timestamps are monotonically increasing, so we can find the cutoff point)
|
||||||
if (timestamps[i] >= cutoff) {
|
let keepStart = 0;
|
||||||
result.push(timestamps[i]);
|
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. */
|
/** Periodically sweep the entire memory store to free memory. */
|
||||||
@@ -58,16 +61,16 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter {
|
|||||||
const cutoff = now - windowMs;
|
const cutoff = now - windowMs;
|
||||||
|
|
||||||
for (const [key, timestamps] of store) {
|
for (const [key, timestamps] of store) {
|
||||||
const pruned: number[] = [];
|
// Find first valid index (timestamps are sorted ascending)
|
||||||
for (let i = 0; i < timestamps.length; i++) {
|
let keepStart = 0;
|
||||||
if (timestamps[i] >= cutoff) {
|
while (keepStart < timestamps.length && timestamps[keepStart]! < cutoff) {
|
||||||
pruned.push(timestamps[i]);
|
keepStart++;
|
||||||
}
|
|
||||||
}
|
}
|
||||||
if (pruned.length === 0) {
|
if (keepStart === 0) continue; // all valid, skip
|
||||||
|
if (keepStart >= timestamps.length) {
|
||||||
store.delete(key);
|
store.delete(key);
|
||||||
} else {
|
} else {
|
||||||
store.set(key, pruned);
|
store.set(key, timestamps.slice(keepStart));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user