Merge pull request #7348 from cloudflare/jmorrell/add-span-status

[WO-1581] Enable span status and renaming

Author: jmorrell-cloudflare

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

@@ -23,6 +23,7 @@ export function createInstrumentationState() {
     invocationPromises: [],
     invocations: new Map(),
     spans: new Map(),
+    spanUpdates: [],
   };
 }
 
@@ -85,6 +86,10 @@ export function createTailStreamHandler(state) {
           state.spans.set(spanKey, span);
           break;
         }
+        case 'spanUpdate': {
+          state.spanUpdates.push(event);
+          break;
+        }
         case 'outcome':
           invocation.outcome = event.event;
           resolveFn();
@@ -259,6 +264,7 @@ export function createTailStreamCollector() {
   const tailStream = createTailStreamHandler(state);
 
   const spans = state.spans;
+  const spanUpdates = state.spanUpdates;
   const invocations = state.invocations;
   const invocationPromises = state.invocationPromises;
   const waitForCompletion = () => {
@@ -270,6 +276,7 @@ export function createTailStreamCollector() {
     waitForCompletion,
     invocations,
     spans,
+    spanUpdates,
   };
 }
 

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

@@ -22,6 +22,14 @@ export const validateSpans = {
     const allSpans = collector.spans.values();
     const spansByTest = groupSpansBy(allSpans, 'test');
     const invocations = [...collector.invocations.values()];
+    const truncatedUpdatedName = `updated-${'x'.repeat(56)}`;
+    const spanUpdatesFor = (spanKey) =>
+      collector.spanUpdates
+        .filter(
+          (update) =>
+            `${update.invocationId}#${update.spanContext.spanId}` === spanKey
+        )
+        .map((update) => update.event.info);
     const rootAttributes = invocations.flatMap(
       (invocation) => invocation.attributes
     );
@@ -39,6 +47,8 @@ export const validateSpans = {
         expectedSpan: 'undefined-attr-op',
       },
       { test: 'setAttributes', expectedSpan: 'set-attributes-op' },
+      { test: 'setStatus', expectedSpan: 'status-error-op' },
+      { test: 'updateName', expectedSpan: 'update-name-original' },
       { test: 'publicImportTracing', expectedSpan: 'public-import-op' },
       {
         test: 'publicImportStartActiveSpan',
@@ -66,6 +76,66 @@ export const validateSpans = {
       assert(span.closed, `${test}: Span '${expectedSpan}' should be closed`);
     }
 
+    {
+      const statusSpans = spansByTest.get('setStatus') || [];
+      const errorSpan = statusSpans.find((s) => s.name === 'status-error-op');
+      const okSpan = statusSpans.find((s) => s.name === 'status-ok-op');
+
+      const entries = [...collector.spans.entries()];
+      const errorKey = entries.find(([, span]) => span === errorSpan)[0];
+      const okKey = entries.find(([, span]) => span === okSpan)[0];
+      assert.deepStrictEqual(spanUpdatesFor(errorKey), [
+        {
+          type: 'status',
+          status: { code: 'error', message: 'first error' },
+        },
+        {
+          type: 'status',
+          status: { code: 'error', message: 'second error' },
+        },
+        { type: 'status', status: { code: 'unset' } },
+      ]);
+      assert.deepStrictEqual(spanUpdatesFor(okKey), [
+        {
+          type: 'status',
+          status: { code: 'error', message: 'temporary error' },
+        },
+        { type: 'status', status: { code: 'ok' } },
+        {
+          type: 'status',
+          status: { code: 'error', message: 'error after ok' },
+        },
+      ]);
+    }
+
+    {
+      const [[spanKey]] = [...collector.spans.entries()].filter(
+        ([, span]) => span.test === 'updateName'
+      );
+      assert.deepStrictEqual(spanUpdatesFor(spanKey), [
+        { type: 'name', name: 'update-name-intermediate' },
+        { type: 'name', name: truncatedUpdatedName },
+      ]);
+    }
+
+    {
+      const invocation = invocations.find((candidate) =>
+        candidate.attributes.some(
+          ({ name, value }) =>
+            name === 'test' && value === 'updateInvocationSpan'
+        )
+      );
+      assert(invocation, 'updateInvocationSpan: invocation present');
+      const spanKey = `${invocation.invocationId}#${invocation.rootSpanId}`;
+      assert.deepStrictEqual(spanUpdatesFor(spanKey), [
+        { type: 'name', name: 'updated-invocation' },
+        {
+          type: 'status',
+          status: { code: 'error', message: 'invocation error' },
+        },
+      ]);
+    }
+
     // setAttributeUndefined should NOT have a 'skipped' attribute recorded.
     {
       const span = (spansByTest.get('setAttributeUndefined') || []).find(

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

@@ -224,6 +224,53 @@ export const setAttributes = {
   },
 };
 
+export const setStatus = {
+  async test(ctrl, env, ctx) {
+    const errorSpan = publicTracing.startSpan('status-error-op');
+    errorSpan.setAttribute('test', 'setStatus');
+    assert.strictEqual(
+      errorSpan.setStatus({ code: 'error', message: 'first error' }),
+      errorSpan
+    );
+    errorSpan.setStatus({ code: 'error', message: 'second error' });
+    errorSpan.setStatus({ code: 'unset' });
+    errorSpan.end();
+    // All span mutations are no-ops after end().
+    errorSpan.setStatus({ code: 'ok' });
+
+    const okSpan = publicTracing.startSpan('status-ok-op');
+    okSpan.setAttribute('test', 'setStatus');
+    okSpan.setStatus({ code: 'error', message: 'temporary error' });
+    okSpan.setStatus({ code: 'ok', message: 'not retained' });
+    okSpan.setStatus({ code: 'error', message: 'error after ok' });
+    okSpan.end();
+  },
+};
+
+export const updateName = {
+  async test() {
+    const span = publicTracing.startSpan('update-name-original');
+    span.setAttribute('test', 'updateName');
+    assert.strictEqual(span.updateName('update-name-intermediate'), span);
+    span.updateName(`updated-${'x'.repeat(100)}`);
+    span.end();
+    span.updateName('update-name-after-end');
+  },
+};
+
+export const updateInvocationSpan = {
+  async test() {
+    const span = publicTracing.getActiveSpan();
+    assert(span);
+    span.setAttribute('test', 'updateInvocationSpan');
+    assert.strictEqual(span.updateName('updated-invocation'), span);
+    assert.strictEqual(
+      span.setStatus({ code: 'error', message: 'invocation error' }),
+      span
+    );
+  },
+};
+
 // Verify that nested withSpan calls produce correctly nested spans. This exercises the
 // AsyncContextFrame push path in enterSpan: the inner span should be parented on the
 // outer span.

src/cloudflare/internal/tracing.d.ts

@@ -33,6 +33,14 @@ interface ExceptionWithName {
 type Exception =
   ExceptionWithCode | ExceptionWithMessage | ExceptionWithName | string;
 
+type TracingSpanStatusCode = 'unset' | 'ok' | 'error';
+
+interface TracingSpanStatus {
+  code: TracingSpanStatusCode;
+  /** A developer-facing error message. Ignored unless code is "error". */
+  message?: string;
+}
+
 declare class Span {
   // Returns true if this span will be recorded to the tracing system. False when the
   // current async context is not being traced, or when the span has already been submitted.
@@ -48,6 +56,13 @@ declare class Span {
   // Records an exception event on the span. Calls after the span has ended are ignored.
   recordException(exception: Exception): void;
 
+  // Changes the span name. Calls after the span has ended are ignored.
+  updateName(name: string): this;
+
+  // Sets the span status. Calls after the span has ended are ignored. Messages are retained only
+  // for errors.
+  setStatus(status: TracingSpanStatus): this;
+
   // Ends the span and submits its attributes to the tracing system. Idempotent.
   end(): void;
 }

src/workerd/api/tracing.c++

@@ -125,6 +125,14 @@ class UserSpanState final: public SpanState {
     return workerd::SpanParent(builder);
   }
 
+  void updateName(kj::ConstString operationName) override {
+    builder.setOperationName(kj::mv(operationName));
+  }
+
+  void setStatus(tracing::SpanStatus status) override {
+    builder.setStatus(kj::mv(status));
+  }
+
  protected:
   bool canRecordAttributes() override {
     return builder.isObserved();
@@ -199,6 +207,27 @@ class InvocationSpanState final: public SpanState {
     return parent.addRef();
   }
 
+  void updateName(kj::ConstString operationName) override {
+    KJ_IF_SOME(valueContext, context) {
+      KJ_IF_SOME(valueTracer, tracer) {
+        valueTracer->runIfAlive([&](BaseTracer& tracer) {
+          tracer.addSpanUpdate(
+              valueContext.getSpanId(), tracing::SpanUpdate(kj::mv(operationName)));
+        });
+      }
+    }
+  }
+
+  void setStatus(tracing::SpanStatus status) override {
+    KJ_IF_SOME(valueContext, context) {
+      KJ_IF_SOME(valueTracer, tracer) {
+        valueTracer->runIfAlive([&](BaseTracer& tracer) {
+          tracer.addSpanUpdate(valueContext.getSpanId(), tracing::SpanUpdate(kj::mv(status)));
+        });
+      }
+    }
+  }
+
  protected:
   bool canRecordAttributes() override {
     return getIsTraced();
@@ -260,6 +289,10 @@ class NoopSpanState final: public SpanState {
     return workerd::SpanParent(nullptr);
   }
 
+  void updateName(kj::ConstString) override {}
+
+  void setStatus(tracing::SpanStatus) override {}
+
  protected:
   bool canRecordAttributes() override {
     return false;
@@ -365,6 +398,55 @@ void Span::recordException(
   }
 }
 
+jsg::Ref<Span> Span::updateName(jsg::Lock& js, kj::String operationName) {
+  if (operationName.size() > MAX_USER_OPERATION_NAME_BYTES) {
+    operationName = kj::str(operationName.first(MAX_USER_OPERATION_NAME_BYTES));
+  }
+  auto name = kj::ConstString(kj::mv(operationName));
+  KJ_SWITCH_ONEOF(state) {
+    KJ_CASE_ONEOF(s, kj::Own<SpanState>) {
+      s->updateName(kj::mv(name));
+    }
+    KJ_CASE_ONEOF(s, IoOwn<SpanState>) {
+      s->updateName(kj::mv(name));
+    }
+  }
+  return JSG_THIS;
+}
+
+jsg::Ref<Span> Span::setStatus(jsg::Lock& js, TracingSpanStatus status) {
+  tracing::SpanStatusCode code;
+  if (status.code == "unset") {
+    code = tracing::SpanStatusCode::UNSET;
+  } else if (status.code == "ok") {
+    code = tracing::SpanStatusCode::OK;
+  } else if (status.code == "error") {
+    code = tracing::SpanStatusCode::ERROR;
+  } else {
+    JSG_FAIL_REQUIRE(TypeError, "Span status code must be 'unset', 'ok', or 'error'.");
+  }
+
+  kj::Maybe<kj::ConstString> message;
+  if (code == tracing::SpanStatusCode::ERROR) {
+    KJ_IF_SOME(value, status.message) {
+      if (value.size() > MAX_SPAN_BYTES) {
+        value = kj::str(value.first(MAX_SPAN_BYTES));
+      }
+      message = kj::ConstString(kj::mv(value));
+    }
+  }
+  tracing::SpanStatus internalStatus(code, kj::mv(message));
+  KJ_SWITCH_ONEOF(state) {
+    KJ_CASE_ONEOF(s, kj::Own<SpanState>) {
+      s->setStatus(kj::mv(internalStatus));
+    }
+    KJ_CASE_ONEOF(s, IoOwn<SpanState>) {
+      s->setStatus(kj::mv(internalStatus));
+    }
+  }
+  return JSG_THIS;
+}
+
 void Span::end() {
   KJ_SWITCH_ONEOF(state) {
     KJ_CASE_ONEOF(s, kj::Own<SpanState>) {

src/workerd/api/tracing.h

@@ -19,7 +19,7 @@ class Tracing;  // Forward decl; defined further down after user_tracing::Span.
 // the surrounding workerd::api namespace.
 namespace workerd::api::user_tracing {
 
-// Max length of a user-supplied operation name in `ctx.tracing.enterSpan(name, ...)`.
+// Max length of a user-supplied span operation name.
 // Longer names are truncated at the API surface so the limit holds for every downstream
 // SpanSubmitter. Span names identify operations, not carry data; the bound is tight on
 // purpose.
@@ -39,6 +39,19 @@ struct ExceptionData {
   JSG_STRUCT(code, name, message, stack);
 };
 
+struct TracingSpanStatus {
+  kj::String code;
+  jsg::Optional<kj::String> message;
+
+  JSG_STRUCT(code, message);
+  JSG_STRUCT_TS_DEFINE(type TracingSpanStatusCode = "unset" | "ok" | "error");
+  JSG_STRUCT_TS_OVERRIDE({
+    code: TracingSpanStatusCode;
+    /** A developer-facing error message. Ignored unless code is "error". */
+    message?: string;
+  });
+};
+
 // Polymorphic state behind the JS Span wrapper. Concrete states represent recording user spans and
 // no-op spans, while sharing JS-side attribute byte-limit enforcement.
 class SpanState: public kj::Refcounted {
@@ -56,6 +69,9 @@ class SpanState: public kj::Refcounted {
   // ended or has no observer. Used by Tracing methods to push onto the AsyncContextFrame.
   virtual workerd::SpanParent makeSpanParent() = 0;
 
+  virtual void updateName(kj::ConstString operationName) = 0;
+  virtual void setStatus(tracing::SpanStatus status) = 0;
+
   // Sets a single attribute on the span. If value is kj::none, the attribute is not set.
   void setAttribute(kj::String key, kj::Maybe<TagValue> maybeValue);
 
@@ -106,6 +122,13 @@ class Span: public jsg::Object {
   void recordException(
       jsg::Lock& js, jsg::Value exception, const jsg::TypeHandler<ExceptionData>& exceptionHandler);
 
+  // Changes the span name. Calls after the span has ended are ignored.
+  jsg::Ref<Span> updateName(jsg::Lock& js, kj::String operationName);
+
+  // Sets the span status. Calls after the span has ended are ignored. Messages are retained only
+  // for errors.
+  jsg::Ref<Span> setStatus(jsg::Lock& js, TracingSpanStatus status);
+
   // Ends the span and submits its content to the tracing system. Idempotent.
   void end();
 
@@ -115,6 +138,8 @@ class Span: public jsg::Object {
     JSG_METHOD(setAttribute);
     JSG_METHOD(setAttributes);
     JSG_METHOD(recordException);
+    JSG_METHOD(updateName);
+    JSG_METHOD(setStatus);
     JSG_METHOD(end);
 
     JSG_TS_OVERRIDE({
@@ -126,6 +151,8 @@ class Span: public jsg::Object {
         | { code: string | number; name?: string; message?: string; stack?: string }
         | { code?: string | number; name: string; message?: string; stack?: string }
         | { code?: string | number; name?: string; message: string; stack?: string }): void;
+      updateName(name: string): this;
+      setStatus(status: TracingSpanStatus): this;
     });
   }
 
@@ -242,4 +269,5 @@ kj::Own<jsg::modules::ModuleBundle> getInternalTracingModuleBundle(auto featureF
 }  // namespace workerd::api
 
 #define EW_TRACING_ISOLATE_TYPES                                                                   \
-  api::Tracing, api::user_tracing::Span, api::user_tracing::ExceptionData
+  api::Tracing, api::user_tracing::Span, api::user_tracing::ExceptionData,                         \
+      api::user_tracing::TracingSpanStatus