1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;

/// Wraps an actix message with the current Span's context.
/// This lets us trace execution across several actix Actors.
#[derive(actix::Message, Debug)]
#[rtype(result = "<T as actix::Message>::Result")]
pub struct WithSpanContext<T: actix::Message> {
    pub msg: T,
    pub context: opentelemetry::Context,
}

impl<T: actix::Message> WithSpanContext<T> {
    pub fn new(msg: T) -> Self {
        Self { msg, context: Span::current().context() }
    }
}

/// Allows easily attaching the current span's context to a Message.
pub trait WithSpanContextExt: actix::Message {
    fn with_span_context(self) -> WithSpanContext<Self>
    where
        Self: Sized,
    {
        WithSpanContext::<Self>::new(self)
    }
}
impl<T: actix::Message> WithSpanContextExt for T {}