feat(streams): Add stream splitting (#503)

We removed stream multiplexing in #349 because we no longer had a use
for it. With some of the plans we have for our logger, we once again
find ourselves in a position where we'd like to differentiate traffic
over a single stream. Multiplexing was pretty complex, so we'd prefer
not to go back to it if possible. With these requirements in mind, this
PR introduces "stream splitting" via the function `split()` exported
from `@ocap/streams`.

The idea is simple: pass `split()` a duplex stream and a number of type
predicates, and receive an array of streams in return. Each stream
represents the subset of the parent stream's traffic satisfying the
corresponding predicate. Not very efficient for large numbers of
channels or complex predicates, but our usage should be simple enough.

Author: rekmarks

packages/extension/src/background.ts

@@ -1,10 +1,11 @@
 import { isJsonRpcResponse } from '@metamask/utils';
-import type { Json } from '@metamask/utils';
+import type { JsonRpcResponse } from '@metamask/utils';
 import { kernelMethodSpecs } from '@ocap/kernel/rpc';
 import { Logger } from '@ocap/logger';
 import { RpcClient } from '@ocap/rpc-methods';
 import { ChromeRuntimeDuplexStream } from '@ocap/streams/browser';
 import { delay } from '@ocap/utils';
+import type { JsonRpcCall } from '@ocap/utils';
 
 const OFFSCREEN_DOCUMENT_PATH = '/offscreen.html';
 
@@ -25,11 +26,10 @@ async function main(): Promise<void> {
   // Without this delay, sending messages via the chrome.runtime API can fail.
   await delay(50);
 
-  const offscreenStream = await ChromeRuntimeDuplexStream.make(
-    chrome.runtime,
-    'background',
-    'offscreen',
-  );
+  const offscreenStream = await ChromeRuntimeDuplexStream.make<
+    JsonRpcResponse,
+    JsonRpcCall
+  >(chrome.runtime, 'background', 'offscreen', isJsonRpcResponse);
 
   const rpcClient = new RpcClient(
     kernelMethodSpecs,
@@ -50,7 +50,8 @@ async function main(): Promise<void> {
       value: ping,
     },
     sendMessage: {
-      value: async (message: Json) => await offscreenStream.write(message),
+      value: async (message: JsonRpcCall) =>
+        await offscreenStream.write(message),
     },
   });
   harden(globalThis.kernel);
@@ -60,14 +61,8 @@ async function main(): Promise<void> {
     ping().catch(logger.error);
   });
 
-  // Handle replies from the offscreen document
-  for await (const message of offscreenStream) {
-    if (!isJsonRpcResponse(message)) {
-      logger.error('Background received unexpected message', message);
-      continue;
-    }
-    rpcClient.handleResponse(message.id as string, message);
-  }
-
+  await offscreenStream.drain(async (message) =>
+    rpcClient.handleResponse(message.id as string, message),
+  );
   throw new Error('Offscreen connection closed unexpectedly');
 }

packages/streams/src/browser/index.test.ts

@@ -16,6 +16,7 @@ describe('index', () => {
       'PostMessageWriter',
       'initializeMessageChannel',
       'receiveMessagePort',
+      'split',
     ]);
   });
 });

packages/streams/src/browser/index.ts

@@ -23,3 +23,4 @@ export type {
   PostMessageEnvelope,
   PostMessageTarget,
 } from './PostMessageStream.ts';
+export { split } from '../split.ts';

packages/streams/src/index.test.ts

@@ -8,6 +8,7 @@ describe('index', () => {
       'NodeWorkerDuplexStream',
       'NodeWorkerReader',
       'NodeWorkerWriter',
+      'split',
     ]);
   });
 });

packages/streams/src/index.ts

@@ -5,3 +5,4 @@ export {
   NodeWorkerWriter,
   NodeWorkerDuplexStream,
 } from './node/NodeWorkerStream.ts';
+export { split } from './split.ts';

packages/streams/src/split.test.ts

@@ -0,0 +1,199 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { split } from './split.ts';
+import { TestDuplexStream } from '../test/stream-mocks.ts';
+
+describe('split', () => {
+  it('should forward values to the correct sub-stream', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    await stream.receiveInput('a');
+    await stream.receiveInput('b');
+    await stream.end();
+
+    expect(await stream.next()).toStrictEqual({ done: true, value: undefined });
+    expect(await streamA.next()).toStrictEqual({ done: false, value: 'a' });
+    expect(await streamB.next()).toStrictEqual({ done: false, value: 'b' });
+    expect(await streamA.next()).toStrictEqual({
+      done: true,
+      value: undefined,
+    });
+    expect(await streamB.next()).toStrictEqual({
+      done: true,
+      value: undefined,
+    });
+  });
+
+  it('should end the sub-streams when the parent stream ends', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    await stream.end();
+
+    expect(await stream.next()).toStrictEqual({ done: true, value: undefined });
+    expect(await streamA.next()).toStrictEqual({
+      done: true,
+      value: undefined,
+    });
+    expect(await streamB.next()).toStrictEqual({
+      done: true,
+      value: undefined,
+    });
+  });
+
+  it('should end all streams when a sub-stream ends', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    await streamA.return();
+
+    expect(await stream.next()).toStrictEqual({ done: true, value: undefined });
+    expect(await streamA.next()).toStrictEqual({
+      done: true,
+      value: undefined,
+    });
+    expect(await streamB.next()).toStrictEqual({
+      done: true,
+      value: undefined,
+    });
+  });
+
+  it('should forward errors', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    const nextA = streamA.next();
+    const nextB = streamB.next();
+    await stream.throw(new Error('test'));
+
+    expect(await stream.next()).toStrictEqual({ done: true, value: undefined });
+    await expect(nextA).rejects.toThrow('test');
+    await expect(nextB).rejects.toThrow('test');
+  });
+
+  it('should end all streams when a sub-stream errors', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    const nextA = streamA.next();
+    const nextB = streamB.next();
+    await streamA.throw(new Error('test'));
+
+    // We can't observe the error from the parent stream because it's being read
+    // in the body of split().
+    expect(await stream.next()).toStrictEqual({ done: true, value: undefined });
+    await expect(nextA).rejects.toThrow('test');
+    await expect(nextB).rejects.toThrow('test');
+  });
+
+  it('should error if no predicates match', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    const nextA = streamA.next();
+    const nextB = streamB.next();
+    await stream.receiveInput('c');
+
+    expect(await stream.next()).toStrictEqual({ done: true, value: undefined });
+    await expect(nextA).rejects.toThrow(
+      'Failed to match any predicate for value: "c"',
+    );
+    await expect(nextB).rejects.toThrow(
+      'Failed to match any predicate for value: "c"',
+    );
+  });
+
+  it('should allow writing to the sub-streams', async () => {
+    const dispatch = vi.fn();
+    const stream = await TestDuplexStream.make<string>(dispatch);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    await streamA.write('a');
+    await streamB.write('b');
+    await stream.end();
+
+    // The important thing is that the calls are sequential
+    expect(dispatch).toHaveBeenNthCalledWith(3, 'a');
+    expect(dispatch).toHaveBeenNthCalledWith(4, 'b');
+  });
+
+  it('should allow draining sub-streams', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    await streamA.write('a');
+    await streamB.write('b');
+    await stream.end();
+
+    await streamA.drain((value) => {
+      expect(value).toBe('a');
+    });
+    await streamB.drain((value) => {
+      expect(value).toBe('b');
+    });
+  });
+
+  it('should allow iterating over the sub-streams', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    await streamA.write('a');
+    await streamB.write('b');
+    await stream.end();
+
+    for await (const value of streamA) {
+      expect(value).toBe('a');
+    }
+    for await (const value of streamB) {
+      expect(value).toBe('b');
+    }
+  });
+
+  it('should allow piping the sub-streams to a sink', async () => {
+    const stream = await TestDuplexStream.make<string>(() => undefined);
+    const sinkDispatch = vi.fn();
+    const sink = await TestDuplexStream.make<string>(sinkDispatch);
+    const [streamA, streamB] = split(
+      stream,
+      (value) => value === 'a',
+      (value) => value === 'b',
+    );
+    await stream.receiveInput('a');
+    await stream.receiveInput('b');
+    await stream.end();
+    await streamA.pipe(sink);
+    await streamB.pipe(sink);
+    await sink.end();
+
+    expect(sinkDispatch).toHaveBeenNthCalledWith(3, 'a');
+    expect(sinkDispatch).toHaveBeenNthCalledWith(4, 'b');
+  });
+});