Introduce routing metadata as a span property

src/lib.rs

@@ -71,12 +71,29 @@ pub type Result<T> = std::result::Result<T, Error>;
 mod tests {
     use super::*;
     use crate::sampler::AllSampler;
-    use crate::span::{FinishedSpan, Span};
+    use crate::span::{FinishedSpan, RoutingMetadata, Span};
     use crate::tag::{StdTag, Tag, TagValue};
     use std::sync::atomic::{AtomicI64, Ordering};
+    use std::sync::Arc;
     use std::thread;
     use std::time::Duration;
 
+    #[derive(Debug)]
+    struct TestRouting {
+        zone_id: u64,
+        account_id: u64,
+    }
+
+    impl RoutingMetadata for TestRouting {
+        fn group_key(&self) -> String {
+            format!("{}|{}", self.zone_id, self.account_id)
+        }
+
+        fn encode(&self) -> String {
+            format!("zone={};account={}", self.zone_id, self.account_id)
+        }
+    }
+
     #[tokio::test]
     async fn it_works() {
         let (tracer, mut span_rx) = Tracer::new(AllSampler);
@@ -174,4 +191,71 @@ mod tests {
         fn trait_check<T: Send + Sync>() {}
         trait_check::<Span<()>>();
     }
+
+    #[test]
+    fn routing_inherited_by_children() {
+        let routing = Arc::new(TestRouting {
+            zone_id: 7,
+            account_id: 99,
+        });
+        let expected_group_key = routing.group_key();
+        let expected_encoded = routing.encode();
+
+        let (tracer, mut span_rx) = Tracer::new(AllSampler);
+        {
+            // Routing is set once, at construction, on the root only.
+            let parent = tracer
+                .span("parent")
+                .routing(routing.clone())
+                .start_with_state(());
+            let child = parent.child("child", |s| s.start_with_state(()));
+            let grandchild = child.child("grandchild", |s| s.start_with_state(()));
+
+            drop(grandchild);
+            drop(child);
+            drop(parent);
+        }
+
+        // Every span in the tree carries the same routing, without it being
+        // passed to each child.
+        for expected in ["grandchild", "child", "parent"] {
+            let span = span_rx.try_recv().unwrap();
+            assert_eq!(span.operation_name(), expected);
+
+            let r = span
+                .routing()
+                .unwrap_or_else(|| panic!("{expected} should inherit routing"));
+            assert_eq!(r.group_key(), expected_group_key);
+            assert_eq!(r.encode(), expected_encoded);
+        }
+    }
+
+    #[test]
+    fn routing_last_set_wins() {
+        let second = Arc::new(TestRouting {
+            zone_id: 2,
+            account_id: 2,
+        });
+        let expected_group_key = second.group_key();
+        let expected_encoded = second.encode();
+
+        let (tracer, mut span_rx) = Tracer::new(AllSampler);
+        {
+            // If the builder is called more than once, the last value wins.
+            let span = tracer
+                .span("span")
+                .routing(Arc::new(TestRouting {
+                    zone_id: 1,
+                    account_id: 1,
+                }))
+                .routing(second.clone())
+                .start_with_state(());
+            drop(span);
+        }
+
+        let span = span_rx.try_recv().unwrap();
+        let r = span.routing().expect("routing should be set");
+        assert_eq!(r.group_key(), expected_group_key);
+        assert_eq!(r.encode(), expected_encoded);
+    }
 }

src/span.rs

@@ -106,6 +106,17 @@ impl<T, F: Fn(&mut Span<T>) + Send + Sync + 'static> From<F> for FinishSpanCallb
     }
 }
 
+/// Caller-defined routing attached to a span, inherited by children, and read by
+/// the exporter to route the resulting [`FinishedSpan`].
+pub trait RoutingMetadata: fmt::Debug + Send + Sync {
+    /// Batching key: spans sharing it are exported together. Must differ
+    /// whenever [`encode`](RoutingMetadata::encode) would.
+    fn group_key(&self) -> String;
+
+    /// The value the exporter transmits for this routing.
+    fn encode(&self) -> String;
+}
+
 /// Span.
 ///
 /// When this span is dropped, it will be converted to `FinishedSpan` and
@@ -281,8 +292,9 @@ impl<T> Span<T> {
 
     /// Starts a `ChildOf` span if this span is sampled.
     ///
-    /// The child will inherit this span's finish callback, if it has one. To avoid
-    /// this kind of inheritance, you can use `span.handle().child(...)` instead.
+    /// The child will inherit this span's finish callback and routing metadata,
+    /// if set. To avoid this kind of inheritance, you can use
+    /// `span.handle().child(...)` instead.
     pub fn child<N, F>(&self, operation_name: N, f: F) -> Span<T>
     where
         N: Into<Cow<'static, str>>,
@@ -293,6 +305,9 @@ impl<T> Span<T> {
             if let Some(finish_cb) = self.0.as_ref().and_then(|s| s.finish_cb.clone()) {
                 opts = opts.finish_callback(finish_cb);
             }
+            if let Some(routing) = self.0.as_ref().and_then(|s| s.routing.clone()) {
+                opts = opts.routing(routing);
+            }
             f(opts)
         })
     }
@@ -319,6 +334,7 @@ impl<T> Span<T> {
             context,
             finish_cb: opts.finish_cb,
             span_tx: opts.span_tx.clone(),
+            routing: opts.routing,
         };
         Span(Some(inner))
     }
@@ -338,6 +354,7 @@ impl<T> Drop for Span<T> {
                 tags: inner.tags,
                 logs: inner.logs,
                 context: inner.context,
+                routing: inner.routing,
             };
             inner.span_tx.0.consume_span(finished);
         }
@@ -360,6 +377,7 @@ struct SpanInner<T> {
     context: SpanContext<T>,
     finish_cb: Option<FinishSpanCallback<T>>,
     span_tx: SharedSpanConsumer<T>,
+    routing: Option<Arc<dyn RoutingMetadata>>,
 }
 
 /// Finished span.
@@ -372,6 +390,7 @@ pub struct FinishedSpan<T> {
     tags: Vec<Tag>,
     logs: Vec<Log>,
     context: SpanContext<T>,
+    routing: Option<Arc<dyn RoutingMetadata>>,
 }
 impl<T> FinishedSpan<T> {
     /// Returns the operation name of this span.
@@ -408,6 +427,11 @@ impl<T> FinishedSpan<T> {
     pub fn context(&self) -> &SpanContext<T> {
         &self.context
     }
+
+    /// Returns the routing metadata of this span, if any was set.
+    pub fn routing(&self) -> Option<&dyn RoutingMetadata> {
+        self.routing.as_deref()
+    }
 }
 
 /// Span context.
@@ -601,6 +625,7 @@ pub struct StartSpanOptions<'a, S: 'a, T: 'a> {
     references: Vec<SpanReference<T>>,
     baggage_items: Vec<BaggageItem>,
     finish_cb: Option<FinishSpanCallback<T>>,
+    routing: Option<Arc<dyn RoutingMetadata>>,
     span_tx: &'a SharedSpanConsumer<T>,
     sampler: &'a S,
 }
@@ -629,6 +654,17 @@ where
         self
     }
 
+    /// Attaches routing metadata to the span being built.
+    ///
+    /// The metadata is carried into the resulting [`FinishedSpan`] for the
+    /// exporter to read, and is inherited by child spans (see [`Span::child`]).
+    /// This is the only way to set routing — there is no live setter, so routing
+    /// cannot be changed or cleared after the span has started.
+    pub fn routing(mut self, routing: Arc<dyn RoutingMetadata>) -> Self {
+        self.routing = Some(routing);
+        self
+    }
+
     /// Adds the `ChildOf` reference to this span.
     pub fn child_of<C>(mut self, context: &C) -> Self
     where
@@ -696,6 +732,7 @@ where
             references: Vec::new(),
             baggage_items: Vec::new(),
             finish_cb: None,
+            routing: None,
             span_tx,
             sampler,
         }