apalis 1.0.0-rc.8

Simple, extensible and multithreaded background task processing for Rust
Documentation 
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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

use apalis_core::task::metadata::TracingContext;
use opentelemetry::{global, trace::TraceContextExt as _};
use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt as _;

/// OpenTelemetry-focused wrapper around [`TracingContext`].
///
/// This type provides a clearer API surface for users that explicitly work with
/// OpenTelemetry context propagation, while preserving `apalis-core` compatibility.
#[derive(Debug, Default, Clone)]
pub struct OtelTraceContext(pub(crate) TracingContext);

impl OtelTraceContext {
    /// Capture context from the currently active tracing span.
    #[must_use]
    pub fn current() -> Self {
        let mut carrier = HeaderCarrier::default();
        let current = Span::current().context();
        global::get_text_map_propagator(|propagator| {
            propagator.inject_context(&current, &mut carrier);
        });

        Self::from(carrier.into_tracing_context())
    }

    /// Restore this context as the OpenTelemetry parent of the provided span.
    pub(crate) fn restore(&self, span: &tracing::Span) {
        let Some(carrier) = HeaderCarrier::from_tracing_context(&self.0) else {
            return;
        };
        let parent_context =
            global::get_text_map_propagator(|propagator| propagator.extract(&carrier));

        if parent_context.span().span_context().is_valid() {
            let _ = span.set_parent(parent_context);
        }
    }
}

impl From<TracingContext> for OtelTraceContext {
    fn from(value: TracingContext) -> Self {
        Self(value)
    }
}

impl From<OtelTraceContext> for TracingContext {
    fn from(value: OtelTraceContext) -> Self {
        value.0
    }
}

#[derive(Debug, Default, Clone)]
struct HeaderCarrier {
    traceparent: String,
    tracestate: Option<String>,
}

impl HeaderCarrier {
    fn from_tracing_context(context: &TracingContext) -> Option<Self> {
        let trace_id = context.trace_id().as_deref()?;
        let span_id = context.span_id().as_deref()?;
        let trace_flags = context.trace_flags().unwrap_or_default();

        Some(Self {
            traceparent: format!("00-{trace_id}-{span_id}-{trace_flags:02x}"),
            tracestate: context.trace_state().clone(),
        })
    }

    fn into_tracing_context(self) -> TracingContext {
        let mut context = TracingContext::new();
        if let Some((trace_id, span_id, trace_flags)) = parse_traceparent(&self.traceparent) {
            context = context
                .with_trace_id(trace_id)
                .with_span_id(span_id)
                .with_trace_flags(trace_flags);
        }
        if let Some(tracestate) = self.tracestate {
            context = context.with_trace_state(tracestate);
        }
        context
    }
}

impl opentelemetry::propagation::Extractor for HeaderCarrier {
    fn get(&self, key: &str) -> Option<&str> {
        if key.eq_ignore_ascii_case("traceparent") {
            return Some(self.traceparent.as_str());
        }
        if key.eq_ignore_ascii_case("tracestate") {
            return self.tracestate.as_deref();
        }
        None
    }

    fn keys(&self) -> Vec<&str> {
        let mut keys = vec!["traceparent"];
        if self.tracestate.is_some() {
            keys.push("tracestate");
        }
        keys
    }
}

impl opentelemetry::propagation::Injector for HeaderCarrier {
    fn set(&mut self, key: &str, value: String) {
        if key.eq_ignore_ascii_case("traceparent") {
            self.traceparent = value;
        } else if key.eq_ignore_ascii_case("tracestate") {
            self.tracestate = Some(value);
        }
    }
}

fn parse_traceparent(value: &str) -> Option<(String, String, u8)> {
    let mut parts = value.split('-');
    let version = parts.next()?;
    let trace_id = parts.next()?;
    let span_id = parts.next()?;
    let trace_flags = parts.next()?;
    if parts.next().is_some() {
        return None;
    }

    if !is_hex_with_len(version, 2)
        || !is_hex_with_len(trace_id, 32)
        || !is_hex_with_len(span_id, 16)
        || !is_hex_with_len(trace_flags, 2)
    {
        return None;
    }

    let trace_flags = u8::from_str_radix(trace_flags, 16).ok()?;
    Some((trace_id.to_owned(), span_id.to_owned(), trace_flags))
}

fn is_hex_with_len(value: &str, expected_len: usize) -> bool {
    value.len() == expected_len && value.as_bytes().iter().all(u8::is_ascii_hexdigit)
}