From dc2e97ca9bd83d04d310c70d27e59257a980e752 Mon Sep 17 00:00:00 2001
From: Ermal Kaleci <ermalkaleci@gmail.com>
Date: Fri, 24 Nov 2023 15:07:32 +0100
Subject: [PATCH] set attr error.message

---
 src/extensions/client/mod.rs | 18 ++++++++----------
 src/middlewares/mod.rs       |  5 +----
 src/server.rs                | 20 ++++++++------------
 src/utils/mod.rs             | 18 ++++++++++++++++--
 4 files changed, 33 insertions(+), 28 deletions(-)

diff --git a/src/extensions/client/mod.rs b/src/extensions/client/mod.rs
index 9a9ad0b..af27525 100644
--- a/src/extensions/client/mod.rs
+++ b/src/extensions/client/mod.rs
@@ -393,15 +393,14 @@ impl Client {
 
             let result = rx.await.map_err(errors::internal_error)?.map_err(errors::map_error);
 
-            opentelemetry::trace::get_active_span(|span| match result.as_ref() {
+            match result.as_ref() {
                 Ok(_) => {
-                    span.set_status(opentelemetry::trace::Status::Ok);
+                    TRACER.span_ok();
                 }
                 Err(err) => {
-                    span.set_status(opentelemetry::trace::Status::error(err.to_string()));
-                    span.record_error(&err);
+                    TRACER.span_error(format!("{}", err));
                 }
-            });
+            }
 
             result
         }
@@ -430,15 +429,14 @@ impl Client {
 
             let result = rx.await.map_err(errors::failed)?;
 
-            opentelemetry::trace::get_active_span(|span| match result.as_ref() {
+            match result.as_ref() {
                 Ok(_) => {
-                    span.set_status(opentelemetry::trace::Status::Ok);
+                    TRACER.span_ok();
                 }
                 Err(err) => {
-                    span.set_status(opentelemetry::trace::Status::error(err.to_string()));
-                    span.record_error(&err);
+                    TRACER.span_error(format!("{}", err));
                 }
-            });
+            };
 
             result
         }
diff --git a/src/middlewares/mod.rs b/src/middlewares/mod.rs
index 52461f0..93e2228 100644
--- a/src/middlewares/mod.rs
+++ b/src/middlewares/mod.rs
@@ -154,10 +154,7 @@ impl<Request: Debug + Send + 'static, Result: Send + 'static> Middlewares<Reques
         tokio::select! {
             _ = sleep => {
                 tracing::error!("middlewares timeout: {req}");
-                opentelemetry::trace::get_active_span(|span| {
-                    span.set_status(opentelemetry::trace::Status::error("middlewares timeout"));
-                    span.record_error(&jsonrpsee::core::Error::Custom("middlewares timeout".to_string()));
-                });
+                TRACER.span_error("middlewares timeout");
                 task_handle.abort();
             }
             _ = &mut task_handle => {
diff --git a/src/server.rs b/src/server.rs
index b5941c3..b01c0a1 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -92,15 +92,12 @@ pub async fn build(config: Config) -> anyhow::Result<SubwayServerHandle> {
                             .await
                             .map_err(|_| errors::map_error(jsonrpsee::core::Error::RequestTimeout))?;
 
-                        opentelemetry::trace::get_active_span(|span| match result.as_ref() {
-                            Ok(_) => {
-                                span.set_status(opentelemetry::trace::Status::Ok);
-                            }
+                        match result.as_ref() {
+                            Ok(_) => tracer.span_ok(),
                             Err(err) => {
-                                span.set_status(opentelemetry::trace::Status::error(err.to_string()));
-                                span.record_error(&err);
+                                tracer.span_error(format!("{}", err));
                             }
-                        });
+                        };
 
                         result
                     }
@@ -163,15 +160,14 @@ pub async fn build(config: Config) -> anyhow::Result<SubwayServerHandle> {
                                 .await
                                 .map_err(|_| errors::map_error(jsonrpsee::core::Error::RequestTimeout))?;
 
-                            opentelemetry::trace::get_active_span(|span| match result.as_ref() {
+                            match result.as_ref() {
                                 Ok(_) => {
-                                    span.set_status(opentelemetry::trace::Status::Ok);
+                                    tracer.span_ok();
                                 }
                                 Err(err) => {
-                                    span.set_status(opentelemetry::trace::Status::error(format!("{:?}", err)));
-                                    span.record_error(&jsonrpsee::core::Error::Custom(format!("{:?}", err)));
+                                    tracer.span_error(format!("{:?}", err));
                                 }
-                            });
+                            };
 
                             result
                         }
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index d08f9ce..8d4d6a4 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -41,8 +41,8 @@ pub mod errors {
 pub mod telemetry {
     use opentelemetry::{
         global::{self, BoxedSpan},
-        trace::{TraceContextExt, Tracer as _},
-        Context,
+        trace::{get_active_span, Status, TraceContextExt, Tracer as _},
+        Context, KeyValue,
     };
     use std::borrow::Cow;
 
@@ -62,5 +62,19 @@ pub mod telemetry {
             let span = self.span(span_name);
             Context::current_with_span(span)
         }
+
+        pub fn span_ok(&self) {
+            get_active_span(|span| {
+                span.set_status(Status::Ok);
+            });
+        }
+
+        pub fn span_error(&self, err: impl Into<Cow<'static, str>>) {
+            get_active_span(|span| {
+                let err_msg = err.into();
+                span.set_status(Status::error(format!("{}", err_msg)));
+                span.set_attribute(KeyValue::new("error.message", err_msg));
+            });
+        }
     }
 }