logo
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

//! # UDP Jaeger Agent Client
use crate::exporter::runtime::JaegerTraceRuntime;
use crate::exporter::thrift::{
    agent::{self, TAgentSyncClient},
    jaeger,
};
use crate::exporter::transport::{TBufferChannel, TNoopChannel};
use std::fmt;
use std::net::{ToSocketAddrs, UdpSocket};
use thrift::{
    protocol::{TCompactInputProtocol, TCompactOutputProtocol},
    transport::{ReadHalf, TIoChannel, WriteHalf},
};

/// The max size of UDP packet we want to send, synced with jaeger-agent
const UDP_PACKET_MAX_LENGTH: usize = 65_000;

struct BufferClient {
    buffer: ReadHalf<TBufferChannel>,
    client: agent::AgentSyncClient<
        TCompactInputProtocol<TNoopChannel>,
        TCompactOutputProtocol<WriteHalf<TBufferChannel>>,
    >,
}

impl fmt::Debug for BufferClient {
    /// Debug info
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("BufferClient")
            .field("buffer", &self.buffer)
            .field("client", &"AgentSyncClient")
            .finish()
    }
}

/// `AgentSyncClientUDP` implements a version of the `TAgentSyncClient`
/// interface over UDP.
#[derive(Debug)]
pub(crate) struct AgentSyncClientUdp {
    conn: UdpSocket,
    buffer_client: BufferClient,
    max_packet_size: usize,
}

impl AgentSyncClientUdp {
    /// Create a new UDP agent client
    pub(crate) fn new<T: ToSocketAddrs>(
        host_port: T,
        max_packet_size: Option<usize>,
    ) -> thrift::Result<Self> {
        let max_packet_size = max_packet_size.unwrap_or(UDP_PACKET_MAX_LENGTH);
        let (buffer, write) = TBufferChannel::with_capacity(max_packet_size).split()?;
        let client = agent::AgentSyncClient::new(
            TCompactInputProtocol::new(TNoopChannel),
            TCompactOutputProtocol::new(write),
        );

        let conn = UdpSocket::bind("0.0.0.0:0")?;
        conn.connect(host_port)?;

        Ok(AgentSyncClientUdp {
            conn,
            buffer_client: BufferClient { buffer, client },
            max_packet_size,
        })
    }

    /// Emit standard Jaeger batch
    pub(crate) fn emit_batch(&mut self, batch: jaeger::Batch) -> thrift::Result<()> {
        // Write payload to buffer
        self.buffer_client.client.emit_batch(batch)?;
        let payload = self.buffer_client.buffer.take_bytes();

        if payload.len() > self.max_packet_size {
            return Err(thrift::ProtocolError::new(
                thrift::ProtocolErrorKind::SizeLimit,
                format!(
                    "jaeger exporter payload size of {} bytes over max UDP packet size of {} bytes. Try setting a smaller batch size.",
                    payload.len(),
                    self.max_packet_size,
                ),
            )
            .into());
        }

        self.conn.send(&payload)?;

        Ok(())
    }
}

/// `AgentAsyncClientUDP` implements an async version of the `TAgentSyncClient`
/// interface over UDP.
#[derive(Debug)]
pub(crate) struct AgentAsyncClientUdp<R: JaegerTraceRuntime> {
    runtime: R,
    conn: <R as JaegerTraceRuntime>::Socket,
    buffer_client: BufferClient,
    max_packet_size: usize,
}

impl<R: JaegerTraceRuntime> AgentAsyncClientUdp<R> {
    /// Create a new UDP agent client
    pub(crate) fn new<T: ToSocketAddrs>(
        host_port: T,
        max_packet_size: Option<usize>,
        runtime: R,
    ) -> thrift::Result<Self> {
        let max_packet_size = max_packet_size.unwrap_or(UDP_PACKET_MAX_LENGTH);
        let (buffer, write) = TBufferChannel::with_capacity(max_packet_size).split()?;
        let client = agent::AgentSyncClient::new(
            TCompactInputProtocol::new(TNoopChannel),
            TCompactOutputProtocol::new(write),
        );

        let conn = runtime.create_socket(host_port)?;

        Ok(AgentAsyncClientUdp {
            runtime,
            conn,
            buffer_client: BufferClient { buffer, client },
            max_packet_size,
        })
    }

    /// Emit standard Jaeger batch
    pub(crate) async fn emit_batch(&mut self, batch: jaeger::Batch) -> thrift::Result<()> {
        // Write payload to buffer
        self.buffer_client.client.emit_batch(batch)?;
        let payload = self.buffer_client.buffer.take_bytes();

        if payload.len() > self.max_packet_size {
            return Err(thrift::ProtocolError::new(
                thrift::ProtocolErrorKind::SizeLimit,
                format!(
                    "jaeger exporter payload size of {} bytes over max UDP packet size of {} bytes. Try setting a smaller batch size.",
                    payload.len(),
                    self.max_packet_size,
                ),
            )
            .into());
        }

        // Write async to socket, reading from buffer
        self.runtime.write_to_socket(&self.conn, payload).await?;

        Ok(())
    }
}