@senzops/apm-node 1.2.7 → 1.2.8

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@senzops/apm-node",
-  "version": "1.2.7",
+  "version": "1.2.8",
   "description": "Universal APM SDK for Senzor",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",

package/src/core/context.ts

@@ -5,17 +5,41 @@ interface IStorage<T> {
   getStore(): T | undefined;
 }
 
+/**
+ * Async-safe fallback when AsyncLocalStorage is unavailable.
+ *
+ * Not concurrency-safe across truly parallel requests in a single isolate,
+ * but correct for sequential and async/await patterns (e.g. Cloudflare Workers
+ * where each request gets its own execution context).
+ */
 class NaiveStorage<T> implements IStorage<T> {
   private store: T | undefined;
 
   run<R>(store: T, callback: (...args: any[]) => R, ...args: any[]): R {
     const prev = this.store;
     this.store = store;
+
+    let result: R;
     try {
-      return callback(...args);
-    } finally {
+      result = callback(...args);
+    } catch (err) {
       this.store = prev;
+      throw err;
+    }
+
+    // If the callback returned a thenable (async handler), defer the restore
+    // until the promise settles so Context.current() works across awaits.
+    if (result != null && typeof (result as any).then === 'function') {
+      const promise = (result as any).then(
+        (val: any) => { this.store = prev; return val; },
+        (err: any) => { this.store = prev; throw err; }
+      );
+      return promise as R;
     }
+
+    // Sync callback — restore immediately.
+    this.store = prev;
+    return result;
   }
 
   getStore(): T | undefined {
@@ -23,29 +47,67 @@ class NaiveStorage<T> implements IStorage<T> {
   }
 }
 
-const resolveStorage = <T>(): IStorage<T> => {
+/**
+ * Resolve the best available async context storage.
+ *
+ * Uses lazy re-resolution: if the initial attempt (at module evaluation time)
+ * falls back to NaiveStorage, subsequent calls to resolveStorage() will retry.
+ * This handles runtimes where AsyncLocalStorage becomes available after module
+ * init (e.g. some Workers configurations).
+ */
+const tryResolveALS = <T>(): IStorage<T> | null => {
+  // 1. Check globalThis (Cloudflare Workers nodejs_compat_v2, Bun, Deno)
   if (typeof globalThis !== 'undefined' && (globalThis as any).AsyncLocalStorage) {
     return new (globalThis as any).AsyncLocalStorage();
   }
 
+  // 2. Node.js CJS require
   try {
     if (typeof require !== 'undefined') {
-      const { AsyncLocalStorage } = require('node:async_hooks');
-      if (AsyncLocalStorage) return new AsyncLocalStorage();
+      const mod = require('node:async_hooks');
+      if (mod?.AsyncLocalStorage) return new mod.AsyncLocalStorage();
     }
   } catch {}
 
   try {
     if (typeof require !== 'undefined') {
-      const { AsyncLocalStorage } = require('async_hooks');
-      if (AsyncLocalStorage) return new AsyncLocalStorage();
+      const mod = require('async_hooks');
+      if (mod?.AsyncLocalStorage) return new mod.AsyncLocalStorage();
     }
   } catch {}
 
-  return new NaiveStorage<T>();
+  return null;
 };
 
-export const storage = resolveStorage<ActiveTrace>();
+class LazyStorage<T> implements IStorage<T> {
+  private inner: IStorage<T>;
+  private resolved = false;
+
+  constructor() {
+    this.inner = tryResolveALS<T>() || new NaiveStorage<T>();
+    this.resolved = !(this.inner instanceof NaiveStorage);
+  }
+
+  private ensureResolved() {
+    if (this.resolved) return;
+    const als = tryResolveALS<T>();
+    if (als) {
+      this.inner = als;
+      this.resolved = true;
+    }
+  }
+
+  run<R>(store: T, callback: (...args: any[]) => R, ...args: any[]): R {
+    this.ensureResolved();
+    return this.inner.run(store, callback, ...args);
+  }
+
+  getStore(): T | undefined {
+    return this.inner.getStore();
+  }
+}
+
+export const storage = new LazyStorage<ActiveTrace>();
 
 export const Context = {
   run: <T>(trace: ActiveTrace, fn: () => T): T => {