fix: allow aborting drain waiting

Add onDrain method that can be aborted

Author: achingbrain

doc/migrations/v2.0.0-v3.0.0.md

@@ -84,9 +84,9 @@ and `sink` properties, we now listen for `message` events and call `.send` with
 `Uint8Array`/`Uint8ArrayList` values.
 
 We can detect when write backpressure needs to be applied by `.send` returning
-`false` (we should await a `drain` event when this happens). and read
-backpressure can be explicitly applied by calling the new `.pause` and `.resume`
-methods.
+`false` (we should await a `drain` event or for the promise returned from
+`onDrain()` to resolve when this happens). and read backpressure can be
+explicitly applied by calling the new `.pause` and `.resume` methods.
 
 > [!CAUTION]
 > If no `message` event handler is added, streams will buffer incoming data
@@ -179,7 +179,8 @@ When streams close they emit a `close` event. This event has an `error: Error` p
 
 ### Write backpressure
 
-You can use [p-event](https://www.npmjs.com/package/p-event) or [race-event](https://www.npmjs.com/package/race-event) to pause writing due to backpressure:
+When `.send()` returns false, the promise returned from `onDrain` will resolve
+when the stream can accept new data:
 
 ```ts
 import { createLibp2p } from 'libp2p'
@@ -198,6 +199,37 @@ const bufs = [
   // a lot of data
 ]
 
+for (const buf of bufs) {
+  if (!stream.send(buf)) {
+    await stream.onDrain({
+      signal: AbortSignal.timeout(5_000)
+    })
+  }
+}
+```
+
+Alternatively you can use [p-event](https://www.npmjs.com/package/p-event) or
+[race-event](https://www.npmjs.com/package/race-event) to wait for the stream to
+emit a `drain` event before continuing:
+
+```ts
+import { createLibp2p } from 'libp2p'
+import { peerIdFromString } from '@libp2p/peer-id'
+import { pEvent } from 'p-event'
+
+const node = createLibp2p({
+  // libp2p config here
+})
+
+const remotePeer = peerIdFromString('123Foo...')
+const stream = await node.dialProtocol(remotePeer, '/echo/1.0.0', {
+  signal: AbortSignal.timeout(5_000)
+})
+
+const bufs = [
+  // a lot of data
+]
+
 for (const buf of bufs) {
   if (!stream.send(buf)) {
     await pEvent(stream, 'drain', {

packages/interface/src/message-stream.ts

@@ -196,4 +196,10 @@ export interface MessageStream<Timeline extends MessageStreamTimeline = MessageS
    * emitted before any other queued data.
    */
   unshift (data: Uint8Array | Uint8ArrayList): void
+
+  /**
+   * Returns a promise that resolves when the stream can accept new data or
+   * rejects if the stream is closed or reset before this occurs.
+   */
+  onDrain (options?: AbortOptions): Promise<void>
 }

packages/utils/src/abstract-message-stream.ts

@@ -1,6 +1,8 @@
 import { StreamResetError, TypedEventEmitter, StreamMessageEvent, StreamBufferError, StreamResetEvent, StreamAbortEvent, StreamCloseEvent, StreamStateError } from '@libp2p/interface'
 import { pushable } from 'it-pushable'
+import { raceSignal } from 'race-signal'
 import { Uint8ArrayList } from 'uint8arraylist'
+import { StreamClosedError } from './errors.ts'
 import type { MessageStreamEvents, MessageStreamStatus, MessageStream, AbortOptions, MessageStreamTimeline, MessageStreamDirection, EventHandler, StreamOptions, MessageStreamReadStatus, MessageStreamWriteStatus } from '@libp2p/interface'
 import type { Logger } from '@libp2p/logger'
 
@@ -63,6 +65,8 @@ export abstract class AbstractMessageStream<Timeline extends MessageStreamTimeli
   protected readonly writeBuffer: Uint8ArrayList
   protected sendingData: boolean
 
+  private onDrainPromise?: PromiseWithResolvers<void>
+
   constructor (init: MessageStreamInit) {
     super()
 
@@ -96,8 +100,15 @@ export abstract class AbstractMessageStream<Timeline extends MessageStreamTimeli
         this.writableNeedsDrain = false
         this.processSendQueue()
       }
+
+      this.onDrainPromise?.resolve()
     }
     this.addEventListener('drain', continueSendingOnDrain)
+
+    const rejectOnDrainOnClose = (evt: StreamCloseEvent): void => {
+      this.onDrainPromise?.reject(evt.error ?? new StreamClosedError())
+    }
+    this.addEventListener('close', rejectOnDrainOnClose)
   }
 
   get readBufferLength (): number {
@@ -108,6 +119,18 @@ export abstract class AbstractMessageStream<Timeline extends MessageStreamTimeli
     return this.writeBuffer.byteLength
   }
 
+  async onDrain (options?: AbortOptions): Promise<void> {
+    if (this.writableNeedsDrain !== true) {
+      return Promise.resolve()
+    }
+
+    if (this.onDrainPromise == null) {
+      this.onDrainPromise = Promise.withResolvers()
+    }
+
+    return raceSignal(this.onDrainPromise.promise, options?.signal)
+  }
+
   async * [Symbol.asyncIterator] (): AsyncGenerator<Uint8Array | Uint8ArrayList> {
     if (this.readStatus !== 'readable' && this.readStatus !== 'paused') {
       return

packages/utils/src/errors.ts

@@ -37,3 +37,8 @@ export class MaxEarlyStreamsError extends Error {
   static name = 'MaxEarlyStreamsError'
   name = 'MaxEarlyStreamsError'
 }
+
+export class StreamClosedError extends Error {
+  static name = 'StreamClosedError'
+  name = 'StreamClosedError'
+}

packages/utils/test/stream-utils-test.spec.ts

@@ -313,4 +313,82 @@ describe('stream-pair', () => {
       Uint8Array.from([4, 5, 6, 7, 0, 1, 2, 3])
     )
   })
+
+  it('should return a promise from onDrain', async () => {
+    const [outgoing] = await streamPair({
+      delay: 100
+    })
+
+    // should resolve immediately when backpressure is not being applied
+    await expect(outgoing.onDrain()).to.eventually.be.undefined()
+
+    while (true) {
+      if (outgoing.send(Uint8Array.from([0, 1, 2, 3])) === false) {
+        break
+      }
+    }
+
+    await expect(outgoing.onDrain()).to.eventually.be.undefined()
+  })
+
+  it('should abort a promise from onDrain via signal', async () => {
+    const [outgoing] = await streamPair({
+      delay: 100
+    })
+
+    // should resolve immediately when backpressure is not being applied
+    await expect(outgoing.onDrain()).to.eventually.be.undefined()
+
+    while (true) {
+      if (outgoing.send(Uint8Array.from([0, 1, 2, 3])) === false) {
+        break
+      }
+    }
+
+    const controller = new AbortController()
+    controller.abort()
+
+    await expect(outgoing.onDrain({
+      signal: controller.signal
+    })).to.eventually.be.rejected
+      .with.property('name', 'AbortError')
+  })
+
+  it('should abort a promise from onDrain via abort', async () => {
+    const [outgoing] = await streamPair({
+      delay: 100
+    })
+
+    while (true) {
+      if (outgoing.send(Uint8Array.from([0, 1, 2, 3])) === false) {
+        break
+      }
+    }
+
+    await expect(
+      Promise.all([
+        outgoing.onDrain(),
+        // eslint-disable-next-line @typescript-eslint/await-thenable
+        outgoing.abort(new Error('Abort!'))
+      ])
+    ).to.eventually.be.rejected
+      .with.property('message', 'Abort!')
+  })
+
+  it('should abort a promise from onDrain via reset', async () => {
+    const [outgoing, incoming] = await streamPair({
+      delay: 100
+    })
+
+    while (true) {
+      if (outgoing.send(Uint8Array.from([0, 1, 2, 3])) === false) {
+        break
+      }
+    }
+
+    incoming.abort(new Error('Reset!'))
+
+    await expect(outgoing.onDrain()).to.eventually.be.rejected
+      .with.property('name', 'StreamResetError')
+  })
 })