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_core::accounting::{MemoryBounds, MemoryBoundsBuilder};
13use saluki_core::runtime;
14use saluki_core::{
15    components::{encoders::*, BuildContext},
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: BuildContext) -> Result<Box<dyn Encoder + Send>, GenericError> {
95        let metrics_builder = MetricsBuilder::from_component_context(context.component_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        // Run our request builder task on the worker pool.
162        //
163        // The request builder task ignores the shutdown signal on purpose: it drains its incoming event buffer channel
164        // until the channel closes, which is what guarantees every buffered metric is encoded and dispatched.
165        let request_builder_fut = run_request_builder(stats_rb, telemetry, events_rx, payloads_tx, flush_timeout);
166        runtime::worker("request_builder", request_builder_fut)
167            .on_runtime(context.topology_context().global_thread_pool().clone())
168            .spawn();
169
170        health.mark_ready();
171        debug!("Datadog APM Stats encoder started.");
172
173        loop {
174            select! {
175                biased;
176
177                _ = health.live() => continue,
178                maybe_payload = payloads_rx.recv() => match maybe_payload {
179                    Some(payload) => {
180                        if let Err(e) = context.dispatcher().dispatch(payload).await {
181                            error!("Failed to dispatch payload: {}", e);
182                        }
183                    }
184                    None => break,
185                },
186                maybe_event_buffer = context.events().next() => match maybe_event_buffer {
187                    Some(event_buffer) => events_tx.send(event_buffer).await
188                        .error_context("Failed to send event buffer to request builder task.")?,
189                    None => break,
190                },
191            }
192        }
193
194        drop(events_tx);
195
196        // Continue draining the payloads receiver until it is closed.
197        while let Some(payload) = payloads_rx.recv().await {
198            if let Err(e) = context.dispatcher().dispatch(payload).await {
199                error!("Failed to dispatch payload: {}", e);
200            }
201        }
202
203        // Draining `payloads_rx` to completion already implies the request builder finished: it owns the only sender,
204        // so the channel only closes once that child's future has run to completion (or been dropped).
205        debug!("Datadog APM Stats encoder stopped.");
206
207        Ok(())
208    }
209}
210
211async fn run_request_builder(
212    mut stats_request_builder: RequestBuilder<StatsEndpointEncoder>, telemetry: ComponentTelemetry,
213    mut events_rx: Receiver<EventsBuffer>, payloads_tx: Sender<PayloadsBuffer>, flush_timeout: std::time::Duration,
214) -> Result<(), GenericError> {
215    let mut pending_flush = false;
216    let pending_flush_timeout = sleep(flush_timeout);
217    pin!(pending_flush_timeout);
218
219    loop {
220        select! {
221            Some(event_buffer) = events_rx.recv() => {
222                for event in event_buffer {
223                    let trace_stats = match event.try_into_trace_stats() {
224                        Some(stats) => stats,
225                        None => continue,
226                    };
227
228                    // Encode the stats. If we get it back, that means the current request is full, and we need to
229                    // flush it before we can try to encode the stats again.
230                    let stats_to_retry = match stats_request_builder.encode(trace_stats).await {
231                        Ok(None) => continue,
232                        Ok(Some(stats)) => stats,
233                        Err(e) => {
234                            error!(error = %e, "Failed to encode stats.");
235                            telemetry.events_dropped_encoder().increment(1);
236                            continue;
237                        }
238                    };
239
240                    let maybe_requests = stats_request_builder.flush().await;
241                    if maybe_requests.is_empty() {
242                        panic!("builder told us to flush, but gave us nothing");
243                    }
244
245                    for maybe_request in maybe_requests {
246                        match maybe_request {
247                            Ok((events, _data_points, request)) => {
248                                let payload_meta = PayloadMetadata::from_event_count(events);
249                                let http_payload = HttpPayload::new(payload_meta, request);
250                                let payload = Payload::Http(http_payload);
251
252                                payloads_tx.send(payload).await
253                                    .map_err(|_| generic_error!("Failed to send payload to encoder."))?;
254                            },
255                            Err(e) => if e.is_recoverable() {
256                                continue;
257                            } else {
258                                return Err(GenericError::from(e).context("Failed to flush request."));
259                            }
260                        }
261                    }
262
263                    if let Err(e) = stats_request_builder.encode(stats_to_retry).await {
264                        error!(error = %e, "Failed to encode stats.");
265                        telemetry.events_dropped_encoder().increment(1);
266                    }
267                }
268
269                debug!("Processed event buffer.");
270
271                // If we're not already pending a flush, we'll start the countdown.
272                if !pending_flush {
273                    pending_flush_timeout.as_mut().reset(tokio::time::Instant::now() + flush_timeout);
274                    pending_flush = true;
275                }
276            },
277            _ = &mut pending_flush_timeout, if pending_flush => {
278                debug!("Flushing pending request(s).");
279
280                pending_flush = false;
281
282                let maybe_stats_requests = stats_request_builder.flush().await;
283                for maybe_request in maybe_stats_requests {
284                    match maybe_request {
285                        Ok((events, _data_points, request)) => {
286                            let payload_meta = PayloadMetadata::from_event_count(events);
287                            let http_payload = HttpPayload::new(payload_meta, request);
288                            let payload = Payload::Http(http_payload);
289
290                            payloads_tx.send(payload).await
291                                .map_err(|_| generic_error!("Failed to send payload to encoder."))?;
292                        },
293                        Err(e) => if e.is_recoverable() {
294                            continue;
295                        } else {
296                            return Err(GenericError::from(e).context("Failed to flush request."));
297                        }
298                    }
299                }
300
301                debug!("All flushed requests sent to I/O task. Waiting for next event buffer...");
302            },
303
304            else => break,
305        }
306    }
307
308    Ok(())
309}
310
311#[derive(Debug)]
312struct StatsEndpointEncoder {
313    agent_hostname: MetaString,
314    agent_version: MetaString,
315    agent_env: MetaString,
316}
317
318impl StatsEndpointEncoder {
319    fn new(agent_hostname: MetaString, agent_version: MetaString, agent_env: MetaString) -> Self {
320        Self {
321            agent_hostname,
322            agent_version,
323            agent_env,
324        }
325    }
326
327    fn to_proto_stats_payload(&self, stats: &TraceStats) -> ProtoStatsPayload {
328        let mut payload = ProtoStatsPayload::new();
329        payload.set_agentHostname(self.agent_hostname.to_string());
330        payload.set_agentEnv(self.agent_env.to_string());
331        payload.set_agentVersion(self.agent_version.to_string());
332        payload.set_clientComputed(false);
333        payload.set_splitPayload(false);
334        payload.set_stats(stats.stats().iter().map(convert_client_stats_payload).collect());
335
336        payload
337    }
338}
339
340fn convert_client_stats_payload(client_payload: &ClientStatsPayload) -> ProtoClientStatsPayload {
341    let mut proto_client = ProtoClientStatsPayload::new();
342    proto_client.set_hostname(client_payload.hostname().to_string());
343    proto_client.set_env(client_payload.env().to_string());
344    proto_client.set_version(client_payload.version().to_string());
345    proto_client.set_lang(client_payload.lang().to_string());
346    proto_client.set_tracerVersion(client_payload.tracer_version().to_string());
347    proto_client.set_runtimeID(client_payload.runtime_id().to_string());
348    proto_client.set_sequence(client_payload.sequence());
349    proto_client.set_agentAggregation(client_payload.agent_aggregation().to_string());
350    proto_client.set_service(client_payload.service().to_string());
351    proto_client.set_containerID(client_payload.container_id().to_string());
352    proto_client.set_tags(client_payload.tags().into_iter().map(|s| s.to_string()).collect());
353    proto_client.set_git_commit_sha(client_payload.git_commit_sha().to_string());
354    proto_client.set_image_tag(client_payload.image_tag().to_string());
355    proto_client.set_process_tags_hash(client_payload.process_tags_hash());
356    proto_client.set_process_tags(client_payload.process_tags().to_string());
357    proto_client.set_stats(client_payload.stats().iter().map(convert_client_stats_bucket).collect());
358    proto_client
359}
360
361fn convert_client_stats_bucket(bucket: &ClientStatsBucket) -> ProtoClientStatsBucket {
362    let mut proto_bucket = ProtoClientStatsBucket::new();
363    proto_bucket.set_start(bucket.start());
364    proto_bucket.set_duration(bucket.duration());
365    proto_bucket.set_agentTimeShift(bucket.agent_time_shift());
366    proto_bucket.set_stats(bucket.stats().iter().map(convert_client_grouped_stats).collect());
367    proto_bucket
368}
369
370fn convert_client_grouped_stats(grouped: &ClientGroupedStats) -> ProtoClientGroupedStats {
371    let mut proto_grouped = ProtoClientGroupedStats::new();
372    proto_grouped.set_service(grouped.service().to_string());
373    proto_grouped.set_name(grouped.name().to_string());
374    proto_grouped.set_resource(grouped.resource().to_string());
375    proto_grouped.set_HTTP_status_code(grouped.http_status_code());
376    proto_grouped.set_type(grouped.span_type().to_string());
377    proto_grouped.set_DB_type(grouped.db_type().to_string());
378    proto_grouped.set_hits(grouped.hits());
379    proto_grouped.set_errors(grouped.errors());
380    proto_grouped.set_duration(grouped.duration());
381    proto_grouped.set_okSummary(grouped.ok_summary().to_vec());
382    proto_grouped.set_errorSummary(grouped.error_summary().to_vec());
383    proto_grouped.set_synthetics(grouped.synthetics());
384    proto_grouped.set_topLevelHits(grouped.top_level_hits());
385    proto_grouped.set_span_kind(grouped.span_kind().to_string());
386    proto_grouped.set_peer_tags(grouped.peer_tags().iter().map(|s| s.to_string()).collect());
387    proto_grouped.set_is_trace_root(match grouped.is_trace_root() {
388        None => Trilean::NOT_SET,
389        Some(true) => Trilean::TRUE,
390        Some(false) => Trilean::FALSE,
391    });
392    proto_grouped.set_GRPC_status_code(grouped.grpc_status_code().to_string());
393    proto_grouped.set_HTTP_method(grouped.http_method().to_string());
394    proto_grouped.set_HTTP_endpoint(grouped.http_endpoint().to_string());
395    proto_grouped
396}
397
398/// Error type for stats encoding.
399#[derive(Debug)]
400pub struct StatsEncodeError(rmp_serde::encode::Error);
401
402impl std::fmt::Display for StatsEncodeError {
403    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
404        write!(f, "failed to encode stats as MessagePack: {}", self.0)
405    }
406}
407
408impl std::error::Error for StatsEncodeError {
409    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
410        Some(&self.0)
411    }
412}
413
414impl EndpointEncoder for StatsEndpointEncoder {
415    type Input = TraceStats;
416    type EncodeError = StatsEncodeError;
417
418    fn encoder_name() -> &'static str {
419        "stats"
420    }
421
422    fn compressed_size_limit(&self) -> usize {
423        DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT
424    }
425
426    fn uncompressed_size_limit(&self) -> usize {
427        DEFAULT_INTAKE_UNCOMPRESSED_SIZE_LIMIT
428    }
429
430    fn encode(&mut self, stats: &Self::Input, buffer: &mut Vec<u8>) -> Result<(), Self::EncodeError> {
431        let payload = self.to_proto_stats_payload(stats);
432        rmp_serde::encode::write_named(buffer, &payload).map_err(StatsEncodeError)?;
433        Ok(())
434    }
435
436    fn endpoint_uri(&self) -> Uri {
437        PathAndQuery::from_static("/api/v0.2/stats").into()
438    }
439
440    fn endpoint_method(&self) -> Method {
441        Method::POST
442    }
443
444    fn content_type(&self) -> HeaderValue {
445        CONTENT_TYPE_MSGPACK.clone()
446    }
447}
saluki_components/encoders/datadog/stats/mod.rs

saluki_components/encoders/datadog/stats/
mod.rs
