[o11y] Fix tail worker return reporting for fetch event exceptions

We still need to report the fetch return code following an exception.

Author: fhanau

src/workerd/io/tracer.c++

@@ -437,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
@@ -460,16 +481,15 @@ 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::setReturn(
-    kj::Date timestamp, kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo) {
+    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.
@@ -480,9 +500,10 @@ void WorkerTracer::setReturn(
   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);
+        timestamp.orDefault([&]() { return getTime(); }));
   }
 
   // Add fetch response info for legacy tail worker

src/workerd/io/tracer.h

@@ -124,8 +124,8 @@ class BaseTracer: public kj::Refcounted {
   // 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::Date time, kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo = kj::none) = 0;
+  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.
@@ -146,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);
 
@@ -203,8 +208,8 @@ class WorkerTracer final: public BaseTracer {
   // Set a worker-level tag/attribute to be provided in the onset event.
   void setWorkerAttribute(kj::ConstString key, Span::TagValue value);
 
-  void setReturn(
-      kj::Date time, kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo = kj::none) override;
+  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,

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

@@ -261,7 +261,9 @@ kj::Promise<void> WorkerEntrypoint::request(kj::HttpMethod method,
 
   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.
+  // 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()) {
@@ -401,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));
@@ -467,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 &&
@@ -480,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) {
@@ -503,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;