diff --git a/.changeset/fix-abandoned-response-body-activity.md b/.changeset/fix-abandoned-response-body-activity.md new file mode 100644 index 0000000..b3ac915 --- /dev/null +++ b/.changeset/fix-abandoned-response-body-activity.md @@ -0,0 +1,9 @@ +--- +'@cloudflare/containers': patch +--- + +Fix containers never sleeping after a client abandons a response body. `containerFetch` counted a +proxied response as in flight until its body had been fully piped. A body nobody read stalled on +backpressure, so the in-flight counter never dropped back to zero and every alarm renewed +`sleepAfter`. The request now leaves the in-flight count as soon as the container answers, and +bytes flowing through the response body renew the activity timeout instead. diff --git a/.changeset/fix-aborted-request-inflight-leak.md b/.changeset/fix-aborted-request-inflight-leak.md new file mode 100644 index 0000000..932d3da --- /dev/null +++ b/.changeset/fix-aborted-request-inflight-leak.md @@ -0,0 +1,10 @@ +--- +'@cloudflare/containers': patch +--- + +Fix containers never sleeping after a client aborts a request the container has not answered. +`containerFetch` counted a request as in flight until the proxied fetch settled. If the client +disconnected while the container was still working and that fetch never settled, the count stayed +above zero and every alarm renewed `sleepAfter`. The request's abort signal now releases the count +too. The runtime fires that signal on client disconnect when the `enable_request_signal` +compatibility flag is set. diff --git a/src/lib/container.ts b/src/lib/container.ts index bde983c..1dba33f 100644 --- a/src/lib/container.ts +++ b/src/lib/container.ts @@ -37,6 +37,7 @@ const PING_TIMEOUT_MS = 5000; const DEFAULT_SLEEP_AFTER = '10m'; // Default sleep after inactivity time const INSTANCE_POLL_INTERVAL_MS = 300; // Default interval for polling container state +const PROXY_BODY_READ_BYTES = 256 * 1024; // Buffer size for each read of a proxied response body // Timeout for getting container instance and launching a VM // Time to find an instance, attach a DO, call start, but NOT @@ -1056,6 +1057,64 @@ export class Container extends DurableObject { this.sleepAfterMs = Date.now() + timeoutInMs; } + /** + * Forward a container response body to the client and renew the activity timeout as bytes move. + * + * The renewal needs to see the bytes, so the body has to pass through JavaScript, and that is + * where proxying gets expensive. A `TransformStream` with a transform callback hands JavaScript + * 4KB at a time. Reading with a BYOB reader into one reused buffer instead makes each trip a + * memcpy plus a promise, and the buffer size decides how many trips a body takes. The output + * side is an `IdentityTransformStream`, which the runtime pumps natively, so the bytes only + * cross into JavaScript once. + */ + private proxyResponseBody( + body: ReadableStream, + headers: Headers + ): ReadableStream { + // Size the buffer to the body when its length is known, so a small reply does not allocate + // the full read buffer. + const contentLength = Number(headers.get('Content-Length')); + const bufferSize = + contentLength > 0 ? Math.min(contentLength, PROXY_BODY_READ_BYTES) : PROXY_BODY_READ_BYTES; + const { readable, writable } = new IdentityTransformStream(); + void this.pumpResponseBody(body, writable, bufferSize); + return readable; + } + + private async pumpResponseBody( + body: ReadableStream, + writable: WritableStream, + bufferSize: number + ) { + const writer = writable.getWriter(); + const reader = body.getReader({ mode: 'byob' }); + let buffer: ArrayBufferLike = new ArrayBuffer(bufferSize); + + try { + for (;;) { + const { done, value }: ReadableStreamReadResult = await reader.read( + new Uint8Array(buffer) + ); + if (done) { + break; + } + this.renewActivityTimeout(); + // `IdentityTransformStream` resolves a write only once the client has taken the bytes, so + // a body nobody reads parks here without renewing anything, and the container can sleep. + await writer.write(value); + // A BYOB read detaches the buffer it was given and returns it as the result's backing + // store. Taking it back means one allocation for the whole body. + buffer = value.buffer; + } + await writer.close(); + } catch (e) { + // Either side failing tears down the other. A client cancel errors the writable and rejects + // the pending write, and an error from the container rejects the read. Cancelling the source + // stops the container producing into a body nobody will receive. + await Promise.allSettled([writer.abort(e), reader.cancel(e)]); + } + } + /** * Decrement the inflight request counter. * When the counter transitions to 0, renew the activity timeout so the @@ -1214,6 +1273,28 @@ export class Container extends DurableObject { this.inflightRequests++; + // Whichever of these happens first releases the count, and the rest are no-ops: the container + // answers, the proxied fetch throws, the WebSocket closes or errors, the client aborts. More + // than one can fire. A WebSocket sends both close and error, and an abort can race the answer. + const { signal } = request; + let settled = false; + const settleInflight = () => { + if (!settled) { + settled = true; + signal.removeEventListener('abort', settleInflight); + this.decrementInflight(); + } + }; + + // If the client disconnects and the proxied fetch never settles, this listener is the only + // thing that releases the count. The runtime only fires abort on an incoming request's signal + // when the `enable_request_signal` compatibility flag is on. Without it, this never runs. + if (signal.aborted) { + settleInflight(); + } else { + signal.addEventListener('abort', settleInflight); + } + try { // Renew the activity timeout whenever a request is proxied this.renewActivityTimeout(); @@ -1224,16 +1305,6 @@ export class Container extends DurableObject { const containerWs = res.webSocket; const [client, server] = Object.values(new WebSocketPair()) as [WebSocket, WebSocket]; - // Guard to ensure we only decrement inflight once per WebSocket, - // since both close and error events can fire. - let settled = false; - const settleInflight = () => { - if (!settled) { - settled = true; - this.decrementInflight(); - } - }; - // Accept both WebSocket ends containerWs.accept(); server.accept(); @@ -1290,19 +1361,20 @@ export class Container extends DurableObject { return new Response(null, { status: res.status, webSocket: client, headers: res.headers }); } - if (res.body !== null) { - const { readable, writable } = new IdentityTransformStream(); - res.body?.pipeTo(writable).finally(() => { - this.decrementInflight(); - }); + // A response is no longer in flight once the container has answered. Waiting for the body + // to finish would pin `inflightRequests` above zero whenever nobody reads it, since the + // stream stalls on backpressure. Bytes flowing through the body renew the timeout below. + // WebSockets keep the pin until close because the runtime guarantees a close/error event. + // Body completion has no such guarantee. + settleInflight(); - return new Response(readable, res); + if (res.body !== null) { + return new Response(this.proxyResponseBody(res.body, res.headers), res); } - this.decrementInflight(); return res; } catch (e) { - this.decrementInflight(); + settleInflight(); if (!(e instanceof Error)) { throw e; diff --git a/src/tests/container.test.ts b/src/tests/container.test.ts index 08f48a7..b7ac803 100644 --- a/src/tests/container.test.ts +++ b/src/tests/container.test.ts @@ -1,7 +1,8 @@ import { describe, expect, test as baseTest, vi } from 'vitest'; import { Container } from '../lib/container'; import { getRandom } from '../lib/utils'; -import { MockWebSocket, test, webSocketPairSpy, type MockCtx } from './fixtures'; +import { setImmediate } from 'node:timers/promises'; +import { MockWebSocket, mockContainerBody, test, webSocketPairSpy, type MockCtx } from './fixtures'; describe('Container', () => { test('should initialize with default values', ({ container }) => { @@ -458,6 +459,202 @@ describe('Container', () => { expect(renewSpy).toHaveBeenCalled(); }); + test('an abandoned response body must not keep the container awake', async ({ + mockCtx, + container, + }) => { + // Only freeze the clock: the pump moves bytes across real microtasks and setImmediate. + vi.useFakeTimers({ toFake: ['Date'] }); + try { + container.sleepAfter = '1s'; + mockCtx.container.running = true; + mockCtx.storage.get.mockResolvedValue({ status: 'healthy', lastChange: Date.now() }); + // An endless body that the container keeps producing and nobody reads. Container bodies + // are byte streams, so the proxy can read them with a BYOB reader. + mockContainerBody( + mockCtx, + new ReadableStream({ + type: 'bytes', + pull(controller) { + controller.enqueue(new Uint8Array(1024)); + }, + }) + ); + + const response = await container.containerFetch(new Request('https://example.com/big')); + + // @ts-expect-error - inflightRequests is private + expect(container.inflightRequests).toBe(0); + + vi.advanceTimersByTime(2_000); + // @ts-expect-error - isActivityExpired is private + expect(container.isActivityExpired()).toBe(true); + + // Reading the body counts as activity again. + await response.body!.getReader().read(); + await setImmediate(); + // @ts-expect-error - isActivityExpired is private + expect(container.isActivityExpired()).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + + test('a proxied response body must arrive intact', async ({ mockCtx, container }) => { + mockCtx.container.running = true; + mockCtx.storage.get.mockResolvedValue({ status: 'healthy', lastChange: Date.now() }); + // Three chunks: two that fit a read buffer and one larger than it, so the pump has to + // hand back a partial view and reuse its buffer. + const chunks = [ + new Uint8Array(3).fill(1), + new Uint8Array(5).fill(2), + new Uint8Array(300 * 1024).fill(3), + ]; + mockContainerBody( + mockCtx, + new ReadableStream({ + type: 'bytes', + pull(controller) { + const chunk = chunks.shift(); + if (chunk) { + controller.enqueue(chunk); + } else { + // Closing a byte source leaves a pending BYOB read hanging until it is answered. + controller.close(); + controller.byobRequest?.respond(0); + } + }, + }) + ); + + const response = await container.containerFetch(new Request('https://example.com/file')); + const bytes = new Uint8Array(await response.arrayBuffer()); + + expect(bytes.byteLength).toBe(3 + 5 + 300 * 1024); + expect(bytes.subarray(0, 3)).toEqual(new Uint8Array(3).fill(1)); + expect(bytes.subarray(3, 8)).toEqual(new Uint8Array(5).fill(2)); + expect(bytes.subarray(8).every(b => b === 3)).toBe(true); + }); + + test('cancelling a proxied response body must cancel the container body', async ({ + mockCtx, + container, + }) => { + mockCtx.container.running = true; + mockCtx.storage.get.mockResolvedValue({ status: 'healthy', lastChange: Date.now() }); + const cancel = vi.fn(); + mockContainerBody( + mockCtx, + new ReadableStream({ + type: 'bytes', + pull(controller) { + controller.enqueue(new Uint8Array(1024)); + }, + cancel, + }) + ); + + const response = await container.containerFetch(new Request('https://example.com/big')); + await response.body!.cancel('client went away'); + await setImmediate(); + + expect(cancel).toHaveBeenCalledWith('client went away'); + }); + + test('a client abort must release the in-flight count when the container never answers', async ({ + mockCtx, + container, + }) => { + vi.useFakeTimers(); + try { + container.sleepAfter = '1s'; + mockCtx.container.running = true; + mockCtx.storage.get.mockResolvedValue({ status: 'healthy', lastChange: Date.now() }); + // The container accepts the connection but never responds. + mockCtx.container.getTcpPort.mockReturnValue({ fetch: vi.fn(() => new Promise(() => {})) }); + + const abort = new AbortController(); + void container.containerFetch( + new Request('https://example.com/slow', { signal: abort.signal }) + ); + await vi.advanceTimersByTimeAsync(0); + + // @ts-expect-error - inflightRequests is private + expect(container.inflightRequests).toBe(1); + vi.advanceTimersByTime(2_000); + // @ts-expect-error - isActivityExpired is private + expect(container.isActivityExpired()).toBe(false); + + abort.abort(); + + // @ts-expect-error - inflightRequests is private + expect(container.inflightRequests).toBe(0); + vi.advanceTimersByTime(2_000); + // @ts-expect-error - isActivityExpired is private + expect(container.isActivityExpired()).toBe(true); + } finally { + vi.useRealTimers(); + } + }); + + test('a request aborted before the container answers must release the count only once', async ({ + mockCtx, + container, + }) => { + mockCtx.container.running = true; + mockCtx.storage.get.mockResolvedValue({ status: 'healthy', lastChange: Date.now() }); + let answer!: (res: unknown) => void; + const answered = new Promise(resolve => { + answer = resolve; + }); + mockCtx.container.getTcpPort.mockReturnValue({ + fetch: vi + .fn() + .mockReturnValueOnce(answered) + .mockReturnValue(new Promise(() => {})), + }); + + const abort = new AbortController(); + const aborted = container.containerFetch( + new Request('https://example.com/aborted', { signal: abort.signal }) + ); + void container.containerFetch(new Request('https://example.com/still-running')); + await new Promise(resolve => setTimeout(resolve, 0)); + + // @ts-expect-error - inflightRequests is private + expect(container.inflightRequests).toBe(2); + + abort.abort(); + // @ts-expect-error - inflightRequests is private + expect(container.inflightRequests).toBe(1); + + // When the container later answers the aborted request, that must not decrement the count + // for the other request. + answer({ status: 200, webSocket: null, headers: new Headers(), body: null }); + await aborted; + // @ts-expect-error - inflightRequests is private + expect(container.inflightRequests).toBe(1); + }); + + test('a request whose signal is already aborted must not be counted as in flight', async ({ + mockCtx, + container, + }) => { + mockCtx.container.running = true; + mockCtx.storage.get.mockResolvedValue({ status: 'healthy', lastChange: Date.now() }); + mockCtx.container.getTcpPort.mockReturnValue({ fetch: vi.fn(() => new Promise(() => {})) }); + + const abort = new AbortController(); + abort.abort(); + void container.containerFetch( + new Request('https://example.com/gone', { signal: abort.signal }) + ); + await new Promise(resolve => setTimeout(resolve, 0)); + + // @ts-expect-error - inflightRequests is private + expect(container.inflightRequests).toBe(0); + }); + test('containerFetch should create a WebSocket connection when requested', async ({ mockCtx, container, diff --git a/src/tests/fixtures.ts b/src/tests/fixtures.ts index c84d6d8..5c2bd88 100644 --- a/src/tests/fixtures.ts +++ b/src/tests/fixtures.ts @@ -94,6 +94,15 @@ export function makeMockCtx() { export type MockCtx = ReturnType; +/** Make the container answer the next request with a 200 carrying `body`. */ +export function mockContainerBody(mockCtx: MockCtx, body: ReadableStream): void { + mockCtx.container.getTcpPort.mockReturnValue({ + fetch: vi + .fn() + .mockResolvedValue({ status: 200, webSocket: null, headers: new Headers(), body }), + }); +} + export const test = baseTest .extend('mockCtx', () => makeMockCtx()) .extend('container', ({ mockCtx }) => { diff --git a/src/tests/setup.ts b/src/tests/setup.ts index 01c0d2f..c64e9ef 100644 --- a/src/tests/setup.ts +++ b/src/tests/setup.ts @@ -18,3 +18,23 @@ vi.mock('cloudflare:workers', () => { WorkerEntrypoint: MockWorkerEntrypoint, }; }); + +// The runtime's `IdentityTransformStream` is a native byte pipe with no equivalent in Node. A +// `TransformStream` gives the tests the same shape: a writable the pump feeds and a readable the +// caller consumes, with backpressure once nobody reads. It copies each chunk on write, as the +// real one does, because once a write settles the pump hands its buffer to the next read, and +// that read detaches the buffer from any view the stream still holds. +if (!('IdentityTransformStream' in globalThis)) { + vi.stubGlobal( + 'IdentityTransformStream', + class IdentityTransformStream extends TransformStream { + constructor() { + super({ + transform(chunk, controller) { + controller.enqueue(chunk.slice()); + }, + }); + } + } + ); +}