[o11y] Consolidate STW instrumentation test tail workers

fhanau · fhanau · commit d4cffdef654d · 2025-10-12T17:48:25.000-04:00
diff --git a/src/workerd/api/BUILD.bazel b/src/workerd/api/BUILD.bazel
@@ -434,6 +434,7 @@ wd_test(
     data = [
         "kv-instrumentation-test.js",
         "kv-test.js",
+        "tests/instrumentation-tail-worker.js",
     ],
 )
 
@@ -449,6 +450,7 @@ wd_test(
     data = [
         "r2-instrumentation-test.js",
         "r2-test.js",
+        "tests/instrumentation-tail-worker.js",
     ],
 )
 
@@ -892,5 +894,6 @@ wd_test(
     data = [
         "actor-kv-test.js",
         "actor-kv-test-tail.js",
+        "tests/instrumentation-tail-worker.js",
     ],
 )
diff --git a/src/workerd/api/actor-kv-test-tail.js b/src/workerd/api/actor-kv-test-tail.js
@@ -1,16 +1,17 @@
 import * as assert from 'node:assert';
+import {
+  invocationPromises,
+  spans,
+  testTailHandler,
+} from 'test:instumentation-tail';
 
-let invocationPromises = [];
-let spans = new Map();
+// Use shared instrumentation test tail worker
+export default testTailHandler;
 
-export default {
+export const test = {
   async test(ctrl, env, ctx) {
     const expected = [
-      {
-        name: 'durable_object_storage_put',
-        closed: true,
-        subrequests: [{ name: 'durable_object_subrequest' }],
-      },
+      { name: 'durable_object_storage_put', closed: true },
       { name: 'durable_object_storage_put', closed: true },
       { name: 'durable_object_storage_get', closed: true },
       { name: 'durable_object_storage_get', closed: true },
@@ -22,6 +23,7 @@ export default {
       { name: 'durable_object_storage_deleteAlarm', closed: true },
       { name: 'durable_object_storage_transaction', closed: true },
       { name: 'durable_object_storage_sync', closed: true },
+      { name: 'durable_object_subrequest', closed: true },
     ];
 
     await Promise.allSettled(invocationPromises);
@@ -31,49 +33,4 @@ export default {
     assert.deepStrictEqual(received, expected);
     return new Response('');
   },
-
-  tailStream(event, env, ctx) {
-    let resolveFn;
-
-    invocationPromises.push(
-      new Promise((resolve, reject) => {
-        resolveFn = resolve;
-      })
-    );
-
-    return (event) => {
-      switch (event.event.type) {
-        case 'spanOpen':
-          if (event.event.name === 'durable_object_subrequest') {
-            let span = spans.get(event.event.spanId);
-            span['subrequests'] = span['subrequests']
-              ? (span['subrequests'].push({ name: { ...event.name } }),
-                span['subrequests'])
-              : [{ name: event.event.name }];
-          } else {
-            // The span ids will change between tests, but Map preserves insertion order
-            spans.set(event.event.spanId, { name: event.event.name });
-          }
-          break;
-        case 'attributes': {
-          let span = spans.get(event.event.spanId);
-          for (let { name, value } of event.event.info) {
-            span[name] = value;
-          }
-          spans.set(event.event.spanId, span);
-          break;
-        }
-        case 'spanClose': {
-          const spanId = event.spanContext.spanId;
-          let span = spans.get(spanId);
-          span['closed'] = true;
-          spans.set(spanId, span);
-          break;
-        }
-        case 'outcome':
-          resolveFn();
-          break;
-      }
-    };
-  },
 };
diff --git a/src/workerd/api/actor-kv-test.wd-test b/src/workerd/api/actor-kv-test.wd-test
@@ -6,6 +6,11 @@ const config :Workerd.Config = (
     (name = "TEST_TMPDIR", disk = (writable = true)),
     (name = "tail", worker = .tailWorker),
   ],
+  extensions = [ (
+    modules = [
+      ( name = "test:instumentation-tail", esModule = embed "tests/instrumentation-tail-worker.js" ),
+    ]
+  ) ]
 );
 
 const mainWorker :Workerd.Worker = (
diff --git a/src/workerd/api/kv-instrumentation-test.js b/src/workerd/api/kv-instrumentation-test.js
@@ -2,73 +2,14 @@
 // Licensed under the Apache 2.0 license found in the LICENSE file or at:
 //     https://opensource.org/licenses/Apache-2.0
 import * as assert from 'node:assert';
+import {
+  invocationPromises,
+  spans,
+  testTailHandler,
+} from 'test:instumentation-tail';
 
-// tailStream is going to be invoked multiple times, but we want to wait
-// to run the test until all executions are done. Collect promises for
-// each
-let invocationPromises = [];
-let spans = new Map();
-
-export default {
-  tailStream(event, env, ctx) {
-    // For each "onset" event, store a promise which we will resolve when
-    // we receive the equivalent "outcome" event
-    let resolveFn;
-    invocationPromises.push(
-      new Promise((resolve, reject) => {
-        resolveFn = resolve;
-      })
-    );
-
-    // Capture the top-level span ID from the onset event
-    const topLevelSpanId = event.event.spanId;
-
-    // Accumulate the span info for easier testing
-    return (event) => {
-      // For spanOpen events, the new span ID is in event.event.spanId
-      // For other events, they reference an existing span via event.spanContext.spanId
-      let spanKey = event.invocationId + event.spanContext.spanId;
-      switch (event.event.type) {
-        case 'spanOpen':
-          // spanOpen creates a new span with ID in event.event.spanId
-          spanKey = event.invocationId + event.event.spanId;
-          spans.set(spanKey, {
-            name: event.event.name,
-          });
-          break;
-        case 'attributes': {
-          // Filter out top-level attributes events (jsRpcSession span)
-          if (topLevelSpanId && event.spanContext.spanId === topLevelSpanId) {
-            // Ignore attributes for the top-level span
-            break;
-          }
-
-          // attributes references an existing span via spanContext.spanId
-          let span = spans.get(spanKey);
-          if (!span) {
-            throw new Error(`Attributes event for unknown span: ${spanKey}`);
-          }
-          for (let { name, value } of event.event.info) {
-            span[name] = value;
-          }
-          break;
-        }
-        case 'spanClose': {
-          // spanClose references an existing span via spanContext.spanId
-          let span = spans.get(spanKey);
-          if (!span) {
-            throw new Error(`SpanClose event for unknown span: ${spanKey}`);
-          }
-          span['closed'] = true;
-          break;
-        }
-        case 'outcome':
-          resolveFn();
-          break;
-      }
-    };
-  },
-};
+// Use shared instrumentation test tail worker
+export default testTailHandler;
 
 export const test = {
   async test() {
diff --git a/src/workerd/api/kv-test.wd-test b/src/workerd/api/kv-test.wd-test
@@ -16,6 +16,11 @@ const unitTests :Workerd.Config = (
     ),
     (name = "tail", worker = .tailWorker, ),
   ],
+  extensions = [ (
+    modules = [
+      ( name = "test:instumentation-tail", esModule = embed "tests/instrumentation-tail-worker.js" ),
+    ]
+  ) ]
 );
 
 const tailWorker :Workerd.Worker = (
diff --git a/src/workerd/api/r2-instrumentation-test.js b/src/workerd/api/r2-instrumentation-test.js
@@ -2,55 +2,14 @@
 // Licensed under the Apache 2.0 license found in the LICENSE file or at:
 //     https://opensource.org/licenses/Apache-2.0
 import * as assert from 'node:assert';
+import {
+  invocationPromises,
+  spans,
+  testTailHandler,
+} from 'test:instumentation-tail';
 
-// tailStream is going to be invoked multiple times, but we want to wait
-// to run the test until all executions are done. Collect promises for
-// each
-let invocationPromises = [];
-let spans = new Map();
-
-export default {
-  tailStream(event, env, ctx) {
-    // For each "onset" event, store a promise which we will resolve when
-    // we receive the equivalent "outcome" event
-    let resolveFn;
-    invocationPromises.push(
-      new Promise((resolve, reject) => {
-        resolveFn = resolve;
-      })
-    );
-
-    // Accumulate the span info for easier testing
-    return (event) => {
-      // span ids are simple counters for tests, but invocation ID allows us to differentiate them
-      let spanKey = event.invocationId + event.spanContext.spanId;
-      switch (event.event.type) {
-        case 'spanOpen':
-          spans.set(event.invocationId + event.event.spanId, {
-            name: event.event.name,
-          });
-          break;
-        case 'attributes': {
-          let span = spans.get(spanKey);
-          for (let { name, value } of event.event.info) {
-            span[name] = value;
-          }
-          spans.set(spanKey, span);
-          break;
-        }
-        case 'spanClose': {
-          let span = spans.get(spanKey);
-          span['closed'] = true;
-          spans.set(spanKey, span);
-          break;
-        }
-        case 'outcome':
-          resolveFn();
-          break;
-      }
-    };
-  },
-};
+// Use shared instrumentation test tail worker
+export default testTailHandler;
 
 export const test = {
   async test() {
diff --git a/src/workerd/api/r2-test.wd-test b/src/workerd/api/r2-test.wd-test
@@ -18,6 +18,11 @@ const unitTests :Workerd.Config = (
     ),
     (name = "tail", worker = .tailWorker, ),
   ],
+  extensions = [ (
+    modules = [
+      ( name = "test:instumentation-tail", esModule = embed "tests/instrumentation-tail-worker.js" ),
+    ]
+  ) ]
 );
 
 const tailWorker :Workerd.Worker = (
diff --git a/src/workerd/api/tests/instrumentation-tail-worker.js b/src/workerd/api/tests/instrumentation-tail-worker.js
@@ -0,0 +1,66 @@
+// tailStream is going to be invoked multiple times, but we want to wait
+// to run the test until all executions are done. Collect promises for
+// each
+export let spans = new Map();
+export let invocationPromises = [];
+
+// tail stream handler function used in several STW instrumentation tests.
+export const testTailHandler = {
+  tailStream(event, env, ctx) {
+    // Capture the top-level span ID from the onset event
+    const topLevelSpanId = event.event.spanId;
+
+    // For each "onset" event, store a promise which we will resolve when
+    // we receive the equivalent "outcome" event
+    let resolveFn;
+    invocationPromises.push(
+      new Promise((resolve, reject) => {
+        resolveFn = resolve;
+      })
+    );
+
+    return (event) => {
+      // For spanOpen events, the new span ID is in event.event.spanId
+      // For other events, they reference an existing span via event.spanContext.spanId
+      let spanKey = event.invocationId + event.spanContext.spanId;
+      switch (event.event.type) {
+        case 'spanOpen':
+          // spanOpen creates a new span with ID in event.event.spanId
+          spanKey = event.invocationId + event.event.spanId;
+          spans.set(spanKey, {
+            name: event.event.name,
+          });
+          break;
+        case 'attributes': {
+          // Filter out top-level attributes events (jsRpcSession span)
+          if (event.spanContext.spanId === topLevelSpanId) {
+            // Ignore attributes for the top-level span
+            break;
+          }
+
+          // attributes references an existing span via spanContext.spanId
+          let span = spans.get(spanKey);
+          if (!span) {
+            throw new Error(`Attributes event for unknown span: ${spanKey}`);
+          }
+          for (let { name, value } of event.event.info) {
+            span[name] = value;
+          }
+          break;
+        }
+        case 'spanClose': {
+          // spanClose references an existing span via spanContext.spanId
+          let span = spans.get(spanKey);
+          if (!span) {
+            throw new Error(`SpanClose event for unknown span: ${spanKey}`);
+          }
+          span['closed'] = true;
+          break;
+        }
+        case 'outcome':
+          resolveFn();
+          break;
+      }
+    };
+  },
+};

Original file line number	Diff line number	Diff line change
`@@ -434,6 +434,7 @@ wd_test(`
`434`	`434`	`data = [`
`435`	`435`	`"kv-instrumentation-test.js",`
`436`	`436`	`"kv-test.js",`
	`437`	`+ "tests/instrumentation-tail-worker.js",`
`437`	`438`	`],`
`438`	`439`	`)`
`439`	`440`
`@@ -449,6 +450,7 @@ wd_test(`
`449`	`450`	`data = [`
`450`	`451`	`"r2-instrumentation-test.js",`
`451`	`452`	`"r2-test.js",`
	`453`	`+ "tests/instrumentation-tail-worker.js",`
`452`	`454`	`],`
`453`	`455`	`)`
`454`	`456`
`@@ -892,5 +894,6 @@ wd_test(`
`892`	`894`	`data = [`
`893`	`895`	`"actor-kv-test.js",`
`894`	`896`	`"actor-kv-test-tail.js",`
	`897`	`+ "tests/instrumentation-tail-worker.js",`
`895`	`898`	`],`
`896`	`899`	`)`