saluki_components/encoders/datadog/stats/
mod.rs

1//! Datadog APM Stats encoder.
2
3use std::time::Duration;
4
5use agent_data_plane_config::{defaults::DEFAULT_TRACE_ENV, domains, shared::SharedConfiguration};
6use async_trait::async_trait;
7use datadog_protos::traces::{
8    ClientGroupedStats as ProtoClientGroupedStats, ClientStatsBucket as ProtoClientStatsBucket,
9    ClientStatsPayload as ProtoClientStatsPayload, StatsPayload as ProtoStatsPayload, Trilean,
10};
11use http::{uri::PathAndQuery, HeaderValue, Method, Uri};
12use saluki_common::task::HandleExt as _;
13use saluki_core::accounting::{MemoryBounds, MemoryBoundsBuilder};
14use saluki_core::{
15    components::{encoders::*, ComponentContext},
16    data_model::{
17        event::{
18            trace_stats::{ClientGroupedStats, ClientStatsBucket, ClientStatsPayload, TraceStats},
19            EventType,
20        },
21        payload::{HttpPayload, Payload, PayloadMetadata, PayloadType},
22    },
23    observability::ComponentMetricsExt as _,
24    topology::{EventsBuffer, PayloadsBuffer},
25};
26use saluki_env::{host::providers::BoxedHostProvider, EnvironmentProvider, HostProvider};
27use saluki_error::{generic_error, ErrorContext as _, GenericError};
28use saluki_io::compression::CompressionScheme;
29use saluki_metrics::MetricsBuilder;
30use stringtheory::MetaString;
31use tokio::{
32    pin, select,
33    sync::mpsc::{self, Receiver, Sender},
34    time::sleep,
35};
36use tracing::{debug, error};
37
38use crate::common::datadog::{
39    io::RB_BUFFER_CHUNK_SIZE,
40    request_builder::{EndpointEncoder, RequestBuilder},
41    telemetry::ComponentTelemetry,
42    DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT, DEFAULT_INTAKE_UNCOMPRESSED_SIZE_LIMIT,
43};
44
45const MAX_STATS_PER_PAYLOAD: usize = 4000;
46static CONTENT_TYPE_MSGPACK: HeaderValue = HeaderValue::from_static("application/msgpack");
47
48/// Configuration for the Datadog APM Stats encoder.
49pub struct DatadogApmStatsEncoderConfiguration {
50    flush_timeout: Duration,
51    agent_hostname: Option<String>,
52    agent_version: String,
53    env: String,
54}
55
56impl DatadogApmStatsEncoderConfiguration {
57    /// Creates a new `DatadogApmStatsEncoderConfiguration` from the resolved configuration.
58    pub fn from_configuration(traces: &domains::traces::Domain, shared: &SharedConfiguration) -> Self {
59        let app_details = saluki_metadata::get_app_details();
60        Self {
61            flush_timeout: shared.metrics_encoding.flush_timeout,
62            agent_hostname: None,
63            agent_version: format!("agent-data-plane/{}", app_details.version().raw()),
64            env: if traces.env.is_empty() {
65                DEFAULT_TRACE_ENV.to_owned()
66            } else {
67                traces.env.clone()
68            },
69        }
70    }
71
72    /// Sets the agent hostname using the environment provider.
73    pub async fn with_environment_provider<E>(mut self, environment_provider: E) -> Result<Self, GenericError>
74    where
75        E: EnvironmentProvider<Host = BoxedHostProvider>,
76    {
77        let host_provider = environment_provider.host();
78        let hostname = host_provider.get_hostname().await?;
79        self.agent_hostname = Some(hostname);
80        Ok(self)
81    }
82}
83
84#[async_trait]
85impl EncoderBuilder for DatadogApmStatsEncoderConfiguration {
86    fn input_event_type(&self) -> EventType {
87        EventType::TraceStats
88    }
89
90    fn output_payload_type(&self) -> PayloadType {
91        PayloadType::Http
92    }
93
94    async fn build(&self, context: ComponentContext) -> Result<Box<dyn Encoder + Send>, GenericError> {
95        let metrics_builder = MetricsBuilder::from_component_context(&context);
96        let telemetry = ComponentTelemetry::from_builder(&metrics_builder);
97        let compression_scheme = CompressionScheme::gzip_default();
98
99        let agent_hostname = MetaString::from(self.agent_hostname.clone().unwrap_or_default());
100        let agent_version = MetaString::from(self.agent_version.clone());
101        let agent_env = MetaString::from(self.env.clone());
102
103        let mut stats_rb = RequestBuilder::new(
104            StatsEndpointEncoder::new(agent_hostname, agent_version, agent_env),
105            compression_scheme,
106            RB_BUFFER_CHUNK_SIZE,
107        )
108        .await?;
109        stats_rb.with_max_inputs_per_payload(MAX_STATS_PER_PAYLOAD);
110
111        let flush_timeout = if self.flush_timeout.is_zero() {
112            // We always give ourselves a minimum flush timeout of 10ms to allow for some very minimal amount of
113            // batching, while still practically flushing things almost immediately.
114            Duration::from_millis(10)
115        } else {
116            self.flush_timeout
117        };
118
119        Ok(Box::new(DatadogStats {
120            stats_rb,
121            telemetry,
122            flush_timeout,
123        }))
124    }
125}
126
127impl MemoryBounds for DatadogApmStatsEncoderConfiguration {
128    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
129        builder
130            .minimum()
131            .with_single_value::<DatadogStats>("component struct")
132            .with_array::<EventsBuffer>("request builder events channel", 8)
133            .with_array::<PayloadsBuffer>("request builder payloads channel", 8);
134
135        builder
136            .firm()
137            .with_array::<TraceStats>("stats split re-encode buffer", MAX_STATS_PER_PAYLOAD);
138    }
139}
140
141pub struct DatadogStats {
142    stats_rb: RequestBuilder<StatsEndpointEncoder>,
143    telemetry: ComponentTelemetry,
144    flush_timeout: Duration,
145}
146
147#[async_trait]
148impl Encoder for DatadogStats {
149    async fn run(mut self: Box<Self>, mut context: EncoderContext) -> Result<(), GenericError> {
150        let Self {
151            stats_rb,
152            telemetry,
153            flush_timeout,
154        } = *self;
155
156        let mut health = context.take_health_handle();
157
158        let (events_tx, events_rx) = mpsc::channel(8);
159        let (payloads_tx, mut payloads_rx) = mpsc::channel(8);
160
161        let request_builder_fut = run_request_builder(stats_rb, telemetry, events_rx, payloads_tx, flush_timeout);
162        let request_builder_handle = context
163            .topology_context()
164            .global_thread_pool()
165            .spawn_traced_named("dd-stats-request-builder", request_builder_fut);
166
167        health.mark_ready();
168        debug!("Datadog APM Stats encoder started.");
169
170        loop {
171            select! {
172                biased;
173
174                _ = health.live() => continue,
175                maybe_payload = payloads_rx.recv() => match maybe_payload {
176                    Some(payload) => {
177                        if let Err(e) = context.dispatcher().dispatch(payload).await {
178                            error!("Failed to dispatch payload: {}", e);
179                        }
180                    }
181                    None => break,
182                },
183                maybe_event_buffer = context.events().next() => match maybe_event_buffer {
184                    Some(event_buffer) => events_tx.send(event_buffer).await
185                        .error_context("Failed to send event buffer to request builder task.")?,
186                    None => break,
187                },
188            }
189        }
190
191        drop(events_tx);
192
193        // Continue draining the payloads receiver until it is closed.
194        while let Some(payload) = payloads_rx.recv().await {
195            if let Err(e) = context.dispatcher().dispatch(payload).await {
196                error!("Failed to dispatch payload: {}", e);
197            }
198        }
199
200        // Request build task should now be stopped.
201        match request_builder_handle.await {
202            Ok(Ok(())) => debug!("Request builder task stopped."),
203            Ok(Err(e)) => error!(error = %e, "Request builder task failed."),
204            Err(e) => error!(error = %e, "Request builder task panicked."),
205        }
206
207        debug!("Datadog APM Stats encoder stopped.");
208
209        Ok(())
210    }
211}
212
213async fn run_request_builder(
214    mut stats_request_builder: RequestBuilder<StatsEndpointEncoder>, telemetry: ComponentTelemetry,
215    mut events_rx: Receiver<EventsBuffer>, payloads_tx: Sender<PayloadsBuffer>, flush_timeout: std::time::Duration,
216) -> Result<(), GenericError> {
217    let mut pending_flush = false;
218    let pending_flush_timeout = sleep(flush_timeout);
219    pin!(pending_flush_timeout);
220
221    loop {
222        select! {
223            Some(event_buffer) = events_rx.recv() => {
224                for event in event_buffer {
225                    let trace_stats = match event.try_into_trace_stats() {
226                        Some(stats) => stats,
227                        None => continue,
228                    };
229
230                    // Encode the stats. If we get it back, that means the current request is full, and we need to
231                    // flush it before we can try to encode the stats again.
232                    let stats_to_retry = match stats_request_builder.encode(trace_stats).await {
233                        Ok(None) => continue,
234                        Ok(Some(stats)) => stats,
235                        Err(e) => {
236                            error!(error = %e, "Failed to encode stats.");
237                            telemetry.events_dropped_encoder().increment(1);
238                            continue;
239                        }
240                    };
241
242                    let maybe_requests = stats_request_builder.flush().await;
243                    if maybe_requests.is_empty() {
244                        panic!("builder told us to flush, but gave us nothing");
245                    }
246
247                    for maybe_request in maybe_requests {
248                        match maybe_request {
249                            Ok((events, _data_points, request)) => {
250                                let payload_meta = PayloadMetadata::from_event_count(events);
251                                let http_payload = HttpPayload::new(payload_meta, request);
252                                let payload = Payload::Http(http_payload);
253
254                                payloads_tx.send(payload).await
255                                    .map_err(|_| generic_error!("Failed to send payload to encoder."))?;
256                            },
257                            Err(e) => if e.is_recoverable() {
258                                continue;
259                            } else {
260                                return Err(GenericError::from(e).context("Failed to flush request."));
261                            }
262                        }
263                    }
264
265                    if let Err(e) = stats_request_builder.encode(stats_to_retry).await {
266                        error!(error = %e, "Failed to encode stats.");
267                        telemetry.events_dropped_encoder().increment(1);
268                    }
269                }
270
271                debug!("Processed event buffer.");
272
273                // If we're not already pending a flush, we'll start the countdown.
274                if !pending_flush {
275                    pending_flush_timeout.as_mut().reset(tokio::time::Instant::now() + flush_timeout);
276                    pending_flush = true;
277                }
278            },
279            _ = &mut pending_flush_timeout, if pending_flush => {
280                debug!("Flushing pending request(s).");
281
282                pending_flush = false;
283
284                let maybe_stats_requests = stats_request_builder.flush().await;
285                for maybe_request in maybe_stats_requests {
286                    match maybe_request {
287                        Ok((events, _data_points, request)) => {
288                            let payload_meta = PayloadMetadata::from_event_count(events);
289                            let http_payload = HttpPayload::new(payload_meta, request);
290                            let payload = Payload::Http(http_payload);
291
292                            payloads_tx.send(payload).await
293                                .map_err(|_| generic_error!("Failed to send payload to encoder."))?;
294                        },
295                        Err(e) => if e.is_recoverable() {
296                            continue;
297                        } else {
298                            return Err(GenericError::from(e).context("Failed to flush request."));
299                        }
300                    }
301                }
302
303                debug!("All flushed requests sent to I/O task. Waiting for next event buffer...");
304            },
305
306            else => break,
307        }
308    }
309
310    Ok(())
311}
312
313#[derive(Debug)]
314struct StatsEndpointEncoder {
315    agent_hostname: MetaString,
316    agent_version: MetaString,
317    agent_env: MetaString,
318}
319
320impl StatsEndpointEncoder {
321    fn new(agent_hostname: MetaString, agent_version: MetaString, agent_env: MetaString) -> Self {
322        Self {
323            agent_hostname,
324            agent_version,
325            agent_env,
326        }
327    }
328
329    fn to_proto_stats_payload(&self, stats: &TraceStats) -> ProtoStatsPayload {
330        let mut payload = ProtoStatsPayload::new();
331        payload.set_agentHostname(self.agent_hostname.to_string());
332        payload.set_agentEnv(self.agent_env.to_string());
333        payload.set_agentVersion(self.agent_version.to_string());
334        payload.set_clientComputed(false);
335        payload.set_splitPayload(false);
336        payload.set_stats(stats.stats().iter().map(convert_client_stats_payload).collect());
337
338        payload
339    }
340}
341
342fn convert_client_stats_payload(client_payload: &ClientStatsPayload) -> ProtoClientStatsPayload {
343    let mut proto_client = ProtoClientStatsPayload::new();
344    proto_client.set_hostname(client_payload.hostname().to_string());
345    proto_client.set_env(client_payload.env().to_string());
346    proto_client.set_version(client_payload.version().to_string());
347    proto_client.set_lang(client_payload.lang().to_string());
348    proto_client.set_tracerVersion(client_payload.tracer_version().to_string());
349    proto_client.set_runtimeID(client_payload.runtime_id().to_string());
350    proto_client.set_sequence(client_payload.sequence());
351    proto_client.set_agentAggregation(client_payload.agent_aggregation().to_string());
352    proto_client.set_service(client_payload.service().to_string());
353    proto_client.set_containerID(client_payload.container_id().to_string());
354    proto_client.set_tags(client_payload.tags().into_iter().map(|s| s.to_string()).collect());
355    proto_client.set_git_commit_sha(client_payload.git_commit_sha().to_string());
356    proto_client.set_image_tag(client_payload.image_tag().to_string());
357    proto_client.set_process_tags_hash(client_payload.process_tags_hash());
358    proto_client.set_process_tags(client_payload.process_tags().to_string());
359    proto_client.set_stats(client_payload.stats().iter().map(convert_client_stats_bucket).collect());
360    proto_client
361}
362
363fn convert_client_stats_bucket(bucket: &ClientStatsBucket) -> ProtoClientStatsBucket {
364    let mut proto_bucket = ProtoClientStatsBucket::new();
365    proto_bucket.set_start(bucket.start());
366    proto_bucket.set_duration(bucket.duration());
367    proto_bucket.set_agentTimeShift(bucket.agent_time_shift());
368    proto_bucket.set_stats(bucket.stats().iter().map(convert_client_grouped_stats).collect());
369    proto_bucket
370}
371
372fn convert_client_grouped_stats(grouped: &ClientGroupedStats) -> ProtoClientGroupedStats {
373    let mut proto_grouped = ProtoClientGroupedStats::new();
374    proto_grouped.set_service(grouped.service().to_string());
375    proto_grouped.set_name(grouped.name().to_string());
376    proto_grouped.set_resource(grouped.resource().to_string());
377    proto_grouped.set_HTTP_status_code(grouped.http_status_code());
378    proto_grouped.set_type(grouped.span_type().to_string());
379    proto_grouped.set_DB_type(grouped.db_type().to_string());
380    proto_grouped.set_hits(grouped.hits());
381    proto_grouped.set_errors(grouped.errors());
382    proto_grouped.set_duration(grouped.duration());
383    proto_grouped.set_okSummary(grouped.ok_summary().to_vec());
384    proto_grouped.set_errorSummary(grouped.error_summary().to_vec());
385    proto_grouped.set_synthetics(grouped.synthetics());
386    proto_grouped.set_topLevelHits(grouped.top_level_hits());
387    proto_grouped.set_span_kind(grouped.span_kind().to_string());
388    proto_grouped.set_peer_tags(grouped.peer_tags().iter().map(|s| s.to_string()).collect());
389    proto_grouped.set_is_trace_root(match grouped.is_trace_root() {
390        None => Trilean::NOT_SET,
391        Some(true) => Trilean::TRUE,
392        Some(false) => Trilean::FALSE,
393    });
394    proto_grouped.set_GRPC_status_code(grouped.grpc_status_code().to_string());
395    proto_grouped.set_HTTP_method(grouped.http_method().to_string());
396    proto_grouped.set_HTTP_endpoint(grouped.http_endpoint().to_string());
397    proto_grouped
398}
399
400/// Error type for stats encoding.
401#[derive(Debug)]
402pub struct StatsEncodeError(rmp_serde::encode::Error);
403
404impl std::fmt::Display for StatsEncodeError {
405    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
406        write!(f, "failed to encode stats as MessagePack: {}", self.0)
407    }
408}
409
410impl std::error::Error for StatsEncodeError {
411    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
412        Some(&self.0)
413    }
414}
415
416impl EndpointEncoder for StatsEndpointEncoder {
417    type Input = TraceStats;
418    type EncodeError = StatsEncodeError;
419
420    fn encoder_name() -> &'static str {
421        "stats"
422    }
423
424    fn compressed_size_limit(&self) -> usize {
425        DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT
426    }
427
428    fn uncompressed_size_limit(&self) -> usize {
429        DEFAULT_INTAKE_UNCOMPRESSED_SIZE_LIMIT
430    }
431
432    fn encode(&mut self, stats: &Self::Input, buffer: &mut Vec<u8>) -> Result<(), Self::EncodeError> {
433        let payload = self.to_proto_stats_payload(stats);
434        rmp_serde::encode::write_named(buffer, &payload).map_err(StatsEncodeError)?;
435        Ok(())
436    }
437
438    fn endpoint_uri(&self) -> Uri {
439        PathAndQuery::from_static("/api/v0.2/stats").into()
440    }
441
442    fn endpoint_method(&self) -> Method {
443        Method::POST
444    }
445
446    fn content_type(&self) -> HeaderValue {
447        CONTENT_TYPE_MSGPACK.clone()
448    }
449}
saluki_components/encoders/datadog/stats/mod.rs

saluki_components/encoders/datadog/stats/
mod.rs
