feat!: add stream muxer close (#254)

BREAKING CHANGE: StreamMuxer now has a `close` method

packages/interface-mocks/src/muxer.ts

@@ -266,6 +266,7 @@ class MockMuxer implements StreamMuxer {
   public name: string
   public protocol: string = '/mock-muxer/1.0.0'
 
+  private readonly closeController: AbortController
   private readonly registryInitiatorStreams: Map<string, MuxedStream>
   private readonly registryRecipientStreams: Map<string, MuxedStream>
   private readonly options: StreamMuxerInit
@@ -279,17 +280,11 @@ class MockMuxer implements StreamMuxer {
     this.registryRecipientStreams = new Map()
     this.log('create muxer')
     this.options = init ?? { direction: 'inbound' }
+    this.closeController = new AbortController()
     // receives data from the muxer at the other end of the stream
     this.source = this.input = pushable({
       onEnd: (err) => {
-        this.log('closing muxed streams')
-        for (const stream of this.streams) {
-          if (err == null) {
-            stream.close()
-          } else {
-            stream.abort(err)
-          }
-        }
+        this.close(err)
       }
     })
 
@@ -303,7 +298,7 @@ class MockMuxer implements StreamMuxer {
   async sink (source: Source<Uint8Array>) {
     try {
       await pipe(
-        source,
+        abortableSource(source, this.closeController.signal),
         (source) => map(source, buf => uint8ArrayToString(buf)),
         ndjson.parse,
         async (source) => {
@@ -366,6 +361,9 @@ class MockMuxer implements StreamMuxer {
   }
 
   newStream (name?: string) {
+    if (this.closeController.signal.aborted) {
+      throw new Error('Muxer already closed')
+    }
     this.log('newStream %s', name)
     const storedStream = this.createStream(name, 'initiator')
     this.registryInitiatorStreams.set(storedStream.stream.id, storedStream)
@@ -399,6 +397,19 @@ class MockMuxer implements StreamMuxer {
 
     return muxedStream
   }
+
+  close (err?: Error): void {
+    if (this.closeController.signal.aborted) return
+    this.log('closing muxed streams')
+
+    if (err == null) {
+      this.streams.forEach(s => s.close())
+    } else {
+      this.streams.forEach(s => s.abort(err))
+    }
+    this.closeController.abort()
+    this.input.end(err)
+  }
 }
 
 class MockMuxerFactory implements StreamMuxerFactory {

packages/interface-stream-muxer-compliance-tests/src/close-test.ts

@@ -68,6 +68,120 @@ export default (common: TestSetup<StreamMuxerFactory>) => {
       expect(dialer.streams).to.have.lengthOf(0)
     })
 
+    it('calling close closes streams', async () => {
+      let openedStreams = 0
+      const expectedStreams = 5
+      const dialerFactory = await common.setup()
+      const dialer = dialerFactory.createStreamMuxer({ direction: 'outbound' })
+
+      // Listener is echo server :)
+      const listenerFactory = await common.setup()
+      const listener = listenerFactory.createStreamMuxer({
+        direction: 'inbound',
+        onIncomingStream: (stream) => {
+          openedStreams++
+          void pipe(stream, stream)
+        }
+      })
+
+      const p = duplexPair<Uint8Array>()
+      void pipe(p[0], dialer, p[0])
+      void pipe(p[1], listener, p[1])
+
+      const streams = Array(expectedStreams).fill(0).map(() => dialer.newStream())
+
+      void Promise.all(
+        streams.map(async stream => {
+          return await pipe(
+            infiniteRandom,
+            stream,
+            drain
+          )
+        })
+      )
+
+      expect(dialer.streams, 'dialer - number of opened streams should match number of calls to newStream').to.have.lengthOf(expectedStreams)
+
+      // Pause, and then close the dialer
+      await delay(50)
+
+      dialer.close()
+
+      expect(openedStreams, 'listener - number of opened streams should match number of calls to newStream').to.have.equal(expectedStreams)
+      expect(dialer.streams, 'all tracked streams should be deleted after the muxer has called close').to.have.lengthOf(0)
+    })
+
+    it('calling close with an error aborts streams', async () => {
+      let openedStreams = 0
+      const expectedStreams = 5
+      const dialerFactory = await common.setup()
+      const dialer = dialerFactory.createStreamMuxer({ direction: 'outbound' })
+
+      // Listener is echo server :)
+      const listenerFactory = await common.setup()
+      const listener = listenerFactory.createStreamMuxer({
+        direction: 'inbound',
+        onIncomingStream: (stream) => {
+          openedStreams++
+          void pipe(stream, stream)
+        }
+      })
+
+      const p = duplexPair<Uint8Array>()
+      void pipe(p[0], dialer, p[0])
+      void pipe(p[1], listener, p[1])
+
+      const streams = Array(expectedStreams).fill(0).map(() => dialer.newStream())
+
+      const streamPipes = streams.map(async stream => {
+        return await pipe(
+          infiniteRandom,
+          stream,
+          drain
+        )
+      })
+
+      expect(dialer.streams, 'dialer - number of opened streams should match number of calls to newStream').to.have.lengthOf(expectedStreams)
+
+      // Pause, and then close the dialer
+      await delay(50)
+
+      // close _with an error_
+      dialer.close(new Error())
+
+      const timeoutError = new Error('timeout')
+      for (const pipe of streamPipes) {
+        try {
+          await Promise.race([
+            pipe,
+            new Promise((_resolve, reject) => setTimeout(() => reject(timeoutError), 20))
+          ])
+          expect.fail('stream pipe with infinite source should never return')
+        } catch (e) {
+          if (e === timeoutError) {
+            expect.fail('expected stream pipe to throw an error after muxer closed with error')
+          }
+        }
+      }
+
+      expect(openedStreams, 'listener - number of opened streams should match number of calls to newStream').to.have.equal(expectedStreams)
+      expect(dialer.streams, 'all tracked streams should be deleted after the muxer has called close').to.have.lengthOf(0)
+    })
+
+    it('calling newStream after close throws an error', async () => {
+      const dialerFactory = await common.setup()
+      const dialer = dialerFactory.createStreamMuxer({ direction: 'outbound' })
+
+      dialer.close()
+
+      try {
+        dialer.newStream()
+        expect.fail('newStream should throw if called after close')
+      } catch (e) {
+        expect(dialer.streams, 'closed muxer should have no streams').to.have.lengthOf(0)
+      }
+    })
+
     it('closing one of the muxed streams doesn\'t close others', async () => {
       const p = duplexPair<Uint8Array>()
       const dialerFactory = await common.setup()

packages/interface-stream-muxer/src/index.ts

@@ -19,6 +19,11 @@ export interface StreamMuxer extends Duplex<Uint8Array> {
    * provided, the id of the stream will be used.
    */
   newStream: (name?: string) => Stream
+
+  /**
+   * Close or abort all tracked streams and stop the muxer
+   */
+  close: (err?: Error) => void
 }
 
 export interface StreamMuxerInit extends AbortOptions {