feat(util-stream): enable buffering of stream chunks

Author: kuhe

.changeset/thirty-walls-pump.md

@@ -0,0 +1,5 @@
+---
+"@smithy/util-stream": minor
+---
+
+utility for buffering stream chunks

packages/util-stream/package.json

@@ -61,6 +61,7 @@
   "browser": {
     "./dist-es/checksum/ChecksumStream": "./dist-es/checksum/ChecksumStream.browser",
     "./dist-es/checksum/createChecksumStream": "./dist-es/checksum/createChecksumStream.browser",
+    "./dist-es/createBufferedReadable": "./dist-es/createBufferedReadableStream",
     "./dist-es/getAwsChunkedEncodingStream": "./dist-es/getAwsChunkedEncodingStream.browser",
     "./dist-es/headStream": "./dist-es/headStream.browser",
     "./dist-es/sdk-stream-mixin": "./dist-es/sdk-stream-mixin.browser",
@@ -73,12 +74,14 @@
     "./dist-es/sdk-stream-mixin": "./dist-es/sdk-stream-mixin.browser",
     "./dist-es/headStream": "./dist-es/headStream.browser",
     "./dist-es/splitStream": "./dist-es/splitStream.browser",
+    "./dist-es/createBufferedReadable": "./dist-es/createBufferedReadableStream",
     "./dist-cjs/checksum/createChecksumStream": "./dist-cjs/checksum/createChecksumStream.browser",
     "./dist-cjs/checksum/ChecksumStream": "./dist-cjs/checksum/ChecksumStream.browser",
     "./dist-cjs/getAwsChunkedEncodingStream": "./dist-cjs/getAwsChunkedEncodingStream.browser",
     "./dist-cjs/sdk-stream-mixin": "./dist-cjs/sdk-stream-mixin.browser",
     "./dist-cjs/headStream": "./dist-cjs/headStream.browser",
-    "./dist-cjs/splitStream": "./dist-cjs/splitStream.browser"
+    "./dist-cjs/splitStream": "./dist-cjs/splitStream.browser",
+    "./dist-cjs/createBufferedReadable": "./dist-cjs/createBufferedReadableStream"
   },
   "homepage": "https://github.com/awslabs/smithy-typescript/tree/main/packages/util-stream",
   "repository": {

packages/util-stream/src/ByteArrayCollector.ts

@@ -0,0 +1,28 @@
+/**
+ * Aggregates byteArrays on demand.
+ * @internal
+ */
+export class ByteArrayCollector {
+  public byteLength = 0;
+  private byteArrays = [] as Uint8Array[];
+
+  public constructor(public readonly ctor: (size: number) => Uint8Array) {}
+
+  public push(byteArray: Uint8Array) {
+    this.byteArrays.push(byteArray);
+    this.byteLength += byteArray.byteLength;
+  }
+
+  public flush() {
+    const aggregation = this.ctor(this.byteLength);
+    let cursor = 0;
+    for (let i = 0; i < this.byteArrays.length; ++i) {
+      const bytes = this.byteArrays[i];
+      aggregation.set(bytes, cursor);
+      cursor += bytes.byteLength;
+    }
+    this.byteArrays = [];
+    this.byteLength = 0;
+    return aggregation;
+  }
+}

packages/util-stream/src/createBufferedReadable.spec.ts

@@ -0,0 +1,76 @@
+import { Readable } from "node:stream";
+import { describe, expect, test as it, vi } from "vitest";
+
+import { createBufferedReadable } from "./createBufferedReadable";
+import { headStream } from "./headStream";
+
+describe("Buffered Readable stream", () => {
+  function stringStream(size: number, chunkSize: number) {
+    async function* generate() {
+      while (size > 0) {
+        yield "a".repeat(chunkSize);
+        size -= chunkSize;
+      }
+    }
+    return Readable.from(generate());
+  }
+  function byteStream(size: number, chunkSize: number) {
+    async function* generate() {
+      while (size > 0) {
+        yield Buffer.from(new Uint8Array(chunkSize));
+        size -= chunkSize;
+      }
+    }
+    return Readable.from(generate());
+  }
+  const logger = {
+    debug: vi.fn(),
+    info: vi.fn(),
+    warn: vi.fn(),
+    error() {},
+  };
+
+  it("should join upstream chunks if they are too small (stringStream)", async () => {
+    const upstream = stringStream(1024, 8);
+    const downstream = createBufferedReadable(upstream, 64);
+
+    let upstreamChunkCount = 0;
+    upstream.on("data", () => {
+      upstreamChunkCount += 1;
+    });
+
+    let downstreamChunkCount = 0;
+    downstream.on("data", () => {
+      downstreamChunkCount += 1;
+    });
+
+    await headStream(downstream, Infinity);
+
+    expect(upstreamChunkCount).toEqual(128);
+    expect(downstreamChunkCount).toEqual(16);
+  });
+
+  it("should join upstream chunks if they are too small (byteStream)", async () => {
+    const upstream = byteStream(1031, 7);
+    const downstream = createBufferedReadable(upstream, 49, logger);
+
+    let upstreamChunkCount = 0;
+    upstream.on("data", () => {
+      upstreamChunkCount += 1;
+    });
+
+    let downstreamChunkCount = 0;
+    downstream.on("data", () => {
+      downstreamChunkCount += 1;
+    });
+
+    await headStream(downstream, Infinity);
+
+    expect(Math.ceil(1031 / 7)).toBe(148);
+    expect(Math.ceil(1031 / 49)).toBe(22);
+
+    expect(upstreamChunkCount).toEqual(148);
+    expect(downstreamChunkCount).toEqual(22);
+    expect(logger.warn).toHaveBeenCalled();
+  });
+});

packages/util-stream/src/createBufferedReadable.ts

@@ -0,0 +1,78 @@
+import type { Logger } from "@smithy/types";
+import { Readable } from "node:stream";
+
+import { ByteArrayCollector } from "./ByteArrayCollector";
+import type { BufferStore, Modes } from "./createBufferedReadableStream";
+import { createBufferedReadableStream, flush, merge, modeOf, sizeOf } from "./createBufferedReadableStream";
+import { isReadableStream } from "./stream-type-check";
+
+/**
+ * @internal
+ * @param upstream - any Readable or ReadableStream.
+ * @param size - byte or character length minimum. Buffering occurs when a chunk fails to meet this value.
+ * @param onBuffer - for emitting warnings when buffering occurs.
+ * @returns another stream of the same data and stream class, but buffers chunks until
+ * the minimum size is met, except for the last chunk.
+ */
+export function createBufferedReadable(upstream: Readable, size: number, logger?: Logger): Readable;
+export function createBufferedReadable(upstream: ReadableStream, size: number, logger?: Logger): ReadableStream;
+export function createBufferedReadable(
+  upstream: Readable | ReadableStream,
+  size: number,
+  logger?: Logger
+): Readable | ReadableStream {
+  if (isReadableStream(upstream)) {
+    return createBufferedReadableStream(upstream, size, logger);
+  }
+  const downstream = new Readable({ read() {} });
+  let streamBufferingLoggedWarning = false;
+  let bytesSeen = 0;
+
+  const buffers = [
+    "",
+    new ByteArrayCollector((size) => new Uint8Array(size)),
+    new ByteArrayCollector((size) => Buffer.from(new Uint8Array(size))),
+  ] as BufferStore;
+  let mode: Modes | -1 = -1;
+
+  upstream.on("data", (chunk) => {
+    const chunkMode = modeOf(chunk);
+    if (mode !== chunkMode) {
+      if (mode >= 0) {
+        downstream.push(flush(buffers, mode));
+      }
+      mode = chunkMode;
+    }
+    if (mode === -1) {
+      downstream.push(chunk);
+      return;
+    }
+
+    const chunkSize = sizeOf(chunk);
+    bytesSeen += chunkSize;
+    const bufferSize = sizeOf(buffers[mode]);
+    if (chunkSize >= size && bufferSize === 0) {
+      // skip writing to the intermediate buffer
+      // because the upstream chunk is already large enough.
+      downstream.push(chunk);
+    } else {
+      // buffer and potentially flush the data downstream.
+      const newSize = merge(buffers, mode, chunk);
+      if (!streamBufferingLoggedWarning && bytesSeen > size * 2) {
+        streamBufferingLoggedWarning = true;
+        logger?.warn(
+          `@smithy/util-stream - stream chunk size ${chunkSize} is below threshold of ${size}, automatically buffering.`
+        );
+      }
+      if (newSize >= size) {
+        downstream.push(flush(buffers, mode));
+      }
+    }
+  });
+  upstream.on("end", () => {
+    downstream.push(flush(buffers, mode));
+    downstream.push(null);
+  });
+
+  return downstream;
+}

packages/util-stream/src/createBufferedReadableStream.browser.spec.ts

@@ -0,0 +1,107 @@
+import { Readable } from "node:stream";
+import { describe, expect, test as it, vi } from "vitest";
+
+import { createBufferedReadable } from "./createBufferedReadableStream";
+
+describe("Buffered ReadableStream", () => {
+  function stringStream(size: number, chunkSize: number) {
+    async function* generate() {
+      while (size > 0) {
+        yield "a".repeat(chunkSize);
+        size -= chunkSize;
+      }
+    }
+    return Readable.toWeb(Readable.from(generate()));
+  }
+  function byteStream(size: number, chunkSize: number) {
+    async function* generate() {
+      while (size > 0) {
+        yield new Uint8Array(chunkSize);
+        size -= chunkSize;
+      }
+    }
+    return Readable.toWeb(Readable.from(generate()));
+  }
+
+  const logger = {
+    debug: vi.fn(),
+    info: vi.fn(),
+    warn: vi.fn(),
+    error() {},
+  };
+
+  it("should join upstream chunks if they are too small (stringStream)", async () => {
+    let upstreamChunkCount = 0;
+    let downstreamChunkCount = 0;
+
+    const upstream = stringStream(1024, 8);
+    const upstreamReader = upstream.getReader();
+
+    const midstream = new ReadableStream({
+      async pull(controller) {
+        const { value, done } = await upstreamReader.read();
+        if (done) {
+          controller.close();
+        } else {
+          expect(value.length).toBe(8);
+          upstreamChunkCount += 1;
+          controller.enqueue(value);
+        }
+      },
+    });
+    const downstream = createBufferedReadable(midstream, 64);
+    const reader = downstream.getReader();
+
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) {
+        break;
+      } else {
+        downstreamChunkCount += 1;
+        expect(value.length).toBe(64);
+      }
+    }
+
+    expect(upstreamChunkCount).toEqual(128);
+    expect(downstreamChunkCount).toEqual(16);
+  });
+
+  it("should join upstream chunks if they are too small (byteStream)", async () => {
+    let upstreamChunkCount = 0;
+    let downstreamChunkCount = 0;
+
+    const upstream = byteStream(1031, 7);
+    const upstreamReader = upstream.getReader();
+
+    const midstream = new ReadableStream({
+      async pull(controller) {
+        const { value, done } = await upstreamReader.read();
+        if (done) {
+          controller.close();
+        } else {
+          expect(value.length).toBe(7);
+          upstreamChunkCount += 1;
+          controller.enqueue(value);
+        }
+      },
+    });
+    const downstream = createBufferedReadable(midstream, 49, logger);
+    const downstreamReader = downstream.getReader();
+
+    while (true) {
+      const { done, value } = await downstreamReader.read();
+      if (done) {
+        break;
+      } else {
+        downstreamChunkCount += 1;
+        if (value.byteLength > 7) {
+          expect(value.byteLength).toBe(49);
+        }
+      }
+    }
+
+    expect(upstreamChunkCount).toEqual(148);
+    expect(downstreamChunkCount).toEqual(22);
+    expect(logger.warn).toHaveBeenCalled();
+  });
+});