Skip to main content

awssdk_instrumentation/lambda/layer/
tracing.rs

1//! Tracing-backend [`Instrumentor`] implementation for Lambda invocation spans.
2
3use opentelemetry::trace::{SpanContext, TraceContextExt, TraceState};
4use opentelemetry_semantic_conventions::attribute as semco;
5
6use tokio::task::futures::TaskLocalFuture;
7use tracing::instrument::Instrumented;
8use tracing::{Instrument, Span};
9use tracing_opentelemetry::OpenTelemetrySpanExt;
10
11use super::{InstrumentedFuture, Instrumentor, utils::XRayTraceHeader};
12
13/// [`Instrumentor`] implementation for the `tracing-backend` feature.
14///
15/// `TracingInstrumentor` creates a `tracing::info_span!` named
16/// `"Lambda runtime invoke"` for each invocation and stores it in a
17/// Tokio task-local so that [`Instrumentor::with_invocation_span`] can
18/// access it from child tasks.
19///
20/// The X-Ray trace context (when present) is set as the OTel parent of the
21/// span via `tracing-opentelemetry`'s [`OpenTelemetrySpanExt::set_parent`].
22///
23/// This type is re-exported as [`DefaultInstrumentor`] when `tracing-backend`
24/// is the active backend.
25///
26/// [`DefaultInstrumentor`]: crate::lambda::layer::DefaultInstrumentor
27/// [`OpenTelemetrySpanExt::set_parent`]: tracing_opentelemetry::OpenTelemetrySpanExt::set_parent
28#[derive(Debug, Clone)]
29pub struct TracingInstrumentor;
30
31/// [`InstrumentedFuture`] impl for `tracing`'s [`Instrumented`] wrapper.
32impl<Fut: Future> InstrumentedFuture for Instrumented<Fut> {
33    type Fut = Self;
34}
35
36tokio::task_local! {
37    /// Task-local holding the current Lambda invocation's tracing span.
38    static INVOCATION_SPAN: Span;
39}
40
41/// Implements [`Instrumentor`] for the `tracing-backend` feature.
42impl Instrumentor for TracingInstrumentor {
43    type IFut<F: Future> = Instrumented<TaskLocalFuture<Span, F>>;
44    type InvocationSpan = Span;
45
46    fn instrument<F: Future>(inner: F, context: super::InvocationContext) -> Self::IFut<F> {
47        let span = tracing::info_span!(
48            "Lambda runtime invoke",
49            otel.kind = "server",
50            { semco::FAAS_TRIGGER } = context.trigger.to_string(),
51            { semco::CLOUD_RESOURCE_ID } = context.function_arn,
52            { semco::FAAS_INVOCATION_ID } = context.request_id,
53            { semco::CLOUD_ACCOUNT_ID } = context.account_id,
54            { semco::FAAS_COLDSTART } = context.is_coldstart,
55            xray_trace_id = tracing::field::Empty,
56        );
57
58        if let Some(XRayTraceHeader {
59            trace_id,
60            parent_id,
61            sampled,
62        }) = context.xray_trace_header
63        {
64            let otel_context = opentelemetry::Context::new().with_remote_span_context(
65                SpanContext::new(trace_id, parent_id, sampled, true, TraceState::NONE),
66            );
67            span.record("xray_trace_id", trace_id.to_string());
68            span.set_parent(otel_context).expect("not yet activated");
69        }
70
71        // Scope the task-local so with_invocation_span can find it
72        let inner = INVOCATION_SPAN.scope(span.clone(), inner);
73
74        {
75            let _guard = span.enter();
76            inner
77        }
78        .instrument(span)
79    }
80
81    fn with_invocation_span(f: impl FnOnce(&mut Self::InvocationSpan)) {
82        INVOCATION_SPAN.with(|span| {
83            let mut span = span.clone();
84            f(&mut span);
85        });
86    }
87
88    fn spawn<F>(future: F) -> tokio::task::JoinHandle<F::Output>
89    where
90        F: Future + Send + 'static,
91        F::Output: Send + 'static,
92    {
93        let span = INVOCATION_SPAN.with(|span| span.clone());
94        tokio::spawn(INVOCATION_SPAN.scope(span, future))
95    }
96}