From c36678adb7f9352772d9be38a361d1c2ee2fde08 Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 11 Aug 2026 17:26:53 +0800 Subject: [PATCH 1/2] fix(replay): bound detached stream memory --- src/app/v1/_lib/proxy/replay/replay-spool.ts | 214 +++++++++--- src/app/v1/_lib/proxy/replay/replay-store.ts | 70 +++- src/app/v1/_lib/proxy/response-handler.ts | 69 ++-- tests/unit/proxy/replay-spool.test.ts | 313 +++++++++++++----- tests/unit/proxy/replay-store.test.ts | 53 ++- ...esponse-handler-client-abort-drain.test.ts | 154 +++++++++ .../response-handler-stream-terminal.test.ts | 1 + 7 files changed, 713 insertions(+), 161 deletions(-) diff --git a/src/app/v1/_lib/proxy/replay/replay-spool.ts b/src/app/v1/_lib/proxy/replay/replay-spool.ts index 11817432e..0d92545b4 100644 --- a/src/app/v1/_lib/proxy/replay/replay-spool.ts +++ b/src/app/v1/_lib/proxy/replay/replay-spool.ts @@ -27,10 +27,44 @@ import { const FLUSH_INTERVAL_MS = 100; const FLUSH_BYTES_THRESHOLD = 64 * 1024; +const MAX_WRITE_BEHIND_BYTES = 512 * 1024; +const LARGE_PAYLOAD_REBUILD_BYTES = 256 * 1024; +const MAX_CONCURRENT_LARGE_PAYLOAD_REBUILDS = 2; const OWNER_HEARTBEAT_INTERVAL_MS = 15_000; const PRE_SPOOL_ABORT_WAIT_MS = 100; let activeSpoolCount = 0; +let activeLargePayloadRebuilds = 0; +const largePayloadRebuildWaiters: Array<() => void> = []; + +type QueuedReplayBatch = { + chunks: string[]; + byteSize: number; +}; + +async function withPayloadRebuildSlot( + byteSize: number, + operation: () => Promise +): Promise { + if (byteSize < LARGE_PAYLOAD_REBUILD_BYTES) return operation(); + + if (activeLargePayloadRebuilds < MAX_CONCURRENT_LARGE_PAYLOAD_REBUILDS) { + activeLargePayloadRebuilds += 1; + } else { + await new Promise((resolve) => largePayloadRebuildWaiters.push(resolve)); + } + + try { + return await operation(); + } finally { + const next = largePayloadRebuildWaiters.shift(); + if (next) { + next(); + } else { + activeLargePayloadRebuilds = Math.max(0, activeLargePayloadRebuilds - 1); + } + } +} export function getActiveReplaySpoolCount(): number { return activeSpoolCount; @@ -39,10 +73,13 @@ export function getActiveReplaySpoolCount(): number { export class ReplaySpool { private readonly store = getReplayStore(); private readonly decoder = new TextDecoder("utf-8"); - private readonly parts: string[] = []; - private readonly queuedBatches = new Set(); + private readonly encoder = new TextEncoder(); + private readonly queuedBatches = new Set(); + private readonly terminalListeners = new Set<() => void>(); private pending: string[] = []; private pendingBytes = 0; + private activeWriteBatch: QueuedReplayBatch | null = null; + private queuedBytes = 0; private totalBytes = 0; private chunkCount = 0; private disabled = false; @@ -69,6 +106,16 @@ export class ReplaySpool { return this.terminal || this.disabled; } + /** 订阅 spool 终态或失效, 供 detached drain 动态收紧资源窗口. */ + onTerminal(listener: () => void): () => void { + if (this.isTerminal) { + listener(); + return () => {}; + } + this.terminalListeners.add(listener); + return () => this.terminalListeners.delete(listener); + } + /** 流热路径同步观察：累积并调度冲刷。 */ observe(chunk: Uint8Array): void { if (this.disabled || this.terminal || chunk.byteLength === 0) return; @@ -82,8 +129,12 @@ export class ReplaySpool { const text = this.decoder.decode(chunk, { stream: true }); if (text.length === 0) return; this.pending.push(text); - this.parts.push(text); - this.pendingBytes += chunk.byteLength; + this.pendingBytes += this.encoder.encode(text).byteLength; + + if (this.exceedsWriteBehindLimit(this.pendingBytes)) { + this.disable("write_behind_limit"); + return; + } if (this.pendingBytes >= FLUSH_BYTES_THRESHOLD) { this.scheduleFlush(0); @@ -117,22 +168,24 @@ export class ReplaySpool { } private enqueueFlush(): void { - const batch = this.pending; - if (batch.length === 0) return; + const chunks = this.pending; + if (chunks.length === 0) return; + const batch: QueuedReplayBatch = { chunks, byteSize: this.pendingBytes }; this.pending = []; this.pendingBytes = 0; - this.queuedBatches.add(batch); + this.trackQueuedBatch(batch); // 续接体自带 try/catch：链永不 rejected；每个 await 之后复查 disabled， // 防止与 disable/halt 竞态时在清理之后又写回 owning meta this.writeChain = this.writeChain.then(async () => { try { if (this.disabled || this.aborting) return; - const expectedChunkCount = this.chunkCount + batch.length; + const expectedChunkCount = this.chunkCount + batch.chunks.length; const appended = await this.store.writeOwned( this.identity.replayId, this.ownerToken, this.buildMeta("owning", { chunkCount: expectedChunkCount }), - batch + this.chunkCount, + batch.chunks ); if (this.disabled || this.aborting) return; if (appended === null) { @@ -144,6 +197,10 @@ export class ReplaySpool { this.halt("owner_lease_lost"); return; } + if (appended === "chunk_count_mismatch" || appended !== expectedChunkCount) { + this.disable("chunk_count_mismatch"); + return; + } this.chunkCount = appended; this.metaWritten = true; } catch (error) { @@ -153,7 +210,7 @@ export class ReplaySpool { }); this.disable("flush_error"); } finally { - this.queuedBatches.delete(batch); + this.releaseQueuedBatch(batch); } }); } @@ -202,7 +259,8 @@ export class ReplaySpool { const chunkCount = await this.store.writeOwned( this.identity.replayId, this.ownerToken, - this.buildMeta("owning") + this.buildMeta("owning"), + 0 ); if (this.disabled || this.aborting) return; if (chunkCount === null) { @@ -213,6 +271,10 @@ export class ReplaySpool { this.halt("owner_lease_lost"); return; } + if (chunkCount === "chunk_count_mismatch" || chunkCount !== 0) { + this.disable("chunk_count_mismatch"); + return; + } this.chunkCount = chunkCount; this.metaWritten = true; } catch (error) { @@ -232,27 +294,37 @@ export class ReplaySpool { async completeAfterBilling(messageRequestId: number | null): Promise { if (this.disabled || this.terminal) return; this.terminal = true; + this.notifyTerminal(); this.clearFlushTimer(); const tail = this.decoder.decode(); if (tail.length > 0) { this.pending.push(tail); - this.parts.push(tail); + this.pendingBytes += this.encoder.encode(tail).byteLength; + } + if (this.exceedsWriteBehindLimit(this.pendingBytes)) { + this.disable("write_behind_limit"); + await this.writeChain; + return; } - const batch = this.pending; + const batch: QueuedReplayBatch = { + chunks: this.pending, + byteSize: this.pendingBytes, + }; this.pending = []; this.pendingBytes = 0; - this.queuedBatches.add(batch); + this.trackQueuedBatch(batch); this.writeChain = this.writeChain.then(async () => { let pgPersisted = false; try { if (this.disabled || this.aborting) return; - const expectedChunkCount = this.chunkCount + batch.length; + const expectedChunkCount = this.chunkCount + batch.chunks.length; const appended = await this.store.writeOwned( this.identity.replayId, this.ownerToken, this.buildMeta("owning", { chunkCount: expectedChunkCount }), - batch + this.chunkCount, + batch.chunks ); if (appended === false) { throw new Error("replay owner lease lost before completion"); @@ -261,23 +333,47 @@ export class ReplaySpool { // 尾批或 owning meta 丢失时热层条目不完整，绝不能置 completed throw new Error("final replay flush failed"); } + if (appended === "chunk_count_mismatch" || appended !== expectedChunkCount) { + throw new Error("replay chunks changed before completion"); + } this.chunkCount = appended; this.metaWritten = true; - const payload = this.takePayload(); - // 先写 PG（持久 payload），再翻 Redis meta 为 completed（热层可服务） - const persistResult = await this.store.persistCompleted({ - replayId: this.identity.replayId, - verifier: this.identity.verifier, - scopeTag: this.identity.scopeTag, - keyId: this.identity.keyId, - userId: this.identity.userId, - format: this.identity.format, - model: this.identity.model, - statusCode: this.statusCode, - headers: this.headers, - payload, - byteSize: this.totalBytes, - sourceMessageRequestId: messageRequestId, + // 尾批已经进入 Redis，后续 limiter 等待与 PG await 不应继续保留其副本。 + this.releaseQueuedBatch(batch, true); + // Redis 是活跃 spool 的唯一长期正文副本. 大 payload 的 fenced 读取, + // 组装和 PG await 共用同一并发槽, 避免终态同秒到达形成 heap 峰值. + const persistResult = await withPayloadRebuildSlot(this.totalBytes, async () => { + if (this.disabled || this.aborting) { + throw new Error("replay spool stopped before payload rebuild"); + } + const chunks = await this.store.readOwnedChunks( + this.identity.replayId, + this.ownerToken, + this.chunkCount + ); + if (chunks === false) { + throw new Error("replay owner lease lost or chunks incomplete before completion"); + } + if (chunks === null) { + throw new Error("final replay payload read failed"); + } + const payload = chunks.join(""); + chunks.length = 0; + // 先写 PG（持久 payload），再翻 Redis meta 为 completed（热层可服务） + return this.store.persistCompleted({ + replayId: this.identity.replayId, + verifier: this.identity.verifier, + scopeTag: this.identity.scopeTag, + keyId: this.identity.keyId, + userId: this.identity.userId, + format: this.identity.format, + model: this.identity.model, + statusCode: this.statusCode, + headers: this.headers, + payload, + byteSize: this.totalBytes, + sourceMessageRequestId: messageRequestId, + }); }); pgPersisted = true; const completed = await this.store.completeOwned( @@ -321,8 +417,7 @@ export class ReplaySpool { ) .catch(() => false); } finally { - this.queuedBatches.delete(batch); - this.clearPayload(); + this.releaseQueuedBatch(batch); this.release(); } }); @@ -337,11 +432,11 @@ export class ReplaySpool { } if (this.terminal) return; this.terminal = true; + this.notifyTerminal(); this.aborting = true; this.clearTimer(); this.pending = []; this.pendingBytes = 0; - this.clearPayload(); this.clearQueuedBatches(); this.abortPromise = this.writeChain.then(async () => { try { @@ -375,9 +470,9 @@ export class ReplaySpool { private teardown(reason: string, deleteEntry: boolean): void { if (this.disabled) return; this.disabled = true; + this.notifyTerminal(); this.clearTimer(); this.pending = []; - this.parts.length = 0; this.pendingBytes = 0; this.clearQueuedBatches(); // 清理顺着 writeChain 串行：与 in-flight append 竞态时绝不出现「删除后又写回」 @@ -435,19 +530,54 @@ export class ReplaySpool { this.clearOwnerHeartbeat(); } - private takePayload(): string { - const payload = this.parts.join(""); - this.clearPayload(); - return payload; + private notifyTerminal(): void { + const listeners = [...this.terminalListeners]; + this.terminalListeners.clear(); + for (const listener of listeners) { + try { + listener(); + } catch (error) { + logger.debug("[ReplaySpool] terminal listener failed", { + error: error instanceof Error ? error.message : String(error), + }); + } + } + } + + private trackQueuedBatch(batch: QueuedReplayBatch): void { + this.queuedBatches.add(batch); + if (this.activeWriteBatch) { + this.queuedBytes += batch.byteSize; + } else { + this.activeWriteBatch = batch; + } + } + + private releaseQueuedBatch(batch: QueuedReplayBatch, clearChunks = false): void { + if (!this.queuedBatches.delete(batch)) return; + if (clearChunks) batch.chunks.length = 0; + if (this.activeWriteBatch === batch) { + this.activeWriteBatch = null; + const next = this.queuedBatches.values().next().value as QueuedReplayBatch | undefined; + if (next) { + this.activeWriteBatch = next; + this.queuedBytes = Math.max(0, this.queuedBytes - next.byteSize); + } + return; + } + this.queuedBytes = Math.max(0, this.queuedBytes - batch.byteSize); } - private clearPayload(): void { - this.parts.length = 0; + private exceedsWriteBehindLimit(additionalBytes: number): boolean { + if (!this.activeWriteBatch) return false; + return this.queuedBytes + additionalBytes > MAX_WRITE_BEHIND_BYTES; } private clearQueuedBatches(): void { - for (const batch of this.queuedBatches) batch.length = 0; + for (const batch of this.queuedBatches) batch.chunks.length = 0; this.queuedBatches.clear(); + this.activeWriteBatch = null; + this.queuedBytes = 0; } } diff --git a/src/app/v1/_lib/proxy/replay/replay-store.ts b/src/app/v1/_lib/proxy/replay/replay-store.ts index 02a6c44fb..f5ad27095 100644 --- a/src/app/v1/_lib/proxy/replay/replay-store.ts +++ b/src/app/v1/_lib/proxy/replay/replay-store.ts @@ -68,16 +68,34 @@ if redis.call('GET', KEYS[1]) ~= ARGV[1] then return -1 end local len = redis.call('LLEN', KEYS[3]) -if #ARGV > 4 then - len = redis.call('RPUSH', KEYS[3], unpack(ARGV, 5)) +if len ~= tonumber(ARGV[4]) then + return -2 +end +if #ARGV > 5 then + len = redis.call('RPUSH', KEYS[3], unpack(ARGV, 6)) if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[3], ARGV[2]) end end -redis.call('SETEX', KEYS[2], ARGV[2], ARGV[4]) +redis.call('SETEX', KEYS[2], ARGV[2], ARGV[5]) redis.call('EXPIRE', KEYS[1], ARGV[3]) return len`; +const LUA_READ_OWNED_CHUNKS = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then + return {-1} +end +local len = redis.call('LLEN', KEYS[2]) +if len ~= tonumber(ARGV[2]) then + return {-2} +end +local chunks = redis.call('LRANGE', KEYS[2], 0, -1) +local result = {1} +for i = 1, #chunks do + result[#result + 1] = chunks[i] +end +return result`; + const LUA_ABORT_OWNED = ` if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 @@ -196,15 +214,18 @@ export class ReplayStore { /** * owner 热层写入：token 校验、chunk 追加、owning meta 更新和租约续期在同一 Lua - * 内完成，避免旧 owner 在租约交接窗口污染新 owner 的 chunks/meta。 - * null 表示 Redis 不可用，false 表示 token 已失效，number 为当前 chunk 总数。 + * 内完成，避免旧 owner 在租约交接窗口污染新 owner 的 chunks/meta。追加前 + * 同时校验 LIST 长度，防止 eviction 或旧 generation 残留产生截断 payload。 + * null 表示 Redis 不可用，false 表示 token 已失效，chunk_count_mismatch 表示 + * LIST 已不再匹配当前 owner 的本地进度，number 为追加后的 chunk 总数。 */ async writeOwned( replayId: string, ownerToken: string, meta: ReplayMeta, + expectedChunkCount: number, values: string[] = [] - ): Promise { + ): Promise { const redis = this.getRawRedis(); if (!redis) return null; try { @@ -217,11 +238,14 @@ export class ReplayStore { ownerToken, resolveReplayTtlSeconds(), OWNER_LEASE_TTL_SECONDS, + expectedChunkCount, JSON.stringify(meta), ...values ); const length = typeof result === "number" ? result : Number(result); - return length === -1 ? false : length; + if (length === -1) return false; + if (length === -2) return "chunk_count_mismatch"; + return length; } catch (error) { logger.debug("[ReplayStore] fenced owner write failed", { replayId: replayId.slice(0, 12), @@ -231,6 +255,38 @@ export class ReplayStore { } } + /** + * 仅当前 owner 可读取完整 chunks 快照. owner token 校验, LIST 长度校验和 + * LRANGE 在同一 Lua 中执行, 避免租约交接或尾批缺失时持久化错误 payload. + * null 表示 Redis 不可用, false 表示 token 已失效或 chunk 数量不完整. + */ + async readOwnedChunks( + replayId: string, + ownerToken: string, + expectedChunkCount: number + ): Promise { + const redis = this.getRawRedis(); + if (!redis) return null; + try { + const result = await redis.eval( + LUA_READ_OWNED_CHUNKS, + 2, + `cch:replay:owner:${replayId}`, + `cch:replay:chunks:${replayId}`, + ownerToken, + expectedChunkCount + ); + if (!Array.isArray(result) || Number(result[0]) !== 1) return false; + return result.slice(1).map(String); + } catch (error) { + logger.debug("[ReplayStore] fenced owner read failed", { + replayId: replayId.slice(0, 12), + error: error instanceof Error ? error.message : String(error), + }); + return null; + } + } + /** 从 offset（0-based）读到当前末尾；Redis 不可用返回 null。 */ async readChunks(replayId: string, fromIndex: number): Promise { return this.chunks.lrangeFrom(replayId, fromIndex); diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index 81e78c2cf..bd9c9b3a4 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -4219,12 +4219,49 @@ export class ProxyResponseHandler { }, idleTimeoutMs); }; let cleanupClientAbortListener = () => {}; + let cleanupReplaySpoolTerminalListener = () => {}; let clientDetachHandled = false; + let clientDetachStartedAt: number | null = null; + const expireClientAbortDrain = () => { + logger.info("ResponseHandler: Client abort drain window exceeded", { + taskId, + providerId: provider.id, + messageId: messageContext.id, + clientAbortDrainTimeoutMs, + }); + + try { + const sessionWithController = session as typeof session & { + responseController?: AbortController; + }; + sessionWithController.responseController?.abort(new Error("client_abort_drain_timeout")); + } catch (e) { + logger.warn("ResponseHandler: Failed to abort upstream after client drain timeout", { + taskId, + providerId: provider.id, + error: e, + }); + } + + const drainTimeoutError = new Error("client_abort_drain_timeout"); + abortController.abort(drainTimeoutError); + responsePump?.cancelSource(drainTimeoutError); + }; + const armClientAbortDrainTimer = () => { + clearClientAbortDrainTimer(); + if (responsePump?.getState() === "closed") return; + const detachedAt = clientDetachStartedAt ?? Date.now(); + const elapsedMs = Math.max(0, Date.now() - detachedAt); + const remainingMs = Math.max(0, clientAbortDrainTimeoutMs - elapsedMs); + clientAbortDrainTimeoutId = setTimeout(expireClientAbortDrain, remainingMs); + clientAbortDrainTimeoutId.unref?.(); + }; const handleClientAbort = (reason?: unknown) => { if (responsePump?.getState() === "closed") return; responsePump?.startDrain(reason ?? "client_detached"); if (clientDetachHandled) return; clientDetachHandled = true; + clientDetachStartedAt = Date.now(); logger.debug("ResponseHandler: Client disconnected, cleaning up", { taskId, providerId: provider.id, @@ -4237,31 +4274,7 @@ export class ProxyResponseHandler { if (!idleTimeoutId) { startIdleTimer(); } - clientAbortDrainTimeoutId = setTimeout(() => { - logger.info("ResponseHandler: Client abort drain window exceeded", { - taskId, - providerId: provider.id, - messageId: messageContext.id, - clientAbortDrainTimeoutMs, - }); - - try { - const sessionWithController = session as typeof session & { - responseController?: AbortController; - }; - sessionWithController.responseController?.abort(new Error("client_abort_drain_timeout")); - } catch (e) { - logger.warn("ResponseHandler: Failed to abort upstream after client drain timeout", { - taskId, - providerId: provider.id, - error: e, - }); - } - - const drainTimeoutError = new Error("client_abort_drain_timeout"); - abortController.abort(drainTimeoutError); - responsePump?.cancelSource(drainTimeoutError); - }, clientAbortDrainTimeoutMs); + armClientAbortDrainTimer(); }; // 统计/结算只保留有界的“头 + 尾”文本快照，避免长流式响应把进程堆撑满。 @@ -4830,6 +4843,10 @@ export class ProxyResponseHandler { } catch { // env 解析失败保持 60s 现状 } + cleanupReplaySpoolTerminalListener = replaySpool.onTerminal(() => { + clientAbortDrainTimeoutMs = Math.min(clientAbortDrainTimeoutMs, CLIENT_ABORT_DRAIN_MAX_MS); + if (clientDetachHandled) armClientAbortDrainTimer(); + }); } const observeChunk = (value: Uint8Array) => { @@ -4910,6 +4927,8 @@ export class ProxyResponseHandler { cleanupResponseControllerAbortListener(); cleanupClientAbortListener(); cleanupClientAbortListener = () => {}; + cleanupReplaySpoolTerminalListener(); + cleanupReplaySpoolTerminalListener = () => {}; clearClientAbortDrainTimer(); clearIdleTimer(); clearResponseTimeoutOnce(); diff --git a/tests/unit/proxy/replay-spool.test.ts b/tests/unit/proxy/replay-spool.test.ts index e366083d2..aa2c30240 100644 --- a/tests/unit/proxy/replay-spool.test.ts +++ b/tests/unit/proxy/replay-spool.test.ts @@ -29,7 +29,29 @@ const envControl = vi.hoisted(() => ({ const storeControl = vi.hoisted(() => { const order: string[] = []; + const ownedChunks: string[] = []; let ownedChunkCount = 0; + const readOwnedChunks = async ( + _replayId: string, + _ownerToken: string, + expectedChunkCount: number + ) => { + order.push("readOwned"); + return ownedChunks.length === expectedChunkCount ? [...ownedChunks] : false; + }; + const writeOwned = async ( + _replayId: string, + _ownerToken: string, + _meta: { status: string }, + expectedChunkCount: number, + values: string[] = [] + ) => { + order.push(`write:${values.join("|")}`); + if (ownedChunkCount !== expectedChunkCount) return "chunk_count_mismatch" as const; + ownedChunks.push(...values); + ownedChunkCount += values.length; + return ownedChunkCount; + }; const store = { appendChunks: vi.fn(async (_replayId: string, values: string[]) => { order.push(`append:${values.join("|")}`); @@ -39,18 +61,8 @@ const storeControl = vi.hoisted(() => { order.push(`meta:${meta.status}`); return true; }), - writeOwned: vi.fn( - async ( - _replayId: string, - _ownerToken: string, - _meta: { status: string }, - values: string[] = [] - ) => { - order.push(`write:${values.join("|")}`); - ownedChunkCount += values.length; - return ownedChunkCount; - } - ), + writeOwned: vi.fn(writeOwned), + readOwnedChunks: vi.fn(readOwnedChunks), completeOwned: vi.fn( async (_replayId: string, _ownerToken: string, meta: { status: string }) => { order.push(`meta:${meta.status}`); @@ -91,6 +103,9 @@ const storeControl = vi.hoisted(() => { store, resetOwnedChunkCount: () => { ownedChunkCount = 0; + ownedChunks.length = 0; + store.writeOwned.mockImplementation(writeOwned); + store.readOwnedChunks.mockImplementation(readOwnedChunks); }, }; }); @@ -156,13 +171,6 @@ async function drainWriteChain(spool: ReplaySpool): Promise { await (spool as unknown as { writeChain: Promise }).writeChain; } -function retainedAsciiPartBytes(spool: ReplaySpool): number { - return (spool as unknown as { parts: string[] }).parts.reduce( - (total, part) => total + part.length, - 0 - ); -} - function makeOwnerSession(): ProxySession { return { replayState: { identity, ownerToken: "owner-token", role: "owner" }, @@ -228,6 +236,7 @@ describe("ReplaySpool：write-behind 批量冲刷", () => { byteSize: 18, heartbeatAt: expect.any(Number), }), + 0, ["data: a\n\n", "data: b\n\n"] ); @@ -245,6 +254,97 @@ describe("ReplaySpool：write-behind 批量冲刷", () => { await spool.abort("test_cleanup"); }); + it("单个大 chunk 作为当前 Redis write 时不被误判为 backlog", async () => { + const spool = makeSpool(); + const chunk = encoder.encode("x".repeat(513 * 1024)); + + spool.observe(chunk); + await drainWriteChain(spool); + + expect(spool.isTerminal).toBe(false); + expect(storeControl.store.writeOwned).toHaveBeenCalledTimes(1); + expect(storeControl.store.writeOwned.mock.calls[0][4]).toEqual(["x".repeat(513 * 1024)]); + + await spool.abort("test_cleanup"); + }); + + it("跨 chunk UTF-8 序列按实际保留文本字节计入 pending", async () => { + const spool = makeSpool(); + + spool.observe(new Uint8Array([0xf0, 0x9f, 0x98])); + spool.observe(new Uint8Array([0x80])); + + expect((spool as unknown as { pendingBytes: number }).pendingBytes).toBe(4); + + await spool.abort("test_cleanup"); + }); + + it("Redis 写入阻塞且 backlog 超过独立上限时 fail-open 关闭 spool", async () => { + let resolveRedis!: (value: number) => void; + storeControl.store.writeOwned.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveRedis = resolve; + }) + ); + const spool = makeSpool(); + const chunk = encoder.encode("x".repeat(128 * 1024)); + + spool.observe(chunk); + await vi.advanceTimersByTimeAsync(0); + expect(storeControl.store.writeOwned).toHaveBeenCalledTimes(1); + + for (let index = 0; index < 5; index += 1) spool.observe(chunk); + + expect(spool.isTerminal).toBe(true); + expect(storeControl.store.abortOwned).not.toHaveBeenCalled(); + + resolveRedis(1); + await drainWriteChain(spool); + + expect(storeControl.store.writeOwned).toHaveBeenCalledTimes(1); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "write_behind_limit" }) + ); + expect(getActiveReplaySpoolCount()).toBe(0); + }); + + it("40 个 7.5 MiB 活跃 spool 不保留完整本地正文", async () => { + storeControl.store.writeOwned.mockImplementation( + async ( + _replayId: string, + _ownerToken: string, + meta: { status: string; chunkCount: number }, + _expectedChunkCount: number, + _values: string[] = [] + ) => meta.chunkCount + ); + const spools = Array.from({ length: 40 }, () => makeSpool()); + const chunk = encoder.encode("x".repeat(128 * 1024)); + + for (let wave = 0; wave < 60; wave += 1) { + for (const spool of spools) spool.observe(chunk); + await vi.advanceTimersByTimeAsync(0); + } + await Promise.all(spools.map(drainWriteChain)); + + for (const spool of spools) { + const memoryState = spool as unknown as { + parts?: string[]; + pendingBytes: number; + queuedBytes: number; + }; + expect(memoryState.parts).toBeUndefined(); + expect(memoryState.pendingBytes).toBe(0); + expect(memoryState.queuedBytes).toBe(0); + expect(spool.isTerminal).toBe(false); + } + + await Promise.all(spools.map((spool) => spool.abort("test_cleanup"))); + }); + it("空 chunk 不触发任何调度", async () => { const spool = makeSpool(); spool.observe(new Uint8Array(0)); @@ -394,6 +494,53 @@ describe("ReplaySpool：超尺寸自失效", () => { }); describe("ReplaySpool：completeAfterBilling 终态屏障", () => { + it("从当前 owner 的完整 Redis chunks 重建 durable payload", async () => { + const spool = makeSpool(); + spool.observe(encoder.encode("data: first\n\n")); + spool.observe(encoder.encode("data: second\n\n")); + + await spool.completeAfterBilling(42); + + expect(storeControl.store.readOwnedChunks).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + 2 + ); + expect(storeControl.store.persistCompleted).toHaveBeenCalledWith( + expect.objectContaining({ payload: "data: first\n\ndata: second\n\n" }) + ); + }); + + it("同一时间最多重建并持久化两个大 payload", async () => { + const payload = "x".repeat(300 * 1024); + const pendingReads: Array<(chunks: string[]) => void> = []; + storeControl.store.writeOwned.mockImplementation( + async (_replayId: string, _ownerToken: string, meta: { chunkCount: number }) => + meta.chunkCount + ); + storeControl.store.readOwnedChunks.mockImplementation( + () => + new Promise((resolve) => { + pendingReads.push(resolve); + }) + ); + const spools = [makeSpool(), makeSpool(), makeSpool()]; + for (const spool of spools) spool.observe(encoder.encode(payload)); + + const completions = spools.map((spool, index) => spool.completeAfterBilling(index + 1)); + await vi.advanceTimersByTimeAsync(0); + + expect(storeControl.store.readOwnedChunks).toHaveBeenCalledTimes(2); + + pendingReads[0]([payload]); + await vi.advanceTimersByTimeAsync(0); + expect(storeControl.store.readOwnedChunks).toHaveBeenCalledTimes(3); + + pendingReads[1]([payload]); + pendingReads[2]([payload]); + await Promise.all(completions); + }); + it("按 fenced 尾批冲刷 -> PG 持久化 -> completed meta 顺序执行", async () => { const spool = makeSpool(200, "text/event-stream; charset=utf-8"); spool.observe(encoder.encode("data: hello \n\n")); @@ -403,6 +550,7 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { expect(storeControl.order).toEqual([ "write:data: hello \n\n|data: world\n\n", + "readOwned", "persist", "meta:completed", "release", @@ -540,61 +688,10 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { expect(storeControl.store.completeOwned).toHaveBeenCalledTimes(1); }); - it("Redis 阻塞期间不提前复制 payload，PG 阻塞期间释放 parts", async () => { - const payloadBytes = 4 * 1024 * 1024; - const chunk = "x".repeat(64 * 1024); - let resolveRedis!: (value: number) => void; - let resolvePersist!: (value: "persisted") => void; - storeControl.store.writeOwned.mockImplementationOnce( - () => - new Promise((resolve) => { - resolveRedis = resolve; - }) - ); - storeControl.store.persistCompleted.mockImplementationOnce( - () => - new Promise<"persisted">((resolve) => { - resolvePersist = resolve; - }) - ); - const spool = makeSpool(); - for (let index = 0; index < 64; index += 1) { - spool.observe(encoder.encode(chunk)); - } - expect(retainedAsciiPartBytes(spool)).toBe(payloadBytes); - - const completion = spool.completeAfterBilling(10); - await vi.advanceTimersByTimeAsync(0); - - expect(storeControl.store.writeOwned).toHaveBeenCalledTimes(1); - expect(storeControl.store.persistCompleted).not.toHaveBeenCalled(); - expect(retainedAsciiPartBytes(spool)).toBe(payloadBytes); - - resolveRedis(1); - await vi.advanceTimersByTimeAsync(0); - - expect(storeControl.store.persistCompleted).toHaveBeenCalledTimes(1); - expect(storeControl.store.persistCompleted).toHaveBeenCalledWith( - expect.objectContaining({ - payload: "x".repeat(payloadBytes), - byteSize: payloadBytes, - }) - ); - expect(retainedAsciiPartBytes(spool)).toBe(0); - - resolvePersist("persisted"); - await completion; - }); - - it("payload 组装失败时封死热层并释放 heartbeat 与并发配额", async () => { + it("fenced payload 读取失败时封死热层并释放 heartbeat 与并发配额", async () => { + storeControl.store.readOwnedChunks.mockRejectedValueOnce(new Error("payload read failed")); const spool = makeSpool(); spool.observe(encoder.encode("data: partial\n\n")); - const parts = (spool as unknown as { parts: unknown[] }).parts; - parts[0] = { - toString: () => { - throw new Error("payload assembly failed"); - }, - }; await spool.completeAfterBilling(10); @@ -610,6 +707,31 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { expect(storeControl.store.renewOwnerLease).not.toHaveBeenCalled(); }); + it("尾批写入 Redis 后在等待 PG 时立即释放本地 batch", async () => { + let resolvePersist!: () => void; + storeControl.store.persistCompleted.mockImplementationOnce( + () => + new Promise((resolve) => { + resolvePersist = resolve; + }) + ); + const spool = makeSpool(); + spool.observe(encoder.encode("data: final\n\n")); + + const completion = spool.completeAfterBilling(11); + await vi.advanceTimersByTimeAsync(0); + + expect(storeControl.store.persistCompleted).toHaveBeenCalledTimes(1); + const finalBatch = storeControl.store.writeOwned.mock.calls.at(-1)?.[4] as string[]; + expect(finalBatch).toEqual([]); + expect( + (spool as unknown as { queuedBatches: Set<{ chunks: string[] }> }).queuedBatches.size + ).toBe(0); + + resolvePersist(); + await completion; + }); + it("跨 chunk 截断的 UTF-8 序列在 complete 时冲刷解码尾部", async () => { const spool = makeSpool(); // "中" (0xE4 0xB8 0xAD) 只送前两字节：observe 阶段解码挂起，complete 时 flush 出替换字符 @@ -622,7 +744,8 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { identity.replayId, "owner-token", expect.objectContaining({ status: "owning", chunkCount: 1 }), - ["\uFFFD"] + 0, + [] ); expect(storeControl.store.persistCompleted).toHaveBeenCalledWith( expect.objectContaining({ payload: "\uFFFD", byteSize: 2 }) @@ -656,13 +779,14 @@ describe("ReplaySpool：abort 终态", () => { expect(getActiveReplaySpoolCount()).toBe(0); }); - it("abort 立即释放已累积的 payload", async () => { + it("abort 后不再读取或持久化 payload", async () => { const spool = makeSpool(); spool.observe(encoder.encode("data: partial\n\n")); await spool.abort("upstream_error"); - expect((spool as unknown as { parts: string[] }).parts).toEqual([]); + expect(storeControl.store.readOwnedChunks).not.toHaveBeenCalled(); + expect(storeControl.store.persistCompleted).not.toHaveBeenCalled(); }); it("Redis flush 阻塞时 abort 立即释放 batch，并在 fenced cleanup 后释放并发配额", async () => { @@ -681,9 +805,13 @@ describe("ReplaySpool：abort 终态", () => { await vi.advanceTimersByTimeAsync(0); expect(storeControl.store.writeOwned).toHaveBeenCalledTimes(1); - const batch = storeControl.store.writeOwned.mock.calls[0][3] as string[]; + const batch = storeControl.store.writeOwned.mock.calls[0][4] as string[]; expect(batch.length).toBeGreaterThan(0); - const queuedBatches = (spool as unknown as { queuedBatches: Set }).queuedBatches; + const queuedBatches = ( + spool as unknown as { + queuedBatches: Set<{ chunks: string[]; byteSize: number }>; + } + ).queuedBatches; expect(queuedBatches.size).toBe(2); let abortSettled = false; @@ -693,7 +821,7 @@ describe("ReplaySpool：abort 终态", () => { await vi.advanceTimersByTimeAsync(0); expect(batch).toEqual([]); - expect([...queuedBatches].every((queuedBatch) => queuedBatch.length === 0)).toBe(true); + expect([...queuedBatches].every((queuedBatch) => queuedBatch.chunks.length === 0)).toBe(true); expect(abortSettled).toBe(false); expect(getActiveReplaySpoolCount()).toBe(1); expect(storeControl.store.abortOwned).not.toHaveBeenCalled(); @@ -791,6 +919,26 @@ describe("ReplaySpool：abort 终态", () => { }); describe("ReplaySpool：isTerminal", () => { + it("disable 与 abort 会同步通知 detached drain listener", async () => { + envControl.maxPayloadBytes = 16; + const disabledListener = vi.fn(); + const disabledSpool = makeSpool(); + disabledSpool.onTerminal(disabledListener); + + disabledSpool.observe(encoder.encode("x".repeat(32))); + expect(disabledListener).toHaveBeenCalledTimes(1); + await drainWriteChain(disabledSpool); + + envControl.maxPayloadBytes = 8 * 1024 * 1024; + const abortedListener = vi.fn(); + const abortedSpool = makeSpool(); + abortedSpool.onTerminal(abortedListener); + + const abortion = abortedSpool.abort("upstream_error"); + expect(abortedListener).toHaveBeenCalledTimes(1); + await abortion; + }); + it("abort 置 terminal，disable（超限）置 disabled，两者均视为终态", async () => { const aborted = makeSpool(); expect(aborted.isTerminal).toBe(false); @@ -857,7 +1005,8 @@ describe("createReplaySpoolIfOwner", () => { "content-type": "application/json; charset=utf-8", "x-provider-request-id": "req-1", }, - }) + }), + 0 ); spool?.observe(encoder.encode('{"ok":true}')); @@ -903,7 +1052,8 @@ describe("createReplaySpoolIfOwner", () => { delivery: "stream", chunkCount: 0, byteSize: 0, - }) + }), + 0 ); await spool?.abort("test_cleanup"); @@ -917,7 +1067,8 @@ describe("createReplaySpoolIfOwner", () => { expect(storeControl.store.writeOwned).toHaveBeenCalledWith( identity.replayId, "owner-token", - expect.objectContaining({ headers: { "content-type": "text/event-stream" } }) + expect.objectContaining({ headers: { "content-type": "text/event-stream" } }), + 0 ); await spool?.abort("test_cleanup"); diff --git a/tests/unit/proxy/replay-store.test.ts b/tests/unit/proxy/replay-store.test.ts index 50e6cdfb3..64cf3b021 100644 --- a/tests/unit/proxy/replay-store.test.ts +++ b/tests/unit/proxy/replay-store.test.ts @@ -160,11 +160,28 @@ function createFakeRedis() { // 按脚本内容分发：fenced owner write / RPUSH+EXPIRE / terminal fencing / lease fencing eval: vi.fn( async (script: string, _numkeys: number, key: string, ...args: (string | number)[]) => { + if (script.includes("LRANGE") && _numkeys === 2) { + const [chunksKey, token, expectedChunkCount] = args; + if (kv.get(key) !== token) return [-1]; + const chunks = lists.get(String(chunksKey)) ?? []; + if (chunks.length !== Number(expectedChunkCount)) return [-2]; + return [1, ...chunks]; + } if (script.includes("LLEN") && _numkeys === 3) { - const [metaKey, chunksKey, token, replayTtl, _ownerTtl, serializedMeta, ...values] = args; + const [ + metaKey, + chunksKey, + token, + replayTtl, + _ownerTtl, + expectedChunkCount, + serializedMeta, + ...values + ] = args; if (kv.get(key) !== token) return -1; const listKey = String(chunksKey); const list = lists.get(listKey) ?? []; + if (list.length !== Number(expectedChunkCount)) return -2; list.push(...values.map(String)); lists.set(listKey, list); if (values.length > 0 && Number(replayTtl) > 0) { @@ -447,9 +464,9 @@ describe("ReplayStore：owner 租约", () => { const owningMeta = makeMeta({ chunkCount: 2, byteSize: 8 }); await store.tryClaimOwner("r1", "tok-a"); - await expect(store.writeOwned("r1", "tok-a", owningMeta, ["part-a", "part-b"])).resolves.toBe( - 2 - ); + await expect( + store.writeOwned("r1", "tok-a", owningMeta, 0, ["part-a", "part-b"]) + ).resolves.toBe(2); expect(currentRedis().eval).toHaveBeenLastCalledWith( expect.stringContaining("LLEN"), @@ -460,6 +477,7 @@ describe("ReplayStore：owner 租约", () => { "tok-a", 600, 45, + 0, JSON.stringify(owningMeta), "part-a", "part-b" @@ -473,10 +491,10 @@ describe("ReplayStore：owner 租约", () => { const store = new ReplayStore(); const currentMeta = makeMeta({ verifier: "new-owner" }); await store.tryClaimOwner("r1", "tok-new"); - await store.writeOwned("r1", "tok-new", currentMeta, ["new-data"]); + await store.writeOwned("r1", "tok-new", currentMeta, 0, ["new-data"]); await expect( - store.writeOwned("r1", "tok-old", makeMeta({ verifier: "stale-owner" }), ["stale-data"]) + store.writeOwned("r1", "tok-old", makeMeta({ verifier: "stale-owner" }), 1, ["stale-data"]) ).resolves.toBe(false); await expect(store.getMeta("r1")).resolves.toEqual(currentMeta); @@ -484,6 +502,29 @@ describe("ReplayStore：owner 租约", () => { expect(currentRedis().kv.get("cch:replay:owner:r1")).toBe("tok-new"); }); + it("readOwnedChunks 只为当前 owner 返回数量完整的原子 chunks 快照", async () => { + const store = new ReplayStore(); + await store.tryClaimOwner("r1", "tok-a"); + await store.writeOwned("r1", "tok-a", makeMeta({ chunkCount: 2 }), 0, ["first", "second"]); + + await expect(store.readOwnedChunks("r1", "tok-a", 2)).resolves.toEqual(["first", "second"]); + await expect(store.readOwnedChunks("r1", "tok-stale", 2)).resolves.toBe(false); + await expect(store.readOwnedChunks("r1", "tok-a", 3)).resolves.toBe(false); + }); + + it("writeOwned 在追加前原子拒绝不匹配的 LIST generation", async () => { + const store = new ReplayStore(); + await store.appendChunks("r1", ["stale-prefix"]); + await store.tryClaimOwner("r1", "tok-a"); + + await expect( + store.writeOwned("r1", "tok-a", makeMeta({ chunkCount: 1 }), 0, ["new-data"]) + ).resolves.toBe("chunk_count_mismatch"); + + await expect(store.getMeta("r1")).resolves.toBeNull(); + await expect(store.readChunks("r1", 0)).resolves.toEqual(["stale-prefix"]); + }); + it("releaseOwner 是 compare-delete：token 不匹配不删，匹配才删", async () => { const store = new ReplayStore(); await store.tryClaimOwner("r1", "tok-a"); diff --git a/tests/unit/proxy/response-handler-client-abort-drain.test.ts b/tests/unit/proxy/response-handler-client-abort-drain.test.ts index e97233285..fdf9a2c8d 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -28,6 +28,54 @@ const asyncTasks: Promise[] = []; const registeredTasks: Array<{ taskType: string; promise: Promise }> = []; const STREAM_STATS_HEAD_BYTES_FOR_TEST = 1024 * 1024; +const replayControl = vi.hoisted(() => { + const listeners = new Set<() => void>(); + const state = { enabled: false, terminal: false, detachedMs: 300_000 }; + const notifyTerminal = () => { + state.terminal = true; + for (const listener of [...listeners]) listener(); + }; + return { + state, + notifyTerminal, + reset() { + state.enabled = false; + state.terminal = false; + state.detachedMs = 300_000; + listeners.clear(); + }, + spool: { + get isTerminal() { + return state.terminal; + }, + observe: vi.fn(), + abort: vi.fn(async () => notifyTerminal()), + completeAfterBilling: vi.fn(async () => notifyTerminal()), + onTerminal(listener: () => void) { + listeners.add(listener); + return () => listeners.delete(listener); + }, + }, + }; +}); + +vi.mock("@/app/v1/_lib/proxy/replay/replay-spool", () => ({ + abortReplayOwnership: vi.fn(async () => undefined), + createReplaySpoolIfOwner: vi.fn(() => (replayControl.state.enabled ? replayControl.spool : null)), + releaseReplayOwnership: vi.fn(async () => undefined), +})); + +vi.mock("@/lib/config/env.schema", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getEnvConfig: () => ({ + ...actual.getEnvConfig(), + REPLAY_MAX_DETACHED_MS: replayControl.state.detachedMs, + }), + }; +}); + vi.mock("@/app/v1/_lib/proxy/response-fixer", () => ({ ResponseFixer: { process: async (_session: unknown, response: Response) => response, @@ -492,12 +540,20 @@ function createControllableIdleTimeoutResponsesSse(): { function createAbortInsensitiveHangingResponsesSse(): { response: Response; close: () => void; + cancelled: Promise; } { let controller: ReadableStreamDefaultController | null = null; + let resolveCancelled!: (reason: unknown) => void; + const cancelled = new Promise((resolve) => { + resolveCancelled = resolve; + }); const stream = new ReadableStream({ start(streamController) { controller = streamController; }, + cancel(reason) { + resolveCancelled(reason); + }, }); return { @@ -505,6 +561,7 @@ function createAbortInsensitiveHangingResponsesSse(): { status: 200, headers: { "content-type": "text/event-stream" }, }), + cancelled, close() { try { controller?.close(); @@ -1041,6 +1098,7 @@ describe("ProxyResponseHandler stream client abort finalization", () => { asyncTasks.splice(0, asyncTasks.length); registeredTasks.splice(0, registeredTasks.length); vi.clearAllMocks(); + replayControl.reset(); vi.mocked(updateMessageRequestDetailsDurably).mockImplementation( async (_id, details, options) => { try { @@ -2807,6 +2865,102 @@ describe("ProxyResponseHandler stream client abort finalization", () => { } }); + it("spool 失效后按断线起点把 detached drain 降级到 60 秒", async () => { + vi.useFakeTimers(); + replayControl.state.enabled = true; + const upstream = createAbortInsensitiveHangingResponsesSse(); + try { + const clientController = new AbortController(); + const responseController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 120_000; + Object.assign(session, { responseController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch(session, upstream.response); + await downstream.body?.cancel("body_cancel_only"); + await vi.advanceTimersByTimeAsync(30_000); + expect(responseController.signal.aborted).toBe(false); + + replayControl.notifyTerminal(); + await vi.advanceTimersByTimeAsync(29_999); + expect(responseController.signal.aborted).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + expect(responseController.signal.aborted).toBe(true); + + const cancelOutcome = await Promise.race([ + upstream.cancelled.then(() => "cancelled" as const), + new Promise<"pending">((resolve) => setImmediate(() => resolve("pending"))), + ]); + expect(cancelOutcome).toBe("cancelled"); + + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + } finally { + upstream.close(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + vi.useRealTimers(); + } + }); + + it("spool 失效不会延长已配置为 10 秒的 detached drain", async () => { + vi.useFakeTimers(); + replayControl.state.enabled = true; + replayControl.state.detachedMs = 10_000; + const upstream = createAbortInsensitiveHangingResponsesSse(); + try { + const clientController = new AbortController(); + const responseController = new AbortController(); + const session = createSession(clientController.signal); + session.provider.streamingIdleTimeoutMs = 120_000; + Object.assign(session, { responseController }); + setDeferredStreamingFinalization(session, { + providerId: 1, + providerName: "avemujica-responses", + providerPriority: 1, + attemptNumber: 1, + totalProvidersAttempted: 1, + isFirstAttempt: true, + isFailoverSuccess: false, + endpointId: 42, + endpointUrl: "https://api.test.invalid/v1", + upstreamStatusCode: 200, + }); + + const downstream = await ProxyResponseHandler.dispatch(session, upstream.response); + await downstream.body?.cancel("body_cancel_only"); + await vi.advanceTimersByTimeAsync(5_000); + replayControl.notifyTerminal(); + await vi.advanceTimersByTimeAsync(4_999); + expect(responseController.signal.aborted).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + expect(responseController.signal.aborted).toBe(true); + await upstream.cancelled; + + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + } finally { + upstream.close(); + const tasks = asyncTasks.splice(0, asyncTasks.length); + await expectAllFulfilled(tasks); + vi.useRealTimers(); + } + }); + it("settles an abort-insensitive source immediately after Provider idle timeout", async () => { vi.useFakeTimers(); try { diff --git a/tests/unit/proxy/response-handler-stream-terminal.test.ts b/tests/unit/proxy/response-handler-stream-terminal.test.ts index 8486af410..5ccda83fc 100644 --- a/tests/unit/proxy/response-handler-stream-terminal.test.ts +++ b/tests/unit/proxy/response-handler-stream-terminal.test.ts @@ -66,6 +66,7 @@ vi.mock("@/app/v1/_lib/proxy/replay/replay-spool", () => ({ abort: mocks.replayAbort, completeAfterBilling: mocks.replayComplete, isTerminal: false, + onTerminal: () => () => {}, observe: mocks.replayObserve, } : null, From defecfcc9a0badfbe73ae3f0148b251b2f342482 Mon Sep 17 00:00:00 2001 From: ding113 Date: Tue, 11 Aug 2026 18:11:41 +0800 Subject: [PATCH 2/2] fix(replay): bound bootstrap backlog and rebuild slot waits Track pending write operations so the write-behind byte limit also applies while the bootstrap owning-meta write is in flight. Add a 30s timeout that removes stalled large-payload rebuild waiters without consuming a future permit. Move replay terminal listener cleanup into the stream processing finally block, replace per-chunk TextEncoder allocations with Buffer.byteLength, and simplify the fenced Redis chunk-read Lua result. --- src/app/v1/_lib/proxy/replay/replay-spool.ts | 50 +++++++++-- src/app/v1/_lib/proxy/replay/replay-store.ts | 7 +- src/app/v1/_lib/proxy/response-handler.ts | 4 +- tests/unit/proxy/replay-spool.test.ts | 89 +++++++++++++++++++ ...esponse-handler-client-abort-drain.test.ts | 34 ++++++- 5 files changed, 166 insertions(+), 18 deletions(-) diff --git a/src/app/v1/_lib/proxy/replay/replay-spool.ts b/src/app/v1/_lib/proxy/replay/replay-spool.ts index 0d92545b4..0db92666b 100644 --- a/src/app/v1/_lib/proxy/replay/replay-spool.ts +++ b/src/app/v1/_lib/proxy/replay/replay-spool.ts @@ -30,12 +30,16 @@ const FLUSH_BYTES_THRESHOLD = 64 * 1024; const MAX_WRITE_BEHIND_BYTES = 512 * 1024; const LARGE_PAYLOAD_REBUILD_BYTES = 256 * 1024; const MAX_CONCURRENT_LARGE_PAYLOAD_REBUILDS = 2; +const PAYLOAD_REBUILD_SLOT_WAIT_MS = 30_000; const OWNER_HEARTBEAT_INTERVAL_MS = 15_000; const PRE_SPOOL_ABORT_WAIT_MS = 100; let activeSpoolCount = 0; let activeLargePayloadRebuilds = 0; -const largePayloadRebuildWaiters: Array<() => void> = []; +const largePayloadRebuildWaiters: Array<{ + resolve: () => void; + timer: ReturnType; +}> = []; type QueuedReplayBatch = { chunks: string[]; @@ -51,7 +55,22 @@ async function withPayloadRebuildSlot( if (activeLargePayloadRebuilds < MAX_CONCURRENT_LARGE_PAYLOAD_REBUILDS) { activeLargePayloadRebuilds += 1; } else { - await new Promise((resolve) => largePayloadRebuildWaiters.push(resolve)); + await new Promise((resolve, reject) => { + const waiter = { + resolve: () => { + clearTimeout(waiter.timer); + resolve(); + }, + timer: setTimeout(() => { + const index = largePayloadRebuildWaiters.indexOf(waiter); + if (index < 0) return; + largePayloadRebuildWaiters.splice(index, 1); + reject(new Error("replay payload rebuild slot wait timed out")); + }, PAYLOAD_REBUILD_SLOT_WAIT_MS), + }; + waiter.timer.unref?.(); + largePayloadRebuildWaiters.push(waiter); + }); } try { @@ -59,7 +78,7 @@ async function withPayloadRebuildSlot( } finally { const next = largePayloadRebuildWaiters.shift(); if (next) { - next(); + next.resolve(); } else { activeLargePayloadRebuilds = Math.max(0, activeLargePayloadRebuilds - 1); } @@ -73,7 +92,6 @@ export function getActiveReplaySpoolCount(): number { export class ReplaySpool { private readonly store = getReplayStore(); private readonly decoder = new TextDecoder("utf-8"); - private readonly encoder = new TextEncoder(); private readonly queuedBatches = new Set(); private readonly terminalListeners = new Set<() => void>(); private pending: string[] = []; @@ -88,6 +106,7 @@ export class ReplaySpool { private ownerHeartbeatTimer: ReturnType | null = null; private ownerHeartbeatInFlight = false; private writeChain: Promise = Promise.resolve(); + private pendingWriteOperations = 0; private metaWritten = false; constructor( @@ -129,7 +148,7 @@ export class ReplaySpool { const text = this.decoder.decode(chunk, { stream: true }); if (text.length === 0) return; this.pending.push(text); - this.pendingBytes += this.encoder.encode(text).byteLength; + this.pendingBytes += Buffer.byteLength(text, "utf8"); if (this.exceedsWriteBehindLimit(this.pendingBytes)) { this.disable("write_behind_limit"); @@ -174,9 +193,11 @@ export class ReplaySpool { this.pending = []; this.pendingBytes = 0; this.trackQueuedBatch(batch); + this.pendingWriteOperations += 1; // 续接体自带 try/catch：链永不 rejected；每个 await 之后复查 disabled， // 防止与 disable/halt 竞态时在清理之后又写回 owning meta this.writeChain = this.writeChain.then(async () => { + this.activateQueuedBatch(batch); try { if (this.disabled || this.aborting) return; const expectedChunkCount = this.chunkCount + batch.chunks.length; @@ -211,6 +232,7 @@ export class ReplaySpool { this.disable("flush_error"); } finally { this.releaseQueuedBatch(batch); + this.pendingWriteOperations = Math.max(0, this.pendingWriteOperations - 1); } }); } @@ -253,6 +275,7 @@ export class ReplaySpool { /** 立即建立 owning meta（handleStream 创建 spool 时调用，供 attach 读者尽早看到状态）。 */ bootstrap(): void { + this.pendingWriteOperations += 1; this.writeChain = this.writeChain.then(async () => { try { if (this.disabled || this.aborting || this.metaWritten) return; @@ -283,6 +306,8 @@ export class ReplaySpool { error: error instanceof Error ? error.message : String(error), }); this.disable("flush_error"); + } finally { + this.pendingWriteOperations = Math.max(0, this.pendingWriteOperations - 1); } }); } @@ -299,7 +324,7 @@ export class ReplaySpool { const tail = this.decoder.decode(); if (tail.length > 0) { this.pending.push(tail); - this.pendingBytes += this.encoder.encode(tail).byteLength; + this.pendingBytes += Buffer.byteLength(tail, "utf8"); } if (this.exceedsWriteBehindLimit(this.pendingBytes)) { this.disable("write_behind_limit"); @@ -313,8 +338,10 @@ export class ReplaySpool { this.pending = []; this.pendingBytes = 0; this.trackQueuedBatch(batch); + this.pendingWriteOperations += 1; this.writeChain = this.writeChain.then(async () => { + this.activateQueuedBatch(batch); let pgPersisted = false; try { if (this.disabled || this.aborting) return; @@ -418,6 +445,7 @@ export class ReplaySpool { .catch(() => false); } finally { this.releaseQueuedBatch(batch); + this.pendingWriteOperations = Math.max(0, this.pendingWriteOperations - 1); this.release(); } }); @@ -546,13 +574,19 @@ export class ReplaySpool { private trackQueuedBatch(batch: QueuedReplayBatch): void { this.queuedBatches.add(batch); - if (this.activeWriteBatch) { + if (this.pendingWriteOperations > 0 || this.activeWriteBatch) { this.queuedBytes += batch.byteSize; } else { this.activeWriteBatch = batch; } } + private activateQueuedBatch(batch: QueuedReplayBatch): void { + if (!this.queuedBatches.has(batch) || this.activeWriteBatch === batch) return; + this.activeWriteBatch = batch; + this.queuedBytes = Math.max(0, this.queuedBytes - batch.byteSize); + } + private releaseQueuedBatch(batch: QueuedReplayBatch, clearChunks = false): void { if (!this.queuedBatches.delete(batch)) return; if (clearChunks) batch.chunks.length = 0; @@ -569,7 +603,7 @@ export class ReplaySpool { } private exceedsWriteBehindLimit(additionalBytes: number): boolean { - if (!this.activeWriteBatch) return false; + if (this.pendingWriteOperations === 0) return false; return this.queuedBytes + additionalBytes > MAX_WRITE_BEHIND_BYTES; } diff --git a/src/app/v1/_lib/proxy/replay/replay-store.ts b/src/app/v1/_lib/proxy/replay/replay-store.ts index f5ad27095..b35c8dc79 100644 --- a/src/app/v1/_lib/proxy/replay/replay-store.ts +++ b/src/app/v1/_lib/proxy/replay/replay-store.ts @@ -90,11 +90,8 @@ if len ~= tonumber(ARGV[2]) then return {-2} end local chunks = redis.call('LRANGE', KEYS[2], 0, -1) -local result = {1} -for i = 1, #chunks do - result[#result + 1] = chunks[i] -end -return result`; +table.insert(chunks, 1, 1) +return chunks`; const LUA_ABORT_OWNED = ` if redis.call('GET', KEYS[1]) ~= ARGV[1] then diff --git a/src/app/v1/_lib/proxy/response-handler.ts b/src/app/v1/_lib/proxy/response-handler.ts index bd9c9b3a4..511fe8fc8 100644 --- a/src/app/v1/_lib/proxy/response-handler.ts +++ b/src/app/v1/_lib/proxy/response-handler.ts @@ -4927,8 +4927,6 @@ export class ProxyResponseHandler { cleanupResponseControllerAbortListener(); cleanupClientAbortListener(); cleanupClientAbortListener = () => {}; - cleanupReplaySpoolTerminalListener(); - cleanupReplaySpoolTerminalListener = () => {}; clearClientAbortDrainTimer(); clearIdleTimer(); clearResponseTimeoutOnce(); @@ -5222,6 +5220,8 @@ export class ProxyResponseHandler { cleanupTaskAbortBinding(); cleanupResponseControllerAbortListener(); cleanupClientAbortListener(); + cleanupReplaySpoolTerminalListener(); + cleanupReplaySpoolTerminalListener = () => {}; clearClientAbortDrainTimer(); clearIdleTimer(); // 清除静默期计时器（防止泄漏） releaseSessionAgent(session); diff --git a/tests/unit/proxy/replay-spool.test.ts b/tests/unit/proxy/replay-spool.test.ts index aa2c30240..8bce4d8ee 100644 --- a/tests/unit/proxy/replay-spool.test.ts +++ b/tests/unit/proxy/replay-spool.test.ts @@ -268,6 +268,32 @@ describe("ReplaySpool：write-behind 批量冲刷", () => { await spool.abort("test_cleanup"); }); + it("bootstrap 写入阻塞时首个响应 batch 仍受 backlog 上限约束", async () => { + let resolveBootstrap!: (value: number) => void; + storeControl.store.writeOwned.mockImplementationOnce( + () => + new Promise((resolve) => { + resolveBootstrap = resolve; + }) + ); + const spool = makeSpool(); + spool.bootstrap(); + await vi.advanceTimersByTimeAsync(0); + + spool.observe(encoder.encode("x".repeat(513 * 1024))); + + expect(spool.isTerminal).toBe(true); + expect(storeControl.store.writeOwned).toHaveBeenCalledTimes(1); + + resolveBootstrap(0); + await drainWriteChain(spool); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "write_behind_limit" }) + ); + }); + it("跨 chunk UTF-8 序列按实际保留文本字节计入 pending", async () => { const spool = makeSpool(); @@ -541,6 +567,52 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { await Promise.all(completions); }); + it("大 payload 等待 rebuild slot 超时后移出队列且不吞掉后续 permit", async () => { + const payload = "x".repeat(300 * 1024); + const pendingReads: Array<(chunks: string[]) => void> = []; + storeControl.store.writeOwned.mockImplementation( + async (_replayId: string, _ownerToken: string, meta: { chunkCount: number }) => + meta.chunkCount + ); + storeControl.store.readOwnedChunks.mockImplementation( + () => + new Promise((resolve) => { + pendingReads.push(resolve); + }) + ); + const [first, second, timedOut] = [makeSpool(), makeSpool(), makeSpool()]; + for (const spool of [first, second, timedOut]) spool.observe(encoder.encode(payload)); + + const firstCompletion = first.completeAfterBilling(1); + const secondCompletion = second.completeAfterBilling(2); + const timedOutCompletion = timedOut.completeAfterBilling(3); + await vi.advanceTimersByTimeAsync(0); + expect(storeControl.store.readOwnedChunks).toHaveBeenCalledTimes(2); + + await vi.advanceTimersByTimeAsync(30_000); + await timedOutCompletion; + expect(storeControl.store.readOwnedChunks).toHaveBeenCalledTimes(2); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "complete_failed" }) + ); + + const fourth = makeSpool(); + fourth.observe(encoder.encode(payload)); + const fourthCompletion = fourth.completeAfterBilling(4); + await vi.advanceTimersByTimeAsync(0); + expect(storeControl.store.readOwnedChunks).toHaveBeenCalledTimes(2); + + pendingReads[0]([payload]); + await vi.advanceTimersByTimeAsync(0); + expect(storeControl.store.readOwnedChunks).toHaveBeenCalledTimes(3); + + pendingReads[1]([payload]); + pendingReads[2]([payload]); + await Promise.all([firstCompletion, secondCompletion, fourthCompletion]); + }); + it("按 fenced 尾批冲刷 -> PG 持久化 -> completed meta 顺序执行", async () => { const spool = makeSpool(200, "text/event-stream; charset=utf-8"); spool.observe(encoder.encode("data: hello \n\n")); @@ -707,6 +779,23 @@ describe("ReplaySpool：completeAfterBilling 终态屏障", () => { expect(storeControl.store.renewOwnerLease).not.toHaveBeenCalled(); }); + it("fenced payload 读取返回 null 时不写 PG 并封死热层", async () => { + storeControl.store.readOwnedChunks.mockResolvedValueOnce(null); + const spool = makeSpool(); + spool.observe(encoder.encode("data: partial\n\n")); + + await spool.completeAfterBilling(10); + + expect(storeControl.store.persistCompleted).not.toHaveBeenCalled(); + expect(storeControl.store.completeOwned).not.toHaveBeenCalled(); + expect(storeControl.store.abortOwned).toHaveBeenCalledWith( + identity.replayId, + "owner-token", + expect.objectContaining({ status: "aborted", abortReason: "complete_failed" }) + ); + expect(getActiveReplaySpoolCount()).toBe(0); + }); + it("尾批写入 Redis 后在等待 PG 时立即释放本地 batch", async () => { let resolvePersist!: () => void; storeControl.store.persistCompleted.mockImplementationOnce( diff --git a/tests/unit/proxy/response-handler-client-abort-drain.test.ts b/tests/unit/proxy/response-handler-client-abort-drain.test.ts index fdf9a2c8d..d06669b28 100644 --- a/tests/unit/proxy/response-handler-client-abort-drain.test.ts +++ b/tests/unit/proxy/response-handler-client-abort-drain.test.ts @@ -30,10 +30,17 @@ const STREAM_STATS_HEAD_BYTES_FOR_TEST = 1024 * 1024; const replayControl = vi.hoisted(() => { const listeners = new Set<() => void>(); - const state = { enabled: false, terminal: false, detachedMs: 300_000 }; + const state = { + enabled: false, + terminal: false, + detachedMs: 300_000, + unsubscribeCalls: 0, + }; const notifyTerminal = () => { state.terminal = true; - for (const listener of [...listeners]) listener(); + const currentListeners = [...listeners]; + listeners.clear(); + for (const listener of currentListeners) listener(); }; return { state, @@ -42,6 +49,7 @@ const replayControl = vi.hoisted(() => { state.enabled = false; state.terminal = false; state.detachedMs = 300_000; + state.unsubscribeCalls = 0; listeners.clear(); }, spool: { @@ -53,7 +61,10 @@ const replayControl = vi.hoisted(() => { completeAfterBilling: vi.fn(async () => notifyTerminal()), onTerminal(listener: () => void) { listeners.add(listener); - return () => listeners.delete(listener); + return () => { + state.unsubscribeCalls += 1; + listeners.delete(listener); + }; }, }, }; @@ -1236,6 +1247,23 @@ describe("ProxyResponseHandler stream client abort finalization", () => { } }); + it("cleans the replay terminal listener when the session agent release hook throws", async () => { + replayControl.state.enabled = true; + const session = createSession(new AbortController().signal); + vi.mocked(session.releaseAgent).mockImplementationOnce(() => { + throw new Error("release agent failed"); + }); + + const downstream = await ProxyResponseHandler.dispatch(session, createResponsesSse()); + await downstream.text(); + const processingTask = getRegisteredTask("stream-processing"); + expect(processingTask).toBeDefined(); + await expect(processingTask).resolves.toBeUndefined(); + await Promise.allSettled(asyncTasks.splice(0, asyncTasks.length)); + + expect(replayControl.state.unsubscribeCalls).toBe(1); + }); + it("copies Buffer-backed stream windows before retaining stats snapshots", () => { const accumulator = new BoundedStreamTextAccumulator(); const headMarker = "head-copy-marker";