[o11y] Add traceId to user trace

src/workerd/api/trace.c++

@@ -82,6 +82,10 @@ 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(
     jsg::Lock& js, const Trace& trace) {
   return KJ_MAP(x, trace.diagnosticChannelEvents) -> jsg::Ref<TraceDiagnosticChannelEvent> {
@@ -203,6 +207,7 @@ TraceItem::TraceItem(jsg::Lock& js, const Trace& trace)
       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),
@@ -288,6 +293,11 @@ 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;
 }

src/workerd/api/trace.h

@@ -151,6 +151,7 @@ class TraceItem final: public jsg::Object {
   jsg::Optional<kj::Array<kj::StringPtr>> getScriptTags();
   kj::StringPtr getExecutionModel();
   kj::ArrayPtr<jsg::Ref<OTelSpan>> getSpans();
+  kj::StringPtr getTraceId();
   kj::StringPtr getOutcome();
 
   uint getCpuTime();
@@ -163,6 +164,7 @@ class TraceItem final: public jsg::Object {
     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);
@@ -191,6 +193,7 @@ class TraceItem final: public jsg::Object {
   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;

src/workerd/io/trace.c++

@@ -126,6 +126,14 @@ kj::String TraceId::toW3C() const {
   return kj::str(s.releaseAsArray());
 }
 
+// Return ID represented as a network-order/big endian hex string.
+kj::String TraceId::toNEHex() const {
+  kj::Vector<char> s(32);
+  addHex(s, __builtin_bswap64(low));
+  addHex(s, __builtin_bswap64(high));
+  return kj::str(s.releaseAsArray());
+}
+
 TraceId TraceId::fromEntropy(kj::Maybe<kj::EntropySource&> entropySource) {
   if (isPredictableModeForTest()) {
     return TraceId(0x2a2a2a2a2a2a2a2a, 0x2a2a2a2a2a2a2a2a);
@@ -186,11 +194,11 @@ kj::Rc<InvocationSpanContext> InvocationSpanContext::newForInvocation(
       kj::mv(traceId), TraceId::fromEntropy(entropySource), 0, kj::mv(parent));
 }
 
-TraceId TraceId::fromCapnp(rpc::InvocationSpanContext::TraceId::Reader reader) {
+TraceId TraceId::fromCapnp(rpc::TraceId::Reader reader) {
   return TraceId(reader.getLow(), reader.getHigh());
 }
 
-void TraceId::toCapnp(rpc::InvocationSpanContext::TraceId::Builder writer) const {
+void TraceId::toCapnp(rpc::TraceId::Builder writer) const {
   writer.setLow(low);
   writer.setHigh(high);
 }
@@ -503,6 +511,11 @@ void Trace::copyTo(rpc::Trace::Builder builder) {
       spans[i].copyTo(list[i]);
     }
   }
+  // Add trace ID, if available.
+  KJ_IF_SOME(t, traceId) {
+    auto traceIdBuilder = builder.initTraceId();
+    t.toCapnp(traceIdBuilder);
+  }
 
   {
     auto list = builder.initExceptions(exceptions.size());
@@ -619,6 +632,10 @@ void Trace::mergeFrom(rpc::Trace::Reader reader, PipelineLogLevel pipelineLogLev
   if (pipelineLogLevel != PipelineLogLevel::NONE) {
     logs.addAll(reader.getLogs());
     spans.addAll(reader.getSpans());
+    // Set traceId, if not set already
+    if (reader.hasSpans() && traceId == kj::none) {
+      traceId = tracing::TraceId::fromCapnp(reader.getTraceId());
+    }
     exceptions.addAll(reader.getExceptions());
     diagnosticChannelEvents.addAll(reader.getDiagnosticChannelEvents());
   }
@@ -892,6 +909,12 @@ void WorkerTracer::addSpan(CompleteSpan&& span) {
   trace->numSpans++;
 }
 
+void WorkerTracer::setTraceId(tracing::TraceId traceId) {
+  if (trace->traceId == kj::none) {
+    trace->traceId = traceId;
+  }
+}
+
 Span::TagValue spanTagClone(const Span::TagValue& tag) {
   KJ_SWITCH_ONEOF(tag) {
     KJ_CASE_ONEOF(str, kj::String) {

src/workerd/io/trace.h

@@ -92,6 +92,9 @@ class TraceId final {
   // Replicates W3C Serialization
   kj::String toW3C() const;
 
+  // Return network order hex representation
+  kj::String toNEHex() const;
+
   // Creates a random Trace Id, optionally usig a given entropy source. If an
   // entropy source is not given, then we fallback to using BoringSSL's RAND_bytes.
   static TraceId fromEntropy(kj::Maybe<kj::EntropySource&> entropy = kj::none);
@@ -112,8 +115,8 @@ class TraceId final {
     return high;
   }
 
-  static TraceId fromCapnp(rpc::InvocationSpanContext::TraceId::Reader reader);
-  void toCapnp(rpc::InvocationSpanContext::TraceId::Builder writer) const;
+  static TraceId fromCapnp(rpc::TraceId::Reader reader);
+  void toCapnp(rpc::TraceId::Builder writer) const;
 
  private:
   uint64_t low = 0;
@@ -491,7 +494,11 @@ class Trace final: public kj::Refcounted {
   kj::Maybe<kj::String> entrypoint;
 
   kj::Vector<tracing::Log> logs;
+
+  // trace ID, if user spans are being recorded.
+  kj::Maybe<tracing::TraceId> traceId;
   kj::Vector<CompleteSpan> spans;
+
   // A request's trace can have multiple exceptions due to separate request/waitUntil tasks.
   kj::Vector<tracing::Exception> exceptions;
 
@@ -609,6 +616,7 @@ class WorkerTracer final: public kj::Refcounted, public BaseTracer {
 
   void addLog(kj::Date timestamp, LogLevel logLevel, kj::String message) override;
   void addSpan(CompleteSpan&& span) override;
+  void setTraceId(tracing::TraceId traceId);
   void addException(kj::Date timestamp,
       kj::String name,
       kj::String message,

src/workerd/io/worker-interface.capnp

@@ -16,11 +16,12 @@ using import "/workerd/io/outcome.capnp".EventOutcome;
 using import "/workerd/io/script-version.capnp".ScriptVersion;
 using import "/workerd/io/trace.capnp".UserSpanData;
 
+struct TraceId {
+  high @0 :UInt64;
+  low @1 :UInt64;
+}
+
 struct InvocationSpanContext {
-  struct TraceId {
-    high @0 :UInt64;
-    low @1 :UInt64;
-  }
   traceId @0 :TraceId;
   invocationId @1 :TraceId;
   spanId @2 :UInt32;
@@ -44,6 +45,7 @@ struct Trace @0x8e8d911203762d34 {
   }
 
   spans @26 :List(UserSpanData);
+  traceId @27 :TraceId;
 
   exceptions @1 :List(Exception);
   struct Exception {