1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163

//! X-Ray daemon client.

use std::env;
use std::net::{SocketAddr, UdpSocket};
use std::sync::Arc;

use serde::Serialize;

use crate::error::{Error, Result};

/// X-Ray client interface.
pub trait Client: Clone + std::fmt::Debug + Send + Sync {
    /// Sends a segment to the xray daemon this client is connected to.
    fn send<S>(&self, data: &S) -> Result<()>
    where
        S: Serialize;
}

/// X-Ray daemon client.
#[derive(Clone, Debug)]
pub struct DaemonClient {
    socket: Arc<UdpSocket>,
}

impl DaemonClient {
    const HEADER: &'static [u8] = br#"{"format": "json", "version": 1}"#;
    const DELIMITER: &'static [u8] = &[b'\n'];

    /// Return a new X-Ray client connected
    /// to the provided `addr`
    pub fn new(addr: SocketAddr) -> Result<Self> {
        let socket = Arc::new(UdpSocket::bind(&[([0, 0, 0, 0], 0).into()][..])?);
        socket.set_nonblocking(true)?;
        socket.connect(addr)?;
        Ok(DaemonClient { socket })
    }

    /// Creates a new X-Ray client from the Lambda environment variable.
    ///
    /// The following environment variable must be set:
    /// - `AWS_XRAY_DAEMON_ADDRESS`: X-Ray daemon address
    ///
    /// Please refer to the [AWS documentation](https://docs.aws.amazon.com/lambda/latest/dg/configuration-envvars.html#configuration-envvars-runtime)
    /// for more details.
    pub fn from_lambda_env() -> Result<Self> {
        let addr: SocketAddr = env::var("AWS_XRAY_DAEMON_ADDRESS")
            .map_err(|_| Error::MissingEnvVar("AWS_XRAY_DAEMON_ADDRESS"))?
            .parse::<SocketAddr>()
            .map_err(|e| Error::BadConfig(format!("invalid X-Ray daemon address: {e}")))?;
        DaemonClient::new(addr)
    }

    #[inline]
    fn packet<S>(data: S) -> Result<Vec<u8>>
    where
        S: Serialize,
    {
        let bytes = serde_json::to_vec(&data)?;
        Ok([Self::HEADER, Self::DELIMITER, &bytes].concat())
    }
}

impl Client for DaemonClient {
    fn send<S>(&self, data: &S) -> Result<()>
    where
        S: Serialize,
    {
        self.socket.send(&Self::packet(data)?)?;
        Ok(())
    }
}

/// Infallible client.
#[derive(Clone, Debug)]
pub enum InfallibleClient<C> {
    /// Operational client.
    Op(C),
    /// Non-operational client.
    Noop,
}

impl<C> InfallibleClient<C> {
    /// Creates a new infallible client from a result of client creation.
    pub fn new<E>(result: std::result::Result<C, E>) -> Self {
        match result {
            Ok(client) => Self::Op(client),
            Err(_) => Self::Noop,
        }
    }
}

impl<C> Client for InfallibleClient<C>
where
    C: Client,
{
    fn send<S>(&self, data: &S) -> Result<()>
    where
        S: Serialize,
    {
        match self {
            Self::Op(client) => client.send(data),
            Self::Noop => Ok(()),
        }
    }
}

/// Conversion into an [`InfallibleClient`].
///
/// This is useful if you want to fall back to a "no-op" client if the creation
/// of a client fails.
///
/// ```
/// use xray_lite::{
///     Context as _,
///     CustomNamespace,
///     DaemonClient,
///     IntoInfallibleClient as _,
///     SubsegmentContext,
/// };
///
/// let client = DaemonClient::from_lambda_env().into_infallible();
/// # std::env::set_var("_X_AMZN_TRACE_ID", "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1");
/// let context = SubsegmentContext::from_lambda_env(client).unwrap();
/// let _session = context.enter_subsegment(CustomNamespace::new("readme.example"));
/// ```
pub trait IntoInfallibleClient {
    /// Client type.
    type Client: Client;

    /// Converts a value into an [`InfallibleClient`].
    fn into_infallible(self) -> InfallibleClient<Self::Client>;
}

impl<C> IntoInfallibleClient for std::result::Result<C, Error>
where
    C: Client,
{
    type Client = C;

    fn into_infallible(self) -> InfallibleClient<C> {
        InfallibleClient::new(self)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn client_prefixes_packets_with_header() {
        assert_eq!(
            DaemonClient::packet(serde_json::json!({
                "foo": "bar"
            }))
            .unwrap(),
            [
                br#"{"format": "json", "version": 1}"# as &[u8],
                &[b'\n'],
                br#"{"foo":"bar"}"#,
            ]
            .concat()
        )
    }
}