Expose flush in compression

Author: etienne-dldc

src/compression/Compression.ts

@@ -4,9 +4,12 @@ import { HttpHeader } from '../core/mod';
 import { ContentEncoding } from './ContentEnconding';
 import { compress } from './compress';
 
+export type TFlush = () => void;
+
 export interface ICompression {
   readonly acceptedEncoding: readonly ContentEncoding[];
-  readonly usedEncoding: null | readonly ContentEncoding[];
+  readonly usedEncoding: null | ContentEncoding;
+  readonly flush: TFlush;
 }
 
 export const CompressionKey = Key.create<ICompression>('Compress');
@@ -23,27 +26,63 @@ export function Compression(): Middleware {
         ? (acceptedEncodingHeader.split(/, ?/) as any)
         : [ContentEncoding.Identity];
 
-    const usedEncoding = selectEncodings(acceptedEncoding);
-    const compressCtx: ICompression = { acceptedEncoding, usedEncoding };
+    const usedEncoding = selectEncoding(acceptedEncoding);
+    const flushDeferred = createDeferredFlush();
+    const compressCtx: ICompression = {
+      acceptedEncoding,
+      usedEncoding,
+      flush: flushDeferred.flush,
+    };
 
     const response = await next(ctx.with(CompressionKey.Provider(compressCtx)));
     if (response === null) {
       // no response = do nothing
       return response;
     }
-    return compress(response, usedEncoding);
+    return compress(response, usedEncoding, flushDeferred.setFlush);
   };
 }
 
-function selectEncodings(acceptedEncoding: Array<ContentEncoding>): Array<ContentEncoding> {
+function selectEncoding(acceptedEncoding: Array<ContentEncoding>): ContentEncoding {
   if (acceptedEncoding.indexOf(ContentEncoding.Brotli) >= 0) {
-    return [ContentEncoding.Brotli];
+    return ContentEncoding.Brotli;
   }
   if (acceptedEncoding.indexOf(ContentEncoding.Gzip) >= 0) {
-    return [ContentEncoding.Gzip];
+    return ContentEncoding.Gzip;
   }
   if (acceptedEncoding.indexOf(ContentEncoding.Deflate) >= 0) {
-    return [ContentEncoding.Deflate];
+    return ContentEncoding.Deflate;
+  }
+  return ContentEncoding.Identity;
+}
+
+interface IDeferredFlush {
+  flush: TFlush;
+  setFlush: (f: TFlush) => void;
+}
+
+function createDeferredFlush(): IDeferredFlush {
+  let flushFn: TFlush | null = null;
+  // flush requested before flush was set
+  let requested = false;
+
+  return {
+    flush,
+    setFlush,
+  };
+
+  function setFlush(f: TFlush) {
+    flushFn = f;
+    if (requested) {
+      flushFn();
+    }
+  }
+
+  function flush() {
+    if (flushFn === null) {
+      requested = true;
+      return;
+    }
+    flushFn();
   }
-  return [ContentEncoding.Identity];
 }

src/compression/compress.ts

@@ -3,22 +3,27 @@ import zlib from 'node:zlib';
 import type { ReadableStream } from 'stream/web';
 import { Headers, Response } from 'undici';
 import { HttpHeader, HttpStatus, ZenResponse } from '../core/mod';
+import type { TFlush } from './Compression';
 import { ContentEncoding } from './ContentEnconding';
 
 /**
  * Compresses the response body with the given encodings.
  */
-export function compress(originalResponse: ZenResponse, encodings: Array<ContentEncoding>): ZenResponse {
+export function compress(
+  originalResponse: ZenResponse,
+  encoding: ContentEncoding,
+  setFlush: (f: TFlush) => void,
+): ZenResponse {
   if (originalResponse.body === null || HttpStatus.isEmpty(originalResponse.status ?? 200)) {
     return originalResponse;
   }
   const bodyStream = new Response(originalResponse.body).body;
-  const body = encodeBodyWithEncodings(bodyStream, encodings);
+  const body = encodeBodyWithEncoding(bodyStream, encoding, setFlush);
 
   return originalResponse
     .withHeaders((prev) => {
       const nextHeaders = new Headers(prev);
-      nextHeaders.set(HttpHeader.ContentEncoding, encodings.join(', '));
+      nextHeaders.set(HttpHeader.ContentEncoding, encoding);
       // remove content length because we no longer know the size of the body
       // (unless we compress first, then send it but that would be quite bad)
       nextHeaders.delete(HttpHeader.ContentLength);
@@ -27,31 +32,28 @@ export function compress(originalResponse: ZenResponse, encodings: Array<Content
     .with(ZenResponse.BodyKey.Provider(body));
 }
 
-function encodeBodyWithEncodings(
+function encodeBodyWithEncoding(
   body: ReadableStream | null,
-  encodings: Array<ContentEncoding>,
+  encoding: ContentEncoding,
+  setFlush: (f: TFlush) => void,
 ): ReadableStream | null {
   if (body === null) {
     return null;
   }
-  let bodyStream: ReadableStream = body;
-
-  encodings.forEach((encoding) => {
-    bodyStream = encodeBodyWithEncoding(bodyStream, encoding);
-  });
-
-  return bodyStream;
-}
-
-function encodeBodyWithEncoding(body: ReadableStream, encoding: ContentEncoding): ReadableStream {
   if (encoding === ContentEncoding.Brotli) {
-    return Readable.toWeb(Readable.fromWeb(body).pipe(zlib.createBrotliCompress()));
+    const enc = zlib.createBrotliCompress();
+    setFlush(enc.flush.bind(enc));
+    return Readable.toWeb(Readable.fromWeb(body).pipe(enc));
   }
   if (encoding === ContentEncoding.Gzip) {
-    return Readable.toWeb(Readable.fromWeb(body).pipe(zlib.createGzip()));
+    const enc = zlib.createGzip();
+    setFlush(enc.flush.bind(enc));
+    return Readable.toWeb(Readable.fromWeb(body).pipe(enc));
   }
   if (encoding === ContentEncoding.Deflate) {
-    return Readable.toWeb(Readable.fromWeb(body).pipe(zlib.createDeflate()));
+    const enc = zlib.createDeflate();
+    setFlush(enc.flush.bind(enc));
+    return Readable.toWeb(Readable.fromWeb(body).pipe(enc));
   }
   return body;
 }

tests/compress.test.ts

@@ -1,5 +1,6 @@
-import { expect, test } from 'vitest';
-import { Compression, compose, createNodeServer, json } from '../src/mod';
+import { assert, expect, test } from 'vitest';
+import { Compression, CompressionKey, ZenResponse, compose, createNodeServer, json } from '../src/mod';
+import { createPushabledAsyncIterable } from './utils/asyncIterable';
 import { mountServer } from './utils/mountServer';
 
 test('gzip', async () => {
@@ -56,3 +57,53 @@ test('deflate', async () => {
   expect(await res.text()).toBe('{"hello":"world"}');
   await close();
 });
+
+test('compress with asyncIterable body', async () => {
+  const server = createNodeServer(
+    compose(Compression(), (ctx) => {
+      const compression = ctx.get(CompressionKey.Consumer);
+      const body = createPushabledAsyncIterable<Uint8Array>();
+      body.push(Uint8Array.from([1, 2, 3]));
+      compression?.flush();
+      setTimeout(() => {
+        body.push(Uint8Array.from([4, 5, 6]));
+        compression?.flush();
+        body.end();
+      }, 100);
+      return ZenResponse.create(body, {
+        headers: { 'content-type': 'application/octet-stream' },
+      });
+    }),
+  );
+  const { close, url, fetch } = await mountServer(server);
+
+  const res = await fetch(url, {
+    headers: { 'accept-encoding': 'deflate' },
+  });
+  assert(res.body);
+  const reader = res.body.getReader();
+  const message = await reader.read();
+  expect(message).toMatchInlineSnapshot(`
+    {
+      "done": false,
+      "value": Uint8Array [
+        1,
+        2,
+        3,
+      ],
+    }
+  `);
+  const message2 = await reader.read();
+  expect(message2).toMatchInlineSnapshot(`
+    {
+      "done": false,
+      "value": Uint8Array [
+        4,
+        5,
+        6,
+      ],
+    }
+  `);
+
+  await close();
+});

tests/utils/asyncIterable.ts

@@ -0,0 +1,61 @@
+export interface IPushabledAsyncIterable<T> extends AsyncIterable<T> {
+  push: (value: T) => void;
+  end: () => void;
+}
+
+export function createPushabledAsyncIterable<T>(): IPushabledAsyncIterable<T> {
+  // A queue of resolve functions waiting for an incoming event which has not yet arrived.
+  let pullQueue: ((value: IteratorResult<T>) => void)[] = [];
+  // A queue of values waiting for next() calls to be made
+  let pushQueue: T[] = [];
+  let running = true;
+
+  return {
+    [Symbol.asyncIterator]: iterator,
+    push,
+    end,
+  };
+
+  function iterator(): AsyncIterator<T> {
+    return {
+      next,
+    };
+  }
+
+  async function next(): Promise<IteratorResult<T>> {
+    if (running === false) {
+      return { done: true, value: undefined };
+    }
+    return new Promise<IteratorResult<T, undefined>>((resolve) => {
+      if (pushQueue.length !== 0) {
+        // Get value from the pushQueue
+        resolve({ done: false, value: pushQueue.shift() as T });
+        return;
+      }
+      pullQueue.push(resolve);
+    });
+  }
+
+  function end() {
+    if (running) {
+      running = false;
+      pullQueue.forEach((resolve) => resolve({ value: undefined, done: true }));
+      pullQueue = [];
+      pushQueue = [];
+    }
+  }
+
+  function push(value: T) {
+    if (running === false) {
+      // do nothing, pullQueue has been or will be emptied by stop()
+      return;
+    }
+    if (pullQueue.length !== 0) {
+      // call next resolve from the pullQueue
+      const resolvedNext = pullQueue.shift()!;
+      resolvedNext({ value, done: false });
+      return;
+    }
+    pushQueue.push(value);
+  }
+}