Implement the tail stream RPC interface (#3228)

* Implement tracing::TailStreamWriter
* Add the streaming-tail-workers autogate
* Implement the tail stream rpc interface
* Update type snapshot

samples/tail-workers/config.capnp

@@ -5,7 +5,10 @@ const tailWorkerExample :Workerd.Config = (
     (name = "main", worker = .helloWorld),
     (name = "log", worker = .logWorker),
   ],
-  sockets = [ ( name = "http", address = "*:8080", http = (), service = "main" ) ]
+  sockets = [ ( name = "http", address = "*:8080", http = (), service = "main" ) ],
+  autogates = [
+    "workerd-autogate-streaming-tail-workers",
+  ],
 );
 
 const helloWorld :Workerd.Worker = (

samples/tail-workers/tail.js

@@ -6,5 +6,8 @@ export default {
   // https://developers.cloudflare.com/workers/observability/logs/tail-workers/
   tail(traces) {
     console.log(traces[0].logs);
-  }
+  },
+  tailStream() {
+    return {};
+  },
 };

src/workerd/api/global-scope.h

@@ -288,6 +288,10 @@ struct ExportedHandler {
   jsg::LenientOptional<jsg::Function<TailHandler>> tail;
   jsg::LenientOptional<jsg::Function<TailHandler>> trace;
 
+  typedef kj::Promise<void> TailStreamHandler(
+      jsg::JsObject obj, jsg::Value env, jsg::Optional<jsg::Ref<ExecutionContext>> ctx);
+  jsg::LenientOptional<jsg::Function<TailStreamHandler>> tailStream;
+
   typedef kj::Promise<void> ScheduledHandler(jsg::Ref<ScheduledController> controller,
       jsg::Value env,
       jsg::Optional<jsg::Ref<ExecutionContext>> ctx);
@@ -319,6 +323,7 @@ struct ExportedHandler {
   JSG_STRUCT(fetch,
       tail,
       trace,
+      tailStream,
       scheduled,
       alarm,
       test,
@@ -335,6 +340,7 @@ struct ExportedHandler {
     type ExportedHandlerFetchHandler<Env = unknown, CfHostMetadata = unknown> = (request: Request<CfHostMetadata, IncomingRequestCfProperties<CfHostMetadata>>, env: Env, ctx: ExecutionContext) => Response | Promise<Response>;
     type ExportedHandlerTailHandler<Env = unknown> = (events: TraceItem[], env: Env, ctx: ExecutionContext) => void | Promise<void>;
     type ExportedHandlerTraceHandler<Env = unknown> = (traces: TraceItem[], env: Env, ctx: ExecutionContext) => void | Promise<void>;
+    type ExportedHandlerTailStreamHandler<Env = unknown> = (event : TailStream.TailEvent, env: Env, ctx: ExecutionContext) => TailStream.TailEventHandlerType | Promise<TailStream.TailEventHandlerType>;
     type ExportedHandlerScheduledHandler<Env = unknown> = (controller: ScheduledController, env: Env, ctx: ExecutionContext) => void | Promise<void>;
     type ExportedHandlerQueueHandler<Env = unknown, Message = unknown> = (batch: MessageBatch<Message>, env: Env, ctx: ExecutionContext) => void | Promise<void>;
     type ExportedHandlerTestHandler<Env = unknown> = (controller: TestController, env: Env, ctx: ExecutionContext) => void | Promise<void>;
@@ -344,6 +350,7 @@ struct ExportedHandler {
     fetch?: ExportedHandlerFetchHandler<Env, CfHostMetadata>;
     tail?: ExportedHandlerTailHandler<Env>;
     trace?: ExportedHandlerTraceHandler<Env>;
+    tailStream?: ExportedHandlerTailStreamHandler<Env>;
     scheduled?: ExportedHandlerScheduledHandler<Env>;
     alarm: never;
     webSocketMessage: never;

src/workerd/io/BUILD.bazel

@@ -36,6 +36,17 @@ wd_cc_library(
     ],
 )
 
+wd_cc_library(
+    name = "trace-stream",
+    srcs = ["trace-stream.c++"],
+    hdrs = ["trace-stream.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":io",
+        "//src/workerd/util:completion-membrane",
+    ],
+)
+
 wd_cc_library(
     name = "io",
     # HACK: Currently, the `io` and `api` packages are interdependent. We fold all the sources
@@ -106,6 +117,7 @@ wd_cc_library(
         "//src/workerd/api:url",
         "//src/workerd/jsg",
         "//src/workerd/util:autogate",
+        "//src/workerd/util:completion-membrane",
         "//src/workerd/util:sqlite",
         "//src/workerd/util:thread-scopes",
         "//src/workerd/util:uuid",

src/workerd/io/io-context.h

@@ -788,6 +788,10 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler
   SpanParent getCurrentTraceSpan();
   SpanParent getCurrentUserTraceSpan();
 
+  tracing::InvocationSpanContext& getInvocationSpanContext() {
+    return getCurrentIncomingRequest().invocationSpanContext;
+  }
+
   // Returns a builder for recording tracing spans (or a no-op builder if tracing is inactive).
   // If called while the JS lock is held, uses the trace information from the current async
   // context, if available.

src/workerd/io/observer.h

@@ -18,6 +18,7 @@
 
 namespace workerd {
 
+class IoContext;
 class WorkerInterface;
 class LimitEnforcer;
 class TimerChannel;
@@ -127,6 +128,20 @@ 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(IoContext& ioContext, tracing::TailEvent::Event&& event) {
+    reportTailEvent(ioContext, [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(
+      IoContext& ioContext, kj::FunctionParam<tracing::TailEvent::Event()> fn) {}
+
   virtual kj::Own<void> addedContextTask() {
     return kj::Own<void>();
   }