fix: clone vs child

Author: prestwich

src/router.rs

@@ -337,10 +337,10 @@ where
 
         for req in inbound.iter() {
             let req = req.map(|req| {
-                // Cloning here resets the `TracingInfo`, which means each
+                // This resets the `TracingInfo`, which means each
                 // ctx has a separate span with similar metadata.
-                let ctx = ctx.clone();
-                let request_span = ctx.init_request_span(self, Some(&batch_span)).clone();
+                let ctx = ctx.child_ctx(self, Some(&batch_span));
+                let request_span = ctx.span().clone();
 
                 // Several span fields are populated in `HandlerArgs::new`.
                 let args = HandlerArgs::new(ctx, req);

src/routes/ctx.rs

@@ -30,7 +30,7 @@ impl From<SendError<WriteItem>> for NotifyError {
 
 /// Tracing information for OpenTelemetry. This struct is used to store
 /// information about the current request that can be used for tracing.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 #[non_exhaustive]
 pub struct TracingInfo {
     /// The OpenTelemetry service name.
@@ -43,16 +43,6 @@ pub struct TracingInfo {
     span: OnceLock<tracing::Span>,
 }
 
-impl Clone for TracingInfo {
-    fn clone(&self) -> Self {
-        Self {
-            service: self.service,
-            context: self.context.clone(),
-            span: OnceLock::new(),
-        }
-    }
-}
-
 impl TracingInfo {
     /// Create a new tracing info with the given service name and no context.
     #[allow(dead_code)] // used in some features
@@ -76,6 +66,43 @@ impl TracingInfo {
         }
     }
 
+    fn make_span<S>(
+        &self,
+        router: &Router<S>,
+        with_notifications: bool,
+        parent: Option<&tracing::Span>,
+    ) -> tracing::Span
+    where
+        S: Clone + Send + Sync + 'static,
+    {
+        let span = info_span!(
+            parent: parent.and_then(|p| p.id()),
+            "AjjRequest",
+            "otel.kind" = "server",
+            "rpc.system" = "jsonrpc",
+            "rpc.jsonrpc.version" = "2.0",
+            "rpc.service" = router.service_name(),
+            notifications_enabled = with_notifications,
+            "trace_id" = ::tracing::field::Empty,
+            "otel.name" = ::tracing::field::Empty,
+            "otel.status_code" = ::tracing::field::Empty,
+            "rpc.jsonrpc.request_id" = ::tracing::field::Empty,
+            "rpc.jsonrpc.error_code" = ::tracing::field::Empty,
+            "rpc.jsonrpc.error_message" = ::tracing::field::Empty,
+            "rpc.method" = ::tracing::field::Empty,
+            params = ::tracing::field::Empty,
+        );
+        if let Some(context) = &self.context {
+            let _ = span.set_parent(context.clone());
+
+            span.record(
+                "trace_id",
+                context.span().span_context().trace_id().to_string(),
+            );
+        }
+        span
+    }
+
     /// Create a request span for a handler invocation.
     fn init_request_span<S>(
         &self,
@@ -89,35 +116,23 @@ impl TracingInfo {
         // This span is populated with as much detail as possible, and then
         // given to the Request. It will be populated with request-specific
         // details (e.g. method) during request setup.
-        self.span.get_or_init(|| {
-            let span = info_span!(
-                parent: parent.and_then(|p| p.id()),
-                "AjjRequest",
-                "otel.kind" = "server",
-                "rpc.system" = "jsonrpc",
-                "rpc.jsonrpc.version" = "2.0",
-                "rpc.service" = router.service_name(),
-                notifications_enabled = with_notifications,
-                "trace_id" = ::tracing::field::Empty,
-                "otel.name" = ::tracing::field::Empty,
-                "otel.status_code" = ::tracing::field::Empty,
-                "rpc.jsonrpc.request_id" = ::tracing::field::Empty,
-                "rpc.jsonrpc.error_code" = ::tracing::field::Empty,
-                "rpc.jsonrpc.error_message" = ::tracing::field::Empty,
-                "rpc.method" = ::tracing::field::Empty,
-                params = ::tracing::field::Empty,
-            );
-            if let Some(context) = &self.context {
-                let _ = span.set_parent(context.clone());
-
-                span.record(
-                    "trace_id",
-                    context.span().span_context().trace_id().to_string(),
-                );
-            }
+        self.span
+            .get_or_init(|| self.make_span(router, with_notifications, parent))
+    }
 
-            span
-        })
+    /// Create a child tracing info for a new handler context.
+    pub fn child<S: Clone + Send + Sync + 'static>(
+        &self,
+        router: &Router<S>,
+        with_notifications: bool,
+        parent: Option<&tracing::Span>,
+    ) -> Self {
+        let span = self.make_span(router, with_notifications, parent);
+        Self {
+            service: self.service,
+            context: self.context.clone(),
+            span: OnceLock::from(span),
+        }
     }
 
     /// Get a reference to the tracing span for this request.
@@ -186,6 +201,24 @@ impl HandlerCtx {
         }
     }
 
+    /// Create a child handler context for a new handler invocation.
+    ///
+    /// This is used when handling batch requests, to give each handler
+    /// its own context.
+    pub fn child_ctx<S: Clone + Send + Sync + 'static>(
+        &self,
+        router: &Router<S>,
+        parent: Option<&tracing::Span>,
+    ) -> Self {
+        Self {
+            notifications: self.notifications.clone(),
+            tasks: self.tasks.clone(),
+            tracing: self
+                .tracing
+                .child(router, self.notifications_enabled(), parent),
+        }
+    }
+
     /// Get a reference to the tracing information for this handler context.
     pub const fn tracing_info(&self) -> &TracingInfo {
         &self.tracing

tests/common/mod.rs

@@ -104,13 +104,17 @@ async fn test_ping<T: TestClient>(client: &mut T) {
 
 /// basic tests of the test router
 pub async fn basic_tests<T: TestClient>(client: &mut T) {
-    test_ping(client).await;
+    timeout(Duration::from_secs(10), async move {
+        test_ping(client).await;
 
-    test_double(client).await;
+        test_double(client).await;
 
-    test_call_me_later(client).await;
+        test_call_me_later(client).await;
 
-    test_notification(client).await;
+        test_notification(client).await;
+    })
+    .await
+    .unwrap();
 }
 
 /// Test when a full batch request fails to parse
@@ -205,11 +209,15 @@ async fn batch_contains_only_notifications<T: TestClient>(client: &mut T) {
 
 /// Test of batch requests
 pub async fn batch_tests<T: TestClient>(client: &mut T) {
-    batch_parse_fail(client).await;
+    timeout(Duration::from_secs(10), async move {
+        batch_parse_fail(client).await;
 
-    batch_single_req_parse_fail(client).await;
+        batch_single_req_parse_fail(client).await;
 
-    batch_contains_notification(client).await;
+        batch_contains_notification(client).await;
 
-    batch_contains_only_notifications(client).await;
+        batch_contains_only_notifications(client).await;
+    })
+    .await
+    .unwrap();
 }