saluki_components/forwarders/datadog/

mod.rs

use std::collections::HashMap;

use agent_data_plane_config::shared::{Secrets, SharedConfiguration};
use agent_data_plane_config::Live;
use async_trait::async_trait;
use http::Uri;
use saluki_common::buf::FrozenChunkedBytesBuffer;
use saluki_core::accounting::{MemoryBounds, MemoryBoundsBuilder, UsageExpr};
use saluki_core::{
    components::{forwarders::*, BuildContext},
    data_model::payload::{PayloadMetadata, PayloadType},
    observability::ComponentMetricsExt as _,
};
use saluki_error::GenericError;
use saluki_metrics::MetricsBuilder;
use stringtheory::MetaString;
use tokio::select;
use tracing::debug;

use crate::common::datadog::{
    api_key::{ApiKeyView, LiveApiKeys},
    config::ForwarderConfiguration,
    endpoints::SingleDestination,
    io::{LiveForwarderConfiguration, TransactionForwarder},
    protocol::MetricsPayloadInfo,
    telemetry::ComponentTelemetry,
    transaction::{Metadata, Transaction},
    validation::ValidationReadiness,
    DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT, METRICS_SERIES_V3_BETA_PATH, METRICS_SERIES_V3_PATH,
    METRICS_SKETCHES_V3_PATH,
};

/// Datadog forwarder.
///
/// Forwards Datadog-specific payloads to the Datadog platform. Handles the standard Datadog Agent configuration,
/// in terms of specifying additional endpoints, adding the necessary HTTP request headers for authentication,
/// identification, and more.
pub struct DatadogForwarderConfiguration {
    /// Forwarder configuration settings.
    ///
    /// See [`ForwarderConfiguration`] for more information about the available settings.
    forwarder_config: ForwarderConfiguration,

    /// Live views the endpoints take their API keys from.
    api_keys: LiveApiKeys,

    /// Live view of the secrets settings the retry policy's gate reads.
    secrets: Live<Secrets>,
}

impl DatadogForwarderConfiguration {
    /// Creates a new `DatadogForwarderConfiguration` from the resolved shared configuration.
    ///
    /// `api_key` and `additional_endpoints` are the live views the endpoints take their API keys from, so that a key an
    /// operator rotates while the process runs reaches the next request. `secrets` is the live view the retry policy's
    /// gate reads to decide whether a rejected key is worth retrying.
    pub fn from_configuration(
        shared: &SharedConfiguration, api_key: Live<String>, additional_endpoints: Live<HashMap<String, Vec<String>>>,
        secrets: Live<Secrets>,
    ) -> Self {
        Self {
            forwarder_config: ForwarderConfiguration::from_configuration(shared),
            api_keys: LiveApiKeys {
                primary: Some(ApiKeyView::Required(api_key)),
                additional: Some(additional_endpoints),
            },
            secrets,
        }
    }

    /// Creates a new `DatadogForwarderConfiguration` that forwards to a single endpoint override.
    ///
    /// The override replaces the configured endpoints entirely, and `api_key_view` is the live view its key comes from.
    /// Multi-Region Failover needs this: it has its own key, not the primary one, and configuration can leave that key
    /// unset. `secrets` serves the same purpose as in [`from_configuration`](Self::from_configuration): it decides
    /// whether the override's rejected key is worth retrying.
    pub fn for_endpoint_override(
        shared: &SharedConfiguration, dd_url: String, api_key: String, api_key_view: Live<Option<String>>,
        secrets: Live<Secrets>,
    ) -> Self {
        let destination = SingleDestination {
            url: dd_url,
            api_key,
            accepts_v3_series: true,
        };

        Self {
            forwarder_config: ForwarderConfiguration::for_single_destination(shared, &destination),
            api_keys: LiveApiKeys {
                primary: Some(ApiKeyView::Optional(api_key_view)),
                additional: None,
            },
            secrets,
        }
    }
}

#[async_trait]
impl ForwarderBuilder for DatadogForwarderConfiguration {
    fn input_payload_type(&self) -> PayloadType {
        PayloadType::Http
    }

    async fn build(&self, context: BuildContext) -> Result<Box<dyn Forwarder + Send>, GenericError> {
        let metrics_builder = MetricsBuilder::from_component_context(context.component_context());
        let telemetry = ComponentTelemetry::from_builder(&metrics_builder);
        let forwarder = TransactionForwarder::from_config(
            context.component_context().clone(),
            self.forwarder_config.clone(),
            LiveForwarderConfiguration {
                api_keys: self.api_keys.clone(),
                secrets: self.secrets.clone(),
            },
            get_dd_endpoint_name,
            telemetry.clone(),
            metrics_builder,
        )?;

        Ok(Box::new(Datadog { forwarder }))
    }
}

impl MemoryBounds for DatadogForwarderConfiguration {
    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
        builder
            .minimum()
            .with_single_value::<Datadog>("component struct")
            .with_array::<Transaction<FrozenChunkedBytesBuffer>>("requests channel", 8);

        builder
            .firm()
            // TODO: This is a little wonky because we're accounting for the firm bound portion of connected encoders here, as this
            // is the only place where we can calculate how many requests we'll hold on to in memory, which is what ultimately influences
            // the firm usage.
            //
            // We're also cheating by knowing what the largest possible payload is that we'll potentially see, based on the limits on the
            // Datadog encoders. This won't necessarily hold up for future sources/encoders, but is good enough for now.
            .with_expr(UsageExpr::sum(
                "in-flight requests",
                UsageExpr::config(
                    "forwarder_retry_queue_payloads_max_size",
                    self.forwarder_config.retry().queue_max_size_bytes() as usize,
                ),
                UsageExpr::product(
                    "high priority queue",
                    UsageExpr::config(
                        "forwarder_high_prio_buffer_size",
                        self.forwarder_config.endpoint_buffer_size(),
                    ),
                    // TODO: The default compressed size limit just so happens to be the biggest one we currently default with on our side,
                    // but it's not clear that this will always be the case.
                    UsageExpr::constant("maximum compressed payload size", DEFAULT_INTAKE_COMPRESSED_SIZE_LIMIT),
                ),
            ));
    }
}

pub struct Datadog {
    forwarder: TransactionForwarder<FrozenChunkedBytesBuffer>,
}

#[async_trait]
impl Forwarder for Datadog {
    async fn run(mut self: Box<Self>, mut context: ForwarderContext) -> Result<(), GenericError> {
        let Self { forwarder } = *self;

        let mut health = context.take_health_handle();

        let mut validation = forwarder.api_key_validator().spawn();

        // Spawn our forwarder task to handle sending requests.
        let forwarder = forwarder.spawn().await;

        debug!("Datadog forwarder started.");

        loop {
            select! {
                _ = health.live() => continue,
                readiness = validation.wait_for_change() => match readiness {
                    ValidationReadiness::Ready => health.mark_ready(),
                    ValidationReadiness::NotReady => health.mark_not_ready(),
                },
                maybe_payload = context.payloads().next() => match maybe_payload {
                    Some(payload) => if let Some(http_payload) = payload.try_into_http_payload() {
                        let (payload_meta, request) = http_payload.into_parts();
                        let transaction_meta = transaction_metadata_from_payload_metadata(&payload_meta);
                        let transaction = Transaction::from_original(transaction_meta, request);

                        forwarder.send_transaction(transaction).await?;
                    }
                    None => break,
                },
            }
        }

        // Shutdown the forwarder gracefully.
        validation.abort();
        forwarder.shutdown().await;

        debug!("Datadog forwarder stopped.");

        Ok(())
    }
}

fn transaction_metadata_from_payload_metadata(payload_meta: &PayloadMetadata) -> Metadata {
    let mut transaction_meta =
        Metadata::from_event_and_data_point_count(payload_meta.event_count(), payload_meta.data_point_count());
    transaction_meta.payload_info = payload_meta.get::<MetricsPayloadInfo>().copied();
    transaction_meta
}

fn get_dd_endpoint_name(uri: &Uri) -> Option<MetaString> {
    match uri.path() {
        "/api/v2/logs" => Some(MetaString::from_static("logs_v2")),
        "/api/v1/series" => Some(MetaString::from_static("series_v1")),
        "/api/v2/series" => Some(MetaString::from_static("series_v2")),
        METRICS_SERIES_V3_PATH => Some(MetaString::from_static("series_v3")),
        METRICS_SERIES_V3_BETA_PATH => Some(MetaString::from_static("series_v3beta")),
        "/api/beta/sketches" => Some(MetaString::from_static("sketches_v2")),
        METRICS_SKETCHES_V3_PATH => Some(MetaString::from_static("sketches_v3")),
        "/api/v1/check_run" => Some(MetaString::from_static("check_run_v1")),
        "/api/v1/events_batch" => Some(MetaString::from_static("events_batch_v1")),
        "/api/v0.2/traces" => Some(MetaString::from_static("traces_v0.2")),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use agent_data_plane_config::SalukiConfiguration;

    use super::*;
    use crate::common::datadog::{
        api_key::ApiKeyRefresher,
        test_util::{shared_configuration, LiveConfiguration},
    };

    #[test]
    fn dd_endpoint_names_map_from_request_path() {
        // Each intake path maps to the telemetry endpoint name the forwarder tags transactions with; the query and
        // authority are ignored (only the path matters), and unknown paths map to `None`.
        let cases: [(Uri, Option<&str>); 12] = [
            (Uri::from_static("/api/v2/logs"), Some("logs_v2")),
            (Uri::from_static("/api/v1/series"), Some("series_v1")),
            (Uri::from_static("/api/v2/series"), Some("series_v2")),
            (Uri::from_static(METRICS_SERIES_V3_PATH), Some("series_v3")),
            (Uri::from_static(METRICS_SERIES_V3_BETA_PATH), Some("series_v3beta")),
            (Uri::from_static("/api/beta/sketches"), Some("sketches_v2")),
            (Uri::from_static(METRICS_SKETCHES_V3_PATH), Some("sketches_v3")),
            (Uri::from_static("/api/v1/check_run"), Some("check_run_v1")),
            (Uri::from_static("/api/v1/events_batch"), Some("events_batch_v1")),
            (Uri::from_static("/api/v0.2/traces"), Some("traces_v0.2")),
            (
                Uri::from_static("https://app.datadoghq.com/api/v2/series"),
                Some("series_v2"),
            ),
            (Uri::from_static("/api/v1/unknown"), None),
        ];

        for (uri, expected) in cases {
            assert_eq!(
                expected,
                get_dd_endpoint_name(&uri).as_deref(),
                "get_dd_endpoint_name({})",
                uri.path()
            );
        }
    }

    #[test]
    fn transaction_metadata_carries_counts_and_metrics_payload_info() {
        // A non-metrics payload propagates the event/data-point counts and leaves `payload_info` as `None`.
        let payload_meta = PayloadMetadata::from_event_and_data_point_count(3, 11);
        let transaction_meta = transaction_metadata_from_payload_metadata(&payload_meta);
        assert_eq!(3, transaction_meta.event_count);
        assert_eq!(11, transaction_meta.data_point_count);
        assert_eq!(None, transaction_meta.payload_info);

        // A metrics payload additionally copies the `MetricsPayloadInfo` extension through unchanged.
        let payload_meta = PayloadMetadata::from_event_and_data_point_count(2, 7).with(MetricsPayloadInfo::v3_series());
        let transaction_meta = transaction_metadata_from_payload_metadata(&payload_meta);
        assert_eq!(2, transaction_meta.event_count);
        assert_eq!(7, transaction_meta.data_point_count);
        assert_eq!(Some(MetricsPayloadInfo::v3_series()), transaction_meta.payload_info);
    }

    #[tokio::test]
    async fn an_endpoint_override_follows_the_failover_key_and_not_the_primary_one() {
        let mut live_config = SalukiConfiguration::default();
        live_config.shared.endpoints.api_key = "primary-api-key".to_string();
        live_config.domains.multi_region_failover.api_key = Some("mrf-api-key".to_string());
        let live = LiveConfiguration::new(live_config.clone());

        let config = DatadogForwarderConfiguration::for_endpoint_override(
            &shared_configuration(),
            "http://mrf.example.test".to_string(),
            "mrf-api-key".to_string(),
            live.live(|config| &config.domains.multi_region_failover.api_key),
            live.live(|config| &config.shared.secrets),
        );

        let endpoints = config
            .forwarder_config
            .build_routable_endpoints()
            .expect("endpoint should resolve");
        assert_eq!(endpoints.len(), 1);
        assert_eq!("mrf-api-key", &*endpoints[0].endpoint().api_key());

        ApiKeyRefresher::new(&endpoints, &config.api_keys)
            .expect("the destination should follow the failover view")
            .spawn();

        // Rotate the primary key first: the override does not follow it.
        live_config.shared.endpoints.api_key = "rotated-primary-api-key".to_string();
        live.store(live_config.clone());

        live_config.domains.multi_region_failover.api_key = Some("rotated-mrf-api-key".to_string());
        live.store(live_config);

        let endpoint = endpoints[0].endpoint();
        tokio::time::timeout(std::time::Duration::from_secs(2), async {
            while &*endpoint.api_key() != "rotated-mrf-api-key" {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("the endpoint override should follow the failover key");
    }
}