perf: speed up middleware routes and whole-body uploads

packages/ts/src/__tests__/e2e/request-bodies.test.ts

@@ -1,12 +1,21 @@
+import { createHash, randomBytes } from 'node:crypto'
 import { describe, expect, it } from 'vitest'
-import { createApp } from '../../index.js'
+import { createApp, UhttpError } from '../../index.js'
 import { bodyApp, KIB, MIB } from '../helpers/apps.js'
-import { chunkedRequest, openConnection, requestHead, requestWithBody, send } from '../helpers/http-client.js'
+import {
+  chunkedRequest,
+  chunkFrame,
+  openConnection,
+  requestHead,
+  requestWithBody,
+  send,
+} from '../helpers/http-client.js'
 import {
   expectClosedAfterReply,
   expectHello,
   piecesOf,
   SMUGGLED_REQUEST,
+  tooLarge,
   withEveryConnectionChoice,
 } from '../helpers/replies.js'
 import { pause, start } from '../helpers/server.js'
@@ -40,6 +49,82 @@ describe('request bodies over the network', () => {
     await expectHello(connection)
   })
 
+  it('reads a 512 KiB body that arrives in pieces while the handler waits on bytes(), and keeps the connection open', async () => {
+    const server = await start(
+      createApp()
+        .post('/digest', async c => {
+          const bytes = await c.req.body.bytes()
+          return c.json({ length: bytes.byteLength, sha256: createHash('sha256').update(bytes).digest('hex') })
+        })
+        .get('/hello', c => c.text('Hello')),
+    )
+    const payload = randomBytes(512 * KIB)
+    const connection = await openConnection(server.port)
+
+    void connection.write(requestHead('POST', '/digest', [`Content-Length: ${payload.byteLength}`]))
+    for (const piece of piecesOf(payload, 16)) {
+      await connection.write(piece)
+      await pause(5)
+    }
+    const response = await connection.response()
+
+    expect(response.status).toBe(200)
+    expect(JSON.parse(response.body.toString())).toEqual({
+      length: payload.byteLength,
+      sha256: createHash('sha256').update(payload).digest('hex'),
+    })
+    await expectHello(connection)
+  })
+
+  it('fails bytes() with REQUEST_ABORTED when the client goes in the middle of the body, and serves the next client', async () => {
+    const failure = Promise.withResolvers<unknown>()
+    const server = await start(
+      createApp()
+        .post('/upload', async c => {
+          try {
+            await c.req.body.bytes()
+          } catch (error) {
+            failure.resolve(error instanceof UhttpError ? error.code : error)
+            throw error
+          }
+          return c.text('read')
+        })
+        .get('/hello', c => c.text('Hello')),
+    )
+    const leaving = await openConnection(server.port)
+
+    await leaving.write(requestHead('POST', '/upload', [`Content-Length: ${256 * KIB}`]))
+    await leaving.write(Buffer.alloc(64 * KIB, 97))
+    await pause(20)
+    leaving.destroy()
+
+    expect(await failure.promise).toBe('REQUEST_ABORTED')
+    await expectHello(await openConnection(server.port))
+  })
+
+  it('answers 413 to a chunked body that passes the limit while the handler waits on bytes(), and closes the connection', async () => {
+    const handled: string[] = []
+    const server = await start(
+      createApp({ maxBodyBytes: 100 }).post('/upload', async c => {
+        const bytes = await c.req.body.bytes()
+        handled.push(String(bytes.byteLength))
+        return c.text('read')
+      }),
+    )
+    const connection = await openConnection(server.port)
+
+    void connection.write(requestHead('POST', '/upload', ['Transfer-Encoding: chunked']))
+    for (let piece = 0; piece < 3; piece++) {
+      await connection.write(chunkFrame(Buffer.alloc(40, 97)))
+      await pause(10)
+    }
+    const response = await connection.response()
+
+    expect([response.status, response.body.toString()]).toEqual([413, tooLarge(100)])
+    expect(handled).toEqual([])
+    await expectClosedAfterReply(connection, response)
+  })
+
   it('reads an empty body from a POST with no length and no chunked encoding', async () => {
     const server = await start(
       createApp()

packages/ts/src/__tests__/e2e/request-context.test.ts

@@ -1,11 +1,24 @@
 import { describe, expect, it, vi } from 'vitest'
-import { createApp } from '../../index.js'
+import { createApp, type Middleware } from '../../index.js'
 import { KIB, MIB } from '../helpers/apps.js'
 import { openConnection, requestHead, requestWithBody, send } from '../helpers/http-client.js'
 import { errorBody, expectClosedAfterReply, expectHello } from '../helpers/replies.js'
 import { memoryLogger, pause, start } from '../helpers/server.js'
 
 const INTERNAL_ERROR = errorBody('INTERNAL_ERROR', 'The request failed.')
+const MIDDLEWARE_CHOICES: { route: string; middleware: Middleware[] }[] = [
+  { route: 'on a route without middleware', middleware: [] },
+  { route: 'behind middleware that passes it on at once', middleware: [(_c, next) => next()] },
+  {
+    route: 'behind middleware that awaits before it passes it on',
+    middleware: [
+      async (_c, next) => {
+        await pause(5)
+        await next()
+      },
+    ],
+  },
+]
 
 describe('the copy from uWebSockets.js', () => {
   it('copies the method, the path, the query, the parameters, the content type, and the length before any await', async () => {
@@ -184,31 +197,35 @@ describe('the query string and the content type, which the package reads when co
     await expectHello(connection)
   })
 
-  it('copies the request for a handler that first reads it after an await, on a route without middleware', async () => {
-    const app = createApp()
-      .get('/later', async c => {
-        await pause(5)
-        const { method, path, route, query, contentType, contentLength } = c.req
-        return c.json({ method, path, route, query, contentType, contentLength: contentLength ?? null })
-      })
-      .get('/hello', c => c.text('Hello'))
-    const server = await start(app)
+  it.each(MIDDLEWARE_CHOICES)(
+    'copies the request for a handler that first reads it after an await, $route',
+    async ({ middleware }) => {
+      const app = createApp()
+        .use(...middleware)
+        .get('/later', async c => {
+          await pause(5)
+          const { method, path, route, query, contentType, contentLength } = c.req
+          return c.json({ method, path, route, query, contentType, contentLength: contentLength ?? null })
+        })
+        .get('/hello', c => c.text('Hello'))
+      const server = await start(app)
 
-    const { connection, response } = await send(
-      server.port,
-      requestHead('GET', '/later?page=3', ['Content-Type: text/csv']),
-    )
+      const { connection, response } = await send(
+        server.port,
+        requestHead('GET', '/later?page=3', ['Content-Type: text/csv']),
+      )
 
-    expect(JSON.parse(response.body.toString())).toEqual({
-      method: 'GET',
-      path: '/later',
-      route: '/later',
-      query: 'page=3',
-      contentType: 'text/csv',
-      contentLength: null,
-    })
-    await expectHello(connection)
-  })
+      expect(JSON.parse(response.body.toString())).toEqual({
+        method: 'GET',
+        path: '/later',
+        route: '/later',
+        query: 'page=3',
+        contentType: 'text/csv',
+        contentLength: null,
+      })
+      await expectHello(connection)
+    },
+  )
 
   it('gives a task that first reads the request after the reply its method, path, and route, and REQUEST_RELEASED for the rest', async () => {
     const outcome = Promise.withResolvers<unknown[]>()
@@ -241,17 +258,21 @@ describe('the query string and the content type, which the package reads when co
     await expectHello(connection)
   })
 
-  it('reads the length of a body that a client sends to a GET route without middleware, and keeps the connection open', async () => {
-    const app = createApp()
-      .get('/sized', c => c.json({ contentLength: c.req.contentLength ?? null, query: c.req.query }))
-      .get('/hello', c => c.text('Hello'))
-    const server = await start(app)
-
-    const { connection, response } = await send(server.port, requestWithBody('GET', '/sized?x=1', [], 'hello'))
-
-    expect(JSON.parse(response.body.toString())).toEqual({ contentLength: 5, query: 'x=1' })
-    await expectHello(connection)
-  })
+  it.each(MIDDLEWARE_CHOICES)(
+    'reads the length of a body that a client sends to a GET route $route, and keeps the connection open',
+    async ({ middleware }) => {
+      const app = createApp()
+        .use(...middleware)
+        .get('/sized', c => c.json({ contentLength: c.req.contentLength ?? null, query: c.req.query }))
+        .get('/hello', c => c.text('Hello'))
+      const server = await start(app)
+
+      const { connection, response } = await send(server.port, requestWithBody('GET', '/sized?x=1', [], 'hello'))
+
+      expect(JSON.parse(response.body.toString())).toEqual({ contentLength: 5, query: 'x=1' })
+      await expectHello(connection)
+    },
+  )
 
   it('copies both for a stream producer that reads them after the handler returned', async () => {
     const app = createApp()

packages/ts/src/body/reader.ts

@@ -6,8 +6,8 @@ import type { BodyState } from '../runtime.js'
 
 const IDENTITY = 'identity'
 
-interface Waiter {
-  readonly resolve: (chunk: Uint8Array | undefined) => void
+interface Waiter<T> {
+  readonly resolve: (value: T) => void
   readonly reject: (error: unknown) => void
 }
 
@@ -113,7 +113,8 @@ export class BodyReader implements BodyState {
   #paused = false
   #flowing = false
   #error: UhttpError | undefined
-  #waiter: Waiter | undefined
+  #waiter: Waiter<Uint8Array | undefined> | undefined
+  #collector: Waiter<Uint8Array> | undefined
   #pauseTimer: ReturnType<typeof setTimeout> | undefined
 
   /**
@@ -130,8 +131,6 @@ export class BodyReader implements BodyState {
     this.#source = source
     this.#ended = empty
     exchange.body = this
-    exchange.onRelease(error => this.fail(error))
-    exchange.onClose(() => this.#closed())
   }
 
   /** This is `true` while the reader queues the chunks that the source delivers. */
@@ -153,6 +152,7 @@ export class BodyReader implements BodyState {
     if (!last) return
     this.#ended = true
     this.#wake()
+    this.#deliver()
     this.#exchange.arrival.ended()
   }
 
@@ -167,8 +167,11 @@ export class BodyReader implements BodyState {
     this.#queue.length = 0
     this.#queuedBytes = 0
     const waiter = this.#waiter
+    const collector = this.#collector
     this.#waiter = undefined
+    this.#collector = undefined
     waiter?.reject(error)
+    collector?.reject(error)
   }
 
   /**
@@ -193,13 +196,32 @@ export class BodyReader implements BodyState {
     })
   }
 
+  /**
+   * Returns the whole body once its last chunk arrives. While the reader collects the body, it
+   * never pauses the source, because it keeps every chunk until the end of the body anyway, and
+   * the body limit of the route still applies. The caller reads the body either through this
+   * method once or through {@link BodyReader.read}, never both.
+   *
+   * @returns The promise resolves to the bytes of the body. It rejects with the error of a failed
+   * reader.
+   */
+  collect(): Promise<Uint8Array> {
+    if (this.#error !== undefined) return Promise.reject(this.#error)
+    this.keepFlowing()
+    if (this.#ended) return Promise.resolve(this.#drained())
+    return new Promise((resolve, reject) => {
+      this.#collector = { resolve, reject }
+    })
+  }
+
   /**
    * Discards the chunks that the handler hasn't read and every chunk that arrives later, and
    * resumes a paused source, so that the rest of the body keeps arriving after the response starts.
    */
   settle(): void {
     if (this.#paused) this.#resume()
     this.#discarding = true
+    if (this.#collector !== undefined) return
     this.#queue.length = 0
     this.#queuedBytes = 0
   }
@@ -240,6 +262,22 @@ export class BodyReader implements BodyState {
     waiter.resolve(chunk)
   }
 
+  #deliver(): void {
+    const collector = this.#collector
+    if (collector === undefined) return
+    this.#collector = undefined
+    collector.resolve(this.#drained())
+  }
+
+  #drained(): Uint8Array {
+    const queue = this.#queue
+    const first = queue[0]
+    const whole = queue.length === 1 && first !== undefined ? first : Buffer.concat(queue, this.#queuedBytes)
+    queue.length = 0
+    this.#queuedBytes = 0
+    return whole
+  }
+
   #pause(): void {
     if (this.#paused || this.#flowing) return
     this.#paused = true
@@ -256,7 +294,11 @@ export class BodyReader implements BodyState {
     this.#source.resume()
   }
 
-  #closed(): void {
+  /**
+   * Stops the pause timer once the client goes or the response ends, and fails a reader whose body
+   * has not ended with `REQUEST_ABORTED`.
+   */
+  close(): void {
     this.#stopTimer()
     if (!this.#ended) this.fail(new UhttpError('REQUEST_ABORTED', 'The response ended before the request body.'))
   }

packages/ts/src/body/request-body.ts

@@ -119,16 +119,6 @@ async function* chunksOf(reader: BodyReader): AsyncGenerator<Uint8Array> {
   }
 }
 
-async function collect(reader: BodyReader): Promise<Uint8Array> {
-  const chunks: Uint8Array[] = []
-  let total = 0
-  for await (const chunk of chunksOf(reader)) {
-    chunks.push(chunk)
-    total += chunk.byteLength
-  }
-  return chunks.length === 1 && chunks[0] !== undefined ? chunks[0] : Buffer.concat(chunks, total)
-}
-
 /**
  * The package creates this body for each request on the first read of `c.req.body`.
  *
@@ -162,12 +152,18 @@ export class ReaderBody implements RequestBody {
     return chunksOf(reader)
   }
 
-  async bytes(): Promise<Uint8Array> {
-    const reader = this.#claimed()
-    if (this.#bytes === undefined) {
-      if (this.#streamed) throw new UhttpError('BODY_ALREADY_READ', 'Your code has already streamed this body.')
-      this.#bytes = collect(reader)
+  bytes(): Promise<Uint8Array> {
+    if (this.#bytes !== undefined) return this.#bytes
+    let reader: BodyReader
+    try {
+      reader = this.#claimed()
+    } catch (error) {
+      return Promise.reject(error)
+    }
+    if (this.#streamed) {
+      return Promise.reject(new UhttpError('BODY_ALREADY_READ', 'Your code has already streamed this body.'))
     }
+    this.#bytes = reader.collect()
     return this.#bytes
   }