xray_lite/
session.rs

1//! Subsegment session management.
2
3use crate::client::Client;
4use crate::header::Header;
5use crate::namespace::Namespace;
6use crate::segment::Subsegment;
7
8/// Subsegment session.
9#[derive(Debug)]
10pub enum SubsegmentSession<C, N>
11where
12    C: Client,
13    N: Namespace + Send + Sync,
14{
15    /// Entered subsegment.
16    Entered {
17        /// X-Ray client.
18        client: C,
19        /// X-Amzn-Trace-Id header.
20        header: Header,
21        /// Subsegment.
22        subsegment: Subsegment,
23        /// Namespace.
24        namespace: N,
25    },
26    /// Failed subsegment.
27    Failed,
28}
29
30impl<C, N> SubsegmentSession<C, N>
31where
32    C: Client,
33    N: Namespace + Send + Sync,
34{
35    pub(crate) fn new(client: C, header: &Header, namespace: N, name_prefix: &str) -> Self {
36        let mut subsegment = Subsegment::begin(
37            header.trace_id.clone(),
38            header.parent_id.clone(),
39            namespace.name(name_prefix),
40        );
41        namespace.update_subsegment(&mut subsegment);
42        match client.send(&subsegment) {
43            Ok(_) => Self::Entered {
44                client,
45                header: header.with_parent_id(subsegment.id.clone()),
46                subsegment,
47                namespace,
48            },
49            Err(_) => Self::Failed,
50        }
51    }
52
53    pub(crate) fn failed() -> Self {
54        Self::Failed
55    }
56
57    /// Returns the `x-amzn-trace-id` header value.
58    pub fn x_amzn_trace_id(&self) -> Option<String> {
59        match self {
60            Self::Entered { header, .. } => Some(header.to_string()),
61            Self::Failed => None,
62        }
63    }
64
65    /// Returns the namespace as a mutable reference.
66    pub fn namespace_mut(&mut self) -> Option<&mut N> {
67        match self {
68            Self::Entered { namespace, .. } => Some(namespace),
69            Self::Failed => None,
70        }
71    }
72}
73
74impl<C, N> Drop for SubsegmentSession<C, N>
75where
76    C: Client,
77    N: Namespace + Send + Sync,
78{
79    fn drop(&mut self) {
80        match self {
81            Self::Entered {
82                client,
83                subsegment,
84                namespace,
85                ..
86            } => {
87                subsegment.end();
88                namespace.update_subsegment(subsegment);
89                let _ = client
90                    .send(subsegment)
91                    .map_err(|e| eprintln!("failed to end subsegment: {e}"));
92            }
93            Self::Failed => (),
94        }
95    }
96}