Skip to main content

saluki_components/forwarders/datadog/
mod.rs

1use async_trait::async_trait;
2use http::Uri;
3use resource_accounting::{MemoryBounds, MemoryBoundsBuilder, UsageExpr};
4use saluki_common::buf::FrozenChunkedBytesBuffer;
5use saluki_config::GenericConfiguration;
6use saluki_core::{
7    components::{forwarders::*, ComponentContext},
8    data_model::payload::{PayloadMetadata, PayloadType},
9    observability::ComponentMetricsExt as _,
10};
11use saluki_error::GenericError;
12use saluki_metrics::MetricsBuilder;
13use stringtheory::MetaString;
14use tokio::select;
15use tracing::debug;
16
17use crate::common::datadog::{
18    config::ForwarderConfiguration,
19    io::TransactionForwarder,
20    protocol::MetricsPayloadInfo,
21    telemetry::ComponentTelemetry,
22    transaction::{Metadata, Transaction},
23    validation::ValidationReadiness,
24    DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT, METRICS_SERIES_V3_BETA_PATH, METRICS_SERIES_V3_PATH,
25    METRICS_SKETCHES_V3_PATH,
26};
27
28/// Datadog forwarder.
29///
30/// Forwards Datadog-specific payloads to the Datadog platform. Handles the standard Datadog Agent configuration,
31/// in terms of specifying additional endpoints, adding the necessary HTTP request headers for authentication,
32/// identification, and more.
33pub struct DatadogForwarderConfiguration {
34    /// Forwarder configuration settings.
35    ///
36    /// See [`ForwarderConfiguration`] for more information about the available settings.
37    forwarder_config: ForwarderConfiguration,
38
39    configuration: Option<GenericConfiguration>,
40}
41
42impl DatadogForwarderConfiguration {
43    /// Creates a new `DatadogForwarderConfiguration` from the given configuration.
44    pub fn from_configuration(config: &GenericConfiguration) -> Result<Self, GenericError> {
45        let forwarder_config = ForwarderConfiguration::from_configuration(config)?;
46        Ok(Self {
47            forwarder_config,
48            configuration: Some(config.clone()),
49        })
50    }
51
52    /// Overrides the default endpoint and refreshes its API key from the given config path.
53    ///
54    /// This is for override endpoints whose API key does not refresh from the top-level `api_key`
55    /// config path, such as Multi-Region Failover.
56    pub fn with_endpoint_override_and_api_key_refresh_config_path(
57        mut self, dd_url: String, api_key: String, api_key_refresh_config_path: &'static str,
58    ) -> Self {
59        self.apply_endpoint_override(dd_url, api_key, api_key_refresh_config_path);
60
61        self
62    }
63
64    fn apply_endpoint_override(&mut self, dd_url: String, api_key: String, api_key_refresh_config_path: &'static str) {
65        // Clear any existing additional endpoints, and set the new DD URL and API key.
66        //
67        // This ensures that the only endpoint we'll send to is this one.
68        let endpoint = self.forwarder_config.endpoint_mut();
69        endpoint.clear_additional_endpoints();
70        endpoint.set_dd_url(dd_url);
71        endpoint.set_api_key(api_key);
72        endpoint.set_api_key_refresh_config_path(api_key_refresh_config_path);
73        self.forwarder_config.clear_opw_metrics_endpoint();
74    }
75}
76
77#[async_trait]
78impl ForwarderBuilder for DatadogForwarderConfiguration {
79    fn input_payload_type(&self) -> PayloadType {
80        PayloadType::Http
81    }
82
83    async fn build(&self, context: ComponentContext) -> Result<Box<dyn Forwarder + Send>, GenericError> {
84        let metrics_builder = MetricsBuilder::from_component_context(&context);
85        let telemetry = ComponentTelemetry::from_builder(&metrics_builder);
86        let forwarder = TransactionForwarder::from_config(
87            context,
88            self.forwarder_config.clone(),
89            self.configuration.clone(),
90            get_dd_endpoint_name,
91            telemetry.clone(),
92            metrics_builder,
93        )?;
94
95        Ok(Box::new(Datadog { forwarder }))
96    }
97}
98
99impl MemoryBounds for DatadogForwarderConfiguration {
100    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
101        builder
102            .minimum()
103            .with_single_value::<Datadog>("component struct")
104            .with_array::<Transaction<FrozenChunkedBytesBuffer>>("requests channel", 8);
105
106        builder
107            .firm()
108            // TODO: This is a little wonky because we're accounting for the firm bound portion of connected encoders here, as this
109            // is the only place where we can calculate how many requests we'll hold on to in memory, which is what ultimately influences
110            // the firm usage.
111            //
112            // We're also cheating by knowing what the largest possible payload is that we'll potentially see, based on the limits on the
113            // Datadog encoders. This won't necessarily hold up for future sources/encoders, but is good enough for now.
114            .with_expr(UsageExpr::sum(
115                "in-flight requests",
116                UsageExpr::config(
117                    "forwarder_retry_queue_payloads_max_size",
118                    self.forwarder_config.retry().queue_max_size_bytes() as usize,
119                ),
120                UsageExpr::product(
121                    "high priority queue",
122                    UsageExpr::config(
123                        "forwarder_high_prio_buffer_size",
124                        self.forwarder_config.endpoint_buffer_size(),
125                    ),
126                    // TODO: The default compressed size limit just so happens to be the biggest one we currently default with on our side,
127                    // but it's not clear that this will always be the case.
128                    UsageExpr::constant("maximum compressed payload size", DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT),
129                ),
130            ));
131    }
132}
133
134pub struct Datadog {
135    forwarder: TransactionForwarder<FrozenChunkedBytesBuffer>,
136}
137
138#[async_trait]
139impl Forwarder for Datadog {
140    async fn run(mut self: Box<Self>, mut context: ForwarderContext) -> Result<(), GenericError> {
141        let Self { forwarder } = *self;
142
143        let mut health = context.take_health_handle();
144
145        let mut validation = forwarder.api_key_validator().spawn();
146
147        // Spawn our forwarder task to handle sending requests.
148        let forwarder = forwarder.spawn().await;
149
150        debug!("Datadog forwarder started.");
151
152        loop {
153            select! {
154                _ = health.live() => continue,
155                readiness = validation.wait_for_change() => match readiness {
156                    ValidationReadiness::Ready => health.mark_ready(),
157                    ValidationReadiness::NotReady => health.mark_not_ready(),
158                },
159                maybe_payload = context.payloads().next() => match maybe_payload {
160                    Some(payload) => if let Some(http_payload) = payload.try_into_http_payload() {
161                        let (payload_meta, request) = http_payload.into_parts();
162                        let transaction_meta = transaction_metadata_from_payload_metadata(&payload_meta);
163                        let transaction = Transaction::from_original(transaction_meta, request);
164
165                        forwarder.send_transaction(transaction).await?;
166                    }
167                    None => break,
168                },
169            }
170        }
171
172        // Shutdown the forwarder gracefully.
173        validation.abort();
174        forwarder.shutdown().await;
175
176        debug!("Datadog forwarder stopped.");
177
178        Ok(())
179    }
180}
181
182fn transaction_metadata_from_payload_metadata(payload_meta: &PayloadMetadata) -> Metadata {
183    let mut transaction_meta =
184        Metadata::from_event_and_data_point_count(payload_meta.event_count(), payload_meta.data_point_count());
185    transaction_meta.payload_info = payload_meta.get::<MetricsPayloadInfo>().copied();
186    transaction_meta
187}
188
189fn get_dd_endpoint_name(uri: &Uri) -> Option<MetaString> {
190    match uri.path() {
191        "/api/v2/logs" => Some(MetaString::from_static("logs_v2")),
192        "/api/v1/series" => Some(MetaString::from_static("series_v1")),
193        "/api/v2/series" => Some(MetaString::from_static("series_v2")),
194        METRICS_SERIES_V3_PATH => Some(MetaString::from_static("series_v3")),
195        METRICS_SERIES_V3_BETA_PATH => Some(MetaString::from_static("series_v3beta")),
196        "/api/beta/sketches" => Some(MetaString::from_static("sketches_v2")),
197        METRICS_SKETCHES_V3_PATH => Some(MetaString::from_static("sketches_v3")),
198        "/api/v1/check_run" => Some(MetaString::from_static("check_run_v1")),
199        "/api/v1/events_batch" => Some(MetaString::from_static("events_batch_v1")),
200        "/api/v0.2/traces" => Some(MetaString::from_static("traces_v0.2")),
201        _ => None,
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use saluki_config::ConfigurationLoader;
208    use serde_json::json;
209
210    use super::*;
211
212    #[tokio::test]
213    async fn endpoint_override_refreshes_from_mrf_api_key() {
214        let (generic_config, sender) = ConfigurationLoader::for_tests(
215            Some(json!({
216                "api_key": "primary-api-key",
217                "multi_region_failover": {
218                    "api_key": "mrf-api-key"
219                }
220            })),
221            None,
222            true,
223        )
224        .await;
225        let sender = sender.expect("dynamic sender should exist");
226        sender
227            .send(saluki_config::dynamic::ConfigUpdate::Snapshot(json!({})))
228            .await
229            .expect("initial dynamic snapshot should be sent");
230        generic_config.ready().await;
231
232        let config = DatadogForwarderConfiguration::from_configuration(&generic_config)
233            .expect("DatadogForwarderConfiguration should parse")
234            .with_endpoint_override_and_api_key_refresh_config_path(
235                "http://mrf.example.test".to_string(),
236                "mrf-api-key".to_string(),
237                "multi_region_failover.api_key",
238            );
239
240        let mut endpoints = config
241            .forwarder_config
242            .build_routable_endpoints(config.configuration.clone())
243            .expect("endpoint should resolve");
244
245        assert_eq!(endpoints.len(), 1);
246        let (_, mut endpoint) = endpoints.pop().unwrap().into_parts();
247        assert_eq!(endpoint.cached_api_key(), "mrf-api-key");
248        assert!(endpoint.has_configuration());
249        assert_eq!(endpoint.api_key(), "mrf-api-key");
250
251        sender
252            .send(saluki_config::dynamic::ConfigUpdate::Partial {
253                key: "api_key".to_string(),
254                value: json!("rotated-primary-api-key"),
255            })
256            .await
257            .expect("primary API key update should be sent");
258        sender
259            .send(saluki_config::dynamic::ConfigUpdate::Partial {
260                key: "multi_region_failover.api_key".to_string(),
261                value: json!("rotated-mrf-api-key"),
262            })
263            .await
264            .expect("MRF API key update should be sent");
265
266        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
267        loop {
268            if endpoint.api_key() == "rotated-mrf-api-key" {
269                break;
270            }
271            assert!(
272                std::time::Instant::now() < deadline,
273                "timed out waiting for endpoint override to refresh from MRF API key"
274            );
275            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
276        }
277    }
278}