[o11y][nfc] Refactor tests, deduplicate tailStream() implementations, update comments

Author: fhanau

src/cloudflare/internal/test/BUILD.bazel

@@ -4,6 +4,5 @@ exports_files(
         "//src/cloudflare/internal/test/d1:__pkg__",
         "//src/cloudflare/internal/test/images:__pkg__",
         "//src/cloudflare/internal/test/tracing:__pkg__",
-        "//src/workerd/api:__pkg__",
     ],
 )

src/cloudflare/internal/test/d1/d1-api-instrumentation-test.js

@@ -1,7 +1,7 @@
 // Copyright (c) 2017-2023 Cloudflare, Inc.
 // 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 {
   createInstrumentationState,
   createTailStreamHandler,

src/cloudflare/internal/test/instrumentation-test-helper.js

@@ -138,60 +138,20 @@ export async function runInstrumentationTest(
  * @returns {Object} An object with methods to handle spans
  */
 export function createTailStreamCollector() {
-  const invocationPromises = [];
-  const spans = new Map();
+  let state = createInstrumentationState();
 
-  const 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) => {
-      let spanKey = `${event.invocationId}#${event.event.spanId || event.spanContext.spanId}`;
-      switch (event.event.type) {
-        case 'spanOpen':
-          // The span ids will change between tests, but Map preserves insertion order
-          spans.set(spanKey, { 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;
-      }
-    };
-  };
-
-  const waitForCompletion = async () => {
-    await Promise.allSettled(invocationPromises);
-  };
+  const tailStream = createTailStreamHandler(state);
 
-  const getSpans = () => {
-    return Array.from(spans.values());
+  let spans = state.spans;
+  let invocationPromises = state.invocationPromises;
+  const waitForCompletion = () => {
+    return Promise.allSettled(invocationPromises);
   };
 
   return {
     tailStream,
     waitForCompletion,
-    getSpans,
+    spans,
   };
 }
 

src/cloudflare/internal/test/tracing/tracing-helpers-instrumentation-test.js

@@ -19,7 +19,7 @@ export const validateSpans = {
     await collector.waitForCompletion();
 
     // Get all spans and prepare for validation
-    const allSpans = collector.getSpans();
+    const allSpans = collector.spans.values();
     const spansByTest = groupSpansBy(allSpans, 'test');
 
     // Define the core tests that validate withSpan behavior

src/workerd/io/trace.c++

@@ -1484,12 +1484,6 @@ CompleteSpan CompleteSpan::clone() const {
   }
   return copy;
 }
-
-kj::String CompleteSpan::toString() const {
-  return kj::str("CompleteSpan: ", operationName,
-      kj::strArray(
-          KJ_MAP(tag, tags) { return kj::str("(", tag.key, ", ", tag.value, ")"); }, ", "));
-}
 }  // namespace tracing
 
 // ======================================================================================

src/workerd/io/trace.h

@@ -899,14 +899,9 @@ inline kj::String truncateScriptId(kj::StringPtr id) {
 // =======================================================================================
 // Span tracing
 //
-// TODO(cleanup): As of now, this aspect of tracing is actually not related to the rest of this
-//   file. Most of this file defines the interface to feed Trace Workers. Span tracing, however,
-//   is currently designed to feed tracing of the Workers Runtime itself for the benefit of the
-//   developers of the runtime.
-//
-//   We might potentially want to give trace workers some access to span tracing as well, but with
-//   that the trace worker and span interfaces should still be largely independent of each other;
-//   separate span tracing into a separate header.
+// TODO(cleanup): (Streaming) tail workers now have access to span tracing as well, but with that
+// the trace worker and span interfaces are still mostly independent of each other; separate span
+// tracing into a separate header.
 
 class SpanBuilder;
 class SpanObserver;

src/workerd/io/worker-entrypoint.c++

@@ -560,9 +560,8 @@ kj::Promise<void> WorkerEntrypoint::prewarm(kj::StringPtr url) {
       kj::mv(KJ_REQUIRE_NONNULL(this->incomingRequest, "prewarm() can only be called once"));
   incomingRequest->getMetrics().setIsPrewarm();
 
-  // TODO(streaming-tail): Should prewarm be reflected in the tail stream somehow?
-
-  // Intentionally don't call incomingRequest->delivered() for prewarm requests.
+  // Intentionally don't call incomingRequest->delivered() for prewarm requests and do not create
+  // an Onset event, prewarm is not being traced.
 
   // TODO(someday): Ideally, middleware workers would forward prewarm() to the next stage. At
   //   present we don't have a good way to decide what stage that is, especially given that we'll