From 737a718467b63e5b9d87d9f00e57e02ba9be8325 Mon Sep 17 00:00:00 2001
From: Luca Palmieri <lpalmieri@truelayer.com>
Date: Wed, 28 Apr 2021 22:19:32 +0100
Subject: [PATCH] Properly extract errors coming from handlers.

---
 src/middleware.rs        | 39 ++++++++++++++++++++++++++++-----------
 src/root_span_builder.rs | 38 ++++++++++++++++++++++----------------
 2 files changed, 50 insertions(+), 27 deletions(-)

diff --git a/src/middleware.rs b/src/middleware.rs
index 48efd5ad4..7a3926068 100644
--- a/src/middleware.rs
+++ b/src/middleware.rs
@@ -1,6 +1,6 @@
 use crate::{DefaultRootSpanBuilder, RequestId, RootSpan, RootSpanBuilder};
 use actix_web::dev::{Service, ServiceRequest, ServiceResponse, Transform};
-use actix_web::{Error, HttpMessage};
+use actix_web::{Error, HttpMessage, ResponseError};
 use futures::future::{ok, Ready};
 use futures::task::{Context, Poll};
 use std::future::Future;
@@ -143,20 +143,37 @@ where
                 RootSpanType::on_request_end(Span::current(), &outcome);
 
                 #[cfg(feature = "emit_event_on_error")]
-                if let Err(error) = &outcome {
-                    let response_error = error.as_response_error();
-                    let status_code = response_error.status_code();
-                    let error_msg_prefix =
-                        "Error encountered while processing the incoming HTTP request";
-                    if status_code.is_client_error() {
-                        tracing::warn!("{}: {:?}", error_msg_prefix, response_error);
-                    } else {
-                        tracing::error!("{}: {:?}", error_msg_prefix, response_error);
-                    }
+                {
+                    emit_event_on_error(&outcome);
                 }
+
                 outcome
             }
             .instrument(root_span),
         )
     }
 }
+
+fn emit_event_on_error<B: 'static>(outcome: &Result<ServiceResponse<B>, actix_web::Error>) {
+    match outcome {
+        Ok(response) => {
+            if let Some(err) = response.response().error() {
+                emit_error_event(err.as_response_error())
+            }
+        }
+        Err(error) => {
+            let response_error = error.as_response_error();
+            emit_error_event(response_error)
+        }
+    }
+}
+
+fn emit_error_event(response_error: &dyn ResponseError) {
+    let status_code = response_error.status_code();
+    let error_msg_prefix = "Error encountered while processing the incoming HTTP request";
+    if status_code.is_client_error() {
+        tracing::warn!("{}: {:?}", error_msg_prefix, response_error);
+    } else {
+        tracing::error!("{}: {:?}", error_msg_prefix, response_error);
+    }
+}
diff --git a/src/root_span_builder.rs b/src/root_span_builder.rs
index 5b9a39156..b14919187 100644
--- a/src/root_span_builder.rs
+++ b/src/root_span_builder.rs
@@ -42,25 +42,31 @@ impl RootSpanBuilder for DefaultRootSpanBuilder {
     fn on_request_end<B>(span: Span, outcome: &Result<ServiceResponse<B>, Error>) {
         match &outcome {
             Ok(response) => {
-                span.record("http.status_code", &response.response().status().as_u16());
-                span.record("otel.status_code", &"OK");
-            }
-            Err(error) => {
-                let response_error = error.as_response_error();
-                span.record(
-                    "exception.message",
-                    &tracing::field::display(response_error),
-                );
-                span.record("exception.details", &tracing::field::debug(response_error));
-                let status_code = response_error.status_code();
-                span.record("http.status_code", &status_code.as_u16());
-
-                if status_code.is_client_error() {
-                    span.record("otel.status_code", &"OK");
+                if let Some(error) = response.response().error() {
+                    handle_error(span, error)
                 } else {
-                    span.record("otel.status_code", &"ERROR");
+                    span.record("http.status_code", &response.response().status().as_u16());
+                    span.record("otel.status_code", &"OK");
                 }
             }
+            Err(error) => handle_error(span, error),
         };
     }
 }
+
+fn handle_error(span: Span, error: &actix_web::Error) {
+    let response_error = error.as_response_error();
+    span.record(
+        "exception.message",
+        &tracing::field::display(response_error),
+    );
+    span.record("exception.details", &tracing::field::debug(response_error));
+    let status_code = response_error.status_code();
+    span.record("http.status_code", &status_code.as_u16());
+
+    if status_code.is_client_error() {
+        span.record("otel.status_code", &"OK");
+    } else {
+        span.record("otel.status_code", &"ERROR");
+    }
+}