lambda_runtime/layers/
trace.rs

1use tower::{Layer, Service};
2use tracing::{instrument::Instrumented, Instrument};
3
4use crate::{Context, LambdaInvocation};
5use lambda_runtime_api_client::BoxError;
6use std::task;
7
8/// Tower middleware to create a tracing span for invocations of the Lambda function.
9#[derive(Default)]
10pub struct TracingLayer {}
11
12impl TracingLayer {
13    /// Create a new tracing layer.
14    pub fn new() -> Self {
15        Self::default()
16    }
17}
18
19impl<S> Layer<S> for TracingLayer {
20    type Service = TracingService<S>;
21
22    fn layer(&self, inner: S) -> Self::Service {
23        TracingService { inner }
24    }
25}
26
27/// Tower service returned by [TracingLayer].
28#[derive(Clone)]
29pub struct TracingService<S> {
30    inner: S,
31}
32
33impl<S> Service<LambdaInvocation> for TracingService<S>
34where
35    S: Service<LambdaInvocation, Response = (), Error = BoxError>,
36{
37    type Response = ();
38    type Error = BoxError;
39    type Future = Instrumented<S::Future>;
40
41    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> task::Poll<Result<(), Self::Error>> {
42        self.inner.poll_ready(cx)
43    }
44
45    fn call(&mut self, req: LambdaInvocation) -> Self::Future {
46        let span = request_span(&req.context);
47        let future = {
48            // Enter the span before calling the inner service
49            // to ensure that it's assigned as parent of the inner spans.
50            let _guard = span.enter();
51            self.inner.call(req)
52        };
53        future.instrument(span)
54    }
55}
56
57/* ------------------------------------------- UTILS ------------------------------------------- */
58
59/// Creates a tracing span for a Lambda request with context information.
60///
61/// This function creates a span that includes the request ID and optionally
62/// the X-Ray trace ID and tenant ID if they are available in the context.
63pub fn request_span(ctx: &Context) -> tracing::Span {
64    match (&ctx.xray_trace_id, &ctx.tenant_id) {
65        (Some(trace_id), Some(tenant_id)) => {
66            tracing::info_span!(
67                "Lambda runtime invoke",
68                requestId = &ctx.request_id,
69                xrayTraceId = trace_id,
70                tenantId = tenant_id
71            )
72        }
73        (Some(trace_id), None) => {
74            tracing::info_span!(
75                "Lambda runtime invoke",
76                requestId = &ctx.request_id,
77                xrayTraceId = trace_id
78            )
79        }
80        (None, Some(tenant_id)) => {
81            tracing::info_span!(
82                "Lambda runtime invoke",
83                requestId = &ctx.request_id,
84                tenantId = tenant_id
85            )
86        }
87        (None, None) => {
88            tracing::info_span!("Lambda runtime invoke", requestId = &ctx.request_id)
89        }
90    }
91}
lambda_runtime/layers/trace.rs

lambda_runtime/layers/
trace.rs
