From e432cdd7fed5dcb8ee804542f4a3769cd0b8474c Mon Sep 17 00:00:00 2001
From: Attila Szegedi <attila.szegedi@datadoghq.com>
Date: Fri, 26 Jun 2026 16:36:56 +0200
Subject: [PATCH] test(time-profiler): stabilize flaky "should update state"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The test was intermittently failing across platforms (e.g. centos-test-26,
darwin-arm builds) with either "No context found" or "Unexpected context".

Two compounding causes:

1. It established the sampling context with a bare setContext(). Under
   useCPED that only takes effect when the current continuation already has
   an AsyncContextFrame; otherwise it silently no-ops. Whether one existed
   depended on what ran before, so the context was often never stored and
   samples came back with an undefined context.

2. It stopped after a single sample. The first sample is taken at profiler
   start (before any context is set) and is skipped during context
   association, so one sample left nothing reliable to attach the context to.

Establish the context the way production code does — runWithContext when
useCPED is enabled, setContext otherwise — and wait for several samples
before stopping. Reproduced on Node 26: the test failed 30/30 runs in
isolation before, and passes 30/30 in isolation and 25/25 in the full file
after.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 ts/test/test-time-profiler.ts | 81 +++++++++++++++++++++++------------
 1 file changed, 54 insertions(+), 27 deletions(-)

diff --git a/ts/test/test-time-profiler.ts b/ts/test/test-time-profiler.ts
index 54c5b856..ede45f7a 100644
--- a/ts/test/test-time-profiler.ts
+++ b/ts/test/test-time-profiler.ts
@@ -69,42 +69,69 @@ describe('Time Profiler', () => {
         useCPED,
       });
       const initialContext: {[key: string]: string} = {};
-      time.setContext(initialContext);
       const kSampleCount = time.constants.kSampleCount;
       const state = time.getState();
       assert.equal(state[kSampleCount], 0, 'Initial state should be 0');
-      const deadline = Date.now() + 1000;
-      while (state[kSampleCount] === 0) {
-        if (Date.now() > deadline) {
-          assert.fail('State did not change');
-        }
-      }
-      assert(state[kSampleCount] >= 1, 'Unexpected number of samples');
 
       let checked = false;
-      initialContext['aaa'] = 'bbb';
 
-      let endTime = 0n;
-      time.stop(false, ({node, context}: GenerateTimeLabelsArgs) => {
-        if (node.name === time.constants.NON_JS_THREADS_FUNCTION_NAME) {
-          return {};
-        }
-        assert.ok(context !== null, 'Context should not be null');
-        if (!endTime) {
-          endTime = BigInt(Date.now()) * 1000n;
+      // Collect several samples while the context is active, then verify the
+      // context was associated with them. Two things make this robust against
+      // the flakiness this test used to exhibit:
+      //   - We establish the context the way production code does: via
+      //     runWithContext when useCPED is enabled, and setContext otherwise.
+      //     A bare setContext under useCPED only takes effect when the current
+      //     continuation already has an AsyncContextFrame, so on its own it
+      //     silently no-ops — which made the result depend on test ordering.
+      //   - We wait for several samples rather than a single one. The first
+      //     sample is taken at profiler start (before any context is set) and
+      //     is skipped during context association, so relying on just one
+      //     sample left nothing reliable to associate the context with.
+      const minSamples = 5;
+      function sampleAndVerify() {
+        const deadline = Date.now() + 5_000;
+        while (state[kSampleCount] < minSamples) {
+          if (Date.now() > deadline) {
+            assert.fail(
+              `Only ${state[kSampleCount]} sample(s) collected, expected ${minSamples}`,
+            );
+          }
         }
 
-        assert.deepEqual(
-          context!.context,
-          initialContext,
-          'Unexpected context',
-        );
+        // Mutate the context object after the samples were taken: the profiler
+        // stores a live reference to it, so the change must be visible below.
+        initialContext['aaa'] = 'bbb';
+
+        let endTime = 0n;
+        time.stop(false, ({node, context}: GenerateTimeLabelsArgs) => {
+          if (node.name === time.constants.NON_JS_THREADS_FUNCTION_NAME) {
+            return {};
+          }
+          assert.ok(context !== null, 'Context should not be null');
+          if (!endTime) {
+            endTime = BigInt(Date.now()) * 1000n;
+          }
+
+          assert.deepEqual(
+            context!.context,
+            initialContext,
+            'Unexpected context',
+          );
+
+          assert.ok(context!.timestamp >= startTime);
+          assert.ok(context!.timestamp <= endTime);
+          checked = true;
+          return {...context!.context};
+        });
+      }
+
+      if (useCPED) {
+        time.runWithContext(initialContext, sampleAndVerify);
+      } else {
+        time.setContext(initialContext);
+        sampleAndVerify();
+      }
 
-        assert.ok(context!.timestamp >= startTime);
-        assert.ok(context!.timestamp <= endTime);
-        checked = true;
-        return {...context!.context};
-      });
       assert(checked, 'No context found');
     });