Merge pull request #75 from MatrixAI/feature-unforced-end

fix ending `QUICConnection` with `force: false`

src/QUICConnection.ts

@@ -27,6 +27,7 @@ import { buildQuicheConfig, minIdleTimeout } from './config';
 import QUICConnectionId from './QUICConnectionId';
 import QUICStream from './QUICStream';
 import { quiche, ConnectionErrorCode } from './native';
+import { Shutdown } from './native/types';
 import * as utils from './utils';
 import * as events from './events';
 import * as errors from './errors';
@@ -616,7 +617,27 @@ class QUICConnection {
       } = {}) {
     this.logger.info(`Stop ${this.constructor.name}`);
     this.stopKeepAliveIntervalTimer();
-    // Closing the connection first to avoid accepting new streams
+
+    // Yield to allow any background processing to settle before proceeding.
+    // This will allow any streams to process buffers before continuing
+    await utils.yieldMicro();
+
+    // Destroy all streams
+    const streamsDestroyP: Array<Promise<void>> = [];
+    for (const quicStream of this.streamMap.values()) {
+      // The reason is only used if `force` is `true`
+      // If `force` is not true, this will gracefully wait for
+      // both readable and writable to gracefully close
+      streamsDestroyP.push(
+        quicStream.destroy({
+          reason: this.errorLast,
+          force: force || this.conn.isDraining() || this.conn.isClosed(),
+        }),
+      );
+    }
+    await Promise.all(streamsDestroyP);
+
+    // Close after processing all streams
     if (!this.conn.isDraining() && !this.conn.isClosed()) {
       // If `this.conn.close` is already called, the connection will be draining,
       // in that case we just skip doing this local close.
@@ -632,20 +653,7 @@ class QUICConnection {
       });
       this.dispatchEvent(new events.EventQUICConnectionError({ detail: e }));
     }
-    // Destroy all streams
-    const streamsDestroyP: Array<Promise<void>> = [];
-    for (const quicStream of this.streamMap.values()) {
-      // The reason is only used if `force` is `true`
-      // If `force` is not true, this will gracefully wait for
-      // both readable and writable to gracefully close
-      streamsDestroyP.push(
-        quicStream.destroy({
-          reason: this.errorLast,
-          force: force || this.conn.isDraining() || this.conn.isClosed(),
-        }),
-      );
-    }
-    await Promise.all(streamsDestroyP);
+
     // Waiting for `closedP` to resolve
     // Only the `this.connTimeoutTimer` will resolve this promise
     await this.closedP;
@@ -940,6 +948,13 @@ class QUICConnection {
     for (const streamId of this.conn.readable() as Iterable<StreamId>) {
       let quicStream = this.streamMap.get(streamId);
       if (quicStream == null) {
+        if (this[running] === false || this[status] === 'stopping') {
+          // We should reject new connections when stopping
+          this.conn.streamShutdown(streamId, Shutdown.Write, 1);
+          this.conn.streamShutdown(streamId, Shutdown.Read, 1);
+          continue;
+        }
+
         quicStream = QUICStream.createQUICStream({
           initiated: 'peer',
           streamId,
@@ -969,6 +984,13 @@ class QUICConnection {
     for (const streamId of this.conn.writable() as Iterable<StreamId>) {
       let quicStream = this.streamMap.get(streamId);
       if (quicStream == null) {
+        if (this[running] === false || this[status] === 'stopping') {
+          // We should reject new connections when stopping
+          this.conn.streamShutdown(streamId, Shutdown.Write, 1);
+          this.conn.streamShutdown(streamId, Shutdown.Read, 1);
+          continue;
+        }
+
         quicStream = QUICStream.createQUICStream({
           initiated: 'peer',
           streamId,
@@ -1132,6 +1154,21 @@ class QUICConnection {
     return quicStream;
   }
 
+  /**
+   * Destroys all active streams without closing the connection.
+   *
+   * If there are no active streams then it will do nothing.
+   * If the connection is stopped with `force: false` then this can be used
+   * to force close any streams `stop` is waiting for to end.
+   *
+   * Destruction will occur in the background.
+   */
+  public destroyStreams(reason?: any) {
+    for (const quicStream of this.streamMap.values()) {
+      quicStream.cancel(reason);
+    }
+  }
+
   /**
    * Starts the keep alive interval timer.
    *

src/utils.ts

@@ -17,6 +17,13 @@ import * as errors from './errors';
 const textEncoder = new TextEncoder();
 const textDecoder = new TextDecoder('utf-8');
 
+/**
+ * Used to yield to the event loop to allow other micro tasks to process
+ */
+async function yieldMicro(): Promise<void> {
+  return await new Promise<void>((r) => queueMicrotask(r));
+}
+
 /**
  * Convert callback-style to promise-style
  * If this is applied to overloaded function
@@ -550,6 +557,7 @@ function isStreamReset(e: Error): number | false {
 export {
   textEncoder,
   textDecoder,
+  yieldMicro,
   promisify,
   promise,
   bufferWrap,

tests/QUICClient.test.ts

@@ -1775,6 +1775,7 @@ describe(QUICClient.name, () => {
 
     // Handling client error event
     const clientErrorProm = promise<never>();
+    void clientErrorProm.p.catch(() => {}); // Ignore unhandled rejection
     client.addEventListener(
       events.EventQUICClientError.name,
       (evt: events.EventQUICClientError) => clientErrorProm.rejectP(evt.detail),
@@ -1783,6 +1784,7 @@ describe(QUICClient.name, () => {
 
     // Handling client destroy event
     const clientDestroyedProm = promise<void>();
+    void clientDestroyedProm.p.catch(() => {}); // Ignore unhandled rejection
     client.addEventListener(
       events.EventQUICClientDestroyed.name,
       () => clientDestroyedProm.resolveP(),