feat: close read and write streams

README.md

@@ -139,11 +139,17 @@ In addition to `sink` and `source` properties, this stream also has the followin
 
 #### `stream.close()`
 
-Closes the stream for **reading**. If iterating over the source of this stream in a `for await of` loop, it will return (exit the loop) after any buffered data has been consumed.
+Closes the stream for **reading** and **writing**. If iterating over the source of this stream in a `for await of` loop, it will return (exit the loop) after any buffered data has been consumed.
+
+#### `stream.closeRead()`
+
+Closes the stream for **reading**, but still allows writing. This should not typically be called by application code, but may be used for one way, push only streams.
 
 This function is called automatically by the muxer when it receives a `CLOSE` message from the remote.
 
-The source will return normally, the sink will continue to consume.
+#### `stream.closeWrite()`
+
+Closes the stream for **writing**, but still allows reading. This is useful for when you only ever wish to read from a stream.
 
 #### `stream.abort([err])`
 

package.json

@@ -46,7 +46,7 @@
     "aegir": "^33.1.0",
     "cborg": "^1.2.1",
     "iso-random-stream": "^2.0.0",
-    "libp2p-interfaces": "^0.10.0",
+    "libp2p-interfaces": "github:filoozom/js-libp2p-interfaces#close-read-write",
     "p-defer": "^3.0.0",
     "random-int": "^2.0.0",
     "streaming-iterables": "^5.0.4",

src/mplex.js

@@ -4,31 +4,41 @@ const pipe = require('it-pipe')
 const pushable = require('it-pushable')
 const log = require('debug')('libp2p:mplex')
 const abortable = require('abortable-iterator')
+const errCode = require('err-code')
 const Coder = require('./coder')
 const restrictSize = require('./restrict-size')
 const { MessageTypes, MessageTypeNames } = require('./message-types')
 const createStream = require('./stream')
 
+/**
+ * @typedef {import('libp2p-interfaces/src/stream-muxer/types').MuxedStream} MuxedStream
+ * @typedef {import('libp2p-interfaces/src/stream-muxer/types').Sink} Sink
+ */
+
+/**
+ * @typedef {object} MplexOptions
+ * @property {(stream: MuxedStream) => void} [onStream] - Called whenever an inbound stream is created
+ * @property {(stream: MuxedStream) => void} [onStreamEnd] - Called whenever a stream ends
+ * @property {AbortSignal} [signal] - An AbortController signal
+ */
+
 class Mplex {
   /**
    * @class
-   * @param {object} options
-   * @param {function(*)} options.onStream - Called whenever an inbound stream is created
-   * @param {function(*)} options.onStreamEnd - Called whenever a stream ends
-   * @param {AbortSignal} options.signal - An AbortController signal
+   * @param {MplexOptions} options
    */
-  constructor (options) {
+  constructor (options = {}) {
     options = options || {}
     options = typeof options === 'function' ? { onStream: options } : options
 
     this._streamId = 0
     this._streams = {
       /**
-       * @type {Map<number, *>} Stream to ids map
+       * @type {Map<number, MuxedStream>} Stream to ids map
        */
       initiators: new Map(),
       /**
-       * @type {Map<number, *>} Stream to ids map
+       * @type {Map<number, MuxedStream>} Stream to ids map
        */
       receivers: new Map()
     }
@@ -58,7 +68,7 @@ class Mplex {
   /**
    * Returns a Map of streams and their ids
    *
-   * @returns {Map<number,*>}
+   * @returns {Map<number, MuxedStream>}
    */
   get streams () {
     // Inbound and Outbound streams may have the same ids, so we need to make those unique
@@ -90,10 +100,10 @@ class Mplex {
    * Called whenever an inbound stream is created
    *
    * @private
-   * @param {*} options
+   * @param {object} options
    * @param {number} options.id
    * @param {string} options.name
-   * @returns {*} A muxed stream
+   * @returns MuxedStream} A muxed stream
    */
   _newReceiverStream ({ id, name }) {
     const registry = this._streams.receivers
@@ -108,25 +118,32 @@ class Mplex {
    * @param {number} options.id
    * @param {string} options.name
    * @param {string} options.type
-   * @param {Map<number, *>} options.registry - A map of streams to their ids
-   * @returns {*} A muxed stream
+   * @param {Map<number, MuxedStream>} options.registry - A map of streams to their ids
+   * @returns {MuxedStream} A muxed stream
    */
   _newStream ({ id, name, type, registry }) {
     if (registry.has(id)) {
       throw new Error(`${type} stream ${id} already exists!`)
     }
+
     log('new %s stream %s %s', type, id, name)
+
     const send = msg => {
+      if (!registry.has(id)) {
+        throw errCode(new Error('the stream is not in the muxer registry, it may have already been closed'), 'ERR_STREAM_DOESNT_EXIST')
+      }
       if (log.enabled) {
         log('%s stream %s %s send', type, id, name, { ...msg, type: MessageTypeNames[msg.type], data: msg.data && msg.data.slice() })
       }
       return this.source.push(msg)
     }
+
     const onEnd = () => {
       log('%s stream %s %s ended', type, id, name)
       registry.delete(id)
       this.onStreamEnd && this.onStreamEnd(stream)
     }
+
     const stream = createStream({ id, name, send, type, onEnd, maxMsgSize: this._options.maxMsgSize })
     registry.set(id, stream)
     return stream
@@ -137,7 +154,7 @@ class Mplex {
    * also have their size restricted. All messages will be varint decoded.
    *
    * @private
-   * @returns {*} Returns an iterable sink
+   * @returns {Sink} Returns an iterable sink
    */
   _createSink () {
     return async source => {
@@ -216,14 +233,17 @@ class Mplex {
     switch (type) {
       case MessageTypes.MESSAGE_INITIATOR:
       case MessageTypes.MESSAGE_RECEIVER:
+        // We got data from the remote, push it into our local stream
         stream.source.push(data)
         break
       case MessageTypes.CLOSE_INITIATOR:
       case MessageTypes.CLOSE_RECEIVER:
-        stream.close()
+        // We should expect no more data from the remote, stop reading
+        stream.closeRead()
         break
       case MessageTypes.RESET_INITIATOR:
       case MessageTypes.RESET_RECEIVER:
+        // Stop reading and writing to the stream immediately
         stream.reset()
         break
       default:

src/stream.js

@@ -8,9 +8,16 @@ const BufferList = require('bl/BufferList')
 const errCode = require('err-code')
 const { MAX_MSG_SIZE } = require('./restrict-size')
 const { InitiatorMessageTypes, ReceiverMessageTypes } = require('./message-types')
+const pDefer = require('p-defer')
 
 const ERR_MPLEX_STREAM_RESET = 'ERR_MPLEX_STREAM_RESET'
 const ERR_MPLEX_STREAM_ABORT = 'ERR_MPLEX_STREAM_ABORT'
+const MPLEX_WRITE_STREAM_CLOSED = 'MPLEX_WRITE_STREAM_CLOSED'
+
+/**
+ * @typedef {import('libp2p-interfaces/src/stream-muxer/types').MuxedStream} MuxedStream
+ * @typedef {import('libp2p-interfaces/src/stream-muxer/types').Sink} Sink
+ */
 
 /**
  * @param {object} options
@@ -20,18 +27,21 @@ const ERR_MPLEX_STREAM_ABORT = 'ERR_MPLEX_STREAM_ABORT'
  * @param {function(Error)} [options.onEnd] - Called whenever the stream ends
  * @param {string} [options.type] - One of ['initiator','receiver']. Defaults to 'initiator'
  * @param {number} [options.maxMsgSize] - Max size of an mplex message in bytes. Writes > size are automatically split. Defaults to 1MB
- * @returns {*} A muxed stream
+ * @returns {MuxedStream} A muxed stream
  */
 module.exports = ({ id, name, send, onEnd = () => {}, type = 'initiator', maxMsgSize = MAX_MSG_SIZE }) => {
   const abortController = new AbortController()
   const resetController = new AbortController()
+  const writeCloseController = new AbortController()
   const Types = type === 'initiator' ? InitiatorMessageTypes : ReceiverMessageTypes
   const externalId = type === 'initiator' ? (`i${id}`) : `r${id}`
 
   name = String(name == null ? id : name)
 
   let sourceEnded = false
   let sinkEnded = false
+  let sinkInProgress = false
+  let sinkClosedDefer
   let endErr
 
   const onSourceEnd = err => {
@@ -54,11 +64,38 @@ module.exports = ({ id, name, send, onEnd = () => {}, type = 'initiator', maxMsg
       stream.timeline.close = Date.now()
       onEnd(endErr)
     }
+    if (sinkClosedDefer) sinkClosedDefer.resolve()
+  }
+
+  const _send = (message) => {
+    if (!sinkEnded) {
+      send(message)
+    }
   }
 
+  /** @type {MuxedStream} */
   const stream = {
+    // Close for both Reading and Writing
+    close: () => Promise.all([
+      stream.closeRead(),
+      stream.closeWrite()
+    ]),
     // Close for reading
-    close: () => stream.source.end(),
+    closeRead: () => stream.source.end(),
+    // Close for writing
+    closeWrite: () => {
+      if (sinkEnded) {
+        return
+      }
+
+      if (sinkInProgress) {
+        sinkClosedDefer = pDefer()
+        writeCloseController.abort()
+        return sinkClosedDefer.promise
+      }
+
+      return stream.sink([])
+    },
     // Close for reading and writing (local error)
     abort: err => {
       log('%s stream %s abort', type, name, err)
@@ -75,41 +112,49 @@ module.exports = ({ id, name, send, onEnd = () => {}, type = 'initiator', maxMsg
       onSinkEnd(err)
     },
     sink: async source => {
+      if (sinkInProgress) {
+        throw errCode(new Error('the sink was already opened'), 'ERR_SINK_ALREADY_OPENED')
+      }
+
+      sinkInProgress = true
       source = abortable(source, [
         { signal: abortController.signal, options: { abortMessage: 'stream aborted', abortCode: ERR_MPLEX_STREAM_ABORT } },
-        { signal: resetController.signal, options: { abortMessage: 'stream reset', abortCode: ERR_MPLEX_STREAM_RESET } }
+        { signal: resetController.signal, options: { abortMessage: 'stream reset', abortCode: ERR_MPLEX_STREAM_RESET } },
+        { signal: writeCloseController.signal, options: { abortMessage: 'write stream closed', abortCode: MPLEX_WRITE_STREAM_CLOSED } }
       ])
 
       if (type === 'initiator') { // If initiator, open a new stream
-        send({ id, type: Types.NEW_STREAM, data: name })
+        _send({ id, type: Types.NEW_STREAM, data: name })
       }
 
       try {
         for await (let data of source) {
           while (data.length) {
             if (data.length <= maxMsgSize) {
-              send({ id, type: Types.MESSAGE, data })
+              _send({ id, type: Types.MESSAGE, data })
               break
             }
             data = BufferList.isBufferList(data) ? data : new BufferList(data)
-            send({ id, type: Types.MESSAGE, data: data.shallowSlice(0, maxMsgSize) })
+            _send({ id, type: Types.MESSAGE, data: data.shallowSlice(0, maxMsgSize) })
             data.consume(maxMsgSize)
           }
         }
       } catch (err) {
-        // Send no more data if this stream was remotely reset
-        if (err.code === ERR_MPLEX_STREAM_RESET) {
-          log('%s stream %s reset', type, name)
-        } else {
-          log('%s stream %s error', type, name, err)
-          send({ id, type: Types.RESET })
-        }
+        if (err.code !== MPLEX_WRITE_STREAM_CLOSED) {
+          // Send no more data if this stream was remotely reset
+          if (err.code === ERR_MPLEX_STREAM_RESET) {
+            log('%s stream %s reset', type, name)
+          } else {
+            log('%s stream %s error', type, name, err)
+            _send({ id, type: Types.RESET })
+          }
 
-        stream.source.end(err)
-        return onSinkEnd(err)
+          stream.source.end(err)
+          return onSinkEnd(err)
+        }
       }
 
-      send({ id, type: Types.CLOSE })
+      _send({ id, type: Types.CLOSE })
       onSinkEnd()
     },
     source: pushable(onSourceEnd),

test/stream.spec.js

@@ -232,7 +232,7 @@ describe('stream', () => {
       map(msg => {
         // when the initiator sends a CLOSE message, we call close
         if (msg.type === MessageTypes.CLOSE_INITIATOR) {
-          receiver.close()
+          receiver.closeRead()
         }
         return msgToBuffer(msg)
       }),