Merge pull request #5185 from cloudflare/felix/092525-return-event-reporting

Reland "EW-9366 Report fetch response earlier"

Author: fhanau

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

@@ -2,6 +2,7 @@
 // Licensed under the Apache 2.0 license found in the LICENSE file or at:
 //     https://opensource.org/licenses/Apache-2.0
 
+#include <workerd/io/io-context.h>
 #include <workerd/io/tracer.h>
 #include <workerd/util/sentry.h>
 #include <workerd/util/thread-scopes.h>
@@ -141,10 +142,6 @@ WorkerTracer::~WorkerTracer() noexcept(false) {
   KJ_IF_SOME(writer, maybeTailStreamWriter) {
     auto& spanContext = KJ_UNWRAP_OR_RETURN(topLevelInvocationSpanContext);
 
-    KJ_IF_SOME(fetchResponseInfo, trace->fetchResponseInfo) {
-      writer->report(spanContext, tracing::Return({fetchResponseInfo.clone()}), completeTime);
-    }
-
     if (isPredictableModeForTest()) {
       writer->report(spanContext,
           tracing::Outcome(trace->outcome, 0 * kj::MILLISECONDS, 0 * kj::MILLISECONDS),
@@ -440,6 +437,27 @@ void WorkerTracer::recordTimestamp(kj::Date timestamp) {
   }
 }
 
+kj::Date BaseTracer::getTime() {
+  auto& weakIoCtx = KJ_ASSERT_NONNULL(weakIoContext);
+  kj::Date timestamp = kj::UNIX_EPOCH;
+  weakIoCtx->runIfAlive([&timestamp](IoContext& context) { timestamp = context.now(); });
+  if (!weakIoCtx->isValid()) {
+    // This can happen if we the IoContext gets destroyed following an exception, but we still need
+    // to report a time for the return event.
+    if (completeTime != kj::UNIX_EPOCH) {
+      timestamp = completeTime;
+    } else {
+      // Otherwise, we can't actually get an end timestamp that makes sense.
+      if (isPredictableModeForTest()) {
+        KJ_FAIL_ASSERT("reported return event without valid IoContext or completeTime");
+      } else {
+        LOG_WARNING_PERIODICALLY("reported return event without valid IoContext or completeTime");
+      }
+    }
+  }
+  return timestamp;
+}
+
 void BaseTracer::adjustSpanTime(CompleteSpan& span) {
   // To report I/O time, we need the IOContext to still be alive.
   // weakIoContext is only none if we are tracing via RPC (in this case span times have already been
@@ -463,26 +481,37 @@ void BaseTracer::adjustSpanTime(CompleteSpan& span) {
         if (isPredictableModeForTest()) {
           KJ_FAIL_ASSERT("reported span after IoContext was deallocated", span.operationName);
         } else {
-          LOG_WARNING_PERIODICALLY(
-              "reported span after IoContext was deallocated", span.operationName);
+          KJ_LOG(WARNING, "reported span after IoContext was deallocated", span.operationName);
         }
       }
     }
   }
 }
 
-void WorkerTracer::setFetchResponseInfo(tracing::FetchResponseInfo&& info) {
+void WorkerTracer::setReturn(
+    kj::Maybe<kj::Date> timestamp, kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo) {
   // Match the behavior of setEventInfo(). Any resolution of the TODO comments
   // in setEventInfo() that are related to this check while probably also affect
   // this function.
   if (pipelineLogLevel == PipelineLogLevel::NONE) {
     return;
   }
 
-  // Note: In the streaming model, fetchResponseInfo is dispatched when the tail worker returns.
-  KJ_REQUIRE(KJ_REQUIRE_NONNULL(trace->eventInfo).is<tracing::FetchEventInfo>());
-  KJ_ASSERT(trace->fetchResponseInfo == kj::none, "setFetchResponseInfo can only be called once");
-  trace->fetchResponseInfo = kj::mv(info);
+  KJ_IF_SOME(writer, maybeTailStreamWriter) {
+    auto& spanContext = KJ_UNWRAP_OR_RETURN(topLevelInvocationSpanContext);
+
+    // Fall back to weak IoContext if no timestamp is available
+    writer->report(spanContext,
+        tracing::Return({fetchResponseInfo.map([](auto& info) { return info.clone(); })}),
+        timestamp.orDefault([&]() { return getTime(); }));
+  }
+
+  // Add fetch response info for legacy tail worker
+  KJ_IF_SOME(info, fetchResponseInfo) {
+    KJ_REQUIRE(KJ_REQUIRE_NONNULL(trace->eventInfo).is<tracing::FetchEventInfo>());
+    KJ_ASSERT(trace->fetchResponseInfo == kj::none, "setFetchResponseInfo can only be called once");
+    trace->fetchResponseInfo = kj::mv(info);
+  }
 }
 
 void BaseTracer::setUserRequestSpan(SpanParent&& span) {

src/workerd/io/tracer.c++

@@ -120,9 +120,12 @@ class BaseTracer: public kj::Refcounted {
   virtual void setEventInfo(
       IoContext::IncomingRequest& incomingRequest, tracing::EventInfo&& info) = 0;
 
-  // Adds info about the response. Must not be called more than once, and only
-  // after passing a FetchEventInfo to setEventInfo().
-  virtual void setFetchResponseInfo(tracing::FetchResponseInfo&& info) = 0;
+  // Sets the return event for Streaming Tail Worker, including fetchResponseInfo (HTTP status code)
+  // if available. Must not be called more than once, and fetchResponseInfo should only be set for
+  // fetch events. For legacy tail worker, there is no distinct return event so we only add
+  // fetchResponseInfo to the trace if present.
+  virtual void setReturn(kj::Maybe<kj::Date> time = kj::none,
+      kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo = kj::none) = 0;
 
   // Reports the outcome event of the worker invocation. For Streaming Tail Worker, this will be the
   // final event, causing the stream to terminate.
@@ -143,6 +146,11 @@ class BaseTracer: public kj::Refcounted {
       const kj::ConstString& methodName) = 0;
 
  protected:
+  // Retrieves the current timestamp. If the IoContext is no longer available, we assume that the
+  // worker must have wrapped up and reported its outcome event, we report completeTime in that case
+  // acordingly.
+  kj::Date getTime();
+
   // helper method for addSpan() implementations
   void adjustSpanTime(CompleteSpan& span);
 
@@ -194,13 +202,15 @@ class WorkerTracer final: public BaseTracer {
   void setEventInfoInternal(
       const tracing::InvocationSpanContext& context, kj::Date timestamp, tracing::EventInfo&& info);
 
-  void setFetchResponseInfo(tracing::FetchResponseInfo&& info) override;
   void setOutcome(EventOutcome outcome, kj::Duration cpuTime, kj::Duration wallTime) override;
   virtual void recordTimestamp(kj::Date timestamp) override;
 
   // Set a worker-level tag/attribute to be provided in the onset event.
   void setWorkerAttribute(kj::ConstString key, Span::TagValue value);
 
+  void setReturn(kj::Maybe<kj::Date> time = kj::none,
+      kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo = kj::none) override;
+
   void setJsRpcInfo(const tracing::InvocationSpanContext& context,
       kj::Date timestamp,
       const kj::ConstString& methodName) override;

src/workerd/io/tracer.h

@@ -134,6 +134,9 @@ class WorkerEntrypoint::ResponseSentTracker final: public kj::HttpService::Respo
   bool isSent() const {
     return sent;
   }
+  uint getHttpResponseStatus() const {
+    return httpResponseStatus;
+  }
 
   kj::Own<kj::AsyncOutputStream> send(uint statusCode,
       kj::StringPtr statusText,
@@ -142,6 +145,7 @@ class WorkerEntrypoint::ResponseSentTracker final: public kj::HttpService::Respo
     TRACE_EVENT(
         "workerd", "WorkerEntrypoint::ResponseSentTracker::send()", "statusCode", statusCode);
     sent = true;
+    httpResponseStatus = statusCode;
     return inner.send(statusCode, statusText, headers, expectedBodySize);
   }
 
@@ -152,6 +156,7 @@ class WorkerEntrypoint::ResponseSentTracker final: public kj::HttpService::Respo
   }
 
  private:
+  uint httpResponseStatus = 0;
   kj::HttpService::Response& inner;
   bool sent = false;
 };
@@ -255,6 +260,11 @@ kj::Promise<void> WorkerEntrypoint::request(kj::HttpMethod method,
   auto wrappedResponse = kj::heap<ResponseSentTracker>(response);
 
   bool isActor = context.getActor() != kj::none;
+  // HACK: Capture workerTracer directly, it's unclear how to acquire the right tracer from context
+  // when we need it (for DOs, IoContext may point to a different WorkerTracer by the time we use
+  // it). The tracer lives as long or longer than the IoContext (based on being co-owned
+  // by IncomingRequest and PipelineTracer) so long enough.
+  kj::Maybe<BaseTracer&> workerTracer;
 
   KJ_IF_SOME(t, incomingRequest->getWorkerTracer()) {
     kj::String cfJson;
@@ -281,6 +291,7 @@ kj::Promise<void> WorkerEntrypoint::request(kj::HttpMethod method,
 
     t.setEventInfo(*incomingRequest,
         tracing::FetchEventInfo(method, kj::str(url), kj::mv(cfJson), kj::mv(traceHeadersArray)));
+    workerTracer = t;
   }
 
   auto metricsForCatch = kj::addRef(incomingRequest->getMetrics());
@@ -313,10 +324,19 @@ kj::Promise<void> WorkerEntrypoint::request(kj::HttpMethod method,
         cfBlobJson, lock,
         lock.getExportedHandler(entrypointName, kj::mv(props), context.getActor()), kj::mv(signal));
   })
-      .then([this](api::DeferredProxy<void> deferredProxy) {
+      .then([this, &context, &wrappedResponse = *wrappedResponse, workerTracer](
+                api::DeferredProxy<void> deferredProxy) {
     TRACE_EVENT("workerd", "WorkerEntrypoint::request() deferred proxy step",
         PERFETTO_FLOW_FROM_POINTER(this));
     proxyTask = kj::mv(deferredProxy.proxyTask);
+    KJ_IF_SOME(t, workerTracer) {
+      auto httpResponseStatus = wrappedResponse.getHttpResponseStatus();
+      if (httpResponseStatus != 0) {
+        t.setReturn(context.now(), tracing::FetchResponseInfo(httpResponseStatus));
+      } else {
+        t.setReturn(context.now());
+      }
+    }
   })
       .catch_([this, &context](kj::Exception&& exception) mutable -> kj::Promise<void> {
     TRACE_EVENT("workerd", "WorkerEntrypoint::request() catch", PERFETTO_FLOW_FROM_POINTER(this));
@@ -383,8 +403,8 @@ kj::Promise<void> WorkerEntrypoint::request(kj::HttpMethod method,
     proxyTask = kj::none;
   }))
       .catch_([this, wrappedResponse = kj::mv(wrappedResponse), isActor, method, url, &headers,
-                  &requestBody, metrics = kj::mv(metricsForCatch)](
-                  kj::Exception&& exception) mutable -> kj::Promise<void> {
+                  &requestBody, metrics = kj::mv(metricsForCatch),
+                  workerTracer](kj::Exception&& exception) mutable -> kj::Promise<void> {
     // Don't return errors to end user.
     TRACE_EVENT("workerd", "WorkerEntrypoint::request() exception",
         PERFETTO_TERMINATING_FLOW_FROM_POINTER(this));
@@ -449,7 +469,7 @@ kj::Promise<void> WorkerEntrypoint::request(kj::HttpMethod method,
         metrics->setFailedOpen(true);
         return promise.attach(kj::mv(service));
       });
-      return promise.catch_([this, wrappedResponse = kj::mv(wrappedResponse),
+      return promise.catch_([this, wrappedResponse = kj::mv(wrappedResponse), workerTracer,
                                 metrics = kj::mv(metrics)](kj::Exception&& e) mutable {
         metrics->setFailedOpen(false);
         if (e.getType() != kj::Exception::Type::DISCONNECTED &&
@@ -462,6 +482,9 @@ kj::Promise<void> WorkerEntrypoint::request(kj::HttpMethod method,
         if (!wrappedResponse->isSent()) {
           kj::HttpHeaders headers(threadContext.getHeaderTable());
           wrappedResponse->send(500, "Internal Server Error", headers, uint64_t(0));
+          KJ_IF_SOME(t, workerTracer) {
+            t.setReturn(kj::none, tracing::FetchResponseInfo(500));
+          }
         }
       });
     } else if (tunnelExceptions) {
@@ -485,6 +508,10 @@ kj::Promise<void> WorkerEntrypoint::request(kj::HttpMethod method,
         } else {
           wrappedResponse->send(500, "Internal Server Error", headers, uint64_t(0));
         }
+        KJ_IF_SOME(t, workerTracer) {
+          t.setReturn(
+              kj::none, tracing::FetchResponseInfo(wrappedResponse->getHttpResponseStatus()));
+        }
       }
 
       return kj::READY_NOW;

src/workerd/io/worker-entrypoint.c++

@@ -1587,9 +1587,6 @@ class RequestObserverWithTracer final: public RequestObserver, public WorkerInte
         auto time = IoContext::current().now();
         t->recordTimestamp(time);
       }
-      if (fetchStatus != 0) {
-        t->setFetchResponseInfo(tracing::FetchResponseInfo(fetchStatus));
-      }
       t->setOutcome(
           outcome, 0 * kj::MILLISECONDS /* cpu time */, 0 * kj::MILLISECONDS /* wall time */);
     }