v3.0.2: run shards rediscovery in separate traces

.coderabbit.yaml

@@ -0,0 +1,6 @@
+reviews:
+  auto_review:
+    enabled: true
+    base_branches:
+      - main
+      - "grok/*"

.vscode/settings.json

@@ -2,5 +2,13 @@
   "cSpell.words": [
     "PGDATABASE",
     "PGPORT"
-  ]
+  ],
+  "editor.codeActionsOnSave": {
+    "source.fixAll.eslint": "explicit"
+  },
+  "[typescript]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
+  },
+  "js/ts.experimental.useTsgo": true,
+  "typescript.native-preview.tsdk": "node_modules/@typescript/native-preview"
 }

package.json

@@ -1,7 +1,7 @@
 {
   "name": "ent-framework",
   "description": "A PostgreSQL graph-database-alike library with microsharding and row-level security",
-  "version": "3.0.1",
+  "version": "3.0.2",
   "license": "MIT",
   "keywords": [
     "postgresql",

src/internal/CachedRefreshedValue.ts

@@ -2,7 +2,7 @@ import { Memoize } from "fast-typescript-memoize";
 import type { DeferredPromise } from "p-defer";
 import pDefer from "p-defer";
 import type { MaybeCallable } from "./misc";
-import { maybeCall, runInVoid } from "./misc";
+import { maybeCall, maybeRunInSeparateTrace, runInVoid } from "./misc";
 
 export interface CachedRefreshedValueOptions<TValue> {
   /** Delay between calling resolver. */
@@ -111,6 +111,16 @@ export class CachedRefreshedValue<TValue> {
 
   @Memoize()
   private async refreshLoop(): Promise<void> {
+    // Whether the upcoming resolverFn() call must run in the caller's async
+    // context (and thus inherit its APM/tracing trace id), or in a fresh,
+    // detached trace. The very 1st refresh is triggered synchronously by the
+    // very 1st cached() caller (e.g. the 1st query), so it intentionally stays
+    // in that caller's trace. Same for refreshes triggered manually via
+    // refreshAndWait() (e.g. Cluster#rediscover()). But refreshes triggered on
+    // the timer (or by a deps change) must NOT pollute the trace of whoever
+    // happened to start the loop - each such refresh gets its own fresh trace.
+    let keepCallerTrace = true;
+
     while (!this.destroyedError) {
       const warningDelayMs = maybeCall(this.options.warningTimeoutMs);
       const depsDelayMs = maybeCall(this.options.deps.delayMs);
@@ -131,22 +141,30 @@ export class CachedRefreshedValue<TValue> {
       let depsPrev: unknown = undefined;
       try {
         this.resolverFnCallCount++;
-        depsPrev = await this.options.deps.handler();
-        this.latestValue = await this.options.resolverFn();
+        await maybeRunInSeparateTrace(keepCallerTrace, async () => {
+          depsPrev = await this.options.deps.handler();
+          this.latestValue = await this.options.resolverFn();
+        });
         const oldNextValue = this.nextValue;
         this.nextValue = pDefer();
-        oldNextValue.resolve(this.latestValue);
+        oldNextValue.resolve(this.latestValue!);
       } catch (e: unknown) {
         this.onError(e, Math.round(performance.now() - startTime));
       } finally {
         clearTimeout(warningTimeout);
       }
 
       // Wait for delayMs. If this.skipDelay() is called, the code unfreezes
-      // immediately. Also, deps are rechecked every depsDelayMs, and if they
-      // change, the code unfreezes too.
+      // immediately, and the next refresh is treated as caller-triggered (so it
+      // keeps the trace). Also, deps are rechecked every depsDelayMs, and if
+      // they change, the code unfreezes too (and the next refresh is treated as
+      // timer-triggered, i.e. it runs in a fresh detached trace).
+      keepCallerTrace = false;
       const delayDefer = pDefer<void>();
-      this.skipDelay = () => delayDefer.resolve();
+      this.skipDelay = () => {
+        keepCallerTrace = true;
+        delayDefer.resolve();
+      };
 
       let depsTimeoutBody: null | (() => void) = () =>
         runInVoid(async () => {

src/internal/__tests__/CachedRefreshedValue.test.ts

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from "async_hooks";
 import delay from "delay";
 import pDefer from "p-defer";
 import waitForExpect from "wait-for-expect";
@@ -329,3 +330,87 @@ test("changes in deps are respected", async () => {
   depsValue = "other";
   await waitForExpect(async () => expect(await cache.cached()).toEqual("two"));
 });
+
+test("timer-triggered refreshes each run in a separate trace", async () => {
+  // Models how an async_hooks based APM tracer derives a "trace id" from the
+  // current async context: if a query runs with an active trace, it joins it;
+  // otherwise (i.e. when the framework detached us into a fresh async context),
+  // the tracer opens a brand-new trace AND pins it to the current context via
+  // enterWith() (so it would "leak" into the next refresh and coalesce all
+  // refreshes into one trace, were the detachment not restoring a clean
+  // baseline on every iteration).
+  const als = new AsyncLocalStorage<{ traceId: string }>();
+  let detachedTraceSeq = 0;
+  const traces: string[] = [];
+  const recordTrace = (): void => {
+    let store = als.getStore();
+    if (!store) {
+      store = { traceId: `detached-trace-${++detachedTraceSeq}` };
+      als.enterWith(store);
+    }
+
+    traces.push(store.traceId);
+  };
+
+  cache = new CachedRefreshedValue({
+    ...OPTIONS,
+    delayMs: 10,
+    resolverFn: async () => {
+      recordTrace();
+      return "value";
+    },
+  });
+
+  // The very 1st refresh is triggered synchronously by this cached() call (akin
+  // to the 1st query), so it must run in (and thus inherit) the caller's trace.
+  await als.run({ traceId: "caller-trace" }, async () => cache.cached());
+
+  // Let several timer-triggered refreshes happen.
+  await waitForExpect(() => expect(traces.length).toBeGreaterThanOrEqual(4));
+
+  expect(traces[0]).toBe("caller-trace");
+
+  // Every subsequent (timer-triggered) refresh ran detached from the caller's
+  // trace, and each one got its own brand-new trace.
+  const timerTraces = traces.slice(1);
+  expect(
+    timerTraces.every((trace) => trace.startsWith("detached-trace-")),
+  ).toBe(true);
+  expect(new Set(timerTraces).size).toBe(timerTraces.length);
+});
+
+test("manual refreshAndWait() keeps the caller's trace", async () => {
+  const als = new AsyncLocalStorage<{ traceId: string }>();
+  let detachedTraceSeq = 0;
+  const traces: string[] = [];
+  const recordTrace = (): void => {
+    let store = als.getStore();
+    if (!store) {
+      store = { traceId: `detached-trace-${++detachedTraceSeq}` };
+      als.enterWith(store);
+    }
+
+    traces.push(store.traceId);
+  };
+
+  cache = new CachedRefreshedValue({
+    ...OPTIONS,
+    // Large delay, so no timer-triggered refresh interferes: only the initial
+    // refresh and the manual refreshAndWait() ones happen.
+    delayMs: 1_000_000,
+    resolverFn: async () => {
+      recordTrace();
+      return "value";
+    },
+  });
+
+  await als.run({ traceId: "caller-trace" }, async () => {
+    await cache.cached(); // initial refresh
+    await cache.refreshAndWait(); // manual rediscovery
+    await cache.refreshAndWait(); // manual rediscovery
+  });
+
+  // Neither the initial nor the manual refreshes were detached into a fresh
+  // trace: they all kept running in the loop's (caller's) trace.
+  expect(traces).toEqual(["caller-trace", "caller-trace", "caller-trace"]);
+});

src/internal/misc.ts

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from "async_hooks";
 import { createHash } from "crypto";
 import { inspect } from "util";
 import compact from "lodash/compact";
@@ -399,6 +400,43 @@ export function runInVoid(
   }
 }
 
+/**
+ * A snapshot of the "clean" async context, captured at module load time, i.e.
+ * BEFORE any request-scoped async context (such as an incoming HTTP request, and
+ * thus the APM/tracing "trace id" that tracers propagate through
+ * AsyncLocalStorage) could possibly have been established. Calling the returned
+ * runner executes a callback within that captured clean context.
+ *
+ * We deliberately use AsyncLocalStorage.snapshot() and NOT AsyncResource: both
+ * can detach an async context, but AsyncLocalStorage (including .snapshot()) is
+ * supported across runtimes (Node, Deno, Bun), whereas AsyncResource is a
+ * non-functional stub in some of them. If .snapshot() is unavailable (very old
+ * runtimes), we degrade gracefully to running the callback inline (no
+ * detachment).
+ */
+const runInRootAsyncContext: <T>(func: () => T) => T =
+  typeof AsyncLocalStorage?.snapshot === "function"
+    ? AsyncLocalStorage.snapshot()
+    : (func) => func();
+
+/**
+ * Runs an async function detached from the async context (and thus from the
+ * APM/tracing trace id, if any) that is active at the call site: the callback
+ * runs in the "clean" context that was captured at module load time instead.
+ * Each call starts its own fresh async subtree rooted in that clean context, so:
+ * - async_hooks-based tracers observe each call as a separate, brand-new trace
+ *   (the tracer sees no active trace and starts a new one), and
+ * - all the async work spawned within a single call shares that one trace.
+ *
+ * This is used for background loops (e.g. Shards rediscovery) which are kicked
+ * off lazily by the very first query: without this, every subsequent timer-
+ * triggered loop iteration would forever remain attached to (and would thus
+ * pollute) the trace id of that very first query.
+ */
+export async function maybeRunInSeparateTrace<T>(keepCallerTrace: boolean, func: () => Promise<T>): Promise<T> {
+  return keepCallerTrace ? func() : runInRootAsyncContext(func);
+}
+
 /**
  * A typesafe-way to invariant the object's key presence and being
  * non-undefined. It is not always working for union types: sometimes it asserts