refactor: improve tracing span usage

Author: prestwich

Cargo.toml

@@ -10,8 +10,8 @@ edition = "2021"
 rust-version = "1.81"
 authors = ["init4", "James Prestwich"]
 license = "MIT OR Apache-2.0"
-homepage = "https://github.com/init4tech/rpc"
-repository = "https://github.com/init4tech/rpc"
+homepage = "https://github.com/init4tech/ajj"
+repository = "https://github.com/init4tech/ajj"
 
 [dependencies]
 bytes = "1.9.0"

src/pubsub/shared.rs

@@ -8,7 +8,7 @@ use serde_json::value::RawValue;
 use tokio::{pin, select, sync::mpsc, task::JoinHandle};
 use tokio_stream::StreamExt;
 use tokio_util::sync::WaitForCancellationFutureOwned;
-use tracing::{debug, debug_span, error, instrument, trace, Instrument};
+use tracing::{debug, debug_span, error, trace, Instrument};
 
 /// Default notification buffer size per task.
 pub const DEFAULT_NOTIFICATION_BUFFER_PER_CLIENT: usize = 16;
@@ -163,8 +163,7 @@ where
     /// and routes them to the router. For each request, a new task is spawned
     /// to handle the request, and given a sender to the [`WriteTask`]. This
     /// ensures that requests can be handled concurrently.
-    #[instrument(name = "RouteTask", skip(self), fields(conn_id = self.conn_id))]
-    pub async fn task_future(self, cancel: WaitForCancellationFutureOwned) {
+    pub(crate) async fn task_future(self, cancel: WaitForCancellationFutureOwned) {
         let RouteTask {
             router,
             mut requests,
@@ -200,27 +199,24 @@ where
                     // enforces the specification.
                     let reqs = InboundData::try_from(item).unwrap_or_default();
 
-                    let span = debug_span!("pubsub request handling", reqs = reqs.len());
-
-                    let ctx =
-                    HandlerCtx::new(
-                        Some(write_task.clone()),
-                        children.clone(),
-                    );
-
-                    let fut = router.handle_request_batch(ctx, reqs);
-                    let write_task = write_task.clone();
-
                     // Acquiring the permit before spawning the task means that
                     // the write task can backpressure the route task. I.e.
                     // if the client stops accepting responses, we do not keep
                     // handling inbound requests.
-                    let Ok(permit) = write_task.reserve_owned().await else {
+                    let Ok(permit) = write_task.clone().reserve_owned().await else {
                         tracing::error!("write task dropped while waiting for permit");
                         break;
                     };
 
+                    let ctx =
+                    HandlerCtx::new(
+                        Some(write_task.clone()),
+                        children.clone(),
+                    );
+
                     // Run the future in a new task.
+                    let fut = router.handle_request_batch(ctx, reqs);
+
                     children.spawn_cancellable(
                         async move {
                             // Send the response to the write task.
@@ -232,7 +228,6 @@ where
                                 );
                             }
                         }
-                        .instrument(span)
                     );
                 }
             }
@@ -277,7 +272,6 @@ impl<T: Listener> WriteTask<T> {
     /// [`ServerShutdown`] struct.
     ///
     /// [`ServerShutdown`]: crate::pubsub::ServerShutdown
-    #[instrument(skip(self), fields(conn_id = self.conn_id))]
     pub(crate) async fn task_future(self) {
         let WriteTask {
             tasks,
@@ -299,8 +293,9 @@ impl<T: Listener> WriteTask<T> {
                         tracing::error!("Json stream has closed");
                         break;
                     };
-                    if let Err(err) = connection.send_json(json).await {
-                        debug!(%err, "Failed to send json");
+                    let span = debug_span!("WriteTask", conn_id = self.conn_id);
+                    if let Err(err) = connection.send_json(json).instrument(span).await {
+                        debug!(%err, conn_id = self.conn_id, "Failed to send json");
                         break;
                     }
                 }

src/router.rs

@@ -10,7 +10,7 @@ use core::fmt;
 use serde_json::value::RawValue;
 use std::{borrow::Cow, collections::BTreeMap, convert::Infallible, sync::Arc, task::Poll};
 use tower::Service;
-use tracing::{debug_span, trace};
+use tracing::debug_span;
 
 /// A JSON-RPC router. This is the top-level type for handling JSON-RPC
 /// requests. It is heavily inspired by the [`axum::Router`] type.
@@ -301,9 +301,7 @@ where
     pub fn call_with_state(&self, args: HandlerArgs, state: S) -> RouteFuture {
         let id = args.req().id_owned();
         let method = args.req().method();
-
-        let span = debug_span!("Router::call_with_state", %method, ?id);
-        trace!(params = args.req().params());
+        let span = debug_span!(parent: None, "Router::call_with_state", %method, ?id);
 
         self.inner.call_with_state(args, state).with_span(span)
     }
@@ -318,14 +316,19 @@ where
         let mut fut = BatchFuture::new_with_capacity(inbound.single(), inbound.len());
         // According to spec, non-parsable requests should still receive a
         // response.
-        for req in inbound.iter() {
+        let span = debug_span!(parent: None, "BatchFuture::poll", futs = fut.len());
+
+        for (batch_idx, req) in inbound.iter().enumerate() {
             let req = req.map(|req| {
+                let span = debug_span!("RouteFuture::poll", batch_idx, method = req.method(), id = ?req.id());
                 let args = HandlerArgs::new(ctx.clone(), req);
-                self.call_with_state(args, state.clone())
+                self.inner
+                    .call_with_state(args, state.clone())
+                    .with_span(span)
             });
             fut.push_parse_result(req);
         }
-        fut
+        fut.with_span(span)
     }
 
     /// Nest this router into a new Axum router, with the specified path and the currently-running

src/routes/future.rs

@@ -118,6 +118,9 @@ pub struct BatchFuture {
     resps: Vec<Box<RawValue>>,
     /// Whether the batch was a single request.
     single: bool,
+
+    /// The span (if any).
+    span: Option<tracing::Span>,
 }
 
 impl fmt::Debug for BatchFuture {
@@ -136,6 +139,15 @@ impl BatchFuture {
             futs: BatchFutureInner::Prepping(Vec::with_capacity(capacity)),
             resps: Vec::with_capacity(capacity),
             single,
+            span: None,
+        }
+    }
+
+    /// Set the span for the future.
+    pub(crate) fn with_span(self, span: tracing::Span) -> Self {
+        Self {
+            span: Some(span),
+            ..self
         }
     }
 
@@ -167,6 +179,11 @@ impl BatchFuture {
             Err(_) => self.push_parse_error(),
         }
     }
+
+    /// Get the number of futures in the batch.
+    pub(crate) fn len(&self) -> usize {
+        self.futs.len()
+    }
 }
 
 impl std::future::Future for BatchFuture {
@@ -185,6 +202,8 @@ impl std::future::Future for BatchFuture {
         }
 
         let this = self.project();
+        let _enter = this.span.as_ref().map(tracing::Span::enter);
+
         let BatchFutureInnerProj::Running(mut futs) = this.futs.project() else {
             unreachable!()
         };

src/routes/handler.rs

@@ -3,6 +3,7 @@ use crate::{
 };
 use serde_json::value::RawValue;
 use std::{convert::Infallible, future::Future, marker::PhantomData, pin::Pin, task};
+use tracing::{debug_span, enabled, trace, Instrument, Level};
 
 macro_rules! convert_result {
     ($res:expr) => {{
@@ -410,7 +411,27 @@ where
 
     fn call(&mut self, args: HandlerArgs) -> Self::Future {
         let this = self.clone();
-        Box::pin(async move { Ok(this.handler.call_with_state(args, this.state.clone()).await) })
+        Box::pin(async move {
+            let notifications_enabled = args.ctx.notifications_enabled();
+
+            let span = debug_span!(
+                "HandlerService::call",
+                notifications_enabled,
+                params = tracing::field::Empty
+            );
+            if enabled!(Level::TRACE) {
+                span.record("params", args.req.params());
+            }
+
+            Ok(this
+                .handler
+                .call_with_state(args, this.state.clone())
+                .instrument(span)
+                .await
+                .inspect(|res| {
+                    trace!(?res, "handler response");
+                }))
+        })
     }
 }
 

src/routes/mod.rs

@@ -20,6 +20,7 @@ use std::{
     task::{Context, Poll},
 };
 use tower::{util::BoxCloneSyncService, Service, ServiceExt};
+use tracing::{debug_span, enabled, Level};
 
 use crate::types::Response;
 
@@ -95,6 +96,14 @@ impl Service<HandlerArgs> for Route {
     }
 
     fn call(&mut self, args: HandlerArgs) -> Self::Future {
+        let span = debug_span!(
+            "Route::call",
+            notifications_enabled = args.ctx().notifications_enabled(),
+            params = tracing::field::Empty,
+        );
+        if enabled!(Level::TRACE) {
+            span.record("params", args.req.params());
+        }
         self.oneshot_inner(args)
     }
 }

src/types/batch.rs

@@ -3,6 +3,7 @@ use bytes::Bytes;
 use serde::Deserialize;
 use serde_json::value::RawValue;
 use std::ops::Range;
+use tracing::{debug, enabled, instrument, Level};
 
 /// UTF-8, partially deserialized JSON-RPC request batch.
 #[derive(Default)]
@@ -52,7 +53,13 @@ impl core::fmt::Debug for InboundData {
 impl TryFrom<Bytes> for InboundData {
     type Error = RequestError;
 
+    #[instrument(level = "debug", skip(bytes), fields(buf_len = bytes.len(), bytes = tracing::field::Empty))]
     fn try_from(bytes: Bytes) -> Result<Self, Self::Error> {
+        if enabled!(Level::TRACE) {
+            tracing::span::Span::current().record("bytes", &format!("0x{:x}", bytes));
+        }
+        debug!("Parsing inbound data");
+
         // Special-case a single request, rejecting invalid JSON.
         if bytes.starts_with(b"{") {
             let rv: &RawValue = serde_json::from_slice(bytes.as_ref())?;

src/types/req.rs

@@ -183,7 +183,7 @@ impl Request {
     pub fn deser_params<'a: 'de, 'de, T: serde::Deserialize<'de>>(
         &'a self,
     ) -> serde_json::Result<T> {
-        serde_json::from_str(self.params())
+        serde_json::from_str(self.params()).inspect_err(|err| tracing::debug!(%err, expected = std::any::type_name::<T>(), "failed to parse params"))
     }
 }