Adds tail tags from caller to tail workers

Author: Caio-Nogueira

src/workerd/api/trace.c++

@@ -111,6 +111,26 @@ jsg::Optional<kj::Array<kj::String>> getTraceScriptTags(const Trace& trace) {
   }
 }
 
+TraceItem::TailAttributeValue getTraceTailAttributeValue(const tracing::Attribute& tag) {
+  KJ_REQUIRE(tag.value.size() == 1, "tail attributes must contain exactly one value");
+
+  KJ_SWITCH_ONEOF(tag.value[0]) {
+    KJ_CASE_ONEOF(boolean, bool) {
+      return TraceItem::TailAttributeValue(boolean);
+    }
+    KJ_CASE_ONEOF(number, double) {
+      return TraceItem::TailAttributeValue(number);
+    }
+    KJ_CASE_ONEOF(integer, int64_t) {
+      return TraceItem::TailAttributeValue(static_cast<double>(integer));
+    }
+    KJ_CASE_ONEOF(string, kj::ConstString) {
+      return TraceItem::TailAttributeValue(kj::str(string));
+    }
+  }
+  KJ_UNREACHABLE;
+}
+
 kj::Own<TraceItem::FetchEventInfo::Request::Detail> getFetchRequestDetail(
     jsg::Lock& js, const Trace& trace, const tracing::FetchEventInfo& eventInfo) {
   const auto getCf = [&]() -> jsg::Optional<jsg::V8Ref<v8::Object>> {
@@ -191,6 +211,8 @@ TraceItem::TraceItem(jsg::Lock& js, const Trace& trace)
       scriptVersion(getTraceScriptVersion(trace)),
       dispatchNamespace(mapCopyString(trace.dispatchNamespace)),
       scriptTags(getTraceScriptTags(trace)),
+      tailAttributes(trace.tailAttributes.map(
+          [](auto& tags) { return KJ_MAP(tag, tags) { return tag.clone(); }; })),
       durableObjectId(mapCopyString(trace.durableObjectId)),
       executionModel(kj::str(trace.executionModel)),
       outcome(kj::str(trace.outcome)),
@@ -270,6 +292,19 @@ jsg::Optional<kj::Array<kj::StringPtr>> TraceItem::getScriptTags() {
       [](kj::Array<kj::String>& tags) { return KJ_MAP(t, tags) -> kj::StringPtr { return t; }; });
 }
 
+jsg::Optional<jsg::Dict<TraceItem::TailAttributeValue>> TraceItem::getTailAttributes() {
+  return tailAttributes.map([](kj::Array<tracing::Attribute>& tags) {
+    return jsg::Dict<TraceItem::TailAttributeValue>{
+      .fields = KJ_MAP(tag, tags) {
+        return jsg::Dict<TraceItem::TailAttributeValue>::Field{
+          .name = kj::str(tag.name),
+          .value = getTraceTailAttributeValue(tag),
+        };
+      },
+    };
+  });
+}
+
 jsg::Optional<kj::StringPtr> TraceItem::getDurableObjectId() {
   return durableObjectId.map([](auto& id) -> kj::StringPtr { return id; });
 }
@@ -732,6 +767,29 @@ void TraceItem::visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
       tracker.trackField("scriptTag", tag);
     }
   }
+  KJ_IF_SOME(tags, tailAttributes) {
+    for (const auto& tag: tags) {
+      tracker.trackFieldWithSize(
+          "tailAttributeName", tag.name.size());
+      for (const auto& value: tag.value) {
+        KJ_SWITCH_ONEOF(value) {
+          KJ_CASE_ONEOF(boolean, bool) {
+            tracker.trackFieldWithSize("tailAttributeValue", sizeof(boolean));
+          }
+          KJ_CASE_ONEOF(number, double) {
+            tracker.trackFieldWithSize("tailAttributeValue", sizeof(number));
+          }
+          KJ_CASE_ONEOF(integer, int64_t) {
+            tracker.trackFieldWithSize("tailAttributeValue", sizeof(integer));
+          }
+          KJ_CASE_ONEOF(string, kj::ConstString) {
+            tracker.trackFieldWithSize(
+                "tailAttributeValue", string.size());
+          }
+        }
+      }
+    }
+  }
   tracker.trackField("outcome", outcome);
 }
 

src/workerd/api/trace.h

@@ -69,6 +69,8 @@ struct ScriptVersion {
 
 class TraceItem final: public jsg::Object {
  public:
+  using TailAttributeValue = kj::OneOf<bool, double, kj::String>;
+
   class FetchEventInfo;
   class JsRpcEventInfo;
   class ScheduledEventInfo;
@@ -101,6 +103,7 @@ class TraceItem final: public jsg::Object {
   jsg::Optional<ScriptVersion> getScriptVersion();
   jsg::Optional<kj::StringPtr> getDispatchNamespace();
   jsg::Optional<kj::Array<kj::StringPtr>> getScriptTags();
+  jsg::Optional<jsg::Dict<TailAttributeValue>> getTailAttributes();
   jsg::Optional<kj::StringPtr> getDurableObjectId();
   kj::StringPtr getExecutionModel();
   kj::StringPtr getOutcome();
@@ -120,6 +123,7 @@ class TraceItem final: public jsg::Object {
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(scriptVersion, getScriptVersion);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(dispatchNamespace, getDispatchNamespace);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(scriptTags, getScriptTags);
+    JSG_LAZY_READONLY_INSTANCE_PROPERTY(tailAttributes, getTailAttributes);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(durableObjectId, getDurableObjectId);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(outcome, getOutcome);
     JSG_LAZY_READONLY_INSTANCE_PROPERTY(executionModel, getExecutionModel);
@@ -141,6 +145,7 @@ class TraceItem final: public jsg::Object {
   kj::Maybe<ScriptVersion> scriptVersion;
   kj::Maybe<kj::String> dispatchNamespace;
   jsg::Optional<kj::Array<kj::String>> scriptTags;
+  kj::Maybe<kj::Array<tracing::Attribute>> tailAttributes;
   kj::Maybe<kj::String> durableObjectId;
   kj::String executionModel;
   kj::String outcome;

src/workerd/io/trace.c++

@@ -753,6 +753,13 @@ void Trace::copyTo(rpc::Trace::Builder builder) const {
     }
   }
 
+  KJ_IF_SOME(tags, tailAttributes) {
+    auto list = builder.initTailAttributes(tags.size());
+    for (auto i: kj::indices(tags)) {
+      tags[i].copyTo(list[i]);
+    }
+  }
+
   KJ_IF_SOME(e, entrypoint) {
     builder.setEntrypoint(e);
   }
@@ -859,6 +866,10 @@ void Trace::mergeFrom(rpc::Trace::Reader reader, PipelineLogLevel pipelineLogLev
     scriptTags = KJ_MAP(tag, tags) { return kj::str(tag); };
   }
 
+  if (auto tags = reader.getTailAttributes(); tags.size() > 0) {
+    tailAttributes = KJ_MAP(tag, tags) { return tracing::Attribute(tag); };
+  }
+
   if (reader.hasEntrypoint()) {
     entrypoint = kj::str(reader.getEntrypoint());
   }

src/workerd/io/trace.h

@@ -869,6 +869,7 @@ class Trace final: public kj::Refcounted {
   kj::Maybe<kj::String> dispatchNamespace;
   kj::Maybe<kj::String> scriptId;
   kj::Array<kj::String> scriptTags;
+  kj::Maybe<kj::Array<tracing::Attribute>> tailAttributes;
   kj::Maybe<kj::String> entrypoint;
   kj::Maybe<kj::String> durableObjectId;
 

src/workerd/io/tracer.c++

@@ -21,6 +21,24 @@ namespace {
 // TODO(streaming-tail): Add a clear indicator for events being truncated based on MAX_TRACE_BYTES
 // so that developers can understand why this happens.
 static constexpr size_t MAX_TRACE_BYTES = 256 * 1024;
+
+tracing::Attribute::Value cloneAttributeValue(const tracing::Attribute::Value& value) {
+  KJ_SWITCH_ONEOF(value) {
+    KJ_CASE_ONEOF(boolean, bool) {
+      return tracing::Attribute::Value(boolean);
+    }
+    KJ_CASE_ONEOF(number, double) {
+      return tracing::Attribute::Value(number);
+    }
+    KJ_CASE_ONEOF(integer, int64_t) {
+      return tracing::Attribute::Value(integer);
+    }
+    KJ_CASE_ONEOF(string, kj::ConstString) {
+      return tracing::Attribute::Value(string.clone());
+    }
+  }
+  KJ_UNREACHABLE;
+}
 }  // namespace
 
 kj::Promise<kj::Own<Trace>> WorkerTracer::onComplete() {
@@ -34,11 +52,24 @@ kj::Promise<kj::Own<Trace>> WorkerTracer::onComplete() {
 WorkerTracer::WorkerTracer(kj::Maybe<kj::Rc<kj::Refcounted>> parentPipeline,
     kj::Own<Trace> trace,
     PipelineLogLevel pipelineLogLevel,
+    kj::Maybe<kj::Array<tracing::Attribute>> tailAttributes,
     kj::Maybe<kj::Own<tracing::TailStreamWriter>> maybeTailStreamWriter)
     : pipelineLogLevel(pipelineLogLevel),
       trace(kj::mv(trace)),
       parentPipeline(kj::mv(parentPipeline)),
-      maybeTailStreamWriter(kj::mv(maybeTailStreamWriter)) {}
+      maybeTailStreamWriter(kj::mv(maybeTailStreamWriter)) {
+  KJ_IF_SOME(tags, tailAttributes) {
+    if (tags.size() == 0) {
+      tailAttributes = kj::none;
+    } else {
+      for (auto& tag: tags) {
+        KJ_REQUIRE(tag.value.size() == 1, "tail attributes must contain exactly one value");
+        setWorkerAttribute(tag.name.clone(), cloneAttributeValue(tag.value[0]));
+      }
+    }
+  }
+  this->trace->tailAttributes = kj::mv(tailAttributes);
+}
 
 WorkerTracer::~WorkerTracer() noexcept(false) {
   // Report the outcome event, which should have been delivered by now.

src/workerd/io/tracer.h

@@ -113,6 +113,7 @@ class WorkerTracer final: public BaseTracer {
   explicit WorkerTracer(kj::Maybe<kj::Rc<kj::Refcounted>> parentPipeline,
       kj::Own<Trace> trace,
       PipelineLogLevel pipelineLogLevel,
+      kj::Maybe<kj::Array<tracing::Attribute>> tailAttributes,
       kj::Maybe<kj::Own<tracing::TailStreamWriter>> maybeTailStreamWriter);
   virtual ~WorkerTracer() noexcept(false);
   KJ_DISALLOW_COPY_AND_MOVE(WorkerTracer);

src/workerd/io/worker-interface.capnp

@@ -168,6 +168,7 @@ struct Trace @0x8e8d911203762d34 {
 
   entrypoint @22 :Text;
   durableObjectId @27 :Text;
+  tailAttributes @28 :List(Attribute);
 
   diagnosticChannelEvents @17 :List(DiagnosticChannelEvent);
   struct DiagnosticChannelEvent {

src/workerd/server/server.c++

@@ -2157,8 +2157,8 @@ class Server::WorkerService final: public Service,
           kj::none /* scriptVersion */, kj::none /* dispatchNamespace */, kj::none /* scriptId */,
           nullptr /* scriptTags */, mapCopyString(entrypointName), executionModel,
           kj::none /* durableObjectId */);
-      kj::Own<WorkerTracer> tracer = kj::refcounted<WorkerTracer>(
-          kj::none, kj::mv(trace), PipelineLogLevel::FULL, kj::mv(tailStreamWriter));
+      kj::Own<WorkerTracer> tracer = kj::refcounted<WorkerTracer>(kj::none, kj::mv(trace),
+          PipelineLogLevel::FULL, kj::none, kj::mv(tailStreamWriter));
 
       // When the tracer is complete, deliver traces to any buffered tail workers. We end up
       // creating two references to the WorkerTracer, one held by the observer and one that will be