fix: close webrtc streams without data loss

packages/interface/package.json

@@ -164,11 +164,15 @@
     "it-stream-types": "^2.0.1",
     "multiformats": "^12.0.1",
     "p-defer": "^4.0.0",
+    "race-signal": "^1.0.0",
     "uint8arraylist": "^2.4.3"
   },
   "devDependencies": {
     "@types/sinon": "^10.0.15",
     "aegir": "^40.0.8",
+    "delay": "^6.0.0",
+    "it-all": "^3.0.3",
+    "it-drain": "^3.0.3",
     "sinon": "^16.0.0",
     "sinon-ts": "^1.0.0"
   }

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

@@ -1,12 +1,14 @@
 import { abortableSource } from 'abortable-iterator'
 import { type Pushable, pushable } from 'it-pushable'
 import defer, { type DeferredPromise } from 'p-defer'
+import { raceSignal } from 'race-signal'
 import { Uint8ArrayList } from 'uint8arraylist'
 import { CodeError } from '../errors.js'
 import type { Direction, ReadStatus, Stream, StreamStatus, StreamTimeline, WriteStatus } from '../connection/index.js'
 import type { AbortOptions } from '../index.js'
 import type { Source } from 'it-stream-types'
 
+// copied from @libp2p/logger to break a circular dependency
 interface Logger {
   (formatter: any, ...args: any[]): void
   error: (formatter: any, ...args: any[]) => void
@@ -16,6 +18,7 @@ interface Logger {
 
 const ERR_STREAM_RESET = 'ERR_STREAM_RESET'
 const ERR_SINK_INVALID_STATE = 'ERR_SINK_INVALID_STATE'
+const DEFAULT_SEND_CLOSE_WRITE_TIMEOUT = 5000
 
 export interface AbstractStreamInit {
   /**
@@ -68,6 +71,12 @@ export interface AbstractStreamInit {
    * connection when closing the writable end of the stream. (default: 500)
    */
   closeTimeout?: number
+
+  /**
+   * After the stream sink has closed, a limit on how long it takes to send
+   * a close-write message to the remote peer.
+   */
+  sendCloseWriteTimeout?: number
 }
 
 function isPromise (res?: any): res is Promise<void> {
@@ -94,6 +103,7 @@ export abstract class AbstractStream implements Stream {
   private readonly onCloseWrite?: () => void
   private readonly onReset?: () => void
   private readonly onAbort?: (err: Error) => void
+  private readonly sendCloseWriteTimeout: number
 
   protected readonly log: Logger
 
@@ -113,6 +123,7 @@ export abstract class AbstractStream implements Stream {
     this.timeline = {
       open: Date.now()
     }
+    this.sendCloseWriteTimeout = init.sendCloseWriteTimeout ?? DEFAULT_SEND_CLOSE_WRITE_TIMEOUT
 
     this.onEnd = init.onEnd
     this.onCloseRead = init?.onCloseRead
@@ -128,7 +139,6 @@ export abstract class AbstractStream implements Stream {
           this.log.trace('source ended')
         }
 
-        this.readStatus = 'closed'
         this.onSourceEnd(err)
       }
     })
@@ -173,11 +183,19 @@ export abstract class AbstractStream implements Stream {
         }
       }
 
-      this.log.trace('sink finished reading from source')
-      this.writeStatus = 'done'
+      this.log.trace('sink finished reading from source, write status is "%s"', this.writeStatus)
+
+      if (this.writeStatus === 'writing') {
+        this.writeStatus = 'closing'
+
+        this.log.trace('send close write to remote')
+        await this.sendCloseWrite({
+          signal: AbortSignal.timeout(this.sendCloseWriteTimeout)
+        })
+
+        this.writeStatus = 'closed'
+      }
 
-      this.log.trace('sink calling closeWrite')
-      await this.closeWrite(options)
       this.onSinkEnd()
     } catch (err: any) {
       this.log.trace('sink ended with error, calling abort with error', err)
@@ -196,6 +214,7 @@ export abstract class AbstractStream implements Stream {
     }
 
     this.timeline.closeRead = Date.now()
+    this.readStatus = 'closed'
 
     if (err != null && this.endErr == null) {
       this.endErr = err
@@ -207,6 +226,10 @@ export abstract class AbstractStream implements Stream {
       this.log.trace('source and sink ended')
       this.timeline.close = Date.now()
 
+      if (this.status !== 'aborted' && this.status !== 'reset') {
+        this.status = 'closed'
+      }
+
       if (this.onEnd != null) {
         this.onEnd(this.endErr)
       }
@@ -221,6 +244,7 @@ export abstract class AbstractStream implements Stream {
     }
 
     this.timeline.closeWrite = Date.now()
+    this.writeStatus = 'closed'
 
     if (err != null && this.endErr == null) {
       this.endErr = err
@@ -232,6 +256,10 @@ export abstract class AbstractStream implements Stream {
       this.log.trace('sink and source ended')
       this.timeline.close = Date.now()
 
+      if (this.status !== 'aborted' && this.status !== 'reset') {
+        this.status = 'closed'
+      }
+
       if (this.onEnd != null) {
         this.onEnd(this.endErr)
       }
@@ -266,16 +294,16 @@ export abstract class AbstractStream implements Stream {
     const readStatus = this.readStatus
     this.readStatus = 'closing'
 
-    if (readStatus === 'ready') {
-      this.log.trace('ending internal source queue')
-      this.streamSource.end()
-    }
-
     if (this.status !== 'reset' && this.status !== 'aborted' && this.timeline.closeRead == null) {
       this.log.trace('send close read to remote')
       await this.sendCloseRead(options)
     }
 
+    if (readStatus === 'ready') {
+      this.log.trace('ending internal source queue')
+      this.streamSource.end()
+    }
+
     this.log.trace('closed readable end of stream')
   }
 
@@ -286,33 +314,26 @@ export abstract class AbstractStream implements Stream {
 
     this.log.trace('closing writable end of stream with starting write status "%s"', this.writeStatus)
 
-    const writeStatus = this.writeStatus
-
     if (this.writeStatus === 'ready') {
       this.log.trace('sink was never sunk, sink an empty array')
-      await this.sink([])
-    }
 
-    this.writeStatus = 'closing'
+      await raceSignal(this.sink([]), options.signal)
+    }
 
-    if (writeStatus === 'writing') {
+    if (this.writeStatus === 'writing') {
       // stop reading from the source passed to `.sink` in the microtask queue
       // - this lets any data queued by the user in the current tick get read
       // before we exit
       await new Promise((resolve, reject) => {
         queueMicrotask(() => {
           this.log.trace('aborting source passed to .sink')
           this.sinkController.abort()
-          this.sinkEnd.promise.then(resolve, reject)
+          raceSignal(this.sinkEnd.promise, options.signal)
+            .then(resolve, reject)
         })
       })
     }
 
-    if (this.status !== 'reset' && this.status !== 'aborted' && this.timeline.closeWrite == null) {
-      this.log.trace('send close write to remote')
-      await this.sendCloseWrite(options)
-    }
-
     this.writeStatus = 'closed'
 
     this.log.trace('closed writable end of stream')
@@ -357,6 +378,7 @@ export abstract class AbstractStream implements Stream {
     const err = new CodeError('stream reset', ERR_STREAM_RESET)
 
     this.status = 'reset'
+    this.timeline.reset = Date.now()
     this._closeSinkAndSource(err)
     this.onReset?.()
   }
@@ -423,7 +445,7 @@ export abstract class AbstractStream implements Stream {
       return
     }
 
-    this.log.trace('muxer destroyed')
+    this.log.trace('stream destroyed')
 
     this._closeSinkAndSource()
   }

packages/interface/test/fixtures/logger.ts

@@ -0,0 +1,16 @@
+// copied from @libp2p/logger to break a circular dependency
+interface Logger {
+  (): void
+  error: () => void
+  trace: () => void
+  enabled: boolean
+}
+
+export function logger (): Logger {
+  const output = (): void => {}
+  output.trace = (): void => {}
+  output.error = (): void => {}
+  output.enabled = false
+
+  return output
+}