xray_lite/
context.rs

1//! Tracing context.
2
3use crate::client::Client;
4use crate::error::Result;
5use crate::header::Header;
6use crate::lambda;
7use crate::namespace::Namespace;
8use crate::session::SubsegmentSession;
9
10/// Context.
11pub trait Context {
12    /// Client type.
13    type Client: Client;
14
15    /// Enters in a new subsegment.
16    ///
17    /// [`SubsegmentSession`] records the end of the subsegment when it is
18    /// dropped.
19    fn enter_subsegment<N>(&self, namespace: N) -> SubsegmentSession<Self::Client, N>
20    where
21        N: Namespace + Send + Sync;
22}
23
24/// Context as a subsegment of an existing segment.
25#[derive(Clone, Debug)]
26pub struct SubsegmentContext<C> {
27    client: C,
28    header: Header,
29    name_prefix: String,
30}
31
32impl<C> SubsegmentContext<C> {
33    /// Creates a new context from the Lambda environment variable.
34    ///
35    /// The following environment variable must be set:
36    /// - `_X_AMZN_TRACE_ID`: AWS X-Ray trace ID
37    ///
38    /// Please refer to the [AWS documentation](https://docs.aws.amazon.com/lambda/latest/dg/configuration-envvars.html#configuration-envvars-runtime)
39    /// for more details.
40    pub fn from_lambda_env(client: C) -> Result<Self> {
41        let header = lambda::header()?;
42        Ok(Self {
43            client,
44            header,
45            name_prefix: "".to_string(),
46        })
47    }
48
49    /// Updates the context with a given name prefix.
50    ///
51    /// The name prefix is prepended to the name of every custom subsegment.
52    /// Only subsegments associated with
53    /// [`CustomNamespace`][crate::namespace::CustomNamespace] are affected.
54    pub fn with_name_prefix(self, prefix: impl Into<String>) -> Self {
55        Self {
56            client: self.client,
57            header: self.header,
58            name_prefix: prefix.into(),
59        }
60    }
61}
62
63impl<C> Context for SubsegmentContext<C>
64where
65    C: Client,
66{
67    type Client = C;
68
69    fn enter_subsegment<N>(&self, namespace: N) -> SubsegmentSession<Self::Client, N>
70    where
71        N: Namespace + Send + Sync,
72    {
73        SubsegmentSession::new(
74            self.client.clone(),
75            &self.header,
76            namespace,
77            &self.name_prefix,
78        )
79    }
80}
81
82/// Infallible context.
83///
84/// This context is useful if you want to fall back to "no-op" when creation of
85/// the underlying context fails.
86pub enum InfallibleContext<T> {
87    /// Operational context.
88    Op(T),
89    /// Non-operational context.
90    Noop,
91}
92
93impl<T> InfallibleContext<T>
94where
95    T: Context,
96{
97    /// Constructs from a result of the underlying context creation.
98    pub fn new<E>(result: std::result::Result<T, E>) -> Self {
99        match result {
100            Ok(context) => Self::Op(context),
101            Err(_) => Self::Noop,
102        }
103    }
104}
105
106impl<T> Context for InfallibleContext<T>
107where
108    T: Context,
109{
110    type Client = T::Client;
111
112    fn enter_subsegment<N>(&self, namespace: N) -> SubsegmentSession<Self::Client, N>
113    where
114        N: Namespace + Send + Sync,
115    {
116        match self {
117            Self::Op(context) => context.enter_subsegment(namespace),
118            Self::Noop => SubsegmentSession::failed(),
119        }
120    }
121}
122
123impl<T> Clone for InfallibleContext<T>
124where
125    T: Clone,
126{
127    fn clone(&self) -> Self {
128        match self {
129            Self::Op(context) => Self::Op(context.clone()),
130            Self::Noop => Self::Noop,
131        }
132    }
133}
134
135impl<T> std::fmt::Debug for InfallibleContext<T>
136where
137    T: std::fmt::Debug,
138{
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        match self {
141            Self::Op(context) => write!(f, "InfallibleContext::Op({:?})", context),
142            Self::Noop => write!(f, "InfallibleContext::Noop"),
143        }
144    }
145}
146
147/// Conversion into an infallible context.
148///
149/// You can convert `Result<Context, _>` into an infallible context by using
150/// this trait.`
151///
152/// ```
153/// use xray_lite::{
154///     Context as _,
155///     DaemonClient,
156///     IntoInfallibleContext as _,
157///     CustomNamespace,
158///     SubsegmentContext,
159/// };
160///
161/// fn main() {
162///     # std::env::set_var("AWS_XRAY_DAEMON_ADDRESS", "127.0.0.1:2000");
163///     let client = DaemonClient::from_lambda_env().unwrap();
164///     let context = SubsegmentContext::from_lambda_env(client).into_infallible();
165///     let _session = context.enter_subsegment(CustomNamespace::new("readme.example"));
166/// }
167/// ```
168pub trait IntoInfallibleContext {
169    /// Underlying context type.
170    type Context: Context;
171
172    /// Converts into an infallible context.
173    fn into_infallible(self) -> InfallibleContext<Self::Context>;
174}
175
176impl<T, E> IntoInfallibleContext for std::result::Result<T, E>
177where
178    T: Context,
179{
180    type Context = T;
181
182    fn into_infallible(self) -> InfallibleContext<Self::Context> {
183        InfallibleContext::new(self)
184    }
185}