[o11y] Add traceId to user trace

src/workerd/api/trace.c++

@@ -15,6 +15,7 @@
 #include <workerd/util/uuid.h>
 
 #include <capnp/schema.h>
+#include <kj/encoding.h>
 
 namespace workerd::api {
 
@@ -74,15 +75,15 @@ jsg::V8Ref<v8::Object> getTraceLogMessage(jsg::Lock& js, const tracing::Log& log
 }
 
 kj::Array<jsg::Ref<TraceLog>> getTraceLogs(jsg::Lock& js, const Trace& trace) {
-  auto builder = kj::heapArrayBuilder<jsg::Ref<TraceLog>>(trace.logs.size() + trace.spans.size());
-  for (auto i: kj::indices(trace.logs)) {
-    builder.add(jsg::alloc<TraceLog>(js, trace, trace.logs[i]));
-  }
-  // Add spans represented as logs to the logs object.
-  for (auto i: kj::indices(trace.spans)) {
-    builder.add(jsg::alloc<TraceLog>(js, trace, trace.spans[i]));
-  }
-  return builder.finish();
+  return KJ_MAP(x, trace.logs) -> jsg::Ref<TraceLog> { return jsg::alloc<TraceLog>(js, trace, x); };
+}
+
+kj::Array<jsg::Ref<OTelSpan>> getTraceSpans(const Trace& trace) {
+  return KJ_MAP(x, trace.spans) -> jsg::Ref<OTelSpan> { return jsg::alloc<OTelSpan>(x); };
+}
+
+kj::Maybe<kj::String> getTraceIdStr(const Trace& trace) {
+  return trace.traceId.map([](const auto& traceId) { return traceId.toNEHex(); });
 }
 
 kj::Array<jsg::Ref<TraceDiagnosticChannelEvent>> getTraceDiagnosticChannelEvents(
@@ -205,6 +206,8 @@ TraceItem::TraceItem(jsg::Lock& js, const Trace& trace)
       dispatchNamespace(trace.dispatchNamespace.map([](auto& ns) { return kj::str(ns); })),
       scriptTags(getTraceScriptTags(trace)),
       executionModel(enumToStr(trace.executionModel)),
+      spans(getTraceSpans(trace)),
+      traceId(getTraceIdStr(trace)),
       outcome(enumToStr(trace.outcome)),
       cpuTime(trace.cpuTime / kj::MILLISECONDS),
       wallTime(trace.wallTime / kj::MILLISECONDS),
@@ -286,6 +289,15 @@ kj::StringPtr TraceItem::getExecutionModel() {
   return executionModel;
 }
 
+kj::ArrayPtr<jsg::Ref<OTelSpan>> TraceItem::getSpans() {
+  return spans;
+}
+
+kj::StringPtr TraceItem::getTraceId() {
+  // TODO(o11y): Handle this better
+  return traceId.orDefault(kj::str());
+}
+
 kj::StringPtr TraceItem::getOutcome() {
   return outcome;
 }
@@ -549,6 +561,47 @@ bool TraceItem::HibernatableWebSocketEventInfo::Close::getWasClean() {
   return eventInfo.wasClean;
 }
 
+kj::StringPtr OTelSpan::getOperation() {
+  return operation;
+}
+
+kj::Date OTelSpan::getStartTime() {
+  return startTime;
+}
+
+kj::StringPtr OTelSpan::getSpanID() {
+  return spanId;
+}
+kj::StringPtr OTelSpan::getParentSpanID() {
+  return parentSpanId;
+}
+
+kj::Date OTelSpan::getEndTime() {
+  return endTime;
+}
+
+kj::ArrayPtr<OTelSpanTag> OTelSpan::getTags() {
+  return tags;
+}
+
+OTelSpan::OTelSpan(const CompleteSpan& span)
+    : operation(kj::str(span.operationName)),
+      startTime(span.startTime),
+      endTime(span.endTime),
+      tags(kj::heapArray<OTelSpanTag>(span.tags.size())) {
+  // IDs are represented as network-order hex strings.
+  uint64_t netSpanId = __builtin_bswap64(span.spanId);
+  uint64_t netParentSpanId = __builtin_bswap64(span.parentSpanId);
+  spanId = kj::encodeHex(kj::ArrayPtr<byte>((kj::byte*)&netSpanId, sizeof(uint64_t)));
+  parentSpanId = kj::encodeHex(kj::ArrayPtr<byte>((kj::byte*)&netParentSpanId, sizeof(uint64_t)));
+  uint32_t i = 0;
+  for (auto& tag: span.tags) {
+    tags[i].key = kj::str(tag.key);
+    tags[i].value = spanTagClone(tag.value);
+    i++;
+  }
+}
+
 TraceLog::TraceLog(jsg::Lock& js, const Trace& trace, const tracing::Log& log)
     : timestamp(getTraceLogTimestamp(log)),
       level(getTraceLogLevel(log)),

src/workerd/api/trace.h

@@ -66,6 +66,54 @@ struct ScriptVersion {
   }
 };
 
+struct OTelSpanTag final: public jsg::Object {
+  kj::String key;
+  kj::OneOf<bool, int64_t, double, kj::String> value;
+  JSG_STRUCT(key, value);
+};
+
+// OpenTelemetry-compatible span data exposed as part of the trace. Loosely based on https://github.com/open-telemetry/opentelemetry-js/blob/v1.28.0/experimental/packages/otlp-transformer/src/trace/types.ts#L64
+class OTelSpan final: public jsg::Object {
+ public:
+  OTelSpan(const CompleteSpan& span);
+  kj::StringPtr getSpanID();
+  kj::StringPtr getParentSpanID();
+  kj::StringPtr getOperation();
+  kj::ArrayPtr<OTelSpanTag> getTags();
+  kj::Date getStartTime();
+  kj::Date getEndTime();
+
+  JSG_RESOURCE_TYPE(OTelSpan) {
+    JSG_LAZY_READONLY_INSTANCE_PROPERTY(spanId, getSpanID);
+    JSG_LAZY_READONLY_INSTANCE_PROPERTY(parentSpanId, getParentSpanID);
+    JSG_LAZY_READONLY_INSTANCE_PROPERTY(operation, getOperation);
+    JSG_LAZY_READONLY_INSTANCE_PROPERTY(tags, getTags);
+    JSG_LAZY_READONLY_INSTANCE_PROPERTY(startTime, getStartTime);
+    JSG_LAZY_READONLY_INSTANCE_PROPERTY(endTime, getEndTime);
+  }
+
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("operation", operation);
+    for (const OTelSpanTag& tag: tags) {
+      tracker.trackField("key", tag.key);
+      KJ_SWITCH_ONEOF(tag.value) {
+        KJ_CASE_ONEOF(str, kj::String) {
+          tracker.trackField("value", str);
+        }
+        KJ_CASE_ONEOF_DEFAULT break;
+      }
+    }
+  }
+
+ private:
+  kj::String spanId;
+  kj::String parentSpanId;
+  kj::String operation;
+  kj::Date startTime;
+  kj::Date endTime;
+  kj::Array<OTelSpanTag> tags;
+};
+
 class TraceItem final: public jsg::Object {
  public:
   class FetchEventInfo;
@@ -102,16 +150,22 @@ class TraceItem final: public jsg::Object {
   jsg::Optional<kj::StringPtr> getDispatchNamespace();
   jsg::Optional<kj::Array<kj::StringPtr>> getScriptTags();
   kj::StringPtr getExecutionModel();
+  kj::ArrayPtr<jsg::Ref<OTelSpan>> getSpans();
+  kj::StringPtr getTraceId();
   kj::StringPtr getOutcome();
 
   uint getCpuTime();
   uint getWallTime();
   bool getTruncated();
 
-  JSG_RESOURCE_TYPE(TraceItem) {
+  JSG_RESOURCE_TYPE(TraceItem, CompatibilityFlags::Reader flags) {
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(event, getEvent);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(eventTimestamp, getEventTimestamp);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(logs, getLogs);
+    if (flags.getTailWorkerUserSpans()) {
+      JSG_LAZY_READONLY_INSTANCE_PROPERTY(spans, getSpans);
+      JSG_LAZY_READONLY_INSTANCE_PROPERTY(traceId, getTraceId);
+    }
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(exceptions, getExceptions);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(diagnosticsChannelEvents, getDiagnosticChannelEvents);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(scriptName, getScriptName);
@@ -138,6 +192,8 @@ class TraceItem final: public jsg::Object {
   kj::Maybe<kj::String> dispatchNamespace;
   jsg::Optional<kj::Array<kj::String>> scriptTags;
   kj::String executionModel;
+  kj::Array<jsg::Ref<OTelSpan>> spans;
+  kj::Maybe<kj::String> traceId;
   kj::String outcome;
   uint cpuTime;
   uint wallTime;
@@ -646,8 +702,9 @@ class TraceCustomEventImpl final: public WorkerInterface::CustomEvent {
       api::TraceItem::HibernatableWebSocketEventInfo,                                              \
       api::TraceItem::HibernatableWebSocketEventInfo::Message,                                     \
       api::TraceItem::HibernatableWebSocketEventInfo::Close,                                       \
-      api::TraceItem::HibernatableWebSocketEventInfo::Error, api::TraceLog, api::TraceException,   \
-      api::TraceDiagnosticChannelEvent, api::TraceMetrics, api::UnsafeTraceMetrics
+      api::TraceItem::HibernatableWebSocketEventInfo::Error, api::TraceLog, api::OTelSpan,         \
+      api::OTelSpanTag, api::TraceException, api::TraceDiagnosticChannelEvent, api::TraceMetrics,  \
+      api::UnsafeTraceMetrics
 // The list of trace.h types that are added to worker.c++'s JSG_DECLARE_ISOLATE_TYPE
 
 }  // namespace workerd::api

src/workerd/io/BUILD.bazel

@@ -153,6 +153,7 @@ wd_cc_library(
     ],
     visibility = ["//visibility:public"],
     deps = [
+        ":trace_capnp",
         ":worker-interface_capnp",
         "//src/workerd/jsg:memory-tracker",
         "//src/workerd/util",
@@ -259,6 +260,7 @@ wd_capnp_library(
     deps = [
         ":outcome_capnp",
         ":script-version_capnp",
+        ":trace_capnp",
         "@capnp-cpp//src/capnp/compat:http-over-capnp_capnp",
     ],
 )
@@ -269,6 +271,8 @@ wd_capnp_library(src = "outcome.capnp")
 
 wd_capnp_library(src = "script-version.capnp")
 
+wd_capnp_library(src = "trace.capnp")
+
 wd_capnp_library(src = "compatibility-date.capnp")
 
 wd_capnp_library(src = "features.capnp")

src/workerd/io/compatibility-date.capnp

@@ -668,4 +668,9 @@ struct CompatibilityFlags @0x8f8c1b68151b6cef {
   # A bug in the original implementation of TransformStream failed to apply backpressure
   # correctly. The fix, however, can break existing implementations that don't account
   # for the bug so we need to put the fix behind a compat flag.
+
+  # Experimental support for exporting user spans to tail worker.
+  tailWorkerUserSpans @69 :Bool
+      $compatEnableFlag("tail_worker_user_spans")
+      $experimental;
 }