fix: agent start not exiting then stopping Polykey

The life-cycle of the `WorkerManager` is now managed by the `PolykeyAgent`.

Author: tegefaulkes

src/PolykeyAgent.ts

@@ -10,6 +10,7 @@ import process from 'process';
 import Logger from '@matrixai/logger';
 import { DB } from '@matrixai/db';
 import { CreateDestroyStartStop } from '@matrixai/async-init/dist/CreateDestroyStartStop';
+import { WorkerManager } from '@matrixai/workers';
 import * as networkUtils from './network/utils';
 import KeyRing from './keys/KeyRing';
 import CertManager from './keys/CertManager';
@@ -37,6 +38,7 @@ import * as errors from './errors';
 import * as utils from './utils';
 import * as keysUtils from './keys/utils';
 import * as nodesUtils from './nodes/utils';
+import * as workersUtils from './workers/utils';
 import TaskManager from './tasks/TaskManager';
 
 type NetworkConfig = {
@@ -81,6 +83,7 @@ class PolykeyAgent {
     proxyConfig = {},
     nodeConnectionManagerConfig = {},
     seedNodes = {},
+    workers,
     // Optional dependencies
     status,
     schema,
@@ -134,6 +137,7 @@ class PolykeyAgent {
     };
     networkConfig?: NetworkConfig;
     seedNodes?: SeedNodes;
+    workers?: number;
     status?: Status;
     schema?: Schema;
     keyRing?: KeyRing;
@@ -455,6 +459,7 @@ class PolykeyAgent {
     await pkAgent.start({
       password,
       networkConfig,
+      workers,
       fresh,
     });
     logger.info(`Created ${this.name}`);
@@ -485,6 +490,7 @@ class PolykeyAgent {
   public readonly events: EventBus;
   public readonly fs: FileSystem;
   public readonly logger: Logger;
+  protected workerManager: PolykeyWorkerManagerInterface | undefined;
 
   constructor({
     nodePath,
@@ -566,10 +572,12 @@ class PolykeyAgent {
   public async start({
     password,
     networkConfig = {},
+    workers,
     fresh = false,
   }: {
     password: string;
     networkConfig?: NetworkConfig;
+    workers?: number;
     fresh?: boolean;
   }) {
     try {
@@ -748,6 +756,15 @@ class PolykeyAgent {
       await this.notificationsManager.start({ fresh });
       await this.sessionManager.start({ fresh });
       await this.taskManager.startProcessing();
+      if (workers != null) {
+        this.workerManager = await workersUtils.createWorkerManager({
+          // 0 means max workers
+          cores: workers === 0 ? undefined : workers,
+          logger: this.logger.getChild(WorkerManager.name),
+        });
+        this.vaultManager.setWorkerManager(this.workerManager);
+        this.db.setWorkerManager(this.workerManager);
+      }
       await this.status.finishStart({
         pid: process.pid,
         nodeId: this.keyRing.getNodeId(),
@@ -786,6 +803,9 @@ class PolykeyAgent {
       await this.db?.stop();
       await this.keyRing?.stop();
       await this.schema?.stop();
+      this.vaultManager.unsetWorkerManager();
+      this.db.unsetWorkerManager();
+      await this.workerManager?.destroy();
       await this.status?.stop({});
       throw e;
     }
@@ -819,6 +839,9 @@ class PolykeyAgent {
     await this.db.stop();
     await this.keyRing.stop();
     await this.schema.stop();
+    this.vaultManager.unsetWorkerManager();
+    this.db.unsetWorkerManager();
+    await this.workerManager?.destroy();
     await this.status.stop({});
     this.logger.info(`Stopped ${this.constructor.name}`);
   }
@@ -870,16 +893,6 @@ class PolykeyAgent {
     await this.schema.destroy();
     this.logger.info(`Destroyed ${this.constructor.name}`);
   }
-
-  public setWorkerManager(workerManager: PolykeyWorkerManagerInterface) {
-    this.db.setWorkerManager(workerManager);
-    this.vaultManager.setWorkerManager(workerManager);
-  }
-
-  public unsetWorkerManager() {
-    this.db.unsetWorkerManager();
-    this.vaultManager.unsetWorkerManager();
-  }
 }
 
 export default PolykeyAgent;

src/bin/agent/CommandStart.ts

@@ -6,7 +6,6 @@ import type {
 } from '../types';
 import type PolykeyAgent from '../../PolykeyAgent';
 import type { RecoveryCode } from '../../keys/types';
-import type { PolykeyWorkerManagerInterface } from '../../workers/types';
 import path from 'path';
 import childProcess from 'child_process';
 import process from 'process';
@@ -46,9 +45,6 @@ class CommandStart extends CommandPolykey {
       options.clientPort =
         options.clientPort ?? config.defaults.networkConfig.clientPort;
       const { default: PolykeyAgent } = await import('../../PolykeyAgent');
-      const { WorkerManager, utils: workersUtils } = await import(
-        '../../workers'
-      );
       const nodesUtils = await import('../../nodes/utils');
       const keysUtils = await import('../../keys/utils/index');
       let password: string | undefined;
@@ -113,6 +109,7 @@ class CommandStart extends CommandPolykey {
           proxyPort: options.proxyPort,
         },
         seedNodes: seedNodes_,
+        workers: options.workers,
         fresh: options.fresh,
       };
       let statusLiveData: AgentStatusLiveData;
@@ -212,11 +209,7 @@ class CommandStart extends CommandPolykey {
         process.title = 'polykey-agent';
         // eslint-disable-next-line prefer-const
         let pkAgent: PolykeyAgent;
-        // eslint-disable-next-line prefer-const
-        let workerManager: PolykeyWorkerManagerInterface;
         this.exitHandlers.handlers.push(async () => {
-          pkAgent?.unsetWorkerManager();
-          await workerManager?.destroy();
           await pkAgent?.stop();
         });
         try {
@@ -231,13 +224,6 @@ class CommandStart extends CommandPolykey {
           }
           throw e;
         }
-        if (options.workers !== 0) {
-          workerManager = await workersUtils.createWorkerManager({
-            cores: options.workers,
-            logger: this.logger.getChild(WorkerManager.name),
-          });
-          pkAgent.setWorkerManager(workerManager);
-        }
         recoveryCodeOut = pkAgent.keyRing.recoveryCode;
         statusLiveData = {
           pid: process.pid,

src/bin/utils/options.ts

@@ -138,10 +138,10 @@ const network = new commander.Option(
 
 const workers = new commander.Option(
   '-w --workers <count>',
-  'Number of workers to use, defaults to number of cores with `all`, 0 means no multi-threading',
+  'Number of workers to use, defaults to number of cores with `all`, 0 means all cores, `false`|`null`|`none`|`no` means no multi-threading',
 )
   .argParser(binParsers.parseCoreCount)
-  .default(undefined);
+  .default(0, 'all');
 
 const pullVault = new commander.Option(
   '-pv, --pull-vault <pullVaultNameOrId>',

src/bin/utils/parsers.ts

@@ -72,10 +72,17 @@ const parseProviderIdList = validateParserToArgListParser(
 );
 
 function parseCoreCount(v: string): number | undefined {
-  if (v === 'all') {
-    return undefined;
+  switch (v) {
+    case 'all':
+      return 0;
+    case 'none':
+    case 'no':
+    case 'false':
+    case 'null':
+      return undefined;
+    default:
+      return parseInt(v);
   }
-  return parseInt(v);
 }
 
 function parseSecretPath(secretPath: string): [string, string, string?] {

src/workers/errors.ts

@@ -0,0 +1,10 @@
+import { ErrorPolykey, sysexits } from '../errors';
+
+class ErrorWorkers<T> extends ErrorPolykey<T> {}
+
+class ErrorWorkersInvalidCores<T> extends ErrorWorkers<T> {
+  static description = 'specified cores must be positive or 0';
+  exitCode = sysexits.USAGE;
+}
+
+export { ErrorWorkersInvalidCores };

src/workers/utils.ts

@@ -3,6 +3,7 @@ import type { PolykeyWorkerManagerInterface } from './types';
 import type Logger from '@matrixai/logger';
 import { WorkerManager } from '@matrixai/workers';
 import { spawn, Worker } from 'threads';
+import * as workerErrors from './errors';
 
 async function createWorkerManager({
   cores,
@@ -11,6 +12,9 @@ async function createWorkerManager({
   cores?: number;
   logger?: Logger;
 }): Promise<PolykeyWorkerManagerInterface> {
+  if (cores != null && (cores < 0 || isNaN(cores))) {
+    throw new workerErrors.ErrorWorkersInvalidCores();
+  }
   return await WorkerManager.createWorkerManager<PolykeyWorkerModule>({
     workerFactory: () => spawn(new Worker('./polykeyWorker')),
     cores,