feat: enhance otel implementation by setting more span info on request receipt

Author: prestwich

Cargo.toml

@@ -5,7 +5,7 @@ description = "Simple, modern, ergonomic JSON-RPC 2.0 router built with tower an
 keywords = ["json-rpc", "jsonrpc", "json"]
 categories = ["web-programming::http-server", "web-programming::websocket"]
 
-version = "0.3.4"
+version = "0.4.0"
 edition = "2021"
 rust-version = "1.81"
 authors = ["init4", "James Prestwich"]

src/axum.rs

@@ -1,15 +1,23 @@
 use crate::{
     types::{InboundData, Response},
-    HandlerCtx, TaskSet,
+    HandlerCtx, TaskSet, TracingInfo,
 };
 use axum::{
     extract::FromRequest,
     http::{header, HeaderValue},
     response::IntoResponse,
 };
 use bytes::Bytes;
-use std::{future::Future, pin::Pin};
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        atomic::{AtomicU32, Ordering},
+        Arc,
+    },
+};
 use tokio::runtime::Handle;
+use tracing::{debug, debug_span};
 
 /// A wrapper around an [`Router`] that implements the
 /// [`axum::handler::Handler`] trait. This struct is an implementation detail
@@ -21,14 +29,22 @@ use tokio::runtime::Handle;
 #[derive(Debug, Clone)]
 pub(crate) struct IntoAxum<S> {
     pub(crate) router: crate::Router<S>,
+
     pub(crate) task_set: TaskSet,
+
+    /// Counter for OTEL messages received.
+    pub(crate) rx_msg_id: Arc<AtomicU32>,
+    /// Counter for OTEL messages sent.
+    pub(crate) tx_msg_id: Arc<AtomicU32>,
 }
 
 impl<S> From<crate::Router<S>> for IntoAxum<S> {
     fn from(router: crate::Router<S>) -> Self {
         Self {
             router,
             task_set: Default::default(),
+            rx_msg_id: Arc::new(AtomicU32::new(0)),
+            tx_msg_id: Arc::new(AtomicU32::new(0)),
         }
     }
 }
@@ -39,12 +55,35 @@ impl<S> IntoAxum<S> {
         Self {
             router,
             task_set: handle.into(),
+            rx_msg_id: Arc::new(AtomicU32::new(0)),
+            tx_msg_id: Arc::new(AtomicU32::new(0)),
         }
     }
+}
 
-    /// Get a new context, built from the task set.
+impl<S> IntoAxum<S>
+where
+    S: Clone + Send + Sync + 'static,
+{
     fn ctx(&self) -> HandlerCtx {
-        self.task_set.clone().into()
+        let request_span = debug_span!(
+            "ajj.IntoAxum::call",
+            "otel.kind" = "server",
+            "rpc.system" = "jsonrpc",
+            "rpc.jsonrpc.version" = "2.0",
+            "rpc.service" = self.router.service_name(),
+            notifications_enabled = false,
+            params = tracing::field::Empty
+        );
+
+        HandlerCtx::new(
+            None,
+            self.task_set.clone(),
+            TracingInfo {
+                service: self.router.service_name(),
+                request_span,
+            },
+        )
     }
 }
 
@@ -56,25 +95,44 @@ where
 
     fn call(self, req: axum::extract::Request, state: S) -> Self::Future {
         Box::pin(async move {
+            let ctx = self.ctx();
+
             let Ok(bytes) = Bytes::from_request(req, &state).await else {
                 return Box::<str>::from(Response::parse_error()).into_response();
             };
 
-            // If the inbound data is not currently parsable, we
-            // send an empty one it to the router, as the router enforces
-            // the specification.
-            let req = InboundData::try_from(bytes).unwrap_or_default();
+            // https://github.com/open-telemetry/semantic-conventions/blob/main/docs/rpc/rpc-spans.md#message-event
+            let req = ctx.span().in_scope(|| {
+                debug!(
+                    "rpc.message.id" = self.rx_msg_id.fetch_add(1, Ordering::Relaxed),
+                    "rpc.message.type" = "received",
+                    "rpc.message.uncompressed_size" = bytes.len(),
+                    "Received request"
+                );
+
+                // If the inbound data is not currently parsable, we
+                // send an empty one it to the router, as the router enforces
+                // the specification.
+                InboundData::try_from(bytes).unwrap_or_default()
+            });
 
-            if let Some(response) = self
-                .router
-                .call_batch_with_state(self.ctx(), req, state)
-                .await
-            {
+            let span = ctx.span().clone();
+            if let Some(response) = self.router.call_batch_with_state(ctx, req, state).await {
                 let headers = [(
                     header::CONTENT_TYPE,
                     HeaderValue::from_static(mime::APPLICATION_JSON.as_ref()),
                 )];
+
                 let body = Box::<str>::from(response);
+
+                span.in_scope(|| {
+                    debug!(
+                        "rpc.message.id" = self.tx_msg_id.fetch_add(1, Ordering::Relaxed),
+                        "rpc.message.type" = "received",
+                        "Received request"
+                    );
+                });
+
                 (headers, body).into_response()
             } else {
                 ().into_response()

src/lib.rs

@@ -170,7 +170,8 @@ pub use pubsub::ReadJsonStream;
 
 mod routes;
 pub use routes::{
-    BatchFuture, Handler, HandlerArgs, HandlerCtx, NotifyError, Params, RouteFuture, State,
+    BatchFuture, Handler, HandlerArgs, HandlerCtx, NotifyError, TracingInfo, Params, RouteFuture,
+    State,
 };
 pub(crate) use routes::{BoxedIntoRoute, ErasedIntoRoute, Method, Route};
 
@@ -206,7 +207,8 @@ pub(crate) mod test_utils {
 mod test {
 
     use crate::{
-        router::RouterInner, routes::HandlerArgs, test_utils::assert_rv_eq, ResponsePayload,
+        router::RouterInner, routes::HandlerArgs, test_utils::assert_rv_eq, HandlerCtx,
+        ResponsePayload,
     };
     use bytes::Bytes;
     use serde_json::value::RawValue;
@@ -232,7 +234,7 @@ mod test {
         let res = router
             .call_with_state(
                 HandlerArgs {
-                    ctx: Default::default(),
+                    ctx: HandlerCtx::mock(),
                     req: req.try_into().unwrap(),
                 },
                 (),
@@ -251,7 +253,7 @@ mod test {
         let res2 = router
             .call_with_state(
                 HandlerArgs {
-                    ctx: Default::default(),
+                    ctx: HandlerCtx::mock(),
                     req: req2.try_into().unwrap(),
                 },
                 (),

src/pubsub/mod.rs

@@ -95,6 +95,7 @@ mod ipc;
 pub use ipc::ReadJsonStream;
 
 mod shared;
+pub(crate) use shared::WriteItem;
 pub use shared::{ConnectionId, DEFAULT_NOTIFICATION_BUFFER_PER_CLIENT};
 
 mod shutdown;

src/pubsub/shared.rs

@@ -1,11 +1,14 @@
 use crate::{
     pubsub::{In, JsonSink, Listener, Out},
     types::InboundData,
-    HandlerCtx, TaskSet,
+    HandlerCtx, TaskSet, TracingInfo,
 };
 use core::fmt;
 use serde_json::value::RawValue;
-use std::sync::{atomic::AtomicU64, Arc};
+use std::sync::{
+    atomic::{AtomicU32, AtomicU64, Ordering},
+    Arc,
+};
 use tokio::{pin, runtime::Handle, select, sync::mpsc, task::JoinHandle};
 use tokio_stream::StreamExt;
 use tokio_util::sync::WaitForCancellationFutureOwned;
@@ -105,8 +108,7 @@ impl ConnectionManager {
 
     /// Increment the connection ID counter and return an unused ID.
     fn next_id(&self) -> ConnectionId {
-        self.next_id
-            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
+        self.next_id.fetch_add(1, Ordering::Relaxed)
     }
 
     /// Get a clone of the router.
@@ -131,13 +133,15 @@ impl ConnectionManager {
             write_task: tx,
             requests,
             tasks: tasks.clone(),
+            rx_msg_id: Arc::new(AtomicU32::new(0)),
         };
 
         let wt = WriteTask {
             tasks,
             conn_id,
-            json: rx,
+            items: rx,
             connection,
+            tx_msg_id: Arc::new(AtomicU32::new(0)),
         };
 
         (rt, wt)
@@ -168,11 +172,14 @@ struct RouteTask<T: crate::pubsub::Listener> {
     /// Connection ID for the connection serviced by this task.
     pub(crate) conn_id: ConnectionId,
     /// Sender to the write task.
-    pub(crate) write_task: mpsc::Sender<Box<RawValue>>,
+    pub(crate) write_task: mpsc::Sender<WriteItem>,
     /// Stream of requests.
     pub(crate) requests: In<T>,
     /// The task set for this connection
     pub(crate) tasks: TaskSet,
+
+    /// Counter for OTEL messages received.
+    pub(crate) rx_msg_id: Arc<AtomicU32>,
 }
 
 impl<T: crate::pubsub::Listener> fmt::Debug for RouteTask<T> {
@@ -199,6 +206,7 @@ where
             mut requests,
             write_task,
             tasks,
+            rx_msg_id,
             ..
         } = self;
 
@@ -224,6 +232,8 @@ where
                         break;
                     };
 
+                    let item_bytes = item.len();
+
                     // If the inbound data is not currently parsable, we
                     // send an empty one it to the router, as the router
                     // enforces the specification.
@@ -234,16 +244,38 @@ where
                     // if the client stops accepting responses, we do not keep
                     // handling inbound requests.
                     let Ok(permit) = write_task.clone().reserve_owned().await else {
-                        tracing::error!("write task dropped while waiting for permit");
+                        error!("write task dropped while waiting for permit");
                         break;
                     };
 
+                    let tracing = TracingInfo { service: router.service_name(), request_span: debug_span!(
+                        parent: None,
+                        "ajj.pubsub.RouteTask::call",
+                        "otel.kind" = "server",
+                        "rpc.system" = "jsonrpc",
+                        "rpc.jsonrpc.version" = "2.0",
+                        "rpc.service" = router.service_name(),
+                        notifications_enabled = true,
+                        params = tracing::field::Empty
+                    ) };
+
                     let ctx =
                     HandlerCtx::new(
                         Some(write_task.clone()),
                         children.clone(),
+                        tracing,
                     );
 
+                    let span = ctx.span().clone();
+                    span.in_scope(|| {
+                        debug!(
+                            "rpc.message.id" = rx_msg_id.fetch_add(1, Ordering::Relaxed),
+                            "rpc.message.type" = "received",
+                            "rpc.message.uncompressed_size" = item_bytes,
+                            "Received request"
+                        );
+                    });
+
                     // Run the future in a new task.
                     let fut = router.handle_request_batch(ctx, reqs);
 
@@ -252,9 +284,9 @@ where
                             // Send the response to the write task.
                             // we don't care if the receiver has gone away,
                             // as the task is done regardless.
-                            if let Some(rv) = fut.await {
+                            if let Some(json) = fut.await {
                                 let _ = permit.send(
-                                    rv
+                                    WriteItem { span, json }
                                 );
                             }
                         }
@@ -275,6 +307,13 @@ where
     }
 }
 
+/// An item to be written to an outbound JSON pubsub stream.
+#[derive(Debug, Clone)]
+pub(crate) struct WriteItem {
+    pub(crate) span: tracing::Span,
+    pub(crate) json: Box<RawValue>,
+}
+
 /// The Write Task is responsible for writing JSON to the outbound connection.
 struct WriteTask<T: Listener> {
     /// Task set
@@ -287,10 +326,13 @@ struct WriteTask<T: Listener> {
     ///
     /// Dropping this channel will cause the associated [`RouteTask`] to
     /// shutdown.
-    pub(crate) json: mpsc::Receiver<Box<RawValue>>,
+    pub(crate) items: mpsc::Receiver<WriteItem>,
 
     /// Outbound connections.
     pub(crate) connection: Out<T>,
+
+    /// Counter for OTEL messages sent.
+    pub(crate) tx_msg_id: Arc<AtomicU32>,
 }
 
 impl<T: Listener> WriteTask<T> {
@@ -305,8 +347,9 @@ impl<T: Listener> WriteTask<T> {
     pub(crate) async fn task_future(self) {
         let WriteTask {
             tasks,
-            mut json,
+            mut items,
             mut connection,
+            tx_msg_id,
             ..
         } = self;
 
@@ -318,12 +361,20 @@ impl<T: Listener> WriteTask<T> {
                     debug!("Shutdown signal received");
                     break;
                 }
-                json = json.recv() => {
-                    let Some(json) = json else {
+                item = items.recv() => {
+                    let Some(WriteItem { span, json }) = item else {
                         tracing::error!("Json stream has closed");
                         break;
                     };
-                    let span = debug_span!("WriteTask", conn_id = self.conn_id);
+                    span.record("conn_id", self.conn_id);
+                    span.in_scope(|| {
+                        debug!(
+                            "rpc.message.id" = tx_msg_id.fetch_add(1, Ordering::Relaxed),
+                            "rpc.message.type" = "sent",
+                            "Sending response"
+                        );
+                    });
+
                     if let Err(err) = connection.send_json(json).instrument(span).await {
                         debug!(%err, conn_id = self.conn_id, "Failed to send json");
                         break;