Support streaming tail workers over rpc

As it stands workers tailed by streaming tail workers have to know they are tailing to a streaming tail worker.
This is a limitation that should be refactored away in the future but for the moment we can't depend on hasHandler to infer this over rpc.

Author: danlapid

samples/tail-workers/config.capnp

@@ -17,6 +17,7 @@ const helloWorld :Workerd.Worker = (
   ],
   compatibilityDate = "2024-10-14",
   tails = ["log"],
+  streamingTails = ["log"],
 );
 
 const logWorker :Workerd.Worker = (

src/workerd/api/tail-worker-test-invalid.wd-test

@@ -13,7 +13,7 @@ const unitTests :Workerd.Config = (
         ],
         compatibilityDate = "2023-08-01",
         compatibilityFlags = ["nodejs_compat", "service_binding_extra_handlers", "cache_option_disabled"],
-        tails = ["log"],
+        streamingTails = ["log"],
       ),
     ),
     # tail worker with tests
@@ -30,4 +30,4 @@ const unitTests :Workerd.Config = (
   autogates = [
     "workerd-autogate-streaming-tail-workers",
   ],
-);
+);

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

@@ -14,7 +14,7 @@ const unitTests :Workerd.Config = (
         ],
         compatibilityDate = "2023-08-01",
         compatibilityFlags = ["nodejs_compat", "service_binding_extra_handlers", "cache_option_disabled"],
-        tails = ["log"],
+        streamingTails = ["log"],
       ),
     ),
     ( name = "queue-test",
@@ -28,7 +28,7 @@ const unitTests :Workerd.Config = (
         ],
         compatibilityDate = "2023-07-24",
         compatibilityFlags = ["nodejs_compat", "service_binding_extra_handlers"],
-        tails = ["log"],
+        streamingTails = ["log"],
       )
     ),
     (name = "alarms", worker = .alarmsWorker),
@@ -66,7 +66,8 @@ const alarmsWorker :Workerd.Worker = (
     (name = "ns", durableObjectNamespace = "DurableObjectExample"),
   ],
   # tailed by the main tail worker and the dummy legacy tail worker, to get traces for it too.
-  tails = ["log", "legacy"],
+  tails = ["legacy"],
+  streamingTails = ["log"],
 );
 
 const hiberWorker :Workerd.Worker = (
@@ -86,7 +87,7 @@ const hiberWorker :Workerd.Worker = (
   bindings = [
     (name = "ns", durableObjectNamespace = "DurableObjectExample"),
   ],
-  tails = ["log"],
+  streamingTails = ["log"],
 );
 
 const logWorker :Workerd.Worker = (
@@ -111,5 +112,5 @@ const logLegacy :Workerd.Worker = (
     (name = "worker", esModule = embed "tail-worker-test-dummy.js")
   ],
   compatibilityDate = "2024-10-14",
-  tails = ["log"],
-);
+  streamingTails = ["log"],
+);

src/workerd/server/server.c++

@@ -1631,6 +1631,7 @@ class Server::WorkerService final: public Service,
     kj::Maybe<kj::Own<SqliteDatabase::Vfs>> actorStorage;
     AlarmScheduler& alarmScheduler;
     kj::Array<kj::Own<Service>> tails;
+    kj::Array<kj::Own<Service>> streamingTails;
   };
   using LinkCallback = kj::Function<LinkedIoChannels(WorkerService&)>;
   using AbortActorsCallback = kj::Function<void()>;
@@ -1760,25 +1761,20 @@ class Server::WorkerService final: public Service,
 
     kj::Array<kj::Own<WorkerInterface>> legacyTailWorkers = nullptr;
     kj::Array<kj::Own<WorkerInterface>> streamingTailWorkers = nullptr;
-    // If streaming tail workers is enabled, then we will initialize two lists:
-    // one with services that only export the tail or trace handler (legacy tail
-    // workers) and one that exports the tailStream handler. We'll check tailStreams
-    // first.
+    legacyTailWorkers = KJ_MAP(service, channels.tails) -> kj::Own<WorkerInterface> {
+      // Caution here... if the tail worker ends up have a circular dependency
+      // on the worker we'll end up with an infinite loop trying to initialize.
+      // We can test this directly but it's more difficult to test indirect
+      // loops (dependency of dependency, etc). Here we're just going to keep
+      // it simple and just check the direct dependency.
+      // If service refers to an EntrypointService, we need to compare with the underlying
+      // WorkerService to match this.
+      KJ_ASSERT(service->service() != this, "A worker currently cannot log to itself");
+      return service->startRequest({});
+    };
+
     if (util::Autogate::isEnabled(util::AutogateKey::STREAMING_TAIL_WORKERS)) {
-      kj::Vector<kj::Own<WorkerInterface>> legacyList;
-      kj::Vector<kj::Own<WorkerInterface>> streamingList;
-      for (auto& service: channels.tails) {
-        KJ_ASSERT(service->service() != this, "A worker currently cannot log to itself");
-        if (service->hasHandler("tailStream"_kj)) {
-          streamingList.add(service->startRequest({}));
-        } else if (service->hasHandler("tail") || service->hasHandler("trace")) {
-          legacyList.add(service->startRequest({}));
-        }
-      }
-      legacyTailWorkers = legacyList.releaseAsArray();
-      streamingTailWorkers = streamingList.releaseAsArray();
-    } else {
-      legacyTailWorkers = KJ_MAP(service, channels.tails) -> kj::Own<WorkerInterface> {
+      streamingTailWorkers = KJ_MAP(service, channels.streamingTails) -> kj::Own<WorkerInterface> {
         // Caution here... if the tail worker ends up have a circular dependency
         // on the worker we'll end up with an infinite loop trying to initialize.
         // We can test this directly but it's more difficult to test indirect
@@ -3475,6 +3471,10 @@ kj::Own<Server::Service> Server::makeWorker(kj::StringPtr name,
       return lookupService(tail, kj::str("Worker \"", name, "\"'s tails"));
     };
 
+    result.streamingTails = KJ_MAP(streamingTails, conf.getStreamingTails()) {
+      return lookupService(streamingTails, kj::str("Worker \"", name, "\"'s streaming tails"));
+    };
+
     return result;
   };
 

src/workerd/server/workerd.capnp

@@ -643,6 +643,10 @@ struct Worker {
   tails @14 :List(ServiceDesignator);
   # List of tail worker services that should receive tail events for this worker.
   # See: https://developers.cloudflare.com/workers/observability/logs/tail-workers/
+
+  streamingTails @15 :List(ServiceDesignator);
+  # List of streaming tail worker services that should receive tail events for this worker.
+  # NOTE: This will be deleted in a future refactor, do not depend on this.
 }
 
 struct ExternalServer {