From 5f73f9728e041d3ce3bf809e7cf35fab8bd4da9d Mon Sep 17 00:00:00 2001 From: Erik Marks Date: Fri, 25 Oct 2024 16:49:44 -0700 Subject: [PATCH 01/10] feat(streams): Add stream multiplexer Also tidy up various asynchronous operations in base stream classes. These were causing promises returned from the stream multiplexer to fulfill erratically. --- packages/streams/src/BaseDuplexStream.test.ts | 26 +- packages/streams/src/BaseDuplexStream.ts | 9 +- packages/streams/src/BaseStream.test.ts | 38 +-- packages/streams/src/BaseStream.ts | 23 +- .../streams/src/ChromeRuntimeStream.test.ts | 18 ++ packages/streams/src/ChromeRuntimeStream.ts | 4 +- packages/streams/src/MessagePortStream.ts | 2 +- packages/streams/src/PostMessageStream.ts | 2 +- .../streams/src/StreamMultiplexer.test.ts | 247 +++++++++++++++++ packages/streams/src/StreamMultiplexer.ts | 260 ++++++++++++++++++ packages/streams/src/index.test.ts | 2 + packages/streams/src/index.ts | 2 + packages/streams/test/stream-mocks.ts | 7 +- 13 files changed, 596 insertions(+), 44 deletions(-) create mode 100644 packages/streams/src/StreamMultiplexer.test.ts create mode 100644 packages/streams/src/StreamMultiplexer.ts diff --git a/packages/streams/src/BaseDuplexStream.test.ts b/packages/streams/src/BaseDuplexStream.test.ts index 91eaf0d08..e97888651 100644 --- a/packages/streams/src/BaseDuplexStream.test.ts +++ b/packages/streams/src/BaseDuplexStream.test.ts @@ -26,7 +26,7 @@ describe('BaseDuplexStream', () => { it('resolves the synchronization promise when receiving an ACK', async () => { const duplexStream = new TestDuplexStream(() => undefined); - duplexStream.receiveInput(makeAck()); + await duplexStream.receiveInput(makeAck()); expect(await duplexStream.completeSynchronization()).toBeUndefined(); }); @@ -37,7 +37,7 @@ describe('BaseDuplexStream', () => { throw error; }); - duplexStream.receiveInput(makeSyn()); + await duplexStream.receiveInput(makeSyn()); await delay(10); expect(onDispatch).toHaveBeenCalledTimes(2); expect(onDispatch).toHaveBeenNthCalledWith(2, makeAck()); @@ -47,8 +47,8 @@ describe('BaseDuplexStream', () => { const duplexStream = new TestDuplexStream(() => undefined); duplexStream.synchronize().catch(() => undefined); - duplexStream.receiveInput(makeSyn()); - duplexStream.receiveInput(makeSyn()); + await duplexStream.receiveInput(makeSyn()); + await duplexStream.receiveInput(makeSyn()); await expect(duplexStream.next()).rejects.toThrow( 'Received duplicate SYN message during synchronization', ); @@ -58,7 +58,7 @@ describe('BaseDuplexStream', () => { const duplexStream = new TestDuplexStream(() => undefined); duplexStream.synchronize().catch(() => undefined); - duplexStream.receiveInput(42); + await duplexStream.receiveInput(42); await expect(duplexStream.next()).rejects.toThrow( `Received unexpected message during synchronization: ${stringify({ done: false, @@ -71,8 +71,8 @@ describe('BaseDuplexStream', () => { const duplexStream = new TestDuplexStream(() => undefined); duplexStream.synchronize().catch(() => undefined); - duplexStream.receiveInput(makeSyn()); - duplexStream.receiveInput(makeSyn()); + await duplexStream.receiveInput(makeSyn()); + await duplexStream.receiveInput(makeSyn()); await expect(duplexStream.next()).rejects.toThrow( 'Received duplicate SYN message during synchronization', ); @@ -85,8 +85,8 @@ describe('BaseDuplexStream', () => { const duplexStream = new TestDuplexStream(() => undefined); duplexStream.synchronize().catch(() => undefined); - duplexStream.receiveInput(makeSyn()); - duplexStream.receiveInput(makeSyn()); + await duplexStream.receiveInput(makeSyn()); + await duplexStream.receiveInput(makeSyn()); await expect(duplexStream.write(42)).rejects.toThrow( 'Received duplicate SYN message during synchronization', ); @@ -149,7 +149,7 @@ describe('BaseDuplexStream', () => { const duplexStream = await TestDuplexStream.make(() => undefined); const message = 42; - duplexStream.receiveInput(message); + await duplexStream.receiveInput(message); expect(await duplexStream.next()).toStrictEqual(makePendingResult(message)); }); @@ -170,7 +170,9 @@ describe('BaseDuplexStream', () => { const duplexStream = await TestDuplexStream.make(() => undefined); const messages = [1, 2, 3]; - messages.forEach((message) => duplexStream.receiveInput(message)); + await Promise.all( + messages.map(async (message) => duplexStream.receiveInput(message)), + ); await duplexStream.return(); let index = 0; @@ -241,7 +243,7 @@ describe('BaseDuplexStream', () => { readerOnEnd, }); - duplexStream.receiveInput(makeDoneResult()); + await duplexStream.receiveInput(makeDoneResult()); expect(readerOnEnd).toHaveBeenCalledOnce(); }); diff --git a/packages/streams/src/BaseDuplexStream.ts b/packages/streams/src/BaseDuplexStream.ts index 486ba3c4c..1789bd682 100644 --- a/packages/streams/src/BaseDuplexStream.ts +++ b/packages/streams/src/BaseDuplexStream.ts @@ -117,7 +117,7 @@ export abstract class BaseDuplexStream< * * @returns A promise that resolves when the stream is synchronized. */ - protected async synchronize(): Promise { + async synchronize(): Promise { if (this.#synchronizationStatus !== SynchronizationStatus.Idle) { return this.#syncKit.promise; } @@ -225,3 +225,10 @@ export type DuplexStream = Pick< > & { [Symbol.asyncIterator]: () => DuplexStream; }; + +export type SynchronizableDuplexStream< + Read extends Json, + Write extends Json = Read, +> = DuplexStream & { + synchronize(): Promise; +}; diff --git a/packages/streams/src/BaseStream.test.ts b/packages/streams/src/BaseStream.test.ts index 6a7705a10..f6e9fdfa2 100644 --- a/packages/streams/src/BaseStream.test.ts +++ b/packages/streams/src/BaseStream.test.ts @@ -46,7 +46,7 @@ describe('BaseReader', () => { const reader = new TestReader(); const message = 42; - reader.receiveInput(message); + await reader.receiveInput(message); expect(await reader.next()).toStrictEqual(makePendingResult(message)); }); @@ -58,7 +58,7 @@ describe('BaseReader', () => { }); const message = 42; - reader.receiveInput(message); + await reader.receiveInput(message); expect(await reader.next()).toStrictEqual(makePendingResult(message)); expect(validateInput).toHaveBeenCalledWith(message); @@ -69,7 +69,7 @@ describe('BaseReader', () => { const nextP = reader.next(); const message = 42; - reader.receiveInput(message); + await reader.receiveInput(message); expect(await nextP).toStrictEqual(makePendingResult(message)); }); @@ -78,7 +78,9 @@ describe('BaseReader', () => { const reader = new TestReader(); const messages = [1, 2, 3]; - messages.forEach((message) => reader.receiveInput(message)); + await Promise.all( + messages.map(async (message) => reader.receiveInput(message)), + ); let index = 0; for await (const message of reader) { @@ -95,7 +97,7 @@ describe('BaseReader', () => { const reader = new TestReader(); const badMessage = Symbol('foo'); - reader.receiveInput(badMessage); + await reader.receiveInput(badMessage); await expect(reader.next()).rejects.toThrow( 'TestReader: Message cannot be processed (must be JSON-serializable)', @@ -107,7 +109,7 @@ describe('BaseReader', () => { const nextP = reader.next(); const badMessage = Symbol('foo'); - reader.receiveInput(badMessage); + await reader.receiveInput(badMessage); await expect(nextP).rejects.toThrow( 'TestReader: Message cannot be processed (must be JSON-serializable)', @@ -117,7 +119,7 @@ describe('BaseReader', () => { it('throws after receiving marshaled error, before read is enqueued', async () => { const reader = new TestReader(); - reader.receiveInput(makeStreamErrorSignal(new Error('foo'))); + await reader.receiveInput(makeStreamErrorSignal(new Error('foo'))); await expect(reader.next()).rejects.toThrow('foo'); }); @@ -126,7 +128,7 @@ describe('BaseReader', () => { const reader = new TestReader(); const nextP = reader.next(); - reader.receiveInput(makeStreamErrorSignal(new Error('foo'))); + await reader.receiveInput(makeStreamErrorSignal(new Error('foo'))); await expect(nextP).rejects.toThrow('foo'); }); @@ -136,7 +138,7 @@ describe('BaseReader', () => { validateInput: (value) => typeof value === 'number', }); - reader.receiveInput({}); + await reader.receiveInput({}); await expect(reader.next()).rejects.toThrow( 'TestReader: Message failed type validation', @@ -149,7 +151,7 @@ describe('BaseReader', () => { }); const nextP = reader.next(); - reader.receiveInput({}); + await reader.receiveInput({}); await expect(nextP).rejects.toThrow( 'TestReader: Message failed type validation', @@ -159,7 +161,7 @@ describe('BaseReader', () => { it('ends after receiving done signal, before read is enqueued', async () => { const reader = new TestReader(); - reader.receiveInput(makeStreamDoneSignal()); + await reader.receiveInput(makeStreamDoneSignal()); expect(await reader.next()).toStrictEqual(makeDoneResult()); expect(await reader.next()).toStrictEqual(makeDoneResult()); @@ -169,7 +171,7 @@ describe('BaseReader', () => { const reader = new TestReader(); const nextP = reader.next(); - reader.receiveInput(makeStreamDoneSignal()); + await reader.receiveInput(makeStreamDoneSignal()); expect(await nextP).toStrictEqual(makeDoneResult()); expect(await reader.next()).toStrictEqual(makeDoneResult()); @@ -178,7 +180,7 @@ describe('BaseReader', () => { it('enqueues input before returning', async () => { const reader = new TestReader(); - reader.receiveInput(1); + await reader.receiveInput(1); await reader.return(); expect(await reader.next()).toStrictEqual(makePendingResult(1)); @@ -189,7 +191,7 @@ describe('BaseReader', () => { const reader = new TestReader(); await reader.return(); - reader.receiveInput(1); + await reader.receiveInput(1); expect(await reader.next()).toStrictEqual(makeDoneResult()); expect(await reader.next()).toStrictEqual(makeDoneResult()); @@ -198,8 +200,8 @@ describe('BaseReader', () => { it('enqueues input before throwing', async () => { const reader = new TestReader(); - reader.receiveInput(1); - reader.receiveInput(makeStreamErrorSignal(new Error('foo'))); + await reader.receiveInput(1); + await reader.receiveInput(makeStreamErrorSignal(new Error('foo'))); expect(await reader.next()).toStrictEqual(makePendingResult(1)); await expect(reader.next()).rejects.toThrow('foo'); @@ -209,8 +211,8 @@ describe('BaseReader', () => { it('ignores input after throwing', async () => { const reader = new TestReader(); - reader.receiveInput(makeStreamErrorSignal(new Error('foo'))); - reader.receiveInput(1); + await reader.receiveInput(makeStreamErrorSignal(new Error('foo'))); + await reader.receiveInput(1); await expect(reader.next()).rejects.toThrow('foo'); expect(await reader.next()).toStrictEqual(makeDoneResult()); diff --git a/packages/streams/src/BaseStream.ts b/packages/streams/src/BaseStream.ts index 3b2b4156a..a624803a4 100644 --- a/packages/streams/src/BaseStream.ts +++ b/packages/streams/src/BaseStream.ts @@ -34,15 +34,15 @@ const makeStreamBuffer = >() => { if (done) { return; } - done = true; + for (const { resolve, reject } of outputBuffer) { error ? reject(error) : resolve(makeDoneResult() as Value); } outputBuffer.length = 0; }, - hasPendingReads() { + hasPendingReads(): boolean { return outputBuffer.length > 0; }, @@ -52,7 +52,7 @@ const makeStreamBuffer = >() => { * @see `end()` for behavior when the stream ends. * @param value - The value or error to put. */ - put(value: Value | Error) { + put(value: Value | Error): void { if (done) { return; } @@ -65,7 +65,7 @@ const makeStreamBuffer = >() => { inputBuffer.push(value); }, - async get() { + async get(): Promise { if (inputBuffer.length > 0) { const value = inputBuffer.shift() as Value; return value instanceof Error @@ -103,7 +103,7 @@ export type ValidateInput = (input: Json) => input is Read; * A function that receives input from a transport mechanism to a readable stream. * Validates that the input is an {@link IteratorResult}, and throws if it is not. */ -export type ReceiveInput = (input: unknown) => void; +export type ReceiveInput = (input: unknown) => Promise; export type BaseReaderArgs = { name?: string | undefined; @@ -168,7 +168,10 @@ export class BaseReader implements Reader { return this.#receiveInput.bind(this); } - readonly #receiveInput: ReceiveInput = async (input) => { + readonly #receiveInput = async (input: unknown): Promise => { + // eslint-disable-next-line @typescript-eslint/await-thenable + await null; + if (!isDispatchable(input)) { await this.#handleInputError( new Error( @@ -216,8 +219,9 @@ export class BaseReader implements Reader { */ async #end(error?: Error): Promise { this.#buffer.end(error); - await this.#onEnd?.(); + const onEndP = this.#onEnd?.(); this.#onEnd = undefined; + await onEndP; } [Symbol.asyncIterator](): typeof this { @@ -277,7 +281,7 @@ export class BaseWriter implements Writer { readonly #onDispatch: Dispatch; - #onEnd: OnEnd | undefined; + #onEnd?: OnEnd | undefined; /** * Constructs a {@link BaseWriter}. @@ -342,8 +346,9 @@ export class BaseWriter implements Writer { async #end(): Promise { this.#isDone = true; - await this.#onEnd?.(); + const onEndP = this.#onEnd?.(); this.#onEnd = undefined; + await onEndP; } [Symbol.asyncIterator](): typeof this { diff --git a/packages/streams/src/ChromeRuntimeStream.test.ts b/packages/streams/src/ChromeRuntimeStream.test.ts index d4245a97f..b5198b845 100644 --- a/packages/streams/src/ChromeRuntimeStream.test.ts +++ b/packages/streams/src/ChromeRuntimeStream.test.ts @@ -216,6 +216,24 @@ describe('ChromeRuntimeReader', () => { expect(await reader.next()).toStrictEqual(makeDoneResult()); expect(onEnd).toHaveBeenCalledTimes(1); }); + + it('handles errors from onEnd function', async () => { + const { runtime, dispatchRuntimeMessage } = makeRuntime(); + const onEnd = vi.fn(() => { + throw new Error('foo'); + }); + const reader = new ChromeRuntimeReader( + asChromeRuntime(runtime), + ChromeRuntimeStreamTarget.Background, + { onEnd }, + ); + + dispatchRuntimeMessage(makeStreamDoneSignal()); + expect(await reader.next()).toStrictEqual(makeDoneResult()); + expect(onEnd).toHaveBeenCalledTimes(1); + expect(await reader.next()).toStrictEqual(makeDoneResult()); + expect(onEnd).toHaveBeenCalledTimes(1); + }); }); describe.concurrent('ChromeRuntimeWriter', () => { diff --git a/packages/streams/src/ChromeRuntimeStream.ts b/packages/streams/src/ChromeRuntimeStream.ts index f325b3a8d..a330a0de3 100644 --- a/packages/streams/src/ChromeRuntimeStream.ts +++ b/packages/streams/src/ChromeRuntimeStream.ts @@ -123,7 +123,9 @@ export class ChromeRuntimeReader extends BaseReader { return; } - this.#receiveInput(message.payload); + this.#receiveInput(message.payload).catch(async (error) => + this.throw(error), + ); } } harden(ChromeRuntimeReader); diff --git a/packages/streams/src/MessagePortStream.ts b/packages/streams/src/MessagePortStream.ts index 27a398f3e..6bebc5b99 100644 --- a/packages/streams/src/MessagePortStream.ts +++ b/packages/streams/src/MessagePortStream.ts @@ -65,7 +65,7 @@ export class MessagePortReader extends BaseReader { return; } - receiveInput(messageEvent.data); + receiveInput(messageEvent.data).catch(async (error) => this.throw(error)); }; port.addEventListener('message', onMessage); port.start(); diff --git a/packages/streams/src/PostMessageStream.ts b/packages/streams/src/PostMessageStream.ts index d5e554071..814660322 100644 --- a/packages/streams/src/PostMessageStream.ts +++ b/packages/streams/src/PostMessageStream.ts @@ -51,7 +51,7 @@ export class PostMessageReader extends BaseReader { return; } - receiveInput(messageEvent.data); + receiveInput(messageEvent.data).catch(async (error) => this.throw(error)); }; setListener(onMessage); diff --git a/packages/streams/src/StreamMultiplexer.test.ts b/packages/streams/src/StreamMultiplexer.test.ts new file mode 100644 index 000000000..b89f5ee7d --- /dev/null +++ b/packages/streams/src/StreamMultiplexer.test.ts @@ -0,0 +1,247 @@ +import type { Json } from '@metamask/utils'; +import { delay, makePromiseKitMock } from '@ocap/test-utils'; +import { describe, expect, it, vi } from 'vitest'; + +import type { ValidateInput } from './BaseStream.js'; +import { makeChannelParams, StreamMultiplexer } from './StreamMultiplexer.js'; +import type { MultiplexEnvelope } from './StreamMultiplexer.js'; +import { makeDoneResult } from './utils.js'; +import { TestDuplexStream, TestMultiplexer } from '../test/stream-mocks.js'; + +vi.mock('@endo/promise-kit', () => makePromiseKitMock()); + +const isString: ValidateInput = (value) => typeof value === 'string'; + +const isNumber: ValidateInput = (value) => typeof value === 'number'; + +const makeMultiplexer = async ( + duplex?: TestDuplexStream, +): Promise< + [TestMultiplexer, TestDuplexStream] +> => { + // We can't use the async factory for a parameter default + // eslint-disable-next-line no-param-reassign + duplex ??= await TestDuplexStream.make( + () => undefined, + ); + return [new TestMultiplexer(duplex), duplex]; +}; + +const makeEnvelope = (channel: string, payload: Json): MultiplexEnvelope => ({ + channel, + payload, +}); + +describe('StreamMultiplexer', () => { + it('constructs a StreamMultiplexer', () => { + const duplex = new TestDuplexStream( + () => undefined, + ); + const multiplex = new TestMultiplexer(duplex); + expect(multiplex).toBeInstanceOf(StreamMultiplexer); + }); + + describe('addChannels', () => { + it('makes and adds channels', async () => { + const [multiplex] = await makeMultiplexer(); + const [ch1, ch2] = multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('2', (_value: number) => undefined, isNumber), + ); + expect(ch1[Symbol.asyncIterator]()).toBe(ch1); + expect(ch2[Symbol.asyncIterator]()).toBe(ch2); + }); + + it('throws if adding a channel with the same name multiple times', async () => { + const [multiplex] = await makeMultiplexer(); + + expect(() => + multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('1', (_value: string) => undefined, isString), + ), + ).toThrow('Channel "1" already exists'); + }); + + it('throws if adding channels after ending', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + // Add one channel so we can start the multiplexer. + multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + ); + + await Promise.all([multiplex.drainAll(), duplex.return()]); + + expect(() => + multiplex.addChannels( + makeChannelParams('2', (_value: number) => undefined, isNumber), + ), + ).toThrow('TestMultiplexer has already ended'); + }); + }); + + describe('drainAll', () => { + it('throws if draining when there are no channels', async () => { + const [multiplex] = await makeMultiplexer(); + await expect(multiplex.drainAll()).rejects.toThrow( + 'TestMultiplexer has no channels', + ); + }); + + it('forwards input to the correct channel', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + const receiveCh1Input = vi.fn(); + const receiveCh2Input = vi.fn(); + multiplex.addChannels( + makeChannelParams('1', receiveCh1Input, isString), + makeChannelParams('2', receiveCh2Input, isNumber), + ); + multiplex.drainAll().catch((error) => { + throw error; + }); + + await Promise.all([ + duplex.receiveInput(makeEnvelope('1', 'foo')), + duplex.receiveInput(makeEnvelope('2', 42)), + ]); + + await delay(10); + + expect(receiveCh1Input).toHaveBeenCalledWith('foo'); + expect(receiveCh2Input).toHaveBeenCalledWith(42); + }); + + it('ends all streams when the duplex stream returns', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + const [ch1, ch2] = multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('2', (_value: number) => undefined, isNumber), + ); + const drainP = multiplex.drainAll(); + + await duplex.return(); + + expect(await duplex.next()).toStrictEqual(makeDoneResult()); + expect(await ch1.next()).toStrictEqual(makeDoneResult()); + expect(await ch2.next()).toStrictEqual(makeDoneResult()); + expect(await drainP).toBeUndefined(); + }); + + it('ends all streams when any channel returns', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + const [ch1, ch2] = multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('2', (_value: number) => undefined, isNumber), + ); + const drainP = multiplex.drainAll(); + + await ch1.return(); + + expect(await duplex.next()).toStrictEqual(makeDoneResult()); + expect(await ch1.next()).toStrictEqual(makeDoneResult()); + expect(await ch2.next()).toStrictEqual(makeDoneResult()); + expect(await drainP).toBeUndefined(); + }); + + it('ends all streams when the duplex stream throws', async () => { + const onDispatch = vi.fn(); + const [multiplex] = await makeMultiplexer( + await TestDuplexStream.make(onDispatch), + ); + const [ch1, ch2] = multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('2', (_value: number) => undefined, isNumber), + ); + onDispatch.mockImplementationOnce(() => { + throw new Error('foo'); + }); + + const drainP = multiplex.drainAll(); + + await expect(ch1.write('foo')).rejects.toThrow( + 'TestDuplexStream experienced a dispatch failure', + ); + + await expect(drainP).rejects.toThrow( + 'TestDuplexStream experienced a dispatch failure', + ); + expect(await ch1.next()).toStrictEqual(makeDoneResult()); + expect(await ch2.next()).toStrictEqual(makeDoneResult()); + }); + + it('ends all streams when a channel throws', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + const [ch1, ch2] = multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('2', (_value: number) => undefined, isNumber), + ); + + const drainP = multiplex.drainAll(); + + await duplex.receiveInput(makeEnvelope('1', 42)); + + await expect(drainP).rejects.toThrow( + 'TestMultiplexer#1: Message failed type validation', + ); + expect(await ch1.next()).toStrictEqual(makeDoneResult()); + expect(await ch2.next()).toStrictEqual(makeDoneResult()); + }); + + it('ends all streams when receiving a message for a non-existent channel', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + const [ch1, ch2] = multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('2', (_value: number) => undefined, isNumber), + ); + + const drainP = multiplex.drainAll(); + + // There is no channel 3 + await duplex.receiveInput(makeEnvelope('3', 42)); + + await expect(drainP).rejects.toThrow( + 'TestMultiplexer received message for unknown channel: 3', + ); + expect(await ch1.next()).toStrictEqual(makeDoneResult()); + expect(await ch2.next()).toStrictEqual(makeDoneResult()); + }); + }); + + describe('writing', () => { + it('writes channel messages correctly', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + const [ch1, ch2] = multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('2', (_value: number) => undefined, isNumber), + ); + + const writeSpy = vi.spyOn(duplex, 'write'); + + await ch1.write('foo'); + await ch2.write(42); + + expect(writeSpy).toHaveBeenCalledTimes(2); + expect(writeSpy).toHaveBeenNthCalledWith(1, { + channel: '1', + payload: 'foo', + }); + expect(writeSpy).toHaveBeenNthCalledWith(2, { + channel: '2', + payload: 42, + }); + }); + + it('returns done results from channel writes after ending', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + const [ch1, ch2] = multiplex.addChannels( + makeChannelParams('1', (_value: string) => undefined, isString), + makeChannelParams('2', (_value: number) => undefined, isNumber), + ); + + await Promise.all([multiplex.drainAll(), duplex.return()]); + + expect(await ch1.write('foo')).toStrictEqual(makeDoneResult()); + expect(await ch2.write(42)).toStrictEqual(makeDoneResult()); + }); + }); +}); diff --git a/packages/streams/src/StreamMultiplexer.ts b/packages/streams/src/StreamMultiplexer.ts new file mode 100644 index 000000000..58e46b8bb --- /dev/null +++ b/packages/streams/src/StreamMultiplexer.ts @@ -0,0 +1,260 @@ +import type { Json } from '@metamask/utils'; + +import type { + DuplexStream, + SynchronizableDuplexStream, +} from './BaseDuplexStream.js'; +import type { + BaseReaderArgs, + ValidateInput, + ReceiveInput, +} from './BaseStream.js'; +import { BaseReader } from './BaseStream.js'; +import { makeDoneResult } from './utils.js'; + +class ChannelReader extends BaseReader { + // eslint-disable-next-line no-restricted-syntax + private constructor(args: BaseReaderArgs) { + super(args); + } + + static make( + args: BaseReaderArgs, + ): [ChannelReader, ReceiveInput] { + const channel = new ChannelReader(args); + return [channel, channel.getReceiveInput()] as const; + } +} + +type ChannelName = string; + +export type MultiplexEnvelope = { + channel: ChannelName; + payload: Json; +}; + +type HandleRead = (value: Read) => void | Promise; + +export type ChannelParameters = { + channelName: ChannelName; + handleRead: HandleRead; + validateInput: ValidateInput; +} & { _write: Write }; + +export const makeChannelParams = ( + channelName: ChannelName, + handleRead: HandleRead, + validateInput: ValidateInput, +): ChannelParameters => + ({ + channelName, + handleRead, + validateInput, + }) as ChannelParameters; + +type ChannelRecord = { + channelName: ChannelName; + stream: DuplexStream; + handleRead: HandleRead; + receiveInput: ReceiveInput; +}; + +type MappedChannels[]> = + Readonly<{ + [K in keyof Channels]: DuplexStream< + // We need to use `any` for the types to actually be inferred, and our + // usage should ensure that we avoid polluting the outside world. + /* eslint-disable @typescript-eslint/no-explicit-any */ + Channels[K] extends ChannelParameters ? Read : never, + Channels[K] extends ChannelParameters ? Write : never + /* eslint-enable @typescript-eslint/no-explicit-any */ + >; + }>; + +export class StreamMultiplexer { + #isDone: boolean; + + readonly #name: string; + + readonly #channels: Map>; + + readonly #stream: SynchronizableDuplexStream< + MultiplexEnvelope, + MultiplexEnvelope + >; + + constructor( + stream: SynchronizableDuplexStream, + name?: string, + ) { + this.#isDone = false; + this.#channels = new Map(); + this.#name = name ?? this.constructor.name; + this.#stream = stream; + } + + /** + * Starts the multiplexer and drains all of its channels. Waits for the underlying + * duplex stream to be synchronized before reading from it. + * + * @returns A promise resolves when the multiplexer and its channels have ended. + */ + async drainAll(): Promise { + if (this.#channels.size === 0) { + throw new Error(`${this.#name} has no channels`); + } + await this.#stream.synchronize(); + + const promise = Promise.all([ + this.#drain(), + ...Array.from(this.#channels.values()).map( + async ({ stream, handleRead }) => stream.drain(handleRead), + ), + ]).then(async () => this.#end()); + + // Set up cleanup and prevent unhandled rejections. The caller is still expected to + // handle rejections. + promise.catch(async (error) => this.#end(error)); + + return promise; + } + + async #drain(): Promise { + for await (const envelope of this.#stream) { + const channel = this.#channels.get(envelope.channel); + if (channel === undefined) { + await this.#end( + new Error( + `${this.#name} received message for unknown channel: ${envelope.channel}`, + ), + ); + return; + } + await channel.receiveInput(envelope.payload); + } + await this.#end(); + } + + /** + * Adds a set of channels. To avoid messages loss, the underlying duplex stream must not + * be synchronized until all channels have been created. + * + * @param channels - The channels to add. + * @returns The added channels. + */ + // See MappedChannels for why we need to use `any`. + // eslint-disable-next-line @typescript-eslint/no-explicit-any + addChannels[]>( + ...channels: Channels + ): MappedChannels { + if (this.#isDone) { + throw new Error(`${this.#name} has already ended`); + } + + return channels.map((params) => + this.#addChannel(params), + ) as MappedChannels; + } + + #addChannel({ + channelName, + handleRead, + validateInput, + }: ChannelParameters): DuplexStream { + if (this.#channels.has(channelName)) { + throw new Error(`Channel "${channelName}" already exists.`); + } + + const { stream, receiveInput } = this.#makeChannel( + channelName, + validateInput, + ); + + // We downcast some properties in order to store all records in one place. + this.#channels.set(channelName, { + channelName, + handleRead: handleRead as HandleRead, + stream: stream as unknown as DuplexStream, + receiveInput, + }); + + return stream; + } + + #makeChannel( + channelName: ChannelName, + validateInput: ValidateInput, + ): { + stream: DuplexStream; + receiveInput: ReceiveInput; + } { + let isDone = false; + + const [reader, receiveInput] = ChannelReader.make({ + validateInput, + name: `${this.#name}#${channelName}`, + onEnd: async () => { + isDone = true; + await this.#end(); + }, + }); + + const write = async ( + payload: Json, + ): Promise> => { + if (isDone) { + return makeDoneResult(); + } + + const writeP = this.#stream.write({ + channel: channelName, + payload, + }); + writeP.catch(async (error) => { + isDone = true; + await reader.throw(error); + }); + return writeP; + }; + + const drain = async (handler: HandleRead): Promise => { + for await (const value of reader) { + await handler(value); + } + }; + + // Create and return the DuplexStream interface + const stream: DuplexStream = { + next: reader.next.bind(reader), + return: reader.return.bind(reader), + throw: reader.throw.bind(reader), + write, + drain, + [Symbol.asyncIterator]() { + return stream; + }, + }; + + return { stream, receiveInput }; + } + + async #end(error?: Error): Promise { + if (this.#isDone) { + return; + } + this.#isDone = true; + + const end = async ( + stream: DuplexStream, + ): Promise => + error === undefined ? stream.return() : stream.throw(error); + + // eslint-disable-next-line promise/no-promise-in-callback + await Promise.all([ + end(this.#stream), + ...Array.from(this.#channels.values()).map(async (channel) => + end(channel.stream), + ), + ]).catch(() => undefined); + } +} diff --git a/packages/streams/src/index.test.ts b/packages/streams/src/index.test.ts index 33e944809..b72d96dd1 100644 --- a/packages/streams/src/index.test.ts +++ b/packages/streams/src/index.test.ts @@ -15,7 +15,9 @@ describe('index', () => { 'PostMessageDuplexStream', 'PostMessageReader', 'PostMessageWriter', + 'StreamMultiplexer', 'initializeMessageChannel', + 'makeChannelParams', 'makeStreamEnvelopeKit', 'receiveMessagePort', ]); diff --git a/packages/streams/src/index.ts b/packages/streams/src/index.ts index f7e59910f..fc12f4c02 100644 --- a/packages/streams/src/index.ts +++ b/packages/streams/src/index.ts @@ -29,3 +29,5 @@ export type { MakeStreamEnvelopeHandler, StreamEnvelopeKit, } from './envelope-kit.js'; +export { StreamMultiplexer, makeChannelParams } from './StreamMultiplexer.js'; +export type { MultiplexEnvelope } from './StreamMultiplexer.js'; diff --git a/packages/streams/test/stream-mocks.ts b/packages/streams/test/stream-mocks.ts index 76cef0917..c0d470c32 100644 --- a/packages/streams/test/stream-mocks.ts +++ b/packages/streams/test/stream-mocks.ts @@ -9,6 +9,9 @@ import type { BaseWriterArgs, } from '../src/BaseStream.js'; import { BaseReader, BaseWriter } from '../src/BaseStream.js'; +import { StreamMultiplexer } from '../src/StreamMultiplexer.js'; + +export type { MultiplexEnvelope } from '../src/StreamMultiplexer.js'; export class TestReader extends BaseReader { readonly #receiveInput: ReceiveInput; @@ -98,7 +101,7 @@ export class TestDuplexStream< */ async completeSynchronization(): Promise { const syncP = super.synchronize().catch(() => undefined); - this.receiveInput(makeAck()); + await this.receiveInput(makeAck()); return syncP; } @@ -118,3 +121,5 @@ export class TestDuplexStream< return stream; } } + +export class TestMultiplexer extends StreamMultiplexer {} From 1ec3f7ff0a6371ec2dd8f13a38364e56ac8ba61c Mon Sep 17 00:00:00 2001 From: Erik Marks Date: Mon, 4 Nov 2024 11:10:58 +0700 Subject: [PATCH 02/10] refactor(streams): Simplify multiplexer interface Get rid of the `makeChannels()` contrivance in favor of exposing a singular `addChannel()`. Also expose `start()` function in order to enable either "manually" draining channels or using the `drainAll()` method. --- .../streams/src/StreamMultiplexer.test.ts | 148 +++++++++++------- packages/streams/src/StreamMultiplexer.ts | 140 +++++++++-------- packages/streams/src/index.test.ts | 1 - packages/streams/src/index.ts | 2 +- 4 files changed, 161 insertions(+), 130 deletions(-) diff --git a/packages/streams/src/StreamMultiplexer.test.ts b/packages/streams/src/StreamMultiplexer.test.ts index b89f5ee7d..cde195d12 100644 --- a/packages/streams/src/StreamMultiplexer.test.ts +++ b/packages/streams/src/StreamMultiplexer.test.ts @@ -3,7 +3,7 @@ import { delay, makePromiseKitMock } from '@ocap/test-utils'; import { describe, expect, it, vi } from 'vitest'; import type { ValidateInput } from './BaseStream.js'; -import { makeChannelParams, StreamMultiplexer } from './StreamMultiplexer.js'; +import { StreamMultiplexer } from './StreamMultiplexer.js'; import type { MultiplexEnvelope } from './StreamMultiplexer.js'; import { makeDoneResult } from './utils.js'; import { TestDuplexStream, TestMultiplexer } from '../test/stream-mocks.js'; @@ -32,6 +32,8 @@ const makeEnvelope = (channel: string, payload: Json): MultiplexEnvelope => ({ payload, }); +const noop = (_value: unknown): void => undefined; + describe('StreamMultiplexer', () => { it('constructs a StreamMultiplexer', () => { const duplex = new TestDuplexStream( @@ -41,42 +43,85 @@ describe('StreamMultiplexer', () => { expect(multiplex).toBeInstanceOf(StreamMultiplexer); }); - describe('addChannels', () => { + describe('addChannel', () => { it('makes and adds channels', async () => { const [multiplex] = await makeMultiplexer(); - const [ch1, ch2] = multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('2', (_value: number) => undefined, isNumber), - ); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); expect(ch1[Symbol.asyncIterator]()).toBe(ch1); expect(ch2[Symbol.asyncIterator]()).toBe(ch2); }); it('throws if adding a channel with the same name multiple times', async () => { const [multiplex] = await makeMultiplexer(); + multiplex.addChannel('1', isString, noop); + expect(() => multiplex.addChannel('1', isString, noop)).toThrow( + 'Channel "1" already exists', + ); + }); + + it('throws if adding channels after starting', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + multiplex.addChannel('1', isString, noop); + const startP = multiplex.start(); - expect(() => - multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('1', (_value: string) => undefined, isString), - ), - ).toThrow('Channel "1" already exists'); + expect(() => multiplex.addChannel('2', isNumber, noop)).toThrow( + 'Channels must be added before starting the multiplexer', + ); + + await Promise.all([startP, duplex.return()]); }); it('throws if adding channels after ending', async () => { const [multiplex, duplex] = await makeMultiplexer(); // Add one channel so we can start the multiplexer. - multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - ); + multiplex.addChannel('1', isString, noop); await Promise.all([multiplex.drainAll(), duplex.return()]); - expect(() => - multiplex.addChannels( - makeChannelParams('2', (_value: number) => undefined, isNumber), - ), - ).toThrow('TestMultiplexer has already ended'); + expect(() => multiplex.addChannel('2', isNumber, noop)).toThrow( + 'Channels must be added before starting', + ); + }); + }); + + describe('start', () => { + it('is idempotent', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + // Add one channel so we can start the multiplexer. + multiplex.addChannel('1', isString, noop); + const startP = Promise.all([multiplex.start(), multiplex.start()]).then( + () => undefined, + ); + await duplex.return(); + expect(await startP).toBeUndefined(); + }); + + it('enables draining channels separately', async () => { + const [multiplex, duplex] = await makeMultiplexer(); + const ch1Handler = vi.fn(); + const ch2Handler = vi.fn(); + const ch1 = multiplex.addChannel('1', isString, ch1Handler); + const ch2 = multiplex.addChannel('2', isNumber, ch2Handler); + + const startAndDrainP = Promise.all([ + multiplex.start(), + ch1.drain(), + ch2.drain(), + ]); + + await Promise.all([ + duplex.receiveInput(makeEnvelope('1', 'foo')), + duplex.receiveInput(makeEnvelope('2', 42)), + ]); + + await delay(10); + + await duplex.return(); + await startAndDrainP; + + expect(ch1Handler).toHaveBeenCalledWith('foo'); + expect(ch2Handler).toHaveBeenCalledWith(42); }); }); @@ -90,15 +135,11 @@ describe('StreamMultiplexer', () => { it('forwards input to the correct channel', async () => { const [multiplex, duplex] = await makeMultiplexer(); - const receiveCh1Input = vi.fn(); - const receiveCh2Input = vi.fn(); - multiplex.addChannels( - makeChannelParams('1', receiveCh1Input, isString), - makeChannelParams('2', receiveCh2Input, isNumber), - ); - multiplex.drainAll().catch((error) => { - throw error; - }); + const ch1Handler = vi.fn(); + const ch2Handler = vi.fn(); + multiplex.addChannel('1', isString, ch1Handler); + multiplex.addChannel('2', isNumber, ch2Handler); + const drainP = multiplex.drainAll(); await Promise.all([ duplex.receiveInput(makeEnvelope('1', 'foo')), @@ -107,16 +148,17 @@ describe('StreamMultiplexer', () => { await delay(10); - expect(receiveCh1Input).toHaveBeenCalledWith('foo'); - expect(receiveCh2Input).toHaveBeenCalledWith(42); + expect(ch1Handler).toHaveBeenCalledWith('foo'); + expect(ch2Handler).toHaveBeenCalledWith(42); + + await duplex.return(); + await drainP; }); it('ends all streams when the duplex stream returns', async () => { const [multiplex, duplex] = await makeMultiplexer(); - const [ch1, ch2] = multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('2', (_value: number) => undefined, isNumber), - ); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); const drainP = multiplex.drainAll(); await duplex.return(); @@ -129,10 +171,8 @@ describe('StreamMultiplexer', () => { it('ends all streams when any channel returns', async () => { const [multiplex, duplex] = await makeMultiplexer(); - const [ch1, ch2] = multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('2', (_value: number) => undefined, isNumber), - ); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); const drainP = multiplex.drainAll(); await ch1.return(); @@ -148,10 +188,8 @@ describe('StreamMultiplexer', () => { const [multiplex] = await makeMultiplexer( await TestDuplexStream.make(onDispatch), ); - const [ch1, ch2] = multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('2', (_value: number) => undefined, isNumber), - ); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); onDispatch.mockImplementationOnce(() => { throw new Error('foo'); }); @@ -171,10 +209,8 @@ describe('StreamMultiplexer', () => { it('ends all streams when a channel throws', async () => { const [multiplex, duplex] = await makeMultiplexer(); - const [ch1, ch2] = multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('2', (_value: number) => undefined, isNumber), - ); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); const drainP = multiplex.drainAll(); @@ -189,10 +225,8 @@ describe('StreamMultiplexer', () => { it('ends all streams when receiving a message for a non-existent channel', async () => { const [multiplex, duplex] = await makeMultiplexer(); - const [ch1, ch2] = multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('2', (_value: number) => undefined, isNumber), - ); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); const drainP = multiplex.drainAll(); @@ -210,10 +244,8 @@ describe('StreamMultiplexer', () => { describe('writing', () => { it('writes channel messages correctly', async () => { const [multiplex, duplex] = await makeMultiplexer(); - const [ch1, ch2] = multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('2', (_value: number) => undefined, isNumber), - ); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); const writeSpy = vi.spyOn(duplex, 'write'); @@ -233,10 +265,8 @@ describe('StreamMultiplexer', () => { it('returns done results from channel writes after ending', async () => { const [multiplex, duplex] = await makeMultiplexer(); - const [ch1, ch2] = multiplex.addChannels( - makeChannelParams('1', (_value: string) => undefined, isString), - makeChannelParams('2', (_value: number) => undefined, isNumber), - ); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); await Promise.all([multiplex.drainAll(), duplex.return()]); diff --git a/packages/streams/src/StreamMultiplexer.ts b/packages/streams/src/StreamMultiplexer.ts index 58e46b8bb..91ea76764 100644 --- a/packages/streams/src/StreamMultiplexer.ts +++ b/packages/streams/src/StreamMultiplexer.ts @@ -35,44 +35,27 @@ export type MultiplexEnvelope = { type HandleRead = (value: Read) => void | Promise; -export type ChannelParameters = { - channelName: ChannelName; - handleRead: HandleRead; - validateInput: ValidateInput; -} & { _write: Write }; - -export const makeChannelParams = ( - channelName: ChannelName, - handleRead: HandleRead, - validateInput: ValidateInput, -): ChannelParameters => - ({ - channelName, - handleRead, - validateInput, - }) as ChannelParameters; - type ChannelRecord = { channelName: ChannelName; - stream: DuplexStream; - handleRead: HandleRead; + stream: HandledDuplexStream; receiveInput: ReceiveInput; }; -type MappedChannels[]> = - Readonly<{ - [K in keyof Channels]: DuplexStream< - // We need to use `any` for the types to actually be inferred, and our - // usage should ensure that we avoid polluting the outside world. - /* eslint-disable @typescript-eslint/no-explicit-any */ - Channels[K] extends ChannelParameters ? Read : never, - Channels[K] extends ChannelParameters ? Write : never - /* eslint-enable @typescript-eslint/no-explicit-any */ - >; - }>; +type HandledDuplexStream = Omit< + DuplexStream, + 'drain' +> & { + drain: () => Promise; +}; + +enum MultiplexerStatus { + Idle = 0, + Running = 1, + Done = 2, +} export class StreamMultiplexer { - #isDone: boolean; + #status: MultiplexerStatus; readonly #name: string; @@ -83,19 +66,26 @@ export class StreamMultiplexer { MultiplexEnvelope >; + /** + * Creates a new multiplexer over the specified duplex stream. The duplex stream will + * be synchronized by the multiplexer, and **should not** be synchronized by the caller. + * + * @param stream - The underlying duplex stream. + * @param name - The multiplexer name. + */ constructor( stream: SynchronizableDuplexStream, name?: string, ) { - this.#isDone = false; + this.#status = MultiplexerStatus.Idle; this.#channels = new Map(); this.#name = name ?? this.constructor.name; this.#stream = stream; } /** - * Starts the multiplexer and drains all of its channels. Waits for the underlying - * duplex stream to be synchronized before reading from it. + * Starts the multiplexer and drains all of its channels. Use either this method or + * {@link start} to drain the multiplexer. * * @returns A promise resolves when the multiplexer and its channels have ended. */ @@ -103,12 +93,11 @@ export class StreamMultiplexer { if (this.#channels.size === 0) { throw new Error(`${this.#name} has no channels`); } - await this.#stream.synchronize(); const promise = Promise.all([ - this.#drain(), - ...Array.from(this.#channels.values()).map( - async ({ stream, handleRead }) => stream.drain(handleRead), + this.start(), + ...Array.from(this.#channels.values()).map(async ({ stream }) => + stream.drain(), ), ]).then(async () => this.#end()); @@ -119,7 +108,20 @@ export class StreamMultiplexer { return promise; } - async #drain(): Promise { + /** + * Idempotently starts the multiplexer by draining the underlying duplex stream and + * forwarding messages to the appropriate channels. Waits for the underlying duplex + * stream to be synchronized before reading from it. Ends the multiplexer if the duplex + * stream ends. Use either this method or {@link drainAll} to drain the multiplexer. + */ + async start(): Promise { + if (this.#status !== MultiplexerStatus.Idle) { + return; + } + this.#status = MultiplexerStatus.Running; + + await this.#stream.synchronize(); + for await (const envelope of this.#stream) { const channel = this.#channels.get(envelope.channel); if (channel === undefined) { @@ -136,31 +138,22 @@ export class StreamMultiplexer { } /** - * Adds a set of channels. To avoid messages loss, the underlying duplex stream must not - * be synchronized until all channels have been created. + * Adds a channel to the multiplexer. To avoid messages loss, the underlying duplex + * stream must not be synchronized until all channels have been created. * - * @param channels - The channels to add. - * @returns The added channels. + * @param channelName - The channel name. + * @param validateInput - The input validator. + * @param handleRead - The channel stream's drain handler. + * @returns The channel stream. */ - // See MappedChannels for why we need to use `any`. - // eslint-disable-next-line @typescript-eslint/no-explicit-any - addChannels[]>( - ...channels: Channels - ): MappedChannels { - if (this.#isDone) { - throw new Error(`${this.#name} has already ended`); + addChannel( + channelName: ChannelName, + validateInput: ValidateInput, + handleRead: HandleRead, + ): HandledDuplexStream { + if (this.#status !== MultiplexerStatus.Idle) { + throw new Error('Channels must be added before starting the multiplexer'); } - - return channels.map((params) => - this.#addChannel(params), - ) as MappedChannels; - } - - #addChannel({ - channelName, - handleRead, - validateInput, - }: ChannelParameters): DuplexStream { if (this.#channels.has(channelName)) { throw new Error(`Channel "${channelName}" already exists.`); } @@ -168,24 +161,33 @@ export class StreamMultiplexer { const { stream, receiveInput } = this.#makeChannel( channelName, validateInput, + handleRead, ); // We downcast some properties in order to store all records in one place. this.#channels.set(channelName, { channelName, - handleRead: handleRead as HandleRead, - stream: stream as unknown as DuplexStream, + stream: stream as unknown as HandledDuplexStream, receiveInput, }); return stream; } + /** + * Constructs a channel. + * + * @param channelName - The channel name. + * @param validateInput - The input validator. + * @param handleRead - The channel stream's drain handler. + * @returns The channel stream and its receiveInput method. + */ #makeChannel( channelName: ChannelName, validateInput: ValidateInput, + handleRead: HandleRead, ): { - stream: DuplexStream; + stream: HandledDuplexStream; receiveInput: ReceiveInput; } { let isDone = false; @@ -217,14 +219,14 @@ export class StreamMultiplexer { return writeP; }; - const drain = async (handler: HandleRead): Promise => { + const drain = async (): Promise => { for await (const value of reader) { - await handler(value); + await handleRead(value); } }; // Create and return the DuplexStream interface - const stream: DuplexStream = { + const stream: HandledDuplexStream = { next: reader.next.bind(reader), return: reader.return.bind(reader), throw: reader.throw.bind(reader), @@ -239,10 +241,10 @@ export class StreamMultiplexer { } async #end(error?: Error): Promise { - if (this.#isDone) { + if (this.#status === MultiplexerStatus.Done) { return; } - this.#isDone = true; + this.#status = MultiplexerStatus.Done; const end = async ( stream: DuplexStream, diff --git a/packages/streams/src/index.test.ts b/packages/streams/src/index.test.ts index b72d96dd1..9fee2a387 100644 --- a/packages/streams/src/index.test.ts +++ b/packages/streams/src/index.test.ts @@ -17,7 +17,6 @@ describe('index', () => { 'PostMessageWriter', 'StreamMultiplexer', 'initializeMessageChannel', - 'makeChannelParams', 'makeStreamEnvelopeKit', 'receiveMessagePort', ]); diff --git a/packages/streams/src/index.ts b/packages/streams/src/index.ts index fc12f4c02..4b4ac279d 100644 --- a/packages/streams/src/index.ts +++ b/packages/streams/src/index.ts @@ -29,5 +29,5 @@ export type { MakeStreamEnvelopeHandler, StreamEnvelopeKit, } from './envelope-kit.js'; -export { StreamMultiplexer, makeChannelParams } from './StreamMultiplexer.js'; +export { StreamMultiplexer } from './StreamMultiplexer.js'; export type { MultiplexEnvelope } from './StreamMultiplexer.js'; From 4ffa6b77698627324800e8d1fd3d772ede22ea8b Mon Sep 17 00:00:00 2001 From: Erik Marks Date: Mon, 4 Nov 2024 17:04:42 +0700 Subject: [PATCH 03/10] feat(streams): Add concrete StreamMultiplexer implementations --- .../streams/src/ChromeRuntimeStream.test.ts | 37 +++++++++++++++++ packages/streams/src/ChromeRuntimeStream.ts | 21 ++++++++-- .../streams/src/MessagePortStream.test.ts | 29 ++++++++++++++ packages/streams/src/MessagePortStream.ts | 13 ++++-- .../streams/src/PostMessageStream.test.ts | 40 +++++++++++++++++++ packages/streams/src/PostMessageStream.ts | 21 ++++++++-- packages/streams/src/index.test.ts | 3 ++ packages/streams/src/index.ts | 11 +++-- 8 files changed, 162 insertions(+), 13 deletions(-) diff --git a/packages/streams/src/ChromeRuntimeStream.test.ts b/packages/streams/src/ChromeRuntimeStream.test.ts index b5198b845..a0be0a52e 100644 --- a/packages/streams/src/ChromeRuntimeStream.test.ts +++ b/packages/streams/src/ChromeRuntimeStream.test.ts @@ -11,7 +11,9 @@ import { ChromeRuntimeWriter, ChromeRuntimeStreamTarget, ChromeRuntimeDuplexStream, + ChromeRuntimeMultiplexer, } from './ChromeRuntimeStream.js'; +import { StreamMultiplexer } from './StreamMultiplexer.js'; import { makeDoneResult, makePendingResult, @@ -357,3 +359,38 @@ describe.concurrent('ChromeRuntimeDuplexStream', () => { expect(await readP).toStrictEqual(makeDoneResult()); }); }); + +describe('ChromeRuntimeMultiplexer', () => { + it('constructs a ChromeRuntimeMultiplexer', () => { + const multiplexer = new ChromeRuntimeMultiplexer( + asChromeRuntime(makeRuntime().runtime), + ChromeRuntimeStreamTarget.Background, + ChromeRuntimeStreamTarget.Offscreen, + ); + + expect(multiplexer).toBeInstanceOf(StreamMultiplexer); + }); + + it('can create and drain channels', async () => { + const { runtime, dispatchRuntimeMessage } = makeRuntime(); + const multiplexer = new ChromeRuntimeMultiplexer( + asChromeRuntime(runtime), + ChromeRuntimeStreamTarget.Background, + ChromeRuntimeStreamTarget.Offscreen, + ); + const handleRead = vi.fn(); + multiplexer.addChannel( + '1', + (value: unknown): value is number => typeof value === 'number', + handleRead, + ); + + const drainP = multiplexer.drainAll(); + dispatchRuntimeMessage(makeAck()); + dispatchRuntimeMessage({ channel: '1', payload: 42 }); + dispatchRuntimeMessage(makeStreamDoneSignal()); + + await drainP; + expect(handleRead).toHaveBeenCalledWith(42); + }); +}); diff --git a/packages/streams/src/ChromeRuntimeStream.ts b/packages/streams/src/ChromeRuntimeStream.ts index a330a0de3..328d852c2 100644 --- a/packages/streams/src/ChromeRuntimeStream.ts +++ b/packages/streams/src/ChromeRuntimeStream.ts @@ -26,6 +26,7 @@ import type { } from './BaseStream.js'; import { BaseReader, BaseWriter } from './BaseStream.js'; import type { ChromeRuntime, ChromeMessageSender } from './chrome.js'; +import { StreamMultiplexer } from './StreamMultiplexer.js'; import type { Dispatchable } from './utils.js'; export enum ChromeRuntimeStreamTarget { @@ -180,9 +181,7 @@ export class ChromeRuntimeDuplexStream< Write, ChromeRuntimeWriter > { - // Unavoidable exception to our preference for #-private names. - // eslint-disable-next-line no-restricted-syntax - private constructor( + constructor( runtime: ChromeRuntime, localTarget: ChromeRuntimeStreamTarget, remoteTarget: ChromeRuntimeStreamTarget, @@ -237,3 +236,19 @@ export class ChromeRuntimeDuplexStream< } } harden(ChromeRuntimeDuplexStream); + +export class ChromeRuntimeMultiplexer extends StreamMultiplexer { + constructor( + runtime: ChromeRuntime, + localTarget: ChromeRuntimeStreamTarget, + remoteTarget: ChromeRuntimeStreamTarget, + name?: string, + ) { + super( + new ChromeRuntimeDuplexStream(runtime, localTarget, remoteTarget), + name, + ); + harden(this); + } +} +harden(ChromeRuntimeMultiplexer); diff --git a/packages/streams/src/MessagePortStream.test.ts b/packages/streams/src/MessagePortStream.test.ts index 23275d311..ad1cb3b98 100644 --- a/packages/streams/src/MessagePortStream.test.ts +++ b/packages/streams/src/MessagePortStream.test.ts @@ -5,6 +5,7 @@ import { makeAck } from './BaseDuplexStream.js'; import type { ValidateInput } from './BaseStream.js'; import { MessagePortDuplexStream, + MessagePortMultiplexer, MessagePortReader, MessagePortWriter, } from './MessagePortStream.js'; @@ -201,3 +202,31 @@ describe('MessagePortDuplexStream', () => { expect(await readP).toStrictEqual(makeDoneResult()); }); }); + +describe('MessagePortMultiplexer', () => { + it('constructs a MessagePortMultiplexer', () => { + const { port1 } = new MessageChannel(); + const multiplexer = new MessagePortMultiplexer(port1); + + expect(multiplexer).toBeInstanceOf(MessagePortMultiplexer); + }); + + it('can create and drain channels', async () => { + const { port1, port2 } = new MessageChannel(); + const multiplexer = new MessagePortMultiplexer(port1); + const handleRead = vi.fn(); + multiplexer.addChannel( + '1', + (value: unknown): value is number => typeof value === 'number', + handleRead, + ); + + const drainP = multiplexer.drainAll(); + port2.postMessage(makeAck()); + port2.postMessage({ channel: '1', payload: 42 }); + port2.postMessage(makeStreamDoneSignal()); + + await drainP; + expect(handleRead).toHaveBeenCalledWith(42); + }); +}); diff --git a/packages/streams/src/MessagePortStream.ts b/packages/streams/src/MessagePortStream.ts index 6bebc5b99..b574ff89e 100644 --- a/packages/streams/src/MessagePortStream.ts +++ b/packages/streams/src/MessagePortStream.ts @@ -28,6 +28,7 @@ import type { ValidateInput, } from './BaseStream.js'; import { BaseReader, BaseWriter } from './BaseStream.js'; +import { StreamMultiplexer } from './StreamMultiplexer.js'; import type { Dispatchable, OnMessage } from './utils.js'; /** @@ -110,9 +111,7 @@ export class MessagePortDuplexStream< Write, MessagePortWriter > { - // Unavoidable exception to our preference for #-private names. - // eslint-disable-next-line no-restricted-syntax - private constructor(port: MessagePort, validateInput?: ValidateInput) { + constructor(port: MessagePort, validateInput?: ValidateInput) { let writer: MessagePortWriter; // eslint-disable-line prefer-const const reader = new MessagePortReader(port, { name: 'MessagePortDuplexStream', @@ -143,3 +142,11 @@ export class MessagePortDuplexStream< } } harden(MessagePortDuplexStream); + +export class MessagePortMultiplexer extends StreamMultiplexer { + constructor(port: MessagePort, name?: string) { + super(new MessagePortDuplexStream(port), name); + harden(this); + } +} +harden(MessagePortMultiplexer); diff --git a/packages/streams/src/PostMessageStream.test.ts b/packages/streams/src/PostMessageStream.test.ts index 7282bee75..695c86378 100644 --- a/packages/streams/src/PostMessageStream.test.ts +++ b/packages/streams/src/PostMessageStream.test.ts @@ -5,9 +5,11 @@ import { makeAck } from './BaseDuplexStream.js'; import type { ValidateInput } from './BaseStream.js'; import { PostMessageDuplexStream, + PostMessageMultiplexer, PostMessageReader, PostMessageWriter, } from './PostMessageStream.js'; +import { StreamMultiplexer } from './StreamMultiplexer.js'; import type { PostMessage } from './utils.js'; import { makeDoneResult, @@ -213,3 +215,41 @@ describe('PostMessageDuplexStream', () => { expect(await readP).toStrictEqual(makeDoneResult()); }); }); + +describe('PostMessageMultiplexer', () => { + it('constructs a PostMessageMultiplexer', () => { + const { postMessageFn, setListener, removeListener } = + makePostMessageMock(); + const multiplexer = new PostMessageMultiplexer( + postMessageFn, + setListener, + removeListener, + ); + + expect(multiplexer).toBeInstanceOf(StreamMultiplexer); + }); + + it('can create and drain channels', async () => { + const { postMessageFn, setListener, removeListener } = + makePostMessageMock(); + const multiplexer = new PostMessageMultiplexer( + postMessageFn, + setListener, + removeListener, + ); + const handleRead = vi.fn(); + multiplexer.addChannel( + '1', + (value: unknown): value is number => typeof value === 'number', + handleRead, + ); + + const drainP = multiplexer.drainAll(); + postMessageFn(makeAck()); + postMessageFn({ channel: '1', payload: 42 }); + postMessageFn(makeStreamDoneSignal()); + + await drainP; + expect(handleRead).toHaveBeenCalledWith(42); + }); +}); diff --git a/packages/streams/src/PostMessageStream.ts b/packages/streams/src/PostMessageStream.ts index 814660322..bf6a9cd78 100644 --- a/packages/streams/src/PostMessageStream.ts +++ b/packages/streams/src/PostMessageStream.ts @@ -15,6 +15,7 @@ import type { ValidateInput, } from './BaseStream.js'; import { BaseReader, BaseWriter } from './BaseStream.js'; +import { StreamMultiplexer } from './StreamMultiplexer.js'; import type { Dispatchable, OnMessage, PostMessage } from './utils.js'; type SetListener = (onMessage: OnMessage) => void; @@ -97,9 +98,7 @@ export class PostMessageDuplexStream< Write, PostMessageWriter > { - // Unavoidable exception to our preference for #-private names. - // eslint-disable-next-line no-restricted-syntax - private constructor( + constructor( postMessageFn: PostMessage, setListener: SetListener, removeListener: RemoveListener, @@ -139,3 +138,19 @@ export class PostMessageDuplexStream< } } harden(PostMessageDuplexStream); + +export class PostMessageMultiplexer extends StreamMultiplexer { + constructor( + postMessageFn: PostMessage, + setListener: SetListener, + removeListener: RemoveListener, + name?: string, + ) { + super( + new PostMessageDuplexStream(postMessageFn, setListener, removeListener), + name, + ); + harden(this); + } +} +harden(PostMessageMultiplexer); diff --git a/packages/streams/src/index.test.ts b/packages/streams/src/index.test.ts index 9fee2a387..cbacc5921 100644 --- a/packages/streams/src/index.test.ts +++ b/packages/streams/src/index.test.ts @@ -6,13 +6,16 @@ describe('index', () => { it('has the expected exports', () => { expect(Object.keys(indexModule).sort()).toStrictEqual([ 'ChromeRuntimeDuplexStream', + 'ChromeRuntimeMultiplexer', 'ChromeRuntimeReader', 'ChromeRuntimeTarget', 'ChromeRuntimeWriter', 'MessagePortDuplexStream', + 'MessagePortMultiplexer', 'MessagePortReader', 'MessagePortWriter', 'PostMessageDuplexStream', + 'PostMessageMultiplexer', 'PostMessageReader', 'PostMessageWriter', 'StreamMultiplexer', diff --git a/packages/streams/src/index.ts b/packages/streams/src/index.ts index 4b4ac279d..ee2493f1e 100644 --- a/packages/streams/src/index.ts +++ b/packages/streams/src/index.ts @@ -5,21 +5,24 @@ export { export type { Reader, Writer } from './utils.js'; export type { DuplexStream } from './BaseDuplexStream.js'; export { + MessagePortDuplexStream, + MessagePortMultiplexer, MessagePortReader, MessagePortWriter, - MessagePortDuplexStream, } from './MessagePortStream.js'; export type { ChromeRuntime, ChromeMessageSender } from './chrome.d.ts'; export { - ChromeRuntimeReader, - ChromeRuntimeWriter, ChromeRuntimeDuplexStream, + ChromeRuntimeMultiplexer, ChromeRuntimeStreamTarget as ChromeRuntimeTarget, + ChromeRuntimeReader, + ChromeRuntimeWriter, } from './ChromeRuntimeStream.js'; export { + PostMessageDuplexStream, + PostMessageMultiplexer, PostMessageReader, PostMessageWriter, - PostMessageDuplexStream, } from './PostMessageStream.js'; export { makeStreamEnvelopeKit } from './envelope-kit.js'; export type { StreamEnveloper } from './enveloper.js'; From 7f7d9a249683fd8e13703ceebe647d3d7210c60b Mon Sep 17 00:00:00 2001 From: Erik Marks Date: Tue, 5 Nov 2024 13:59:42 +0700 Subject: [PATCH 04/10] feat(streams): Add StreamMultiplexer.return() Also make `synchronize()` method of duplex streams optional. --- packages/streams/src/BaseDuplexStream.ts | 7 -- .../streams/src/StreamMultiplexer.test.ts | 72 +++++++++++-------- packages/streams/src/StreamMultiplexer.ts | 37 ++++++---- packages/streams/src/index.ts | 5 +- packages/streams/test/stream-mocks.ts | 26 ++++++- 5 files changed, 98 insertions(+), 49 deletions(-) diff --git a/packages/streams/src/BaseDuplexStream.ts b/packages/streams/src/BaseDuplexStream.ts index 1789bd682..094e7ae88 100644 --- a/packages/streams/src/BaseDuplexStream.ts +++ b/packages/streams/src/BaseDuplexStream.ts @@ -225,10 +225,3 @@ export type DuplexStream = Pick< > & { [Symbol.asyncIterator]: () => DuplexStream; }; - -export type SynchronizableDuplexStream< - Read extends Json, - Write extends Json = Read, -> = DuplexStream & { - synchronize(): Promise; -}; diff --git a/packages/streams/src/StreamMultiplexer.test.ts b/packages/streams/src/StreamMultiplexer.test.ts index cde195d12..8516079fa 100644 --- a/packages/streams/src/StreamMultiplexer.test.ts +++ b/packages/streams/src/StreamMultiplexer.test.ts @@ -14,19 +14,6 @@ const isString: ValidateInput = (value) => typeof value === 'string'; const isNumber: ValidateInput = (value) => typeof value === 'number'; -const makeMultiplexer = async ( - duplex?: TestDuplexStream, -): Promise< - [TestMultiplexer, TestDuplexStream] -> => { - // We can't use the async factory for a parameter default - // eslint-disable-next-line no-param-reassign - duplex ??= await TestDuplexStream.make( - () => undefined, - ); - return [new TestMultiplexer(duplex), duplex]; -}; - const makeEnvelope = (channel: string, payload: Json): MultiplexEnvelope => ({ channel, payload, @@ -45,7 +32,7 @@ describe('StreamMultiplexer', () => { describe('addChannel', () => { it('makes and adds channels', async () => { - const [multiplex] = await makeMultiplexer(); + const [multiplex] = await TestMultiplexer.make(); const ch1 = multiplex.addChannel('1', isString, noop); const ch2 = multiplex.addChannel('2', isNumber, noop); expect(ch1[Symbol.asyncIterator]()).toBe(ch1); @@ -53,7 +40,7 @@ describe('StreamMultiplexer', () => { }); it('throws if adding a channel with the same name multiple times', async () => { - const [multiplex] = await makeMultiplexer(); + const [multiplex] = await TestMultiplexer.make(); multiplex.addChannel('1', isString, noop); expect(() => multiplex.addChannel('1', isString, noop)).toThrow( 'Channel "1" already exists', @@ -61,7 +48,7 @@ describe('StreamMultiplexer', () => { }); it('throws if adding channels after starting', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); multiplex.addChannel('1', isString, noop); const startP = multiplex.start(); @@ -73,7 +60,7 @@ describe('StreamMultiplexer', () => { }); it('throws if adding channels after ending', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); // Add one channel so we can start the multiplexer. multiplex.addChannel('1', isString, noop); @@ -87,7 +74,7 @@ describe('StreamMultiplexer', () => { describe('start', () => { it('is idempotent', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); // Add one channel so we can start the multiplexer. multiplex.addChannel('1', isString, noop); const startP = Promise.all([multiplex.start(), multiplex.start()]).then( @@ -98,7 +85,7 @@ describe('StreamMultiplexer', () => { }); it('enables draining channels separately', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); const ch1Handler = vi.fn(); const ch2Handler = vi.fn(); const ch1 = multiplex.addChannel('1', isString, ch1Handler); @@ -127,14 +114,14 @@ describe('StreamMultiplexer', () => { describe('drainAll', () => { it('throws if draining when there are no channels', async () => { - const [multiplex] = await makeMultiplexer(); + const [multiplex] = await TestMultiplexer.make(); await expect(multiplex.drainAll()).rejects.toThrow( 'TestMultiplexer has no channels', ); }); it('forwards input to the correct channel', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); const ch1Handler = vi.fn(); const ch2Handler = vi.fn(); multiplex.addChannel('1', isString, ch1Handler); @@ -156,7 +143,7 @@ describe('StreamMultiplexer', () => { }); it('ends all streams when the duplex stream returns', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); const ch1 = multiplex.addChannel('1', isString, noop); const ch2 = multiplex.addChannel('2', isNumber, noop); const drainP = multiplex.drainAll(); @@ -170,7 +157,7 @@ describe('StreamMultiplexer', () => { }); it('ends all streams when any channel returns', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); const ch1 = multiplex.addChannel('1', isString, noop); const ch2 = multiplex.addChannel('2', isNumber, noop); const drainP = multiplex.drainAll(); @@ -185,7 +172,7 @@ describe('StreamMultiplexer', () => { it('ends all streams when the duplex stream throws', async () => { const onDispatch = vi.fn(); - const [multiplex] = await makeMultiplexer( + const [multiplex] = await TestMultiplexer.make( await TestDuplexStream.make(onDispatch), ); const ch1 = multiplex.addChannel('1', isString, noop); @@ -208,7 +195,7 @@ describe('StreamMultiplexer', () => { }); it('ends all streams when a channel throws', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); const ch1 = multiplex.addChannel('1', isString, noop); const ch2 = multiplex.addChannel('2', isNumber, noop); @@ -224,7 +211,7 @@ describe('StreamMultiplexer', () => { }); it('ends all streams when receiving a message for a non-existent channel', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); const ch1 = multiplex.addChannel('1', isString, noop); const ch2 = multiplex.addChannel('2', isNumber, noop); @@ -243,7 +230,7 @@ describe('StreamMultiplexer', () => { describe('writing', () => { it('writes channel messages correctly', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); const ch1 = multiplex.addChannel('1', isString, noop); const ch2 = multiplex.addChannel('2', isNumber, noop); @@ -264,7 +251,7 @@ describe('StreamMultiplexer', () => { }); it('returns done results from channel writes after ending', async () => { - const [multiplex, duplex] = await makeMultiplexer(); + const [multiplex, duplex] = await TestMultiplexer.make(); const ch1 = multiplex.addChannel('1', isString, noop); const ch2 = multiplex.addChannel('2', isNumber, noop); @@ -274,4 +261,33 @@ describe('StreamMultiplexer', () => { expect(await ch2.write(42)).toStrictEqual(makeDoneResult()); }); }); + + describe('return', () => { + it('ends the multiplexer and its channels', async () => { + const [multiplex] = await TestMultiplexer.make(); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); + + await multiplex.return(); + + expect(await ch1.next()).toStrictEqual(makeDoneResult()); + expect(await ch2.next()).toStrictEqual(makeDoneResult()); + }); + + it('is idempotent', async () => { + const [multiplex] = await TestMultiplexer.make(); + const ch1 = multiplex.addChannel('1', isString, noop); + const ch2 = multiplex.addChannel('2', isNumber, noop); + + await multiplex.return(); + + expect(await ch1.next()).toStrictEqual(makeDoneResult()); + expect(await ch2.next()).toStrictEqual(makeDoneResult()); + + await multiplex.return(); + + expect(await ch1.next()).toStrictEqual(makeDoneResult()); + expect(await ch2.next()).toStrictEqual(makeDoneResult()); + }); + }); }); diff --git a/packages/streams/src/StreamMultiplexer.ts b/packages/streams/src/StreamMultiplexer.ts index 91ea76764..669a5407a 100644 --- a/packages/streams/src/StreamMultiplexer.ts +++ b/packages/streams/src/StreamMultiplexer.ts @@ -1,9 +1,6 @@ import type { Json } from '@metamask/utils'; -import type { - DuplexStream, - SynchronizableDuplexStream, -} from './BaseDuplexStream.js'; +import type { DuplexStream } from './BaseDuplexStream.js'; import type { BaseReaderArgs, ValidateInput, @@ -12,6 +9,13 @@ import type { import { BaseReader } from './BaseStream.js'; import { makeDoneResult } from './utils.js'; +type SynchronizableDuplexStream< + Read extends Json, + Write extends Json = Read, +> = DuplexStream & { + synchronize?: () => Promise; +}; + class ChannelReader extends BaseReader { // eslint-disable-next-line no-restricted-syntax private constructor(args: BaseReaderArgs) { @@ -41,7 +45,7 @@ type ChannelRecord = { receiveInput: ReceiveInput; }; -type HandledDuplexStream = Omit< +export type HandledDuplexStream = Omit< DuplexStream, 'drain' > & { @@ -67,8 +71,9 @@ export class StreamMultiplexer { >; /** - * Creates a new multiplexer over the specified duplex stream. The duplex stream will - * be synchronized by the multiplexer, and **should not** be synchronized by the caller. + * Creates a new multiplexer over the specified duplex stream. If the duplex stream + * is synchronizable, it will be synchronized by the multiplexer and **should not** + * be synchronized by the caller. * * @param stream - The underlying duplex stream. * @param name - The multiplexer name. @@ -110,9 +115,11 @@ export class StreamMultiplexer { /** * Idempotently starts the multiplexer by draining the underlying duplex stream and - * forwarding messages to the appropriate channels. Waits for the underlying duplex - * stream to be synchronized before reading from it. Ends the multiplexer if the duplex + * forwarding messages to the appropriate channels. Ends the multiplexer if the duplex * stream ends. Use either this method or {@link drainAll} to drain the multiplexer. + * + * If the duplex stream is synchronizable, it will be synchronized by the multiplexer + * and **should not** be synchronized by the caller. */ async start(): Promise { if (this.#status !== MultiplexerStatus.Idle) { @@ -120,7 +127,7 @@ export class StreamMultiplexer { } this.#status = MultiplexerStatus.Running; - await this.#stream.synchronize(); + await this.#stream.synchronize?.(); for await (const envelope of this.#stream) { const channel = this.#channels.get(envelope.channel); @@ -138,8 +145,7 @@ export class StreamMultiplexer { } /** - * Adds a channel to the multiplexer. To avoid messages loss, the underlying duplex - * stream must not be synchronized until all channels have been created. + * Adds a channel to the multiplexer. * * @param channelName - The channel name. * @param validateInput - The input validator. @@ -240,6 +246,13 @@ export class StreamMultiplexer { return { stream, receiveInput }; } + /** + * Ends the multiplexer and its channels. + */ + async return(): Promise { + await this.#end(); + } + async #end(error?: Error): Promise { if (this.#status === MultiplexerStatus.Done) { return; diff --git a/packages/streams/src/index.ts b/packages/streams/src/index.ts index ee2493f1e..880634d5f 100644 --- a/packages/streams/src/index.ts +++ b/packages/streams/src/index.ts @@ -33,4 +33,7 @@ export type { StreamEnvelopeKit, } from './envelope-kit.js'; export { StreamMultiplexer } from './StreamMultiplexer.js'; -export type { MultiplexEnvelope } from './StreamMultiplexer.js'; +export type { + HandledDuplexStream, + MultiplexEnvelope, +} from './StreamMultiplexer.js'; diff --git a/packages/streams/test/stream-mocks.ts b/packages/streams/test/stream-mocks.ts index c0d470c32..a8558c1de 100644 --- a/packages/streams/test/stream-mocks.ts +++ b/packages/streams/test/stream-mocks.ts @@ -1,6 +1,7 @@ import type { Json } from '@metamask/utils'; import { BaseDuplexStream, makeAck } from '../src/BaseDuplexStream.js'; +import type { DuplexStream } from '../src/BaseDuplexStream.js'; import type { Dispatch, ReceiveInput, @@ -9,6 +10,7 @@ import type { BaseWriterArgs, } from '../src/BaseStream.js'; import { BaseReader, BaseWriter } from '../src/BaseStream.js'; +import type { MultiplexEnvelope } from '../src/StreamMultiplexer.js'; import { StreamMultiplexer } from '../src/StreamMultiplexer.js'; export type { MultiplexEnvelope } from '../src/StreamMultiplexer.js'; @@ -122,4 +124,26 @@ export class TestDuplexStream< } } -export class TestMultiplexer extends StreamMultiplexer {} +export class TestMultiplexer extends StreamMultiplexer { + constructor( + duplex: DuplexStream = new TestDuplexStream( + () => undefined, + ), + ) { + super(duplex); + } + + static async make( + duplex?: TestDuplexStream, + ): Promise< + [TestMultiplexer, TestDuplexStream] + > { + // We can't use the async factory for a parameter default + // eslint-disable-next-line no-param-reassign + duplex ??= await TestDuplexStream.make< + MultiplexEnvelope, + MultiplexEnvelope + >(() => undefined); + return [new TestMultiplexer(duplex), duplex] as const; + } +} From 6a752d6d44753f40eadc4c4a4c5e66741904310f Mon Sep 17 00:00:00 2001 From: Erik Marks Date: Tue, 5 Nov 2024 16:42:09 +0700 Subject: [PATCH 05/10] fix(steams): Enable passing input validators to duplex streams --- packages/streams/src/BaseDuplexStream.test.ts | 12 +++++++++ packages/streams/src/BaseDuplexStream.ts | 25 ++++++++++++++++++- packages/streams/src/ChromeRuntimeStream.ts | 16 +++++++++--- packages/streams/src/MessagePortStream.ts | 11 +++++--- packages/streams/src/PostMessageStream.ts | 16 +++++++++--- packages/streams/src/StreamMultiplexer.ts | 16 ++++++++++++ packages/streams/test/stream-mocks.ts | 8 ++++-- 7 files changed, 89 insertions(+), 15 deletions(-) diff --git a/packages/streams/src/BaseDuplexStream.test.ts b/packages/streams/src/BaseDuplexStream.test.ts index e97888651..a127822bc 100644 --- a/packages/streams/src/BaseDuplexStream.test.ts +++ b/packages/streams/src/BaseDuplexStream.test.ts @@ -166,6 +166,18 @@ describe('BaseDuplexStream', () => { expect(await nextP).toStrictEqual(makePendingResult(message)); }); + it('reads from the reader, with input validation', async () => { + const duplexStream = await TestDuplexStream.make(() => undefined, { + validateInput: (value: unknown): value is number => + typeof value === 'number', + }); + + const message = 42; + await duplexStream.receiveInput(message); + + expect(await duplexStream.next()).toStrictEqual(makePendingResult(message)); + }); + it('drains the reader in order', async () => { const duplexStream = await TestDuplexStream.make(() => undefined); diff --git a/packages/streams/src/BaseDuplexStream.ts b/packages/streams/src/BaseDuplexStream.ts index 094e7ae88..3400297f8 100644 --- a/packages/streams/src/BaseDuplexStream.ts +++ b/packages/streams/src/BaseDuplexStream.ts @@ -5,7 +5,7 @@ import { isObject } from '@metamask/utils'; import type { Json } from '@metamask/utils'; import { stringify } from '@ocap/utils'; -import type { BaseReader, BaseWriter } from './BaseStream.js'; +import type { BaseReader, BaseWriter, ValidateInput } from './BaseStream.js'; import { makeDoneResult } from './utils.js'; export enum DuplexStreamSentinel { @@ -35,6 +35,29 @@ export const makeAck = (): DuplexStreamAck => ({ const isAck = (value: unknown): value is DuplexStreamAck => isObject(value) && value[DuplexStreamSentinel.Ack] === true; +type StreamSignal = DuplexStreamSyn | DuplexStreamAck; + +const isDuplexStreamSignal = (value: unknown): value is StreamSignal => + isSyn(value) || isAck(value); + +/** + * Make a validator for input to a duplex stream. Constructor helper for concrete + * duplex stream implementations. + * + * Validators passed in by consumers must be augmented such that errors aren't + * thrown for {@link StreamSignal} values. + * + * @param validateInput - The validator for the stream's input type. + * @returns A validator for the stream's input type, or `undefined` if no + * validation is desired. + */ +export const makeDuplexStreamInputValidator = ( + validateInput?: ValidateInput, +): ((value: unknown) => value is Read) | undefined => + validateInput && + ((value: unknown): value is Read => + isDuplexStreamSignal(value) || validateInput(value as Json)); + enum SynchronizationStatus { Idle = 0, Pending = 1, diff --git a/packages/streams/src/ChromeRuntimeStream.ts b/packages/streams/src/ChromeRuntimeStream.ts index 328d852c2..e6665fee4 100644 --- a/packages/streams/src/ChromeRuntimeStream.ts +++ b/packages/streams/src/ChromeRuntimeStream.ts @@ -17,7 +17,10 @@ import type { Json } from '@metamask/utils'; import { stringify } from '@ocap/utils'; -import { BaseDuplexStream } from './BaseDuplexStream.js'; +import { + BaseDuplexStream, + makeDuplexStreamInputValidator, +} from './BaseDuplexStream.js'; import type { BaseReaderArgs, ValidateInput, @@ -26,7 +29,7 @@ import type { } from './BaseStream.js'; import { BaseReader, BaseWriter } from './BaseStream.js'; import type { ChromeRuntime, ChromeMessageSender } from './chrome.js'; -import { StreamMultiplexer } from './StreamMultiplexer.js'; +import { isMultiplexEnvelope, StreamMultiplexer } from './StreamMultiplexer.js'; import type { Dispatchable } from './utils.js'; export enum ChromeRuntimeStreamTarget { @@ -194,7 +197,7 @@ export class ChromeRuntimeDuplexStream< remoteTarget, { name: 'ChromeRuntimeDuplexStream', - validateInput, + validateInput: makeDuplexStreamInputValidator(validateInput), onEnd: async () => { await writer.return(); }, @@ -245,7 +248,12 @@ export class ChromeRuntimeMultiplexer extends StreamMultiplexer { name?: string, ) { super( - new ChromeRuntimeDuplexStream(runtime, localTarget, remoteTarget), + new ChromeRuntimeDuplexStream( + runtime, + localTarget, + remoteTarget, + isMultiplexEnvelope, + ), name, ); harden(this); diff --git a/packages/streams/src/MessagePortStream.ts b/packages/streams/src/MessagePortStream.ts index b574ff89e..7d1ff006d 100644 --- a/packages/streams/src/MessagePortStream.ts +++ b/packages/streams/src/MessagePortStream.ts @@ -21,14 +21,17 @@ import type { Json } from '@metamask/utils'; -import { BaseDuplexStream } from './BaseDuplexStream.js'; +import { + BaseDuplexStream, + makeDuplexStreamInputValidator, +} from './BaseDuplexStream.js'; import type { BaseReaderArgs, BaseWriterArgs, ValidateInput, } from './BaseStream.js'; import { BaseReader, BaseWriter } from './BaseStream.js'; -import { StreamMultiplexer } from './StreamMultiplexer.js'; +import { isMultiplexEnvelope, StreamMultiplexer } from './StreamMultiplexer.js'; import type { Dispatchable, OnMessage } from './utils.js'; /** @@ -115,7 +118,7 @@ export class MessagePortDuplexStream< let writer: MessagePortWriter; // eslint-disable-line prefer-const const reader = new MessagePortReader(port, { name: 'MessagePortDuplexStream', - validateInput, + validateInput: makeDuplexStreamInputValidator(validateInput), onEnd: async () => { await writer.return(); }, @@ -145,7 +148,7 @@ harden(MessagePortDuplexStream); export class MessagePortMultiplexer extends StreamMultiplexer { constructor(port: MessagePort, name?: string) { - super(new MessagePortDuplexStream(port), name); + super(new MessagePortDuplexStream(port, isMultiplexEnvelope), name); harden(this); } } diff --git a/packages/streams/src/PostMessageStream.ts b/packages/streams/src/PostMessageStream.ts index bf6a9cd78..282f8f2bd 100644 --- a/packages/streams/src/PostMessageStream.ts +++ b/packages/streams/src/PostMessageStream.ts @@ -8,14 +8,17 @@ import type { Json } from '@metamask/utils'; -import { BaseDuplexStream } from './BaseDuplexStream.js'; +import { + BaseDuplexStream, + makeDuplexStreamInputValidator, +} from './BaseDuplexStream.js'; import type { BaseReaderArgs, BaseWriterArgs, ValidateInput, } from './BaseStream.js'; import { BaseReader, BaseWriter } from './BaseStream.js'; -import { StreamMultiplexer } from './StreamMultiplexer.js'; +import { isMultiplexEnvelope, StreamMultiplexer } from './StreamMultiplexer.js'; import type { Dispatchable, OnMessage, PostMessage } from './utils.js'; type SetListener = (onMessage: OnMessage) => void; @@ -107,7 +110,7 @@ export class PostMessageDuplexStream< let writer: PostMessageWriter; // eslint-disable-line prefer-const const reader = new PostMessageReader(setListener, removeListener, { name: 'PostMessageDuplexStream', - validateInput, + validateInput: makeDuplexStreamInputValidator(validateInput), onEnd: async () => { await writer.return(); }, @@ -147,7 +150,12 @@ export class PostMessageMultiplexer extends StreamMultiplexer { name?: string, ) { super( - new PostMessageDuplexStream(postMessageFn, setListener, removeListener), + new PostMessageDuplexStream( + postMessageFn, + setListener, + removeListener, + isMultiplexEnvelope, + ), name, ); harden(this); diff --git a/packages/streams/src/StreamMultiplexer.ts b/packages/streams/src/StreamMultiplexer.ts index 669a5407a..84039ee7d 100644 --- a/packages/streams/src/StreamMultiplexer.ts +++ b/packages/streams/src/StreamMultiplexer.ts @@ -1,4 +1,5 @@ import type { Json } from '@metamask/utils'; +import { isObject } from '@metamask/utils'; import type { DuplexStream } from './BaseDuplexStream.js'; import type { @@ -37,6 +38,21 @@ export type MultiplexEnvelope = { payload: Json; }; +/** + * Type guard for {@link MultiplexEnvelope}. Only verifies that the `payload` property + * is not `undefined`, assuming that multiplexer channels will be responsible for + * performing further validation. + * + * @param value - The value to check. + * @returns Whether the value is a {@link MultiplexEnvelope}. + */ +export const isMultiplexEnvelope = ( + value: unknown, +): value is MultiplexEnvelope => + isObject(value) && + typeof value.channel === 'string' && + typeof value.payload !== 'undefined'; + type HandleRead = (value: Read) => void | Promise; type ChannelRecord = { diff --git a/packages/streams/test/stream-mocks.ts b/packages/streams/test/stream-mocks.ts index a8558c1de..4c06a20dd 100644 --- a/packages/streams/test/stream-mocks.ts +++ b/packages/streams/test/stream-mocks.ts @@ -1,6 +1,10 @@ import type { Json } from '@metamask/utils'; -import { BaseDuplexStream, makeAck } from '../src/BaseDuplexStream.js'; +import { + BaseDuplexStream, + makeAck, + makeDuplexStreamInputValidator, +} from '../src/BaseDuplexStream.js'; import type { DuplexStream } from '../src/BaseDuplexStream.js'; import type { Dispatch, @@ -78,7 +82,7 @@ export class TestDuplexStream< const reader = new TestReader({ name: 'TestDuplexStream', onEnd: readerOnEnd, - validateInput, + validateInput: makeDuplexStreamInputValidator(validateInput), }); super( reader, From 6370aa9dbf82284913eac3c91a69eb4e0b69be0e Mon Sep 17 00:00:00 2001 From: Erik Marks Date: Wed, 6 Nov 2024 19:35:05 +0700 Subject: [PATCH 06/10] test(streams): Fix test coverage Also fix issues in tests following rebase. --- packages/kernel/src/Supervisor.test.ts | 7 +++++-- packages/kernel/src/Vat.test.ts | 4 ++-- packages/streams/src/BaseDuplexStream.test.ts | 12 ++++-------- .../streams/src/ChromeRuntimeStream.test.ts | 19 +++++++++++++++++++ .../streams/src/MessagePortStream.test.ts | 12 ++++++++++++ .../streams/src/PostMessageStream.test.ts | 15 +++++++++++++++ .../streams/src/StreamMultiplexer.test.ts | 14 ++++++++------ packages/streams/src/StreamMultiplexer.ts | 2 +- 8 files changed, 66 insertions(+), 19 deletions(-) diff --git a/packages/kernel/src/Supervisor.test.ts b/packages/kernel/src/Supervisor.test.ts index cfa34482d..985adceb6 100644 --- a/packages/kernel/src/Supervisor.test.ts +++ b/packages/kernel/src/Supervisor.test.ts @@ -30,7 +30,7 @@ describe('Supervisor', () => { it('throws if the stream throws', async () => { const { supervisor, stream } = await makeSupervisor(); const consoleErrorSpy = vi.spyOn(console, 'error'); - stream.receiveInput(NaN); + await stream.receiveInput(NaN); await delay(10); expect(consoleErrorSpy).toHaveBeenCalledWith( `Unexpected read error from Supervisor "${supervisor.id}"`, @@ -46,7 +46,10 @@ describe('Supervisor', () => { const { stream } = await makeSupervisor(); const consoleErrorSpy = vi.spyOn(console, 'error'); - stream.receiveInput({ type: 'command', payload: { method: 'test' } }); + await stream.receiveInput({ + type: 'command', + payload: { method: 'test' }, + }); await delay(10); expect(consoleErrorSpy).toHaveBeenCalledOnce(); expect(consoleErrorSpy).toHaveBeenCalledWith( diff --git a/packages/kernel/src/Vat.test.ts b/packages/kernel/src/Vat.test.ts index d7d31cab1..c344334bb 100644 --- a/packages/kernel/src/Vat.test.ts +++ b/packages/kernel/src/Vat.test.ts @@ -64,7 +64,7 @@ describe('Vat', () => { vi.spyOn(vat, 'makeCapTp').mockResolvedValueOnce(undefined); await vat.init(); const consoleErrorSpy = vi.spyOn(vat.logger, 'error'); - stream.receiveInput(NaN); + await stream.receiveInput(NaN); await delay(10); expect(consoleErrorSpy).toHaveBeenCalledWith( 'Unexpected read error', @@ -97,7 +97,7 @@ describe('Vat', () => { const handleSpy = vi.spyOn(vat.streamEnvelopeReplyHandler, 'handle'); await vat.init(); const rawMessage = { type: 'command', payload: { method: 'test' } }; - stream.receiveInput(rawMessage); + await stream.receiveInput(rawMessage); await delay(10); expect(handleSpy).toHaveBeenCalledWith(rawMessage); }); diff --git a/packages/streams/src/BaseDuplexStream.test.ts b/packages/streams/src/BaseDuplexStream.test.ts index a127822bc..6424a76b1 100644 --- a/packages/streams/src/BaseDuplexStream.test.ts +++ b/packages/streams/src/BaseDuplexStream.test.ts @@ -158,12 +158,10 @@ describe('BaseDuplexStream', () => { const duplexStream = new TestDuplexStream(() => undefined); const nextP = duplexStream.next(); - const message = 42; await duplexStream.completeSynchronization(); - duplexStream.receiveInput(message); - - expect(await nextP).toStrictEqual(makePendingResult(message)); + await duplexStream.receiveInput(42); + expect(await nextP).toStrictEqual(makePendingResult(42)); }); it('reads from the reader, with input validation', async () => { @@ -172,10 +170,8 @@ describe('BaseDuplexStream', () => { typeof value === 'number', }); - const message = 42; - await duplexStream.receiveInput(message); - - expect(await duplexStream.next()).toStrictEqual(makePendingResult(message)); + await duplexStream.receiveInput(42); + expect(await duplexStream.next()).toStrictEqual(makePendingResult(42)); }); it('drains the reader in order', async () => { diff --git a/packages/streams/src/ChromeRuntimeStream.test.ts b/packages/streams/src/ChromeRuntimeStream.test.ts index a0be0a52e..ea1572641 100644 --- a/packages/streams/src/ChromeRuntimeStream.test.ts +++ b/packages/streams/src/ChromeRuntimeStream.test.ts @@ -116,6 +116,24 @@ describe('ChromeRuntimeReader', () => { expect(validateInput).toHaveBeenCalledWith(message); }); + it('throws if validateInput throws', async () => { + const { runtime, dispatchRuntimeMessage } = makeRuntime(); + const validateInput = (() => { + throw new Error('foo'); + }) as unknown as ValidateInput; + const reader = new ChromeRuntimeReader( + asChromeRuntime(runtime), + ChromeRuntimeStreamTarget.Background, + ChromeRuntimeStreamTarget.Offscreen, + { validateInput }, + ); + + const message = { foo: 'bar' }; + dispatchRuntimeMessage(message); + await expect(reader.next()).rejects.toThrow('foo'); + expect(await reader.next()).toStrictEqual(makeDoneResult()); + }); + it('ignores messages from other extensions', async () => { const { runtime, dispatchRuntimeMessage } = makeRuntime(); const reader = new ChromeRuntimeReader( @@ -227,6 +245,7 @@ describe('ChromeRuntimeReader', () => { const reader = new ChromeRuntimeReader( asChromeRuntime(runtime), ChromeRuntimeStreamTarget.Background, + ChromeRuntimeStreamTarget.Offscreen, { onEnd }, ); diff --git a/packages/streams/src/MessagePortStream.test.ts b/packages/streams/src/MessagePortStream.test.ts index ad1cb3b98..034a248a3 100644 --- a/packages/streams/src/MessagePortStream.test.ts +++ b/packages/streams/src/MessagePortStream.test.ts @@ -52,6 +52,18 @@ describe('MessagePortReader', () => { expect(validateInput).toHaveBeenCalledWith(message); }); + it('throws if validateInput throws', async () => { + const validateInput = (() => { + throw new Error('foo'); + }) as unknown as ValidateInput; + const { port1, port2 } = new MessageChannel(); + const reader = new MessagePortReader(port1, { validateInput }); + + port2.postMessage(42); + await expect(reader.next()).rejects.toThrow('foo'); + expect(await reader.next()).toStrictEqual(makeDoneResult()); + }); + it('closes the port when done', async () => { const { port1, port2 } = new MessageChannel(); const closeSpy = vi.spyOn(port1, 'close'); diff --git a/packages/streams/src/PostMessageStream.test.ts b/packages/streams/src/PostMessageStream.test.ts index 695c86378..d3dce1a20 100644 --- a/packages/streams/src/PostMessageStream.test.ts +++ b/packages/streams/src/PostMessageStream.test.ts @@ -69,6 +69,21 @@ describe('PostMessageReader', () => { expect(validateInput).toHaveBeenCalledWith(message); }); + it('throws if validateInput throws', async () => { + const { postMessageFn, setListener, removeListener } = + makePostMessageMock(); + const validateInput = (() => { + throw new Error('foo'); + }) as unknown as ValidateInput; + const reader = new PostMessageReader(setListener, removeListener, { + validateInput, + }); + + postMessageFn(42); + await expect(reader.next()).rejects.toThrow('foo'); + expect(await reader.next()).toStrictEqual(makeDoneResult()); + }); + it('removes its listener when it ends', async () => { const { postMessageFn, setListener, removeListener, listeners } = makePostMessageMock(); diff --git a/packages/streams/src/StreamMultiplexer.test.ts b/packages/streams/src/StreamMultiplexer.test.ts index 8516079fa..2e739d734 100644 --- a/packages/streams/src/StreamMultiplexer.test.ts +++ b/packages/streams/src/StreamMultiplexer.test.ts @@ -230,21 +230,23 @@ describe('StreamMultiplexer', () => { describe('writing', () => { it('writes channel messages correctly', async () => { - const [multiplex, duplex] = await TestMultiplexer.make(); + const dispatch = vi.fn(); + const duplex = await TestDuplexStream.make< + MultiplexEnvelope, + MultiplexEnvelope + >(dispatch); + const [multiplex] = await TestMultiplexer.make(duplex); const ch1 = multiplex.addChannel('1', isString, noop); const ch2 = multiplex.addChannel('2', isNumber, noop); - const writeSpy = vi.spyOn(duplex, 'write'); - await ch1.write('foo'); await ch2.write(42); - expect(writeSpy).toHaveBeenCalledTimes(2); - expect(writeSpy).toHaveBeenNthCalledWith(1, { + expect(dispatch).toHaveBeenCalledWith({ channel: '1', payload: 'foo', }); - expect(writeSpy).toHaveBeenNthCalledWith(2, { + expect(dispatch).toHaveBeenLastCalledWith({ channel: '2', payload: 42, }); diff --git a/packages/streams/src/StreamMultiplexer.ts b/packages/streams/src/StreamMultiplexer.ts index 84039ee7d..df3d36ecf 100644 --- a/packages/streams/src/StreamMultiplexer.ts +++ b/packages/streams/src/StreamMultiplexer.ts @@ -286,6 +286,6 @@ export class StreamMultiplexer { ...Array.from(this.#channels.values()).map(async (channel) => end(channel.stream), ), - ]).catch(() => undefined); + ]); } } From bac8db29b61ba39f259449737c736aabe7a2062d Mon Sep 17 00:00:00 2001 From: Erik Marks Date: Thu, 7 Nov 2024 00:20:03 +0700 Subject: [PATCH 07/10] docs(streams): Complete multiplexer documentation --- packages/streams/src/StreamMultiplexer.ts | 67 ++++++++++++++++++----- 1 file changed, 52 insertions(+), 15 deletions(-) diff --git a/packages/streams/src/StreamMultiplexer.ts b/packages/streams/src/StreamMultiplexer.ts index df3d36ecf..c643ada55 100644 --- a/packages/streams/src/StreamMultiplexer.ts +++ b/packages/streams/src/StreamMultiplexer.ts @@ -1,3 +1,23 @@ +/** + * This module provides a class and utilities for multiplexing duplex streams. A + * multiplexer is not a stream itself, but rather a wrapper around a duplex stream. + * The multiplexer provides methods for creating "channels" over the underlying stream, + * which are themselves duplex streams and may have a different message type and + * validation logic. + * + * The multiplexer is constructed in an idle state, and must be explicitly "started" + * via the `start()` or `drainAll()` methods. All channels must be added before the + * multiplexer is started. + * + * Starting the multiplexer will synchronize the underlying duplex stream, if it is + * synchronizable. Therefore, in order to prevent message loss, callers **should not** + * synchronize the underlying duplex stream before passing it to the multiplexer. For + * the same reason, the multiplexer will throw if any channels are added after it has + * started. + * + * @module StreamMultiplexer + */ + import type { Json } from '@metamask/utils'; import { isObject } from '@metamask/utils'; @@ -10,6 +30,9 @@ import type { import { BaseReader } from './BaseStream.js'; import { makeDoneResult } from './utils.js'; +/** + * A duplex stream that can maybe be synchronized. + */ type SynchronizableDuplexStream< Read extends Json, Write extends Json = Read, @@ -17,6 +40,9 @@ type SynchronizableDuplexStream< synchronize?: () => Promise; }; +/** + * The read stream implementation for {@link StreamMultiplexer} channels. + */ class ChannelReader extends BaseReader { // eslint-disable-next-line no-restricted-syntax private constructor(args: BaseReaderArgs) { @@ -33,6 +59,10 @@ class ChannelReader extends BaseReader { type ChannelName = string; +/** + * A multiplex envelope. The wrapper for all values passing through the underlying + * duplex stream of a {@link StreamMultiplexer}. + */ export type MultiplexEnvelope = { channel: ChannelName; payload: Json; @@ -55,12 +85,10 @@ export const isMultiplexEnvelope = ( type HandleRead = (value: Read) => void | Promise; -type ChannelRecord = { - channelName: ChannelName; - stream: HandledDuplexStream; - receiveInput: ReceiveInput; -}; - +/** + * A duplex stream whose `drain` method does not accept a callback. We say it is + * "handled" because in practice the callback is bound to the `drain` method. + */ export type HandledDuplexStream = Omit< DuplexStream, 'drain' @@ -74,6 +102,12 @@ enum MultiplexerStatus { Done = 2, } +type ChannelRecord = { + channelName: ChannelName; + stream: HandledDuplexStream; + receiveInput: ReceiveInput; +}; + export class StreamMultiplexer { #status: MultiplexerStatus; @@ -161,11 +195,12 @@ export class StreamMultiplexer { } /** - * Adds a channel to the multiplexer. + * Adds a channel to the multiplexer. Must be called before starting the + * multiplexer. * - * @param channelName - The channel name. - * @param validateInput - The input validator. - * @param handleRead - The channel stream's drain handler. + * @param channelName - The channel name. Must be unique. + * @param validateInput - The channel's input validator. + * @param handleRead - The channel's drain handler. * @returns The channel stream. */ addChannel( @@ -197,12 +232,14 @@ export class StreamMultiplexer { } /** - * Constructs a channel. + * Constructs a channel. Channels are objects that implement the {@link HandledDuplexStream} + * interface. Internally, they are backed up by a {@link ChannelReader} and a wrapped + * write method that forwards messages to the underlying duplex stream. * - * @param channelName - The channel name. - * @param validateInput - The input validator. - * @param handleRead - The channel stream's drain handler. - * @returns The channel stream and its receiveInput method. + * @param channelName - The channel name. Must be unique. + * @param validateInput - The channel's input validator. + * @param handleRead - The channel's drain handler. + * @returns The channel stream and its `receiveInput` method. */ #makeChannel( channelName: ChannelName, From f9d9ade10c8dd0f077f9ff4ddff982b547b8e254 Mon Sep 17 00:00:00 2001 From: Erik Marks Date: Tue, 5 Nov 2024 14:01:33 +0700 Subject: [PATCH 08/10] refactor: Replace stream envelopes with multiplex streams --- packages/extension/src/iframe.ts | 11 +- .../extension/src/kernel/VatWorkerClient.ts | 11 +- packages/kernel/src/Kernel.test.ts | 11 +- packages/kernel/src/Kernel.ts | 3 +- packages/kernel/src/Supervisor.test.ts | 53 +++-- packages/kernel/src/Supervisor.ts | 68 +++--- packages/kernel/src/Vat.test.ts | 95 ++++---- packages/kernel/src/Vat.ts | 79 +++---- packages/kernel/src/index.ts | 1 - packages/kernel/src/messages/vat.ts | 73 ++++-- packages/kernel/src/stream-envelope.test.ts | 160 ------------- packages/kernel/src/stream-envelope.ts | 79 ------- packages/kernel/src/types.ts | 6 +- .../documents/make-stream-envelope-kit.md | 124 ---------- packages/streams/src/envelope-handler.test.ts | 77 ------- packages/streams/src/envelope-handler.ts | 142 ------------ packages/streams/src/envelope-kit.test.ts | 24 -- packages/streams/src/envelope-kit.ts | 95 -------- .../streams/src/envelope-kit.types.test.ts | 214 ------------------ packages/streams/src/envelope.test.ts | 32 --- packages/streams/src/envelope.ts | 35 --- packages/streams/src/enveloper.test.ts | 89 -------- packages/streams/src/enveloper.ts | 79 ------- packages/streams/src/index.test.ts | 2 +- packages/streams/src/index.ts | 10 +- .../streams/test/envelope-kit-fixtures.ts | 34 --- packages/test-utils/src/streams.ts | 5 +- packages/utils/src/types.ts | 6 +- 28 files changed, 223 insertions(+), 1395 deletions(-) delete mode 100644 packages/kernel/src/stream-envelope.test.ts delete mode 100644 packages/kernel/src/stream-envelope.ts delete mode 100644 packages/streams/documents/make-stream-envelope-kit.md delete mode 100644 packages/streams/src/envelope-handler.test.ts delete mode 100644 packages/streams/src/envelope-handler.ts delete mode 100644 packages/streams/src/envelope-kit.test.ts delete mode 100644 packages/streams/src/envelope-kit.ts delete mode 100644 packages/streams/src/envelope-kit.types.test.ts delete mode 100644 packages/streams/src/envelope.test.ts delete mode 100644 packages/streams/src/envelope.ts delete mode 100644 packages/streams/src/enveloper.test.ts delete mode 100644 packages/streams/src/enveloper.ts delete mode 100644 packages/streams/test/envelope-kit-fixtures.ts diff --git a/packages/extension/src/iframe.ts b/packages/extension/src/iframe.ts index e8db30dd0..df7a00aaf 100644 --- a/packages/extension/src/iframe.ts +++ b/packages/extension/src/iframe.ts @@ -1,8 +1,7 @@ import { makeExo } from '@endo/exo'; import { M } from '@endo/patterns'; import { Supervisor } from '@ocap/kernel'; -import type { StreamEnvelope, StreamEnvelopeReply } from '@ocap/kernel'; -import { MessagePortDuplexStream, receiveMessagePort } from '@ocap/streams'; +import { MessagePortMultiplexer, receiveMessagePort } from '@ocap/streams'; main().catch(console.error); @@ -10,12 +9,10 @@ main().catch(console.error); * The main function for the iframe. */ async function main(): Promise { - const stream = await receiveMessagePort( + const multiplexer = await receiveMessagePort( (listener) => addEventListener('message', listener), (listener) => removeEventListener('message', listener), - ).then(async (port) => - MessagePortDuplexStream.make(port), - ); + ).then(async (port) => new MessagePortMultiplexer(port)); const bootstrap = makeExo( 'TheGreatFrangooly', @@ -23,7 +20,7 @@ async function main(): Promise { { whatIsTheGreatFrangooly: () => 'Crowned with Chaos' }, ); - const supervisor = new Supervisor({ id: 'iframe', stream, bootstrap }); + const supervisor = new Supervisor({ id: 'iframe', multiplexer, bootstrap }); console.log(supervisor.evaluate('["Hello", "world!"].join(" ");')); } diff --git a/packages/extension/src/kernel/VatWorkerClient.ts b/packages/extension/src/kernel/VatWorkerClient.ts index 5ad43edde..cb6bbf574 100644 --- a/packages/extension/src/kernel/VatWorkerClient.ts +++ b/packages/extension/src/kernel/VatWorkerClient.ts @@ -7,14 +7,12 @@ import { isVatWorkerServiceCommandReply, } from '@ocap/kernel'; import type { - StreamEnvelope, - StreamEnvelopeReply, VatWorkerService, VatId, VatWorkerServiceCommand, } from '@ocap/kernel'; -import type { DuplexStream } from '@ocap/streams'; -import { MessagePortDuplexStream } from '@ocap/streams'; +import type { DuplexStream, MultiplexEnvelope } from '@ocap/streams'; +import { isMultiplexEnvelope, MessagePortDuplexStream } from '@ocap/streams'; import type { Logger } from '@ocap/utils'; import { makeCounter, makeHandledCallback, makeLogger } from '@ocap/utils'; @@ -77,7 +75,7 @@ export class ExtensionVatWorkerClient implements VatWorkerService { async launch( vatId: VatId, - ): Promise> { + ): Promise> { return this.#sendMessage({ method: VatWorkerServiceCommandMethod.Launch, params: { vatId }, @@ -128,8 +126,9 @@ export class ExtensionVatWorkerClient implements VatWorkerService { return; } promise.resolve( - MessagePortDuplexStream.make( + new MessagePortDuplexStream( port, + isMultiplexEnvelope, ), ); break; diff --git a/packages/kernel/src/Kernel.test.ts b/packages/kernel/src/Kernel.test.ts index 7ade18203..73f219415 100644 --- a/packages/kernel/src/Kernel.test.ts +++ b/packages/kernel/src/Kernel.test.ts @@ -1,7 +1,11 @@ import '@ocap/shims/endoify'; import { VatAlreadyExistsError, VatNotFoundError } from '@ocap/errors'; -import type { MessagePortDuplexStream, DuplexStream } from '@ocap/streams'; +import type { + MessagePortDuplexStream, + DuplexStream, + MultiplexEnvelope, +} from '@ocap/streams'; import type { MockInstance } from 'vitest'; import { describe, it, expect, vi, beforeEach } from 'vitest'; @@ -12,7 +16,6 @@ import type { KernelCommandReply, VatCommand, } from './messages/index.js'; -import type { StreamEnvelope, StreamEnvelopeReply } from './stream-envelope.js'; import type { VatId, VatWorkerService } from './types.js'; import { Vat } from './Vat.js'; import { makeMapKVStore } from '../test/storage.js'; @@ -45,9 +48,7 @@ describe('Kernel', () => { launchWorkerMock = vi .spyOn(mockWorkerService, 'launch') - .mockResolvedValue( - {} as DuplexStream, - ); + .mockResolvedValue({} as DuplexStream); terminateWorkerMock = vi .spyOn(mockWorkerService, 'terminate') .mockResolvedValue(undefined); diff --git a/packages/kernel/src/Kernel.ts b/packages/kernel/src/Kernel.ts index 8e6d7f31e..d1d0c54bd 100644 --- a/packages/kernel/src/Kernel.ts +++ b/packages/kernel/src/Kernel.ts @@ -5,6 +5,7 @@ import { VatNotFoundError, toError, } from '@ocap/errors'; +import { StreamMultiplexer } from '@ocap/streams'; import type { DuplexStream } from '@ocap/streams'; import type { Logger } from '@ocap/utils'; import { makeLogger, stringify } from '@ocap/utils'; @@ -184,7 +185,7 @@ export class Kernel { throw new VatAlreadyExistsError(id); } const stream = await this.#vatWorkerService.launch(id); - const vat = new Vat({ id, stream }); + const vat = new Vat({ id, multiplexer: new StreamMultiplexer(stream) }); this.#vats.set(vat.id, vat); await vat.init(); return vat; diff --git a/packages/kernel/src/Supervisor.test.ts b/packages/kernel/src/Supervisor.test.ts index 985adceb6..eb56e6160 100644 --- a/packages/kernel/src/Supervisor.test.ts +++ b/packages/kernel/src/Supervisor.test.ts @@ -1,22 +1,30 @@ import '@ocap/shims/endoify'; +import type { MultiplexEnvelope } from '@ocap/streams'; import { delay } from '@ocap/test-utils'; -import { TestDuplexStream } from '@ocap/test-utils/streams'; +import { TestDuplexStream, TestMultiplexer } from '@ocap/test-utils/streams'; +import { stringify } from '@ocap/utils'; import { describe, it, expect, vi } from 'vitest'; import { VatCommandMethod } from './messages/index.js'; -import type { StreamEnvelope, StreamEnvelopeReply } from './stream-envelope.js'; -import * as streamEnvelope from './stream-envelope.js'; import { Supervisor } from './Supervisor.js'; -const makeSupervisor = async (): Promise<{ +const makeSupervisor = async ( + handleWrite: (input: unknown) => void | Promise = () => undefined, +): Promise<{ supervisor: Supervisor; - stream: TestDuplexStream; + stream: TestDuplexStream; }> => { const stream = await TestDuplexStream.make< - StreamEnvelope, - StreamEnvelopeReply - >(() => undefined); - return { supervisor: new Supervisor({ id: 'test-id', stream }), stream }; + MultiplexEnvelope, + MultiplexEnvelope + >(handleWrite); + return { + supervisor: new Supervisor({ + id: 'test-id', + multiplexer: new TestMultiplexer(stream), + }), + stream, + }; }; describe('Supervisor', () => { @@ -42,19 +50,25 @@ describe('Supervisor', () => { }); describe('handleMessage', () => { - it('throws if the stream envelope handler throws', async () => { - const { stream } = await makeSupervisor(); + it('throws if receiving an unexpected message', async () => { + const { supervisor, stream } = await makeSupervisor(); const consoleErrorSpy = vi.spyOn(console, 'error'); await stream.receiveInput({ - type: 'command', + channel: 'command', payload: { method: 'test' }, }); await delay(10); expect(consoleErrorSpy).toHaveBeenCalledOnce(); expect(consoleErrorSpy).toHaveBeenCalledWith( - 'Supervisor stream error:', - 'Stream envelope handler received unexpected value', + `Unexpected read error from Supervisor "${supervisor.id}"`, + new Error( + `TestMultiplexer#command: Message failed type validation:\n${stringify( + { + method: 'test', + }, + )}`, + ), ); }); @@ -89,8 +103,8 @@ describe('Supervisor', () => { }); it('handles CapTP messages', async () => { - const { supervisor } = await makeSupervisor(); - const wrapCapTpSpy = vi.spyOn(streamEnvelope, 'wrapCapTp'); + const handleWrite = vi.fn(); + const { supervisor } = await makeSupervisor(handleWrite); await supervisor.handleMessage({ id: 'v0:0', @@ -106,7 +120,7 @@ describe('Supervisor', () => { await delay(10); - const capTpAnswer = { + const capTpPayload = { type: 'CTP_RETURN', epoch: 0, answerID: 'q-1', @@ -115,7 +129,10 @@ describe('Supervisor', () => { slots: [], }, }; - expect(wrapCapTpSpy).toHaveBeenCalledWith(capTpAnswer); + expect(handleWrite).toHaveBeenCalledWith({ + channel: 'capTp', + payload: capTpPayload, + }); }); it('handles Evaluate messages', async () => { diff --git a/packages/kernel/src/Supervisor.ts b/packages/kernel/src/Supervisor.ts index 4ee979ffa..b1eba5185 100644 --- a/packages/kernel/src/Supervisor.ts +++ b/packages/kernel/src/Supervisor.ts @@ -1,6 +1,6 @@ import { makeCapTP } from '@endo/captp'; import { StreamReadError } from '@ocap/errors'; -import type { DuplexStream } from '@ocap/streams'; +import type { HandledDuplexStream, StreamMultiplexer } from '@ocap/streams'; import { stringify } from '@ocap/utils'; import type { @@ -8,24 +8,26 @@ import type { VatCommand, VatCommandReply, } from './messages/index.js'; -import { VatCommandMethod } from './messages/index.js'; -import type { StreamEnvelope, StreamEnvelopeReply } from './stream-envelope.js'; import { - makeStreamEnvelopeHandler, - wrapCapTp, - wrapStreamCommandReply, -} from './stream-envelope.js'; + isCapTpMessage, + isVatCommand, + VatCommandMethod, +} from './messages/index.js'; type SupervisorConstructorProps = { id: string; - stream: DuplexStream; + multiplexer: StreamMultiplexer; bootstrap?: unknown; }; export class Supervisor { readonly id: string; - readonly #stream: DuplexStream; + readonly #multiplexer: StreamMultiplexer; + + readonly #commandStream: HandledDuplexStream; + + readonly #capTpStream: HandledDuplexStream; readonly #defaultCompartment = new Compartment({ URL }); @@ -33,37 +35,36 @@ export class Supervisor { capTp?: ReturnType; - constructor({ id, stream, bootstrap }: SupervisorConstructorProps) { + constructor({ id, multiplexer, bootstrap }: SupervisorConstructorProps) { this.id = id; this.#bootstrap = bootstrap; - this.#stream = stream; - - const streamEnvelopeHandler = makeStreamEnvelopeHandler( - { - command: this.handleMessage.bind(this), - capTp: async (content) => this.capTp?.dispatch(content), - }, - (error) => console.error('Supervisor stream error:', error), + this.#multiplexer = multiplexer; + this.#commandStream = multiplexer.addChannel( + 'command', + isVatCommand, + this.handleMessage.bind(this), + ); + this.#capTpStream = multiplexer.addChannel( + 'capTp', + isCapTpMessage, + // eslint-disable-next-line no-void + async (content): Promise => void this.capTp?.dispatch(content), ); - this.#stream - .drain(async (value) => { - await streamEnvelopeHandler.handle(value); - }) - .catch((error) => { - console.error( - `Unexpected read error from Supervisor "${this.id}"`, - error, - ); - throw new StreamReadError({ supervisorId: this.id }, error); - }); + multiplexer.drainAll().catch((error) => { + console.error( + `Unexpected read error from Supervisor "${this.id}"`, + error, + ); + throw new StreamReadError({ supervisorId: this.id }, error); + }); } /** * Terminates the Supervisor. */ async terminate(): Promise { - await this.#stream.return(); + await this.#multiplexer.return(); } /** @@ -95,7 +96,7 @@ export class Supervisor { this.capTp = makeCapTP( 'iframe', async (content: unknown) => - this.#stream.write(wrapCapTp(content as CapTpMessage)), + this.#capTpStream.write(content as CapTpMessage), this.#bootstrap, ); await this.replyToMessage(id, { @@ -126,10 +127,11 @@ export class Supervisor { * @param payload - The payload to reply with. */ async replyToMessage( - id: VatCommand['id'], + id: VatCommandReply['id'], payload: VatCommandReply['payload'], ): Promise { - await this.#stream.write(wrapStreamCommandReply({ id, payload })); + console.log('replyToMessage', id, payload); + await this.#commandStream.write({ id, payload }); } /** diff --git a/packages/kernel/src/Vat.test.ts b/packages/kernel/src/Vat.test.ts index c344334bb..4385c2f88 100644 --- a/packages/kernel/src/Vat.test.ts +++ b/packages/kernel/src/Vat.test.ts @@ -3,20 +3,15 @@ import { VatCapTpConnectionExistsError, VatCapTpConnectionNotFoundError, } from '@ocap/errors'; +import type { MultiplexEnvelope } from '@ocap/streams'; import { delay, makePromiseKitMock } from '@ocap/test-utils'; -import { TestDuplexStream } from '@ocap/test-utils/streams'; -import { stringify } from '@ocap/utils'; +import { TestDuplexStream, TestMultiplexer } from '@ocap/test-utils/streams'; +import { makeLogger, stringify } from '@ocap/utils'; +import type { Logger } from '@ocap/utils'; import { describe, it, expect, vi } from 'vitest'; import { VatCommandMethod } from './messages/index.js'; -import type { - CapTpMessage, - VatCommand, - VatCommandReply, -} from './messages/index.js'; -import type { StreamEnvelope, StreamEnvelopeReply } from './stream-envelope.js'; -import * as streamEnvelope from './stream-envelope.js'; -import { makeStreamEnvelopeReplyHandler } from './stream-envelope.js'; +import type { VatCommand, VatCommandReply } from './messages/index.js'; import { Vat } from './Vat.js'; vi.mock('@endo/eventual-send', () => ({ @@ -27,15 +22,24 @@ vi.mock('@endo/eventual-send', () => ({ }), })); -const makeVat = async (): Promise<{ +const makeVat = async ( + logger?: Logger, +): Promise<{ vat: Vat; - stream: TestDuplexStream; + stream: TestDuplexStream; }> => { const stream = await TestDuplexStream.make< - StreamEnvelopeReply, - StreamEnvelope + MultiplexEnvelope, + MultiplexEnvelope >(() => undefined); - return { vat: new Vat({ id: 'v0', stream }), stream }; + return { + vat: new Vat({ + id: 'v0', + multiplexer: new TestMultiplexer(stream), + logger, + }), + stream, + }; }; describe('Vat', () => { @@ -63,10 +67,10 @@ describe('Vat', () => { vi.spyOn(vat, 'sendMessage').mockResolvedValueOnce(undefined); vi.spyOn(vat, 'makeCapTp').mockResolvedValueOnce(undefined); await vat.init(); - const consoleErrorSpy = vi.spyOn(vat.logger, 'error'); + const logErrorSpy = vi.spyOn(vat.logger, 'error'); await stream.receiveInput(NaN); await delay(10); - expect(consoleErrorSpy).toHaveBeenCalledWith( + expect(logErrorSpy).toHaveBeenCalledWith( 'Unexpected read error', new Error( 'TestDuplexStream: Message cannot be processed (must be JSON-serializable):\nnull', @@ -89,20 +93,6 @@ describe('Vat', () => { }); }); - describe('#receiveMessages', () => { - it('receives messages correctly', async () => { - const { vat, stream } = await makeVat(); - vi.spyOn(vat, 'sendMessage').mockResolvedValueOnce(undefined); - vi.spyOn(vat, 'makeCapTp').mockResolvedValueOnce(undefined); - const handleSpy = vi.spyOn(vat.streamEnvelopeReplyHandler, 'handle'); - await vat.init(); - const rawMessage = { type: 'command', payload: { method: 'test' } }; - await stream.receiveInput(rawMessage); - await delay(10); - expect(handleSpy).toHaveBeenCalledWith(rawMessage); - }); - }); - describe('handleMessage', () => { it('resolves the payload when the message id exists in unresolvedMessages', async () => { const { vat } = await makeVat(); @@ -120,7 +110,7 @@ describe('Vat', () => { it('logs an error when the message id does not exist in unresolvedMessages', async () => { const { vat } = await makeVat(); - const consoleErrorSpy = vi.spyOn(console, 'error'); + const logErrorSpy = vi.spyOn(vat.logger, 'error'); const nonExistentMessageId = 'v0:9'; const mockPayload: VatCommandReply['payload'] = { @@ -133,10 +123,10 @@ describe('Vat', () => { payload: mockPayload, }); - expect(consoleErrorSpy).toHaveBeenCalledWith( + expect(logErrorSpy).toHaveBeenCalledWith( `No unresolved message with id "${nonExistentMessageId}".`, ); - consoleErrorSpy.mockRestore(); + logErrorSpy.mockRestore(); }); }); @@ -170,18 +160,10 @@ describe('Vat', () => { it('creates a CapTP connection and sends CapTpInit message', async () => { const { vat } = await makeVat(); - // @ts-expect-error - streamEnvelopeReplyHandler is readonly - vat.streamEnvelopeReplyHandler = makeStreamEnvelopeReplyHandler( - {}, - console.warn, - ); const sendMessageMock = vi .spyOn(vat, 'sendMessage') .mockResolvedValueOnce(undefined); await vat.makeCapTp(); - expect( - vat.streamEnvelopeReplyHandler.contentHandlers.capTp, - ).toBeDefined(); expect(sendMessageMock).toHaveBeenCalledWith({ method: VatCommandMethod.CapTpInit, params: null, @@ -189,25 +171,26 @@ describe('Vat', () => { }); it('handles CapTP messages', async () => { - const { vat } = await makeVat(); - vi.spyOn(vat, 'sendMessage').mockResolvedValueOnce(undefined); - const wrapCapTpSpy = vi.spyOn(streamEnvelope, 'wrapCapTp'); - const consoleLogSpy = vi.spyOn(vat.logger, 'log'); + const logger = makeLogger('[test]'); + const logSpy = vi.spyOn(logger, 'log'); + const { vat, stream } = await makeVat(logger); + vi.spyOn(vat, 'sendMessage') + .mockResolvedValueOnce(undefined) + .mockResolvedValueOnce(undefined); - await vat.makeCapTp(); + await vat.init(); - const capTpQuestion = { + const capTpPayload = { type: 'CTP_BOOTSTRAP', epoch: 0, questionID: 'q-1', }; - await vat.streamEnvelopeReplyHandler.contentHandlers.capTp?.( - capTpQuestion as CapTpMessage, - ); + await stream.receiveInput({ channel: 'capTp', payload: capTpPayload }); + await delay(10); - expect(consoleLogSpy).toHaveBeenCalledWith( + expect(logSpy).toHaveBeenCalledWith( 'CapTP from vat', - stringify(capTpQuestion), + stringify(capTpPayload), ); const capTpAnswer = { @@ -219,7 +202,11 @@ describe('Vat', () => { slots: [], }, }; - expect(wrapCapTpSpy).toHaveBeenCalledWith(capTpAnswer); + + expect(logSpy).toHaveBeenLastCalledWith( + 'CapTP to vat', + stringify(capTpAnswer), + ); }); }); diff --git a/packages/kernel/src/Vat.ts b/packages/kernel/src/Vat.ts index 49b88fdf5..aa742e8a7 100644 --- a/packages/kernel/src/Vat.ts +++ b/packages/kernel/src/Vat.ts @@ -7,38 +7,37 @@ import { VatDeletedError, StreamReadError, } from '@ocap/errors'; -import type { DuplexStream, Reader } from '@ocap/streams'; +import type { HandledDuplexStream, StreamMultiplexer } from '@ocap/streams'; import type { Logger } from '@ocap/utils'; import { makeLogger, makeCounter, stringify } from '@ocap/utils'; -import { VatCommandMethod } from './messages/index.js'; +import { + isCapTpMessage, + isVatCommandReply, + VatCommandMethod, +} from './messages/index.js'; import type { CapTpMessage, CapTpPayload, VatCommandReply, VatCommand, } from './messages/index.js'; -import type { - StreamEnvelope, - StreamEnvelopeReply, - StreamEnvelopeReplyHandler, -} from './stream-envelope.js'; -import { - makeStreamEnvelopeReplyHandler, - wrapCapTp, - wrapStreamCommand, -} from './stream-envelope.js'; import type { PromiseCallbacks, VatId } from './types.js'; type VatConstructorProps = { id: VatId; - stream: DuplexStream; + multiplexer: StreamMultiplexer; + logger?: Logger | undefined; }; export class Vat { readonly id: VatConstructorProps['id']; - readonly #stream: VatConstructorProps['stream']; + readonly #multiplexer: StreamMultiplexer; + + readonly #commandStream: HandledDuplexStream; + + readonly #capTpStream: HandledDuplexStream; readonly logger: Logger; @@ -47,18 +46,25 @@ export class Vat { readonly unresolvedMessages: Map = new Map(); - readonly streamEnvelopeReplyHandler: StreamEnvelopeReplyHandler; - capTp?: ReturnType; - constructor({ id, stream }: VatConstructorProps) { + constructor({ id, multiplexer, logger }: VatConstructorProps) { this.id = id; - this.#stream = stream; - this.logger = makeLogger(`[vat ${id}]`); + this.logger = logger ?? makeLogger(`[vat ${id}]`); this.#messageCounter = makeCounter(); - this.streamEnvelopeReplyHandler = makeStreamEnvelopeReplyHandler( - { command: this.handleMessage.bind(this) }, - (error) => console.error('Vat stream error:', error), + this.#multiplexer = multiplexer; + this.#commandStream = multiplexer.addChannel( + 'command', + isVatCommandReply, + this.handleMessage.bind(this), + ); + this.#capTpStream = multiplexer.addChannel( + 'capTp', + isCapTpMessage, + async (content): Promise => { + this.logger.log('CapTP from vat', stringify(content)); + this.capTp?.dispatch(content); + }, ); } @@ -72,7 +78,7 @@ export class Vat { async handleMessage({ id, payload }: VatCommandReply): Promise { const promiseCallbacks = this.unresolvedMessages.get(id); if (promiseCallbacks === undefined) { - console.error(`No unresolved message with id "${id}".`); + this.logger.error(`No unresolved message with id "${id}".`); } else { this.unresolvedMessages.delete(id); promiseCallbacks.resolve(payload.params); @@ -86,7 +92,7 @@ export class Vat { */ async init(): Promise { /* v8 ignore next 4: Not known to be possible. */ - this.#receiveMessages(this.#stream).catch((error) => { + this.#multiplexer.drainAll().catch((error) => { this.logger.error(`Unexpected read error`, error); throw new StreamReadError({ vatId: this.id }, error); }); @@ -97,18 +103,6 @@ export class Vat { return await this.makeCapTp(); } - /** - * Receives messages from a vat. - * - * @param reader - The reader for the messages. - */ - async #receiveMessages(reader: Reader): Promise { - for await (const rawMessage of reader) { - this.logger.debug('Vat received message', rawMessage); - await this.streamEnvelopeReplyHandler.handle(rawMessage); - } - } - /** * Make a CapTP connection. * @@ -119,19 +113,12 @@ export class Vat { throw new VatCapTpConnectionExistsError(this.id); } - // Handle writes here. #receiveMessages() handles reads. const ctp = makeCapTP(this.id, async (content: unknown) => { this.logger.log('CapTP to vat', stringify(content)); - await this.#stream.write(wrapCapTp(content as CapTpMessage)); + await this.#capTpStream.write(content as CapTpMessage); }); this.capTp = ctp; - this.streamEnvelopeReplyHandler.contentHandlers.capTp = async ( - content: CapTpMessage, - ) => { - this.logger.log('CapTP from vat', stringify(content)); - ctp.dispatch(content); - }; return this.sendMessage({ method: VatCommandMethod.CapTpInit, @@ -156,7 +143,7 @@ export class Vat { * Terminates the vat. */ async terminate(): Promise { - await this.#stream.return(); + await this.#multiplexer.return(); // Handle orphaned messages for (const [messageId, promiseCallback] of this.unresolvedMessages) { @@ -176,7 +163,7 @@ export class Vat { const { promise, reject, resolve } = makePromiseKit(); const messageId = this.#nextMessageId(); this.unresolvedMessages.set(messageId, { reject, resolve }); - await this.#stream.write(wrapStreamCommand({ id: messageId, payload })); + await this.#commandStream.write({ id: messageId, payload }); return promise; } diff --git a/packages/kernel/src/index.ts b/packages/kernel/src/index.ts index e6fc58cb9..6d4566618 100644 --- a/packages/kernel/src/index.ts +++ b/packages/kernel/src/index.ts @@ -3,5 +3,4 @@ export { Kernel } from './Kernel.js'; export type { KVStore } from './kernel-store.js'; export { Vat } from './Vat.js'; export { Supervisor } from './Supervisor.js'; -export type { StreamEnvelope, StreamEnvelopeReply } from './stream-envelope.js'; export type { VatId, VatWorkerService } from './types.js'; diff --git a/packages/kernel/src/messages/vat.ts b/packages/kernel/src/messages/vat.ts index 2ba698102..18c1ff58b 100644 --- a/packages/kernel/src/messages/vat.ts +++ b/packages/kernel/src/messages/vat.ts @@ -1,29 +1,58 @@ -import { makeIdentifiedMessageKit, messageType } from './message-kit.js'; -import { vatTestCommand } from './vat-test.js'; +import { + object, + union, + literal, + refine, + string, + is, +} from '@metamask/superstruct'; +import type { Infer } from '@metamask/superstruct'; + import { isVatId } from '../types.js'; import type { VatId } from '../types.js'; -export const vatCommand = { - CapTpInit: messageType( - (send) => send === null, - (reply) => typeof reply === 'string', - ), - - ...vatTestCommand, -}; - -const vatMessageKit = makeIdentifiedMessageKit({ - source: vatCommand, - isMessageId: (value: unknown): value is `${VatId}:${number}` => - typeof value === 'string' && - /^\w+:\d+$/u.test(value) && - isVatId(value.split(':')[0]), +type VatMessageId = `${VatId}:${number}`; + +const isVatMessageId = (value: unknown): value is VatMessageId => + typeof value === 'string' && + /^\w+:\d+$/u.test(value) && + isVatId(value.split(':')[0]); + +export enum VatCommandMethod { + Evaluate = 'evaluate', + Ping = 'ping', + CapTpInit = 'capTpInit', +} + +const VatMessageIdStruct = refine(string(), 'VatMessageId', isVatMessageId); + +const VatCommandStruct = object({ + id: VatMessageIdStruct, + payload: union([ + object({ method: literal(VatCommandMethod.Evaluate), params: string() }), + object({ method: literal(VatCommandMethod.Ping), params: literal(null) }), + object({ + method: literal(VatCommandMethod.CapTpInit), + params: literal(null), + }), + ]), }); -export const VatCommandMethod = vatMessageKit.methods; +export type VatCommand = Infer; + +export type VatCommandReply = Infer; + +const VatCommandReplyStruct = object({ + id: VatMessageIdStruct, + payload: union([ + object({ method: literal(VatCommandMethod.Evaluate), params: string() }), + object({ method: literal(VatCommandMethod.Ping), params: string() }), + object({ method: literal(VatCommandMethod.CapTpInit), params: string() }), + ]), +}); -export type VatCommand = typeof vatMessageKit.send; -export const isVatCommand = vatMessageKit.sendGuard; +export const isVatCommand = (value: unknown): value is VatCommand => + is(value, VatCommandStruct); -export type VatCommandReply = typeof vatMessageKit.reply; -export const isVatCommandReply = vatMessageKit.replyGuard; +export const isVatCommandReply = (value: unknown): value is VatCommandReply => + is(value, VatCommandReplyStruct); diff --git a/packages/kernel/src/stream-envelope.test.ts b/packages/kernel/src/stream-envelope.test.ts deleted file mode 100644 index 626036a0a..000000000 --- a/packages/kernel/src/stream-envelope.test.ts +++ /dev/null @@ -1,160 +0,0 @@ -import '@ocap/shims/endoify'; - -import { stringify } from '@ocap/utils'; -import { describe, it, expect } from 'vitest'; - -import type { - CapTpMessage, - VatCommand, - VatCommandReply, -} from './messages/index.js'; -import { VatCommandMethod } from './messages/index.js'; -import { - wrapCapTp, - wrapStreamCommand, - makeStreamEnvelopeHandler, - wrapStreamCommandReply, - makeStreamEnvelopeReplyHandler, -} from './stream-envelope.js'; - -describe('StreamEnvelopeHandler', () => { - const commandContent: VatCommand = { - id: 'v0:0', - payload: { method: VatCommandMethod.Evaluate, params: '1 + 1' }, - }; - const capTpContent: CapTpMessage = { - type: 'CTP_CALL', - epoch: 0, - // Our assumptions about the form of a CapTpMessage are weak. - foo: 'bar', - }; - - const commandLabel = wrapStreamCommand(commandContent).label; - const capTpLabel = wrapCapTp(capTpContent).label; - - const testEnvelopeHandlers = { - command: async () => commandLabel, - capTp: async () => capTpLabel, - }; - - const testErrorHandler = (problem: unknown, value: unknown): never => { - throw new Error(`TEST ${String(problem)} ${stringify(value)}`); - }; - - it.each` - wrapper | content | label - ${wrapStreamCommand} | ${commandContent} | ${commandLabel} - ${wrapCapTp} | ${capTpContent} | ${capTpLabel} - `( - 'handles valid StreamEnvelope $content', - async ({ wrapper, content, label }) => { - const handler = makeStreamEnvelopeHandler( - testEnvelopeHandlers, - testErrorHandler, - ); - expect(await handler.handle(wrapper(content))).toStrictEqual(label); - }, - ); - - it('routes invalid envelopes to default error handler', async () => { - const handler = makeStreamEnvelopeHandler(testEnvelopeHandlers); - await expect( - // @ts-expect-error label is intentionally unknown - handler.handle({ label: 'unknown', content: [] }), - ).rejects.toThrow(/^Stream envelope handler received unexpected value/u); - }); - - it('routes invalid envelopes to supplied error handler', async () => { - const handler = makeStreamEnvelopeHandler( - testEnvelopeHandlers, - testErrorHandler, - ); - await expect( - // @ts-expect-error label is intentionally unknown - handler.handle({ label: 'unknown', content: [] }), - ).rejects.toThrow( - /^TEST Stream envelope handler received unexpected value/u, - ); - }); - - it('routes valid stream envelopes with an unhandled label to the error handler', async () => { - const handler = makeStreamEnvelopeHandler( - { command: testEnvelopeHandlers.command }, - testErrorHandler, - ); - await expect(handler.handle(wrapCapTp(capTpContent))).rejects.toThrow( - /^TEST Stream envelope handler received an envelope with known but unexpected label/u, - ); - }); -}); - -describe('StreamEnvelopeReplyHandler', () => { - const commandContent: VatCommandReply = { - id: 'v0:0', - payload: { method: VatCommandMethod.Evaluate, params: '2' }, - }; - const capTpContent: CapTpMessage = { - type: 'CTP_CALL', - epoch: 0, - // Our assumptions about the form of a CapTpMessage are weak. - foo: 'bar', - }; - - const commandLabel = wrapStreamCommandReply(commandContent).label; - const capTpLabel = wrapCapTp(capTpContent).label; - - const testEnvelopeHandlers = { - command: async () => commandLabel, - capTp: async () => capTpLabel, - }; - - const testErrorHandler = (problem: unknown, value: unknown): never => { - throw new Error(`TEST ${String(problem)} ${stringify(value)}`); - }; - - it.each` - wrapper | content | label - ${wrapStreamCommandReply} | ${commandContent} | ${commandLabel} - ${wrapCapTp} | ${capTpContent} | ${capTpLabel} - `( - 'handles valid StreamEnvelopeReply $content', - async ({ wrapper, content, label }) => { - const handler = makeStreamEnvelopeReplyHandler( - testEnvelopeHandlers, - testErrorHandler, - ); - expect(await handler.handle(wrapper(content))).toStrictEqual(label); - }, - ); - - it('routes invalid envelopes to default error handler', async () => { - const handler = makeStreamEnvelopeReplyHandler(testEnvelopeHandlers); - await expect( - // @ts-expect-error label is intentionally unknown - handler.handle({ label: 'unknown', content: [] }), - ).rejects.toThrow(/^Stream envelope handler received unexpected value/u); - }); - - it('routes invalid envelopes to supplied error handler', async () => { - const handler = makeStreamEnvelopeReplyHandler( - testEnvelopeHandlers, - testErrorHandler, - ); - await expect( - // @ts-expect-error label is intentionally unknown - handler.handle({ label: 'unknown', content: [] }), - ).rejects.toThrow( - /^TEST Stream envelope handler received unexpected value/u, - ); - }); - - it('routes valid stream envelopes with an unhandled label to the error handler', async () => { - const handler = makeStreamEnvelopeReplyHandler( - { command: testEnvelopeHandlers.command }, - testErrorHandler, - ); - await expect(handler.handle(wrapCapTp(capTpContent))).rejects.toThrow( - /^TEST Stream envelope handler received an envelope with known but unexpected label/u, - ); - }); -}); diff --git a/packages/kernel/src/stream-envelope.ts b/packages/kernel/src/stream-envelope.ts deleted file mode 100644 index fceb6fb5e..000000000 --- a/packages/kernel/src/stream-envelope.ts +++ /dev/null @@ -1,79 +0,0 @@ -import { makeStreamEnvelopeKit } from '@ocap/streams'; -import type { ExtractGuardType } from '@ocap/utils'; - -import { - isCapTpMessage, - isVatCommand, - isVatCommandReply, -} from './messages/index.js'; -import type { - CapTpMessage, - VatCommand, - VatCommandReply, -} from './messages/index.js'; - -// Declare and destructure the envelope kit. - -enum EnvelopeLabel { - Command = 'command', - CapTp = 'capTp', -} - -// makeStreamEnvelopeKit requires an enum of labels but typescript -// doesn't support enums as bounds on template parameters. -// -// See https://github.com/microsoft/TypeScript/issues/30611 -// -// This workaround makes something equivalently type inferenceable. -// eslint-disable-next-line @typescript-eslint/no-unused-vars -const envelopeLabels = Object.values(EnvelopeLabel); - -// For now, this envelope kit is for intial sends only - -const envelopeKit = makeStreamEnvelopeKit< - typeof envelopeLabels, - { - command: VatCommand; - capTp: CapTpMessage; - } ->({ - command: isVatCommand, - capTp: isCapTpMessage, -}); - -export type StreamEnvelope = ExtractGuardType< - typeof envelopeKit.isStreamEnvelope ->; -export type StreamEnvelopeHandler = ReturnType< - typeof envelopeKit.makeStreamEnvelopeHandler ->; - -export const wrapStreamCommand = envelopeKit.streamEnveloper.command.wrap; -export const wrapCapTp = envelopeKit.streamEnveloper.capTp.wrap; -export const { makeStreamEnvelopeHandler } = envelopeKit; - -// For now, a separate envelope kit for replies only - -const streamEnvelopeReplyKit = makeStreamEnvelopeKit< - typeof envelopeLabels, - { - command: VatCommandReply; - capTp: CapTpMessage; - } ->({ - command: isVatCommandReply, - capTp: isCapTpMessage, -}); - -export type StreamEnvelopeReply = ExtractGuardType< - typeof streamEnvelopeReplyKit.isStreamEnvelope ->; -export type StreamEnvelopeReplyHandler = ReturnType< - typeof streamEnvelopeReplyKit.makeStreamEnvelopeHandler ->; - -export const wrapStreamCommandReply = - streamEnvelopeReplyKit.streamEnveloper.command.wrap; -// Note: We don't differentiate between wrapCapTp and wrapCapTpReply -export const { makeStreamEnvelopeHandler: makeStreamEnvelopeReplyHandler } = - streamEnvelopeReplyKit; diff --git a/packages/kernel/src/types.ts b/packages/kernel/src/types.ts index bdf120041..a5f5c51b2 100644 --- a/packages/kernel/src/types.ts +++ b/packages/kernel/src/types.ts @@ -1,7 +1,5 @@ import type { PromiseKit } from '@endo/promise-kit'; -import type { DuplexStream } from '@ocap/streams'; - -import type { StreamEnvelope, StreamEnvelopeReply } from './stream-envelope.js'; +import type { DuplexStream, MultiplexEnvelope } from '@ocap/streams'; export type VatId = `v${number}`; @@ -25,7 +23,7 @@ export type VatWorkerService = { */ launch: ( vatId: VatId, - ) => Promise>; + ) => Promise>; /** * Terminate a worker identified by its vat id. * diff --git a/packages/streams/documents/make-stream-envelope-kit.md b/packages/streams/documents/make-stream-envelope-kit.md deleted file mode 100644 index 2fd579d73..000000000 --- a/packages/streams/documents/make-stream-envelope-kit.md +++ /dev/null @@ -1,124 +0,0 @@ ---- -title: Making a StreamEnvelopeKit -group: Documents -category: Guides ---- - -# makeStreamEnvelopeKit - -### Template parameters must be explicitly declared - -To ensure proper typescript inference behavior, it is necessary to explicitly declare the template parameters when calling `makeStreamEnvelopeKit`. See the [example](#example) below for the recommended declaration pattern. - -### Passing an enum as a template parameter - -Due to a [typescript limitation](https://github.com/microsoft/TypeScript/issues/30611) it is not possible to specify an enum as the expected type of a template parameter. Therefore `makeStreamEnvelopeKit` will accept template parameters which are not within its intended bounds; improperly specified template parameters will result in improper typescript inference behavior. See the [example](#example) below for the recommended declaration pattern. - -### Example declaration - -```ts -import { makeStreamEnvelopeKit } from '@ocap/streams'; - -// Declare the content types. -type FooContent = { - a: number; - b: string; -}; - -type BarContent = { - c: boolean; -}; - -// Specify envelope labels in an enum. -enum EnvelopeLabel { - Foo = 'foo', - Bar = 'bar', -} - -// Create a string[] from the EnvelopeLabel enum. -const labels = Object.values(EnvelopeLabel); - -// Make the StreamEnvelopeKit. -export const myStreamEnvelopeKit = makeStreamEnvelopeKit< - // Pass the EnvelopeLabel enum as `typeof labels`. - typeof labels, - // Specify the content type for each content label. - { - // foo matches the value 'foo' of EnvelopeLabel.Foo - foo: FooContent; - bar: BarContent; - } ->({ - // Specify the type guards for each envelope label. - foo: (value: unknown): value is FooContent => - isObject(value) && - typeof value.a === 'number' && - typeof value.b === 'string', - - // bar matches the value 'bar' of EnvelopeLabel.Bar - bar: (value: unknown): value is BarContent => - isObject(value) && typeof value.c === 'boolean', -}); -``` - -### Enveloper use - -The low level enveloping functionality is available via the included `streamEnveloper` and `isStreamEnvelope`. - -```ts -// Destructure your new envelope kit. -const { streamEnveloper, isStreamEnvelope } = myStreamEnvelopeKit; - -// Wrap some FooContent. -const envelope = streamEnveloper.foo.wrap({ - a: 1, - b: 'one', -}); - -// Protect your assumptions with the supplied type guard. -if (isStreamEnvelope(envelope)) { - // ~~~ Unwrap your envelope right away! ~~~ - const content = streamEnveloper[envelope.label].unwrap(envelope); -} -``` - -### Handler use - -If you know in advance how you plan to handle with your envelopes, you can let a `StreamEnvelopeHandler` do the checking and unwrapping for you. - -```ts -// Destructure the maker from the kit. -const { makeStreamEnvelopeHandler } = myStreamEnvelopeKit; - -// Declare how you want your envelope labels handled. -const streamEnvelopeHandler = makeStreamEnvelopeHandler( - { - // The content type is automatically inferred in the declaration. - foo: async (content) => { - await delay(content.a); - return content.b; - }, - bar: async (content) => (content.c ? 'yes' : 'no'), - }, - // The optional errorHandler can throw or return. - // If unspecified, the default behavior is to throw. - (reason, value) => { - if (reason.match(/unexpected value/u)) { - throw new Error(`[myStreamError] ${reason}`); - } - return ['[myStreamWarning]', reason, value]; - }, -); - -// Read messages from an @ocap/streams Reader. -for await (const newMessage of myStreamReader) { - // And handle the message. - await streamEnvelopeHandler - .handle(newMessage) - // If the errorHandler throws, you can catch it here. - .catch(console.error) - // Otherwise, the promise resolves to the value returned by - // its appropriate content handler, or by the errorHandler. - .then(console.log); -} -``` diff --git a/packages/streams/src/envelope-handler.test.ts b/packages/streams/src/envelope-handler.test.ts deleted file mode 100644 index 2269153b4..000000000 --- a/packages/streams/src/envelope-handler.test.ts +++ /dev/null @@ -1,77 +0,0 @@ -import { describe, expect, it } from 'vitest'; - -import { makeStreamEnvelopeHandler } from './envelope-handler.js'; -import { - barContent, - fooContent, - isStreamEnvelope, - Label, - streamEnveloper, -} from '../test/envelope-kit-fixtures.js'; - -describe('StreamEnvelopeHandler', () => { - const testEnvelopeHandlers = { - foo: async () => Label.Foo, - bar: async () => Label.Bar, - }; - - const testErrorHandler = (problem: unknown): never => { - throw new Error(`TEST ${String(problem)}`); - }; - - it.each` - wrapper | content | label - ${streamEnveloper.foo.wrap} | ${fooContent} | ${Label.Foo} - ${streamEnveloper.bar.wrap} | ${barContent} | ${Label.Bar} - `('handles valid StreamEnvelopes', async ({ wrapper, content, label }) => { - const handler = makeStreamEnvelopeHandler( - streamEnveloper, - isStreamEnvelope, - testEnvelopeHandlers, - testErrorHandler, - ); - console.debug(wrapper(content)); - expect(await handler.handle(wrapper(content))).toStrictEqual(label); - }); - - it('routes invalid envelopes to default error handler', async () => { - const handler = makeStreamEnvelopeHandler( - streamEnveloper, - isStreamEnvelope, - testEnvelopeHandlers, - ); - await expect( - // @ts-expect-error label is intentionally unknown - handler.handle({ label: 'unknown', content: [] }), - ).rejects.toThrow(/^Stream envelope handler received unexpected value/u); - }); - - it('routes invalid envelopes to supplied error handler', async () => { - const handler = makeStreamEnvelopeHandler( - streamEnveloper, - isStreamEnvelope, - testEnvelopeHandlers, - testErrorHandler, - ); - await expect( - // @ts-expect-error label is intentionally unknown - handler.handle({ label: 'unknown', content: [] }), - ).rejects.toThrow( - /^TEST Stream envelope handler received unexpected value/u, - ); - }); - - it('routes valid stream envelopes with an unhandled label to the error handler', async () => { - const handler = makeStreamEnvelopeHandler( - streamEnveloper, - isStreamEnvelope, - { foo: testEnvelopeHandlers.foo }, - testErrorHandler, - ); - await expect( - handler.handle(streamEnveloper.bar.wrap(barContent)), - ).rejects.toThrow( - /^TEST Stream envelope handler received an envelope with known but unexpected label/u, - ); - }); -}); diff --git a/packages/streams/src/envelope-handler.ts b/packages/streams/src/envelope-handler.ts deleted file mode 100644 index 9bbdf2da5..000000000 --- a/packages/streams/src/envelope-handler.ts +++ /dev/null @@ -1,142 +0,0 @@ -import { stringify } from '@ocap/utils'; - -import type { StreamEnvelope } from './envelope.js'; -import type { StreamEnveloper } from './enveloper.js'; -import type { TypeMap } from './utils/generics.js'; - -/** - * A handler for automatically unwrapping stream envelopes and handling their content. - */ -export type StreamEnvelopeHandler< - Labels extends readonly string[], - ContentMap extends TypeMap, -> = { - /** - * Checks an unknown value for envelope labels, applying the label's handler - * if known, and applying the error handler if the label is not handled or if - * the content did not pass the envelope's type guard. - * - * @template Envelope - The type of the envelope. - * @param envelope - The envelope to handle. - * @returns The result of the handler. - */ - handle: >( - envelope: Envelope, - ) => Promise; - /** - * The bag of async content handlers labeled with the {@link EnvelopeLabel} they handle. - */ - contentHandlers: StreamEnvelopeContentHandlerBag; - /** - * The error handler for the stream envelope handler. - */ - errorHandler: StreamEnvelopeErrorHandler; -}; - -/** - * A handler for a specific stream envelope label. - */ -type StreamEnvelopeContentHandler< - EnvelopeLabel extends string, - ContentMap extends TypeMap, - Label extends EnvelopeLabel, -> = (content: ContentMap[Label]) => Promise; - -/** - * An object with {@link EnvelopeLabel} keys mapping to an appropriate {@link StreamEnvelopeContentHandler}. - * If the stream envelope handler encounters a well-formed stream envelope without a defined handler, - * the envelope will be passed to the {@link ErrorHandler}. - */ -export type StreamEnvelopeContentHandlerBag< - Labels extends readonly string[], - ContentMap extends TypeMap, -> = { - [Label in Labels[number]]?: (content: ContentMap[Label]) => Promise; -}; - -/** - * A handler for stream envelope parsing errors. - * If the {@link StreamEnvelopeHandler} encounters an error while parsing the supplied value, - * it will pass the reason and value to the error handler. - */ -export type StreamEnvelopeErrorHandler = ( - reason: string, - value: unknown, -) => unknown; - -/** - * The default handler for stream envelope parsing errors. - * - * @param reason - The reason for the error. - * @param value - The value that caused the error. - */ -const defaultStreamEnvelopeErrorHandler: StreamEnvelopeErrorHandler = ( - reason, - value, -) => { - throw new Error(`${reason} ${stringify(value)}`); -}; - -/** - * Makes a {@link StreamEnvelopeHandler} which handles an unknown value. - * - * If the supplied value is a valid envelope with a defined {@link StreamEnvelopeHandler}, - * the stream envelope handler will return whatever the defined handler returns. - * - * If the stream envelope handler is passed a well-formed stream envelope without a defined handler, - * an explanation and the envelope will be passed to the supplied {@link StreamEnvelopeErrorHandler}. - * - * If the stream envelope handler encounters an error while parsing the supplied value, - * it will pass the reason and value to the supplied {@link StreamEnvelopeErrorHandler}. - * - * If no error handler is supplied, the default error handling behavior is to throw. - * - * @param streamEnveloper - A {@link StreamEnveloper} made with the same Labels. - * @param isStreamEnvelope - A type guard which identifies stream envelopes. - * @param contentHandlers - A bag of async content handlers labeled with the {@link EnvelopeLabel} they handle. - * @param errorHandler - An optional synchronous error handler. - * @returns The stream envelope handler. - */ -export const makeStreamEnvelopeHandler = < - Labels extends readonly string[], - ContentMap extends TypeMap, ->( - streamEnveloper: StreamEnveloper, - isStreamEnvelope: ( - value: unknown, - ) => value is StreamEnvelope, - contentHandlers: StreamEnvelopeContentHandlerBag, - errorHandler: StreamEnvelopeErrorHandler = defaultStreamEnvelopeErrorHandler, -): StreamEnvelopeHandler => { - return { - handle: async (value: unknown) => { - if (!isStreamEnvelope(value)) { - return errorHandler( - 'Stream envelope handler received unexpected value', - value, - ); - } - const envelope = value; - const handler = contentHandlers[envelope.label] as - | StreamEnvelopeContentHandler< - Labels[number], - ContentMap, - typeof envelope.label - > - | undefined; - const enveloper = streamEnveloper[envelope.label]; - if (!handler || !enveloper) { - console.debug(`handler: ${stringify(handler)}`); - console.debug(`enveloper: ${stringify(enveloper)}`); - return errorHandler( - 'Stream envelope handler received an envelope with known but unexpected label', - envelope, - ); - } - const content = enveloper.unwrap(envelope); - return await handler(content); - }, - contentHandlers, - errorHandler, - }; -}; diff --git a/packages/streams/src/envelope-kit.test.ts b/packages/streams/src/envelope-kit.test.ts deleted file mode 100644 index f4eee1019..000000000 --- a/packages/streams/src/envelope-kit.test.ts +++ /dev/null @@ -1,24 +0,0 @@ -import { describe, expect, it } from 'vitest'; - -import { makeStreamEnvelopeKit } from './envelope-kit.js'; -import type { - Bar, - ContentMap, - Foo, - labels, -} from '../test/envelope-kit-fixtures.js'; - -describe('makeStreamEnvelopeKit', () => { - it.each` - property - ${'streamEnveloper'} - ${'isStreamEnvelope'} - ${'makeStreamEnvelopeHandler'} - `('has the expected property: $property', ({ property }) => { - const streamEnvelopeKit = makeStreamEnvelopeKit({ - foo: (value: unknown): value is Foo => true, - bar: (value: unknown): value is Bar => true, - }); - expect(streamEnvelopeKit).toHaveProperty(property); - }); -}); diff --git a/packages/streams/src/envelope-kit.ts b/packages/streams/src/envelope-kit.ts deleted file mode 100644 index ad6dc0d5a..000000000 --- a/packages/streams/src/envelope-kit.ts +++ /dev/null @@ -1,95 +0,0 @@ -import type { - StreamEnvelopeContentHandlerBag, - StreamEnvelopeErrorHandler, - StreamEnvelopeHandler, -} from './envelope-handler.js'; -import { makeStreamEnvelopeHandler as makeHandler } from './envelope-handler.js'; -import type { StreamEnvelope } from './envelope.js'; -import { isLabeled } from './envelope.js'; -import type { - Enveloper, - StreamEnveloper, - StreamEnveloperGuards, -} from './enveloper.js'; -import { makeStreamEnveloper } from './enveloper.js'; -import type { TypeMap } from './utils/generics.js'; - -export type MakeStreamEnvelopeHandler< - Labels extends readonly string[], - ContentMap extends TypeMap, -> = ( - contentHandlers: StreamEnvelopeContentHandlerBag, - errorHandler?: StreamEnvelopeErrorHandler, -) => StreamEnvelopeHandler; - -export type StreamEnvelopeKit< - Labels extends readonly string[], - ContentMap extends TypeMap, -> = { - streamEnveloper: StreamEnveloper; - isStreamEnvelope: ( - value: unknown, - ) => value is StreamEnvelope; - makeStreamEnvelopeHandler: MakeStreamEnvelopeHandler; -}; - -/** - * Make a {@link StreamEnvelopeKit}. - * The template parameters must be explicitly declared. See tutorial for suggested declaration pattern. - * - * @tutorial documents/make-stream-envelope-kit.md - An example showing how to specify the template parameters, including how to pass an enum type as a template parameter. - * @template Labels - An enum of envelope labels. WARNING: if specified improperly, typescript inference fails. See referenced tutorial. - * @template Content - An object type mapping the specified labels to the type of content they label. - * @param guards - An object mapping the specified envelope labels to a type guard of their contents. - * @returns The {@link StreamEnvelopeKit}. - */ -export const makeStreamEnvelopeKit = < - Labels extends string[], - ContentMap extends TypeMap, ->( - guards: StreamEnveloperGuards, -): StreamEnvelopeKit => { - const streamEnveloper = makeStreamEnveloper(guards); - const isStreamEnvelope = ( - value: unknown, - ): value is StreamEnvelope => - isLabeled(value) && - ( - Object.values(streamEnveloper) as Enveloper[] - ).some((enveloper) => enveloper.check(value)); - - /** - * Makes a {@link StreamEnvelopeHandler} which handles an unknown value. - * - * If the supplied value is a valid envelope with a defined {@link StreamEnvelopeHandler}, - * the stream envelope handler will return whatever the defined handler returns. - * - * If the stream envelope handler is passed a well-formed stream envelope without a defined handler, - * an explanation and the envelope will be passed to the supplied {@link StreamEnvelopeErrorHandler}. - * - * If the stream envelope handler encounters an error while parsing the supplied value, - * it will pass the reason and value to the supplied {@link StreamEnvelopeErrorHandler}. - * - * If no error handler is supplied, the default error handling behavior is to throw. - * - * @param contentHandlers - A bag of async content handlers labeled with the {@link EnvelopeLabel} they handle. - * @param errorHandler - An optional synchronous error handler. - * @returns The stream envelope handler. - */ - const makeStreamEnvelopeHandler = ( - contentHandlers: StreamEnvelopeContentHandlerBag, - errorHandler?: StreamEnvelopeErrorHandler, - ): StreamEnvelopeHandler => - makeHandler( - streamEnveloper, - isStreamEnvelope, - contentHandlers, - errorHandler, - ); - - return { - streamEnveloper, - isStreamEnvelope, - makeStreamEnvelopeHandler, - }; -}; diff --git a/packages/streams/src/envelope-kit.types.test.ts b/packages/streams/src/envelope-kit.types.test.ts deleted file mode 100644 index 76ed93812..000000000 --- a/packages/streams/src/envelope-kit.types.test.ts +++ /dev/null @@ -1,214 +0,0 @@ -import { describe, expect, it } from 'vitest'; - -import { makeStreamEnvelopeKit } from './envelope-kit.js'; -import type { - Bar, - ContentMap, - Foo, - labels, -} from '../test/envelope-kit-fixtures.js'; -import { - makeStreamEnvelopeHandler as kitMakeStreamEnvelopeHandler, - isStreamEnvelope, - Label, - streamEnveloper, - fooContent, - barContent, -} from '../test/envelope-kit-fixtures.js'; - -const inferNumber = (value: number): number => value; -const inferString = (value: string): string => value; -const inferBoolean = (value: boolean): boolean => value; - -describe('makeStreamEnvelopeKit', () => { - it('causes a typescript error when supplying typeguard keys not matching the label type', () => { - // @ts-expect-error the bar key is missing - makeStreamEnvelopeKit({ - foo: (value: unknown): value is Foo => true, - }); - makeStreamEnvelopeKit({ - foo: (value: unknown): value is Foo => true, - bar: (value: unknown): value is Bar => true, - // @ts-expect-error the qux key is not included in labels - qux: (value: unknown): value is 'qux' => false, - }); - }); - - describe('kitted makeStreamEnvelopeHandler', () => { - it('provides proper typescript inferences', () => { - // all label arguments are optional - kitMakeStreamEnvelopeHandler({}); - // bar is optional - kitMakeStreamEnvelopeHandler({ - foo: async (content) => { - inferNumber(content.a); - // @ts-expect-error a is not a string - inferString(content.a); - // @ts-expect-error b is not a number - inferNumber(content.b); - inferString(content.b); - // @ts-expect-error c is undefined - value.content.c; - }, - }); - // keys not included in labels are forbidden - kitMakeStreamEnvelopeHandler({ - // @ts-expect-error the qux key is not included in labels - qux: async (content: any) => content, - }); - }); - }); -}); - -describe('isStreamEnvelope', () => { - it('provides proper typescript inferences', () => { - const value: any = null; - if (isStreamEnvelope(value)) { - switch (value.label) { - case Label.Foo: - inferNumber(value.content.a); - // @ts-expect-error a is not a string - inferString(value.content.a); - // @ts-expect-error b is not a number - inferNumber(value.content.b); - inferString(value.content.b); - // @ts-expect-error c is undefined - value.content.c; - break; - case Label.Bar: - // @ts-expect-error a is undefined - value.content.a; - // @ts-expect-error a is undefined - value.content.b; - inferBoolean(value.content.c); - break; - default: // unreachable - // @ts-expect-error label options are exhausted - value.label; - } - } - }); -}); - -describe('StreamEnveloper', () => { - describe('check', () => { - it('provides proper typescript inferences', () => { - const envelope: any = null; - if (streamEnveloper.foo.check(envelope)) { - inferNumber(envelope.content.a); - // @ts-expect-error a is not a string - inferString(envelope.content.a); - // @ts-expect-error b is not a number - inferNumber(envelope.content.b); - inferString(envelope.content.b); - // @ts-expect-error c is not defined - envelope.content.c; - switch (envelope.label) { - case Label.Foo: - // eslint-disable-next-line vitest/no-conditional-expect - expect(envelope.label).toMatch(Label.Foo); - break; - // @ts-expect-error label is Label.Foo - case Label.Bar: // unreachable - // @ts-expect-error label is inferred to be never - envelope.label.length; - break; - default: // unreachable - // @ts-expect-error label is inferred to be never - envelope.label.length; - } - } - - if (streamEnveloper.bar.check(envelope)) { - // @ts-expect-error a is not defined - envelope.content.a; - // @ts-expect-error b is not defined - envelope.content.b; - inferBoolean(envelope.content.c); - switch (envelope.label) { - // @ts-expect-error label is Label.Bar - case Label.Foo: // unreachable - // @ts-expect-error label is inferred to be never - envelope.label.length; - break; - case Label.Bar: - // eslint-disable-next-line vitest/no-conditional-expect - expect(envelope.label).toMatch(Label.Bar); - break; - default: // unreachable - // @ts-expect-error label is inferred to be never - envelope.label.length; - } - } - }); - }); - - describe('wrap', () => { - it('provides proper typescript inferences', () => { - streamEnveloper.foo.wrap(fooContent); - // @ts-expect-error foo rejects barContent - streamEnveloper.foo.wrap(barContent); - // @ts-expect-error bar rejects fooContent - streamEnveloper.bar.wrap(fooContent); - streamEnveloper.bar.wrap(barContent); - }); - }); - - describe('unwrap', () => { - it('provides proper typescript inferences', () => { - const envelope: any = null; - try { - const content = streamEnveloper.foo.unwrap(envelope); - - inferNumber(content.a); - // @ts-expect-error a is not a string - inferString(content.a); - // @ts-expect-error b is not a number - inferNumber(content.b); - inferString(content.b); - // @ts-expect-error c is undefined - content.c; - } catch { - undefined; - } - - try { - // @ts-expect-error envelope was already inferred to be Envelope - content = streamEnveloper.bar.unwrap(envelope); - } catch { - undefined; - } - }); - }); - - describe('label', () => { - it('provides proper typescript inferences', () => { - const fooEnveloper: any = streamEnveloper.foo; - const inferFooEnveloper = ( - enveloper: typeof streamEnveloper.foo, - ): unknown => enveloper; - const inferBarEnveloper = ( - enveloper: typeof streamEnveloper.bar, - ): unknown => enveloper; - - type Enveloper = (typeof streamEnveloper)[keyof typeof streamEnveloper]; - const ambiguousEnveloper = fooEnveloper as Enveloper; - - switch (ambiguousEnveloper.label) { - case Label.Foo: - inferFooEnveloper(ambiguousEnveloper); - // @ts-expect-error label = Label.Foo implies ambiguousEnveloper is a FooEnveloper - inferBarEnveloper(ambiguousEnveloper); - break; - case Label.Bar: - // @ts-expect-error label = Label.Bar implies ambiguousEnveloper is a BarEnveloper - inferFooEnveloper(ambiguousEnveloper); - inferBarEnveloper(ambiguousEnveloper); - break; - default: // unreachable - // @ts-expect-error label options are exhausted - ambiguousEnveloper.label; - } - }); - }); -}); diff --git a/packages/streams/src/envelope.test.ts b/packages/streams/src/envelope.test.ts deleted file mode 100644 index a7ef53db7..000000000 --- a/packages/streams/src/envelope.test.ts +++ /dev/null @@ -1,32 +0,0 @@ -import { describe, expect, it } from 'vitest'; - -import type { Foo } from '../test/envelope-kit-fixtures.js'; -import { - barContent, - fooContent, - isStreamEnvelope, - streamEnveloper, -} from '../test/envelope-kit-fixtures.js'; - -describe('isStreamEnvelope', () => { - it.each` - value - ${streamEnveloper.foo.wrap(fooContent)} - ${streamEnveloper.bar.wrap(barContent)} - `('returns true for valid envelopes: $value', ({ value }) => { - expect(isStreamEnvelope(value)).toBe(true); - }); - - it.each` - value - ${null} - ${true} - ${[]} - ${{}} - ${fooContent} - ${{ id: '0x5012C312312' }} - ${streamEnveloper.foo.wrap(barContent as unknown as Foo)} - `('returns false for invalid values: $value', ({ value }) => { - expect(isStreamEnvelope(value)).toBe(false); - }); -}); diff --git a/packages/streams/src/envelope.ts b/packages/streams/src/envelope.ts deleted file mode 100644 index 5a53f4f2f..000000000 --- a/packages/streams/src/envelope.ts +++ /dev/null @@ -1,35 +0,0 @@ -// Envelope types and type guards. - -import { isObject } from '@metamask/utils'; - -import type { TypeMap } from './utils/generics.js'; - -export type Envelope