fix: fixed monitor usage in QUICConnection

* Related #48

Author: tegefaulkes

src/QUICConnection.ts

@@ -24,6 +24,8 @@ import {
 import Logger from '@matrixai/logger';
 import { Timer } from '@matrixai/timer';
 import { context, timedCancellable } from '@matrixai/contexts/dist/decorators';
+import { withF } from '@matrixai/resources';
+import { utils as contextsUtils } from '@matrixai/contexts';
 import { buildQuicheConfig } from './config';
 import QUICStream from './QUICStream';
 import { quiche } from './native';
@@ -278,27 +280,15 @@ class QUICConnection extends EventTarget {
     };
     ctx.signal.addEventListener('abort', abortHandler);
     const connection = new this(args);
-    const startProm = connection.start().then(async () => {
-      // If this is a server connection, we need to process the packet that created it
-      if (args.type === 'server') {
-        await utils.withMonitor(
-          undefined,
-          connection.lockbox,
-          RWLockWriter,
-          async (mon) => {
-            await mon.withF(connection.lockCode, async (mon) => {
-              await connection.recv(args.data, args.remoteInfo, mon);
-              await connection.send(mon);
-            });
-          },
-        );
-      }
-    });
+    const initialData =
+      args.type === 'server'
+        ? { data: args.data, remoteInfo: args.remoteInfo }
+        : undefined;
     // This ensures that TLS has been established and verified on both sides
     try {
       await Promise.race([
         Promise.all([
-          startProm,
+          connection.start(initialData),
           connection.establishedP,
           connection.secureEstablishedP,
         ]),
@@ -471,12 +461,24 @@ class QUICConnection extends EventTarget {
 
   /**
    * This will set up the connection initiate sending
+   * @param initialData - If the connection is server initiated then the data that initiated it needs to be provided here
    */
-  public async start(): Promise<void> {
+  public async start(initialData?: {
+    data: Uint8Array;
+    remoteInfo: RemoteInfo;
+  }): Promise<void> {
     this.logger.info(`Start ${this.constructor.name}`);
     // Set the connection up
     this.socket.connectionMap.set(this.connectionId, this);
-    await this.send();
+    await withF(
+      [contextsUtils.monitor(this.lockbox, RWLockWriter)],
+      async ([mon]) => {
+        if (initialData != null) {
+          await this.recv(initialData.data, initialData.remoteInfo, mon);
+        }
+        await this.send(mon);
+      },
+    );
     this.logger.info(`Started ${this.constructor.name}`);
   }
 
@@ -665,13 +667,15 @@ class QUICConnection extends EventTarget {
   public async recv(
     data: Uint8Array,
     remoteInfo: RemoteInfo,
-    mon: Monitor<RWLockWriter>,
+    mon?: Monitor<RWLockWriter>,
   ): Promise<void> {
-    if (!mon.isLocked(this.lockCode)) {
-      return mon.withF(this.lockCode, async (mon) => {
-        return this.recv(data, remoteInfo, mon);
-      });
-    }
+    await utils.withMonitor(mon, this.lockbox, RWLockWriter, async (mon) => {
+      if (!mon.isLocked(this.lockCode)) {
+        return mon.withF(this.lockCode, async (mon) => {
+          return this.recv(data, remoteInfo, mon);
+        });
+      }
+    });
 
     try {
       // The remote information may be changed on each receive

src/QUICSocket.ts

@@ -8,6 +8,8 @@ import { running } from '@matrixai/async-init';
 import { StartStop, ready } from '@matrixai/async-init/dist/StartStop';
 import { RWLockWriter } from '@matrixai/async-locks';
 import { status } from '@matrixai/async-init/dist/utils';
+import { withF } from '@matrixai/resources';
+import { utils as contextsUtils } from '@matrixai/contexts';
 import QUICConnectionId from './QUICConnectionId';
 import QUICConnectionMap from './QUICConnectionMap';
 import { quiche } from './native';
@@ -108,11 +110,9 @@ class QUICSocket extends EventTarget {
     // Acquire the conn lock, this ensures mutual exclusion
     // for state changes on the internal connection
     try {
-      await utils.withMonitor(
-        undefined,
-        connection.lockbox,
-        RWLockWriter,
-        async (mon) => {
+      await withF(
+        [contextsUtils.monitor(connection.lockbox, RWLockWriter)],
+        async ([mon]) => {
           await mon.withF(connection.lockCode, async (mon) => {
             // Even if we are `stopping`, the `quiche` library says we need to
             // continue processing any packets.

src/utils.ts

@@ -9,9 +9,11 @@ import type {
 } from './types';
 import type { Connection } from './native';
 import type { LockBox, RWLockWriter } from '@matrixai/async-locks';
+import type { Monitor } from '@matrixai/async-locks';
 import dns from 'dns';
 import { IPv4, IPv6, Validator } from 'ip-num';
-import { Monitor } from '@matrixai/async-locks';
+import { withF } from '@matrixai/resources';
+import { utils as contextsUtils } from '@matrixai/contexts';
 import QUICConnectionId from './QUICConnectionId';
 import * as errors from './errors';
 
@@ -464,17 +466,26 @@ function streamStats(
 `;
 }
 
+/**
+ * Used as a clean way to create a new monitor if it doesn't exist, otherwise uses the existing one.
+ */
 async function withMonitor<T>(
   mon: Monitor<RWLockWriter> | undefined,
   lockBox: LockBox<RWLockWriter>,
   lockConstructor: { new (): RWLockWriter },
   fun: (mon: Monitor<RWLockWriter>) => Promise<T>,
   locksPending?: Map<string, { count: number }>,
 ): Promise<T> {
-  const _mon = mon ?? new Monitor(lockBox, lockConstructor, locksPending);
-  const result = await fun(_mon);
-  if (mon != null) await _mon.unlockAll();
-  return result;
+  if (mon == null) {
+    return await withF(
+      [contextsUtils.monitor(lockBox, lockConstructor, locksPending)],
+      async ([mon]) => {
+        return await fun(mon);
+      },
+    );
+  } else {
+    return await fun(mon);
+  }
 }
 
 export {