Integrating tracer into object lifecycles

npmDepsHash

@@ -1 +1 @@
-sha256-4fkSE29rEC53aecV/DYYwP5Xmd2wDjw6sBwGGaaDjF0=
+sha256-yOUsAbtVeXMdGjR92SvSCFAlPTHxm1xQCipqbdkxb+o=

package.json

@@ -115,7 +115,7 @@
     "postversion": "npm install --package-lock-only --ignore-scripts --silent",
     "dependencies": "node ./scripts/npmDepsHash.mjs",
     "tsx": "tsx",
-    "test": "node ./scripts/test.mjs",
+    "test": "node ./scripts/build.mjs && node ./scripts/test.mjs",
     "lint": "eslint '{src,tests,scripts,benches}/**/*.{js,mjs,ts,mts,jsx,tsx}'",
     "lintfix": "eslint '{src,tests,scripts,benches}/**/*.{js,mjs,ts,mts,jsx,tsx}' --fix",
     "lint-shell": "find ./src ./tests ./scripts -type f -regextype posix-extended -regex '.*\\.(sh)' -exec shellcheck {} +",
@@ -137,7 +137,7 @@
   },
   "devDependencies": {
     "@matrixai/errors": "^2.1.3",
-    "@matrixai/logger": "^4.0.3",
+    "@matrixai/logger": "^4.0.4-alpha.5",
     "@matrixai/exec": "^1.0.3",
     "@fast-check/jest": "^2.1.1",
     "@swc/core": "1.3.82",
@@ -171,5 +171,8 @@
     "tsx": "^3.12.7",
     "typedoc": "^0.24.8",
     "typescript": "^5.1.6"
+  },
+  "overrides": {
+    "@matrixai/logger": "^4.0.4-alpha.5"
   }
 }

src/CommandPolykey.ts

@@ -6,6 +6,7 @@ import Logger, {
   formatting,
   levelToString,
   evalLogDataValue,
+  tracer,
 } from '@matrixai/logger';
 import * as binUtils from './utils/index.js';
 import * as binOptions from './utils/options.js';
@@ -23,6 +24,7 @@ class CommandPolykey extends Command {
   protected logger: Logger = logger;
   protected fs: FileSystem;
   protected exitHandlers: binUtils.ExitHandlers;
+  protected tracerProm: Promise<void> | undefined;
 
   public constructor({
     exitHandlers,
@@ -94,6 +96,21 @@ class CommandPolykey extends Command {
       if (opts.nodePath == null) {
         throw new errors.ErrorPolykeyCLINodePath();
       }
+      // If verbose level has been enabled, then we want to start tracing
+      if (opts.verbose && this.tracerProm == null) {
+        this.tracerProm = (async () => {
+          const fs = await import('node:fs');
+          const spanFile = await fs.promises.open('span.jsonl', 'w');
+          const gen = tracer.streamEvents();
+          for await (const event of gen) {
+            await spanFile.write(JSON.stringify(event) + '\n');
+          }
+          await spanFile.close();
+        })();
+        this.exitHandlers.setTracerProm(this.tracerProm);
+      } else {
+        tracer.disableTracing();
+      }
       await fn(...args);
     });
   }

src/polykey.ts

@@ -23,16 +23,20 @@ process.removeAllListeners('SIGTERM');
 /**
  * Set the main entrypoint filepath.
  * This can be referred to globally.
- * For ESM, change to using `import.meta.url`.
  */
 const filename = fileURLToPath(import.meta.url);
 globalThis.PK_MAIN_EXECUTABLE = filename;
 
+import type { SpanId } from '@matrixai/logger';
+
+let rootSpan: SpanId;
+
 async function polykeyAgentMain(): Promise<number> {
   const {
     default: Logger,
     StreamHandler,
     formatting,
+    tracer,
   } = await import('@matrixai/logger');
   const { default: PolykeyAgent } = await import('polykey/PolykeyAgent.js');
   const { default: ErrorPolykey } = await import('polykey/ErrorPolykey.js');
@@ -51,6 +55,7 @@ async function polykeyAgentMain(): Promise<number> {
   const messageIn = await messageInP;
   const errFormat = messageIn.format === 'json' ? 'json' : 'error';
   exitHandlers.errFormat = errFormat;
+  exitHandlers.setPreExit(() => tracer.endSpan(rootSpan));
   // Set the logger according to the verbosity
   logger.setLevel(messageIn.logLevel);
   // Set the logger formatter according to the format
@@ -245,15 +250,21 @@ async function polykeyMain(argv: Array<string>): Promise<number> {
 }
 
 async function main(argv = process.argv): Promise<number> {
-  if (argv[argv.length - 1] === '--agent-mode') {
-    // This is an internal mode for running `PolykeyAgent` as a child process
-    // This is not supposed to be used directly by the user
-    process.title = 'polykey-agent';
-    return polykeyAgentMain();
-  } else {
-    process.title = 'polykey';
-    return polykeyMain(argv);
-  }
+  const { spanContext } = await import('@matrixai/async-init/utils.js');
+  const { tracer } = await import('@matrixai/logger');
+
+  rootSpan = tracer.startSpan('polykeyRoot');
+  return spanContext.run({ currentSpanId: rootSpan }, async () => {
+      if (argv[argv.length - 1] === '--agent-mode') {
+        // This is an internal mode for running `PolykeyAgent` as a child process
+        // This is not supposed to be used directly by the user
+        process.title = 'polykey-agent';
+        return await polykeyAgentMain();
+      } else {
+        process.title = 'polykey';
+        return await polykeyMain(argv);
+      }
+  });
 }
 
 if (import.meta.url.startsWith('file:')) {

src/utils/ExitHandlers.ts

@@ -1,4 +1,5 @@
 import process from 'node:process';
+import { tracer } from '@matrixai/logger';
 import ErrorPolykey from 'polykey/ErrorPolykey.js';
 import * as binUtils from './utils.js';
 import * as errors from '../errors.js';
@@ -11,6 +12,16 @@ class ExitHandlers {
   public handlers: Array<(signal?: NodeJS.Signals) => Promise<void>>;
   protected _exiting: boolean = false;
   protected _errFormat: 'json' | 'error';
+  protected tracerProm: Promise<void> | undefined;
+  protected preExit: () => void | Promise<void>;
+
+  public setTracerProm(prom: Promise<void>) {
+    this.tracerProm = prom;
+  }
+
+  public setPreExit(f: () => void | Promise<void>) {
+    this.preExit = f;
+  }
 
   /**
    * Handles termination signals
@@ -49,6 +60,11 @@ class ExitHandlers {
     } finally {
       // Uninstall all handlers to prevent signal loop
       this.uninstall();
+      await this.preExit?.();
+      if (this.tracerProm) {
+        tracer.endTracing();
+        await this.tracerProm;
+      }
       // Propagate signal to NodeJS VM handlers
       process.kill(process.pid, signal);
     }
@@ -75,6 +91,11 @@ class ExitHandlers {
     );
     process.exitCode = error.exitCode;
     // Fail fast pattern
+    await this.preExit?.();
+    if (this.tracerProm) {
+      tracer.endTracing();
+      await this.tracerProm;
+    }
     process.exit();
   };
 
@@ -99,6 +120,11 @@ class ExitHandlers {
     );
     process.exitCode = error.exitCode;
     // Fail fast pattern
+    await this.preExit?.();
+    if (this.tracerProm) {
+      tracer.endTracing();
+      await this.tracerProm;
+    }
     process.exit();
   };