chore: added cancellation to discovery background task

Author: aryanjassal

src/nodes/NodeConnection.ts

@@ -169,12 +169,9 @@ class NodeConnection {
       targetHostname,
       tlsConfig,
       connectionKeepAliveIntervalTime,
-      connectionKeepAliveTimeoutTime = config.defaultsSystem
-        .nodesConnectionIdleTimeoutTimeMin,
-      connectionInitialMaxStreamsBidi = config.defaultsSystem
-        .nodesConnectionInitialMaxStreamsBidi,
-      connectionInitialMaxStreamsUni = config.defaultsSystem
-        .nodesConnectionInitialMaxStreamsUni,
+      connectionKeepAliveTimeoutTime,
+      connectionInitialMaxStreamsBidi,
+      connectionInitialMaxStreamsUni,
       quicSocket,
       manifest,
       logger,

src/nodes/NodeConnectionManager.ts

@@ -663,13 +663,10 @@ class NodeConnectionManager {
     const cancelAuthenticationPs: Array<PromiseCancellable<void>> = [];
     const cancelReason = new nodesErrors.ErrorNodeConnectionManagerStopping();
     for (const [nodeIdString] of this.connections) {
-      const destroyP = this.authenticateCancel(nodeIdString, cancelReason).then(
-        async () => {
-          return await this.destroyConnection(
-            IdInternal.fromString<NodeId>(nodeIdString),
-            force,
-          );
-        },
+      this.authenticateCancel(nodeIdString, cancelReason);
+      const destroyP = this.destroyConnection(
+        IdInternal.fromString<NodeId>(nodeIdString),
+        force,
       );
       destroyConnectionPs.push(destroyP);
     }
@@ -1186,7 +1183,7 @@ class NodeConnectionManager {
     const remainingKeys = Object.keys(connectionsEntry.connections);
     if (remainingKeys.length === 0) {
       // Clean up authentication
-      await this.authenticateCancel(
+      this.authenticateCancel(
         targetNodeIdString,
         new nodesErrors.ErrorNodeManagerAuthenticationFailed(
           'Connection destroyed before authentication could complete',
@@ -1817,7 +1814,6 @@ class NodeConnectionManager {
     nodeId: NodeId,
     @context ctx: ContextTimed,
   ): Promise<void> {
-    ctx.signal.throwIfAborted();
     const targetNodeIdString = nodeId.toString() as NodeIdString;
     const connectionsEntry = this.connections.get(targetNodeIdString);
     if (connectionsEntry == null) {
@@ -1827,11 +1823,16 @@ class NodeConnectionManager {
     const abortHandler = () => {
       rejectAbortP(ctx.signal.reason);
     };
-    ctx.signal.addEventListener('abort', abortHandler, { once: true });
+    if (ctx.signal.aborted) {
+      abortHandler();
+    } else {
+      ctx.signal.addEventListener('abort', abortHandler, { once: true });
+    }
     try {
       return await Promise.race([connectionsEntry.authenticatedP, abortP]);
     } catch (e) {
-      // Capture the stacktrace here since knowing where we're waiting for authentication is more useful
+      // Capture the stacktrace here since knowing where we're waiting for
+      // authentication is more useful.
       Error.captureStackTrace(e);
       throw e;
     } finally {
@@ -1890,7 +1891,7 @@ class NodeConnectionManager {
     }
   }
 
-  protected async authenticateCancel(
+  protected authenticateCancel(
     targetNodeIdString: NodeIdString,
     reason: Error,
   ) {

src/nodes/NodeGraph.ts

@@ -295,18 +295,20 @@ class NodeGraph {
   @ready(new nodesErrors.ErrorNodeGraphNotRunning())
   public async *getNodeContacts(
     order: 'asc' | 'desc' = 'asc',
-    tran?: DBTransaction,
+    tran: DBTransaction | undefined,
+    ctx: ContextTimed,
   ): AsyncGenerator<[NodeId, NodeContact]> {
     if (tran == null) {
       // Lambda generators don't grab the `this` context, so we need to bind it
-      const getNodeContacts = (tran) => this.getNodeContacts(order, tran);
+      const getNodeContacts = (tran) => this.getNodeContacts(order, tran, ctx);
       return yield* this.db.withTransactionG(async function* (tran) {
         return yield* getNodeContacts(tran);
       });
     }
     return yield* nodesUtils.collectNodeContacts(
       [...this.nodeGraphBucketsDbPath],
       tran,
+      ctx,
       { reverse: order !== 'asc' },
     );
   }
@@ -663,6 +665,7 @@ class NodeGraph {
       for await (const result of nodesUtils.collectNodeContacts(
         [...this.nodeGraphBucketsDbPath, bucketKey],
         tran,
+        ctx,
         {
           reverse: order !== 'asc',
           limit,
@@ -740,10 +743,18 @@ class NodeGraph {
    * Resets the bucket according to the new node ID.
    * Run this after new node ID is generated via renewal or reset.
    */
+  public async resetBuckets(
+    tran?: DBTransaction,
+    ctx?: ContextTimed,
+  ): Promise<void>;
   @ready(new nodesErrors.ErrorNodeGraphNotRunning())
-  public async resetBuckets(tran?: DBTransaction): Promise<void> {
+  @timedCancellable(true)
+  public async resetBuckets(
+    tran: DBTransaction | undefined,
+    @context ctx: ContextTimed,
+  ): Promise<void> {
     if (tran == null) {
-      return this.db.withTransactionF((tran) => this.resetBuckets(tran));
+      return this.db.withTransactionF((tran) => this.resetBuckets(tran, ctx));
     }
     // Setup new space
     const spaceNew = this.space === '0' ? '1' : '0';
@@ -764,6 +775,7 @@ class NodeGraph {
     for await (const [nodeId, nodeContact] of nodesUtils.collectNodeContacts(
       [...this.nodeGraphBucketsDbPath],
       tran,
+      ctx,
     )) {
       const nodeIdKey = nodesUtils.bucketDbKey(nodeId);
       const nodeIdOwn = this.keyRing.getNodeId();
@@ -965,12 +977,12 @@ class NodeGraph {
       for await (const nodeEntry of nodesUtils.collectNodeContacts(
         this.nodeGraphBucketsDbPath,
         tran,
+        ctx,
         {
           lt: [bucketIdKey, ''],
           limit: remainingLimit,
         },
       )) {
-        ctx.signal.throwIfAborted();
         nodes.push(nodeEntry);
       }
     }
@@ -982,12 +994,12 @@ class NodeGraph {
       for await (const nodeEntry of nodesUtils.collectNodeContacts(
         this.nodeGraphBucketsDbPath,
         tran,
+        ctx,
         {
           gt: [bucketId, ''],
           limit: remainingLimit,
         },
       )) {
-        ctx.signal.throwIfAborted();
         nodes.push(nodeEntry);
       }
     }

src/nodes/NodeManager.ts

@@ -136,15 +136,7 @@ class NodeManager {
     _taskInfo,
     bucketIndex: NodeBucketIndex,
   ) => {
-    // Don't use defaults like this
-    // if a default is to be used
-    // provide it directly
-
-    await this.refreshBucket(
-      bucketIndex,
-      this.connectionConnectTimeoutTime,
-      ctx,
-    );
+    await this.refreshBucket(bucketIndex, undefined, ctx);
     // When completed reschedule the task
     // if refreshBucketDelay is 0 then it's considered disabled
     if (this.refreshBucketDelayTime > 0) {
@@ -722,6 +714,7 @@ class NodeManager {
     }
 
     while (true) {
+      ctx.signal.throwIfAborted();
       const isDone = await nodeConnectionsQueue.withNodeSignal(
         async (nodeIdTarget, nodeIdSignaller) => {
           let nodeConnection: NodeConnection | undefined;
@@ -863,6 +856,7 @@ class NodeManager {
     }
 
     while (true) {
+      ctx.signal.throwIfAborted();
       const isDone = await nodeConnectionsQueue.withNodeDirect(
         async (nodeIdTarget, nodeContact) => {
           if (!this.nodeConnectionManager.hasConnection(nodeIdTarget)) {
@@ -1135,6 +1129,7 @@ class NodeManager {
           );
         // Collecting results
         for await (const result of resultStream) {
+          ctx.signal.throwIfAborted();
           const nodeIdNew = nodesUtils.decodeNodeId(result.nodeId);
           if (nodeIdNew == null) {
             utils.never(`failed to decode NodeId "${result.nodeId}"`);
@@ -1151,6 +1146,7 @@ class NodeManager {
             ctx,
           );
         for await (const { nodeIdEncoded, nodeContact } of resultStream) {
+          ctx.signal.throwIfAborted();
           const nodeId = nodesUtils.decodeNodeId(nodeIdEncoded);
           if (nodeId == null) {
             utils.never(`failed to decode NodeId "${nodeIdEncoded}"`);
@@ -2026,8 +2022,7 @@ class NodeManager {
   @timedCancellable(true)
   public async refreshBucket(
     bucketIndex: NodeBucketIndex,
-    connectionConnectTimeoutTime: number | undefined = this
-      .connectionConnectTimeoutTime,
+    connectionConnectTimeoutTime: number = this.connectionConnectTimeoutTime,
     @context ctx: ContextTimed,
   ): Promise<void> {
     // We need to generate a random nodeId for this bucket

src/nodes/utils.ts

@@ -18,6 +18,7 @@ import type {
   NodesAuthenticateConnectionMessageBasicPublic,
   NodesAuthenticateConnectionMessageNone,
 } from './agent/types';
+import type { ContextTimed } from '@matrixai/contexts';
 import dns from 'dns';
 import { utils as dbUtils } from '@matrixai/db';
 import { IdInternal } from '@matrixai/id';
@@ -743,6 +744,7 @@ const quicServerCrypto: QUICServerCrypto = {
 async function* collectNodeContacts(
   levelPath: LevelPath,
   tran: DBTransaction,
+  ctx: ContextTimed,
   options: {
     reverse?: boolean;
     lt?: LevelPath;
@@ -763,6 +765,7 @@ async function* collectNodeContacts(
     gt: options.gt,
     valueAsBuffer: false,
   })) {
+    ctx.signal.throwIfAborted();
     const { nodeId: nodeIdCurrent, nodeContactAddress } = parseBucketsDbKey([
       ...(options.pathAdjust ?? []),
       ...keyPath,

tests/nodes/NodeGraph.test.ts

@@ -1,3 +1,4 @@
+import type { ContextTimed } from '@matrixai/contexts';
 import type {
   NodeContactAddress,
   NodeContact,
@@ -321,7 +322,12 @@ describe(`${NodeGraph.name} test`, () => {
         await nodeGraph.setNodeContact(nodeId2, nodeContact2);
 
         const results: Array<[NodeId, NodeContact]> = [];
-        for await (const result of nodeGraph.getNodeContacts()) {
+        const abortController = new AbortController();
+        for await (const result of nodeGraph.getNodeContacts(
+          undefined,
+          undefined,
+          { signal: abortController.signal } as ContextTimed,
+        )) {
           results.push(result);
         }
         expect(results.length).toBe(2);