refactor: carry Arc<TracingInfo> on WriteItem; drop WriteTask service_name

src/axum.rs

@@ -70,7 +70,10 @@ where
         HandlerCtx::new(
             None,
             self.task_set.clone(),
-            TracingInfo::new_with_context(self.router.service_name(), parent_context),
+            Arc::new(TracingInfo::new_with_context(
+                self.router.service_name(),
+                parent_context,
+            )),
         )
     }
 }

src/lib.rs

@@ -90,7 +90,7 @@
 //!
 //! Routers can also be served over axum websockets. When both `axum` and
 //! `pubsub` features are enabled, the `pubsub` module provides
-//! [`pubsub::AxumWsCfg`] and the [`pubsub::ajj_websocket`] axum handler. This
+//! `pubsub::AxumWsCfg` and the `pubsub::ajj_websocket` axum handler. This
 //! handler will serve the router over websockets at a specific route. The
 //! router is a property of the `AxumWsCfg` object, and is passed to the
 //! handler via axum's `State` extractor.

src/macros.rs

@@ -61,6 +61,7 @@ macro_rules! unwrap_infallible {
 /// should rely on the histogram metric for size visibility.
 ///
 /// See <https://github.com/open-telemetry/semantic-conventions/blob/d66109ff41e75f49587114e5bff9d101b87f40bd/docs/rpc/rpc-spans.md#events>
+#[cfg(any(feature = "axum", feature = "pubsub"))]
 macro_rules! message_event {
     ($type:literal, service: $service:expr, counter: $counter:expr, bytes: $bytes:expr,) => {{
         let bytes = $bytes;

src/metrics.rs

@@ -1,4 +1,6 @@
-use metrics::{counter, gauge, histogram, Counter, Gauge, Histogram};
+use metrics::{counter, gauge, Counter, Gauge};
+#[cfg(any(feature = "axum", feature = "pubsub"))]
+use metrics::{histogram, Histogram};
 use std::sync::LazyLock;
 
 /// Metric name for counting router calls.
@@ -256,6 +258,7 @@ fn record_completed_call(service_name: &'static str, method: &str) {
 }
 
 /// Get or register a histogram for message sizes on a specific service and direction.
+#[cfg(any(feature = "axum", feature = "pubsub"))]
 fn message_size(service_name: &'static str, direction: &'static str) -> Histogram {
     let _ = &DESCRIBE;
     histogram!(
@@ -266,6 +269,7 @@ fn message_size(service_name: &'static str, direction: &'static str) -> Histogra
 }
 
 /// Record the uncompressed size in bytes of a message sent or received on a service.
+#[cfg(any(feature = "axum", feature = "pubsub"))]
 pub(crate) fn record_message_size(
     service_name: &'static str,
     direction: &'static str,

src/pubsub/mod.rs

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

src/pubsub/shared.rs

@@ -1,10 +1,10 @@
 use crate::{
     pubsub::{In, JsonSink, Listener, Out},
+    routes::WriteItem,
     types::InboundData,
     HandlerCtx, TaskSet, TracingInfo,
 };
 use core::fmt;
-use serde_json::value::RawValue;
 use std::sync::{
     atomic::{AtomicU32, AtomicU64, Ordering},
     Arc,
@@ -148,7 +148,6 @@ impl ConnectionManager {
             items: rx,
             connection,
             tx_msg_id: self.tx_msg_id.clone(),
-            service_name: self.router.service_name(),
         };
 
         (rt, wt)
@@ -259,7 +258,7 @@ where
                      // possible, and then given to the Handler ctx. It
                      // will be populated with request-specific details
                      // (e.g. method) during ctx instantiation.
-                    let tracing = TracingInfo::new(router.service_name());
+                    let tracing = Arc::new(TracingInfo::new(router.service_name()));
 
                     let ctx =
                     HandlerCtx::new(
@@ -269,11 +268,11 @@ where
                     );
                     ctx.init_request_span(&router, None);
 
-                    let span = ctx.span().clone();
-                    span.in_scope(|| {
+                    let tracing_for_reply = Arc::clone(&ctx.tracing);
+                    ctx.span().in_scope(|| {
                         message_event!(
                             @received,
-                            service: router.service_name(),
+                            service: tracing_for_reply.service,
                             counter: &rx_msg_id,
                             bytes: item_bytes,
                         );
@@ -289,7 +288,7 @@ where
                             // as the task is done regardless.
                             if let Some(json) = fut.await {
                                 let _ = permit.send(
-                                    WriteItem { span, json }
+                                    WriteItem { tracing: tracing_for_reply, json }
                                 );
                             }
                         }
@@ -310,13 +309,6 @@ 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
@@ -336,9 +328,6 @@ struct WriteTask<T: Listener> {
 
     /// Counter for OTEL messages sent.
     pub(crate) tx_msg_id: Arc<AtomicU32>,
-
-    /// Service name, used to tag the `ajj.router.message_size_bytes` histogram.
-    pub(crate) service_name: &'static str,
 }
 
 impl<T: Listener> WriteTask<T> {
@@ -356,7 +345,6 @@ impl<T: Listener> WriteTask<T> {
             mut items,
             mut connection,
             tx_msg_id,
-            service_name,
             ..
         } = self;
 
@@ -369,14 +357,15 @@ impl<T: Listener> WriteTask<T> {
                     break;
                 }
                 item = items.recv() => {
-                    let Some(WriteItem { span, json }) = item else {
-                        tracing::error!("Json stream has closed");
+                    let Some(WriteItem { tracing, json }) = item else {
+                        ::tracing::error!("Json stream has closed");
                         break;
                     };
+                    let span = tracing.request_span().clone();
                     span.in_scope(|| {
                         message_event!(
                             @sent,
-                            service: service_name,
+                            service: tracing.service,
                             counter: &tx_msg_id,
                             bytes: json.get().len(),
                         );