Merge pull request #3773 from cloudflare/felix/031325-tail-worker-refactor-p1

[o11y] Consolidate adding trace events for legacy and streaming tail workers

Author: fhanau

src/workerd/api/hibernatable-web-socket.c++

@@ -94,16 +94,12 @@ kj::Promise<WorkerInterface::CustomEvent::Result> HibernatableWebSocketCustomEve
     KJ_UNREACHABLE;
   };
 
+  // TODO(streaming-tail-workers): Support Hibernate and Resume events properly.
   KJ_IF_SOME(t, incomingRequest->getWorkerTracer()) {
-    t.setEventInfo(context.now(), tracing::HibernatableWebSocketEventInfo(getType()));
+    t.setEventInfo(incomingRequest->getContext().getInvocationSpanContext(), context.now(),
+        tracing::HibernatableWebSocketEventInfo(getType()));
   }
 
-  // TODO(streaming-tail-workers): Support Hibernate and Resume events properly.
-  context.getMetrics().reportTailEvent(context.getInvocationSpanContext(), [&] {
-    return tracing::Onset(
-        tracing::HibernatableWebSocketEventInfo(getType()), tracing::Onset::WorkerInfo{}, kj::none);
-  });
-
   auto outcomeObserver = kj::rc<OutcomeObserver>(
       kj::addRef(incomingRequest->getMetrics()), context.getInvocationSpanContext());
 

src/workerd/api/node/diagnostics-channel.c++

@@ -37,7 +37,8 @@ void Channel::publish(jsg::Lock& js, jsg::Value message) {
         Error,
         "Diagnostic events cannot be published with SharedArrayBuffer or "
         "transferred ArrayBuffer instances");
-    tracer.addDiagnosticChannelEvent(context.now(), name.toString(js), kj::mv(tmp.data));
+    tracer.addDiagnosticChannelEvent(
+        context.getInvocationSpanContext(), context.now(), name.toString(js), kj::mv(tmp.data));
   }
 }
 

src/workerd/api/node/util.c++

@@ -274,8 +274,9 @@ namespace {
     // JSG_KJ_EXCEPTION would not give us that, and we only want to incur the cost
     // of creating and capturing the stack when we actually need it.
     auto ex = KJ_ASSERT_NONNULL(js.error(message).tryCast<jsg::JsObject>());
-    tracer.addException(ioContext.now(), ex.get(js, "name"_kj).toString(js),
-        ex.get(js, "message"_kj).toString(js), ex.get(js, "stack"_kj).toString(js));
+    tracer.addException(ioContext.getInvocationSpanContext(), ioContext.now(),
+        ex.get(js, "name"_kj).toString(js), ex.get(js, "message"_kj).toString(js),
+        ex.get(js, "stack"_kj).toString(js));
     ioContext.abort(js.exceptionToKj(ex));
   } else {
     ioContext.abort(JSG_KJ_EXCEPTION(FAILED, Error, kj::mv(message)));

src/workerd/api/queue.c++

@@ -556,14 +556,10 @@ kj::Promise<WorkerInterface::CustomEvent::Result> QueueCustomEventImpl::run(
   }
 
   KJ_IF_SOME(t, incomingRequest->getWorkerTracer()) {
-    t.setEventInfo(context.now(), tracing::QueueEventInfo(kj::str(queueName), batchSize));
+    t.setEventInfo(context.getInvocationSpanContext(), context.now(),
+        tracing::QueueEventInfo(kj::str(queueName), batchSize));
   }
 
-  context.getMetrics().reportTailEvent(context.getInvocationSpanContext(), [&] {
-    return tracing::Onset(tracing::QueueEventInfo(kj::mv(queueName), batchSize),
-        tracing::Onset::WorkerInfo{}, kj::none);
-  });
-
   auto outcomeObserver = kj::rc<OutcomeObserver>(
       kj::addRef(incomingRequest->getMetrics()), context.getInvocationSpanContext());
 

src/workerd/api/tail-worker-test.js

@@ -30,23 +30,23 @@ export const test = {
 
     let expected = [
       // http-test.js: fetch and scheduled events get reported correctly.
-      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"http://placeholder/body-length","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
-      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"http://placeholder/body-length","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"http://placeholder/body-length","cfJson":"","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"http://placeholder/body-length","cfJson":"","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
       '{"type":"onset","info":{"type":"scheduled","scheduledTime":"1970-01-01T00:00:00.000Z","cron":""}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
       '{"type":"onset","info":{"type":"scheduled","scheduledTime":"1970-01-01T00:00:00.000Z","cron":"* * * * 30"}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
-      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://placeholder/not-found","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
-      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://placeholder/web-socket","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://placeholder/not-found","cfJson":"","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://placeholder/web-socket","cfJson":"","headers":[{"name":"upgrade","value":"websocket"}]}}{"type":"exception","name":"Error","message":"The script will never generate a response."}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
 
       // queue-test.js: queue events
-      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/message","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
-      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/message","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
-      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/message","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
-      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/message","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
-      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/batch","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/message","cfJson":"","headers":[{"name":"content-type","value":"application/octet-stream"},{"name":"x-msg-delay-secs","value":"2"},{"name":"x-msg-fmt","value":"text"}]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/message","cfJson":"","headers":[{"name":"content-type","value":"application/octet-stream"},{"name":"x-msg-fmt","value":"bytes"}]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/message","cfJson":"","headers":[{"name":"content-type","value":"application/octet-stream"},{"name":"x-msg-fmt","value":"json"}]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/message","cfJson":"","headers":[{"name":"content-type","value":"application/octet-stream"},{"name":"x-msg-fmt","value":"v8"}]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"POST","url":"https://fake-host/batch","cfJson":"","headers":[{"name":"cf-queue-batch-bytes","value":"31"},{"name":"cf-queue-batch-count","value":"4"},{"name":"cf-queue-largest-msg","value":"13"},{"name":"content-type","value":"application/json"},{"name":"x-msg-delay-secs","value":"2"}]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
       '{"type":"onset","info":{"type":"queue","queueName":"test-queue","batchSize":5}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
 
       // actor-alarms-test.js: alarm events
-      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://foo/test","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://foo/test","cfJson":"","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
       '{"type":"onset","info":{"type":"alarm","scheduledTime":"1970-01-01T00:00:00.000Z"}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
 
       // legacy tail worker, triggered via alarm test. It would appear that these being recorded
@@ -58,8 +58,8 @@ export const test = {
       '{"type":"onset","info":{"type":"trace","traces":[""]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
 
       // tests/websocket-hibernation.js: hibernatableWebSocket events
-      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://example.com/","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
-      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://example.com/hibernation","cfJson":"{}","headers":[]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://example.com/","cfJson":"","headers":[{"name":"upgrade","value":"websocket"}]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
+      '{"type":"onset","info":{"type":"fetch","method":"GET","url":"http://example.com/hibernation","cfJson":"","headers":[{"name":"upgrade","value":"websocket"}]}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
       '{"type":"onset","info":{"type":"hibernatableWebSocket","info":{"type":"message"}}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
       '{"type":"onset","info":{"type":"hibernatableWebSocket","info":{"type":"close","code":1000,"wasClean":true}}}{"type":"outcome","outcome":"ok","cpuTime":0,"wallTime":0}',
     ];

src/workerd/api/trace.c++

@@ -653,12 +653,10 @@ kj::Promise<void> sendTracesToExportedHandler(kj::Own<IoContext::IncomingRequest
   auto& metrics = incomingRequest->getMetrics();
 
   KJ_IF_SOME(t, incomingRequest->getWorkerTracer()) {
-    t.setEventInfo(context.now(), tracing::TraceEventInfo(traces));
+    t.setEventInfo(
+        context.getInvocationSpanContext(), context.now(), tracing::TraceEventInfo(traces));
   }
 
-  metrics.reportTailEvent(context.getInvocationSpanContext(), [&] {
-    return tracing::Onset(tracing::TraceEventInfo(traces), tracing::Onset::WorkerInfo{}, kj::none);
-  });
   auto outcomeObserver = kj::rc<OutcomeObserver>(
       kj::addRef(incomingRequest->getMetrics()), context.getInvocationSpanContext());
 

src/workerd/api/worker-rpc.c++

@@ -1833,12 +1833,9 @@ class EntrypointJsRpcTarget final: public JsRpcTargetBase {
 
   void addTrace(jsg::Lock& js, IoContext& ioctx, kj::StringPtr methodName) override {
     KJ_IF_SOME(t, tracer) {
-      t->setEventInfo(ioctx.now(), tracing::JsRpcEventInfo(kj::str(methodName)));
+      t->setEventInfo(ioctx.getInvocationSpanContext(), ioctx.now(),
+          tracing::JsRpcEventInfo(kj::str(methodName)));
     }
-    ioctx.getMetrics().reportTailEvent(ioctx.getInvocationSpanContext(), [&] {
-      return tracing::Onset(
-          tracing::JsRpcEventInfo(kj::str(methodName)), tracing::Onset::WorkerInfo{}, kj::none);
-    });
   }
 };
 

src/workerd/io/observer.h

@@ -128,21 +128,6 @@ class RequestObserver: public kj::Refcounted {
     return nullptr;
   }
 
-  // If the worker is configured to support streaming tail workers, reportTailEvent
-  // will forward the given event on to the collection of streaming tail workers
-  // that are configured with this observer. Otherwise, this is a non-op.
-  virtual void reportTailEvent(
-      const tracing::InvocationSpanContext& context, tracing::TailEvent::Event&& event) {
-    reportTailEvent(context, [event = kj::mv(event)]() mutable { return kj::mv(event); });
-  }
-
-  // If the worker is configured to support streaming tail workers, reportTailEvent
-  // will forward the event returned by the callback on to the collection of streaming
-  // fail workers that are configured with this observer. The callback will only be
-  // invoked if there are tail workers.
-  virtual void reportTailEvent(const tracing::InvocationSpanContext& context,
-      kj::FunctionParam<tracing::TailEvent::Event()> fn) {}
-
   // Reports the outcome event to any configured streaming tail workers, signalizing that the
   // request has completed and will not produce any more events.
   virtual void reportOutcome(const tracing::InvocationSpanContext& context) {}

src/workerd/io/trace-stream.c++

@@ -926,7 +926,8 @@ kj::Promise<WorkerInterface::CustomEvent::Result> TailStreamCustomEventImpl::run
   incomingRequest->delivered();
 
   KJ_IF_SOME(t, incomingRequest->getWorkerTracer()) {
-    t.setEventInfo(ioContext.now(), TraceEventInfo(kj::Array<TraceEventInfo::TraceItem>(nullptr)));
+    t.setEventInfo(ioContext.getInvocationSpanContext(), ioContext.now(),
+        TraceEventInfo(kj::Array<TraceEventInfo::TraceItem>(nullptr)));
   }
 
   // TODO(streaming-tail): Support instrementation for streaming tail workers themselves – need to
@@ -1114,10 +1115,10 @@ kj::Maybe<kj::Own<tracing::TailStreamWriter>> initializeTailStreamWriter(
 
   auto state = kj::heap<TailStreamWriterState>(kj::mv(streamingTailWorkers), waitUntilTasks);
 
-  return kj::heap<tracing::TailStreamWriter>(
+  return kj::refcounted<tracing::TailStreamWriter>(
       // This lambda is called for every streaming tail event that is reported. We use
       // the TailStreamWriterState for this stream to actually handle the event.
-      // Pay attention to the ownership of state here. The lamba holds a bare
+      // Pay attention to the ownership of state here. The lambda holds a bare
       // reference while the instance is attached to the kj::Own below.
       [&state = *state, &waitUntilTasks](tracing::TailEvent&& event) mutable {
     KJ_SWITCH_ONEOF(state.inner) {