xray_lite/
client.rs

1//! X-Ray daemon client.
2
3use std::env;
4use std::net::{SocketAddr, UdpSocket};
5use std::sync::Arc;
6
7use serde::Serialize;
8
9use crate::error::{Error, Result};
10
11/// X-Ray client interface.
12pub trait Client: Clone + std::fmt::Debug + Send + Sync {
13    /// Sends a segment to the xray daemon this client is connected to.
14    fn send<S>(&self, data: &S) -> Result<()>
15    where
16        S: Serialize;
17}
18
19/// X-Ray daemon client.
20#[derive(Clone, Debug)]
21pub struct DaemonClient {
22    socket: Arc<UdpSocket>,
23}
24
25impl DaemonClient {
26    const HEADER: &'static [u8] = br#"{"format": "json", "version": 1}"#;
27    const DELIMITER: &'static [u8] = &[b'\n'];
28
29    /// Return a new X-Ray client connected
30    /// to the provided `addr`
31    pub fn new(addr: SocketAddr) -> Result<Self> {
32        let socket = Arc::new(UdpSocket::bind(&[([0, 0, 0, 0], 0).into()][..])?);
33        socket.set_nonblocking(true)?;
34        socket.connect(addr)?;
35        Ok(DaemonClient { socket })
36    }
37
38    /// Creates a new X-Ray client from the Lambda environment variable.
39    ///
40    /// The following environment variable must be set:
41    /// - `AWS_XRAY_DAEMON_ADDRESS`: X-Ray daemon address
42    ///
43    /// Please refer to the [AWS documentation](https://docs.aws.amazon.com/lambda/latest/dg/configuration-envvars.html#configuration-envvars-runtime)
44    /// for more details.
45    pub fn from_lambda_env() -> Result<Self> {
46        let addr: SocketAddr = env::var("AWS_XRAY_DAEMON_ADDRESS")
47            .map_err(|_| Error::MissingEnvVar("AWS_XRAY_DAEMON_ADDRESS"))?
48            .parse::<SocketAddr>()
49            .map_err(|e| Error::BadConfig(format!("invalid X-Ray daemon address: {e}")))?;
50        DaemonClient::new(addr)
51    }
52
53    #[inline]
54    fn packet<S>(data: S) -> Result<Vec<u8>>
55    where
56        S: Serialize,
57    {
58        let bytes = serde_json::to_vec(&data)?;
59        Ok([Self::HEADER, Self::DELIMITER, &bytes].concat())
60    }
61}
62
63impl Client for DaemonClient {
64    fn send<S>(&self, data: &S) -> Result<()>
65    where
66        S: Serialize,
67    {
68        self.socket.send(&Self::packet(data)?)?;
69        Ok(())
70    }
71}
72
73/// Infallible client.
74#[derive(Clone, Debug)]
75pub enum InfallibleClient<C> {
76    /// Operational client.
77    Op(C),
78    /// Non-operational client.
79    Noop,
80}
81
82impl<C> InfallibleClient<C> {
83    /// Creates a new infallible client from a result of client creation.
84    pub fn new<E>(result: std::result::Result<C, E>) -> Self {
85        match result {
86            Ok(client) => Self::Op(client),
87            Err(_) => Self::Noop,
88        }
89    }
90}
91
92impl<C> Client for InfallibleClient<C>
93where
94    C: Client,
95{
96    fn send<S>(&self, data: &S) -> Result<()>
97    where
98        S: Serialize,
99    {
100        match self {
101            Self::Op(client) => client.send(data),
102            Self::Noop => Ok(()),
103        }
104    }
105}
106
107/// Conversion into an [`InfallibleClient`].
108///
109/// This is useful if you want to fall back to a "no-op" client if the creation
110/// of a client fails.
111///
112/// ```
113/// use xray_lite::{
114///     Context as _,
115///     CustomNamespace,
116///     DaemonClient,
117///     IntoInfallibleClient as _,
118///     SubsegmentContext,
119/// };
120///
121/// let client = DaemonClient::from_lambda_env().into_infallible();
122/// # std::env::set_var("_X_AMZN_TRACE_ID", "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
123/// let context = SubsegmentContext::from_lambda_env(client).unwrap();
124/// let _session = context.enter_subsegment(CustomNamespace::new("readme.example"));
125/// ```
126pub trait IntoInfallibleClient {
127    /// Client type.
128    type Client: Client;
129
130    /// Converts a value into an [`InfallibleClient`].
131    fn into_infallible(self) -> InfallibleClient<Self::Client>;
132}
133
134impl<C> IntoInfallibleClient for std::result::Result<C, Error>
135where
136    C: Client,
137{
138    type Client = C;
139
140    fn into_infallible(self) -> InfallibleClient<C> {
141        InfallibleClient::new(self)
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    #[test]
149    fn client_prefixes_packets_with_header() {
150        assert_eq!(
151            DaemonClient::packet(serde_json::json!({
152                "foo": "bar"
153            }))
154            .unwrap(),
155            [
156                br#"{"format": "json", "version": 1}"# as &[u8],
157                &[b'\n'],
158                br#"{"foo":"bar"}"#,
159            ]
160            .concat()
161        )
162    }
163}
xray_lite/client.rs

xray_lite/
client.rs
