saluki_context/

origin.rs

#![allow(warnings)]
//! Metric origin.

use std::{fmt, num::NonZeroU32, sync::Arc};

use indexmap::Equivalent;
use saluki_common::hash::hash_single_fast;
use serde::{Deserialize, Serialize};
use stringtheory::MetaString;
use tracing::warn;

use crate::tags::{SharedTagSet, Tag};

/// The cardinality of tags associated with the origin entity.
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(try_from = "String")]
pub enum OriginTagCardinality {
    /// No cardinality.
    ///
    /// This implies that no tags should be added to the metric based on its origin.
    None,

    /// Low cardinality.
    ///
    /// This generally covers tags which are static, or relatively slow to change, and generally results in a small
    /// number of unique values for the given tag key.
    Low,

    /// Orchestrator cardinality.
    ///
    /// This generally covers orchestrator-specific tags, such as Kubernetes pod UID, and lands somewhere between low
    /// and high cardinality.
    Orchestrator,

    /// High cardinality.
    ///
    /// This generally covers tags which frequently change and generally results in a large number of unique values for
    /// the given tag key.
    High,
}

impl OriginTagCardinality {
    /// Returns the string representation of the cardinality.
    pub const fn as_str(&self) -> &'static str {
        match self {
            Self::None => "none",
            Self::Low => "low",
            Self::Orchestrator => "orchestrator",
            Self::High => "high",
        }
    }
}

impl TryFrom<&str> for OriginTagCardinality {
    type Error = String;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        // Exact lowercase match first (zero-cost happy path), then case-insensitive fallback
        // using eq_ignore_ascii_case (also allocation-free) to match the core Datadog Agent,
        // which normalises with strings.ToLower before matching.
        match value {
            "none" => Ok(Self::None),
            "low" => Ok(Self::Low),
            "high" => Ok(Self::High),
            "orch" | "orchestrator" => Ok(Self::Orchestrator),
            other if other.eq_ignore_ascii_case("none") => Ok(Self::None),
            other if other.eq_ignore_ascii_case("low") => Ok(Self::Low),
            other if other.eq_ignore_ascii_case("high") => Ok(Self::High),
            other if other.eq_ignore_ascii_case("orch") || other.eq_ignore_ascii_case("orchestrator") => {
                Ok(Self::Orchestrator)
            }
            other => Err(format!("unknown tag cardinality type '{}'", other)),
        }
    }
}

impl TryFrom<String> for OriginTagCardinality {
    type Error = String;

    fn try_from(value: String) -> Result<Self, Self::Error> {
        Self::try_from(value.as_str())
    }
}

impl fmt::Display for OriginTagCardinality {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::None => write!(f, "none"),
            Self::Low => write!(f, "low"),
            Self::Orchestrator => write!(f, "orchestrator"),
            Self::High => write!(f, "high"),
        }
    }
}

/// A raw representation of an origin.
///
/// Metrics contain metadata about their origin, in terms of the metric's _reason_ for existing: the metric was ingested
/// via DogStatsD, or was generated by an integration, and so on. However, there is also the concept of a metric
/// originating from a particular _entity_, such as a specific Kubernetes container. This relates directly to the
/// specific sender of the metric, which is used to enrich the metric with additional tags describing the origin entity.
///
/// The origin entity will generally be the process ID of the metric sender, or the container ID, both of which are then
/// generally mapped to the relevant information for the metric, such as the orchestrator-level tags for the
/// container/pod/deployment.
#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
pub struct RawOrigin<'a> {
    /// Process ID of the sender.
    process_id: Option<NonZeroU32>,

    /// Local Data of the sender.
    ///
    /// This will typically be either the container ID, or the inode of the container's cgroups controller, or both. It
    /// may or may not have a special prefix that indicates which of the two it's.
    local_data: Option<&'a str>,

    /// Pod UID of the sender.
    ///
    /// This is generally only used in Kubernetes environments to uniquely identify the pod. UIDs are equivalent to UUIDs.
    pod_uid: Option<&'a str>,

    /// Desired cardinality of any tags associated with the entity.
    ///
    /// This controls the cardinality of the tags added to this metric when enriching based on the available entity IDs.
    cardinality: Option<OriginTagCardinality>,

    /// External Data of the sender.
    ///
    /// See [`ExternalData`] for more information.
    external_data: Option<RawExternalData<'a>>,
}

impl<'a> RawOrigin<'a> {
    /// Returns `true` if the origin information is empty.
    pub fn is_empty(&self) -> bool {
        self.process_id.is_none()
            && self.local_data.is_none()
            && self.pod_uid.is_none()
            && self.cardinality.is_none()
            && self.external_data.is_none()
    }

    /// Unsets the process ID of the sender.
    pub fn clear_process_id(&mut self) {
        self.process_id = None;
    }

    /// Sets the process ID of the sender.
    ///
    /// Must be a non-zero value. If the value is zero, it's silently ignored.
    pub fn set_process_id(&mut self, process_id: u32) {
        self.process_id = NonZeroU32::new(process_id);
    }

    /// Returns the process ID of the sender.
    pub fn process_id(&self) -> Option<u32> {
        self.process_id.map(NonZeroU32::get)
    }

    /// Sets the Local Data of the sender.
    pub fn set_local_data(&mut self, local_data: impl Into<Option<&'a str>>) {
        self.local_data = local_data.into();
    }

    /// Returns the Local Data of the sender.
    pub fn local_data(&self) -> Option<&str> {
        self.local_data
    }

    /// Sets the pod UID of the sender.
    pub fn set_pod_uid(&mut self, pod_uid: impl Into<Option<&'a str>>) {
        self.pod_uid = pod_uid.into();
    }

    /// Returns the pod UID of the sender.
    pub fn pod_uid(&self) -> Option<&str> {
        self.pod_uid
    }

    /// Sets the desired cardinality of any tags associated with the entity.
    pub fn set_cardinality(&mut self, cardinality: impl Into<Option<OriginTagCardinality>>) {
        self.cardinality = cardinality.into();
    }

    /// Returns the desired cardinality of any tags associated with the entity.
    pub fn cardinality(&self) -> Option<OriginTagCardinality> {
        self.cardinality.as_ref().copied()
    }

    /// Sets the External Data of the sender.
    pub fn set_external_data(&mut self, external_data: impl Into<Option<&'a str>>) {
        self.external_data = external_data.into().and_then(RawExternalData::try_from_str);
    }

    /// Returns the External Data of the sender.
    pub fn external_data(&self) -> Option<&RawExternalData<'a>> {
        self.external_data.as_ref()
    }
}

impl fmt::Display for RawOrigin<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut has_written = false;

        write!(f, "RawOrigin(")?;

        if let Some(process_id) = self.process_id {
            write!(f, "process_id={}", process_id)?;
        }

        if let Some(local_data) = self.local_data {
            if has_written {
                write!(f, " ")?;
            } else {
                has_written = true;
            }
            write!(f, "local_data={}", local_data)?;
        }

        if let Some(pod_uid) = self.pod_uid {
            if has_written {
                write!(f, " ")?;
            } else {
                has_written = true;
            }
            write!(f, "pod_uid={}", pod_uid)?;
        }

        if let Some(cardinality) = self.cardinality {
            if has_written {
                write!(f, " ")?;
            } else {
                has_written = true;
            }
            write!(f, "cardinality={}", cardinality)?;
        }

        if let Some(external_data) = self.external_data.as_ref() {
            if has_written {
                write!(f, " ")?;
            }
            write!(f, "external_data={}", external_data)?;
        }

        write!(f, ")")
    }
}

/// A resolver for mapping origins to their associated tags.
pub trait OriginTagsResolver: Send + Sync {
    /// Resolves the origin tags for the given raw origin.
    fn resolve_origin_tags(&self, origin: RawOrigin<'_>) -> SharedTagSet;
}

impl<T> OriginTagsResolver for Arc<T>
where
    T: OriginTagsResolver,
{
    fn resolve_origin_tags(&self, origin: RawOrigin<'_>) -> SharedTagSet {
        (**self).resolve_origin_tags(origin)
    }
}

impl<T> OriginTagsResolver for Option<T>
where
    T: OriginTagsResolver,
{
    fn resolve_origin_tags(&self, origin: RawOrigin<'_>) -> SharedTagSet {
        self.as_ref()
            .map(|resolver| resolver.resolve_origin_tags(origin))
            .unwrap_or_else(|| SharedTagSet::default())
    }
}

/// External Data associated with an origin.
///
/// "External Data" is a concept that's used to aid origin detection of workloads running in Kubernetes environments
/// where introspection isn't possible or may return incorrect information. Origin detection generally centers around
/// determining the container where a metric originates from, and then enriching the metric with tags that describe that
/// container, as well as the pod the container is running within, and so on. In some cases, the origin of a metric
/// can't be detected from the outside (such as by using peer credentials over Unix Domain sockets) and can't be
/// detected by the workload itself (such as when running in nested virtualization environments). In these cases, we
/// need a mechanism that allows passing the necessary information to the client, who then passes it on to us, so that
/// we can correctly resolve the origin.
///
/// "External Data" supports this by allowing for an external Kubernetes admission controller to attach specific
/// metadata -- pod UID and container name -- to application pods, which is then read and sent along with metrics. This
/// information is then used during origin detection in order to correlate the container ID of the origin, which is
/// sufficient to allow enriching the metric with container-specific tags.
///
/// # Format
///
/// An External Data string is a comma-separated list of key/value pairs, where each key represents a particular aspect
/// of the workload entity. The following keys are supported:
///
/// - `it-<true/false>`: A boolean value indicating whether the entity is an init container.
/// - `pu-<pod_uid>`: The pod UID associated with the entity.
/// - `cn-<container_name>`: The container name associated with the entity.
///
/// For parsing external data strings without allocating, see [`RawExternalData`].
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct ExternalData {
    pod_uid: MetaString,
    container_name: MetaString,
    init_container: bool,
}

impl ExternalData {
    /// Creates a new `ExternalData` instance.
    pub fn new(pod_uid: MetaString, container_name: MetaString, init_container: bool) -> Self {
        Self {
            pod_uid,
            container_name,
            init_container,
        }
    }

    /// Returns the pod UID.
    pub fn pod_uid(&self) -> &MetaString {
        &self.pod_uid
    }

    /// Returns the container name.
    pub fn container_name(&self) -> &MetaString {
        &self.container_name
    }

    /// Returns `true` if the entity is an init container.
    pub fn is_init_container(&self) -> bool {
        self.init_container
    }
}

impl std::hash::Hash for ExternalData {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        (*self.pod_uid).hash(state);
        (*self.container_name).hash(state);
        self.init_container.hash(state);
    }
}

/// A borrowed representation of [`ExternalData`].
///
/// This can be used to parse external data strings without needing to allocate backing storage for any of the fields,
/// and can be used to look up map entries (such as when using `HashMap`) when the key is [`ExternalData`].
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RawExternalData<'a> {
    pod_uid: &'a str,
    container_name: &'a str,
    init_container: bool,
}

impl<'a> RawExternalData<'a> {
    /// Creates a new `RawExternalData` from a raw string.
    ///
    /// If the external data isn't valid, `None` is returned.
    pub fn try_from_str(raw: &'a str) -> Option<Self> {
        if raw.is_empty() {
            return None;
        }

        let mut data = Self {
            pod_uid: "",
            container_name: "",
            init_container: false,
        };

        let parts = raw.split(',');
        for part in parts {
            if part.len() < 4 {
                // All key/value pairs have a prefix of `xx-` where `xx` is some short code, so we basically can't have
                // any real key/value pair that's less than four characters overall.
                warn!("Parsed external data with invalid key/value pair: {}", part);
                continue;
            }

            let key = &part[0..3];
            let value = &part[3..];

            match key {
                "it-" => data.init_container = value.parse().unwrap_or(false),
                "pu-" => data.pod_uid = value,
                "cn-" => data.container_name = value,
                _ => {
                    // Unknown key, ignore.
                    warn!("Parsed external data with unknown key: {}", key);
                }
            }
        }

        Some(data)
    }
}

impl std::hash::Hash for RawExternalData<'_> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.pod_uid.hash(state);
        self.container_name.hash(state);
        self.init_container.hash(state);
    }
}

impl fmt::Display for RawExternalData<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "pu-{},cn-{},it-{}",
            self.pod_uid, self.container_name, self.init_container
        )
    }
}

impl Equivalent<ExternalData> for RawExternalData<'_> {
    fn equivalent(&self, other: &ExternalData) -> bool {
        self.pod_uid == &*other.pod_uid
            && self.container_name == &*other.container_name
            && self.init_container == other.init_container
    }
}

#[cfg(test)]
mod tests {
    use std::{
        collections::hash_map::DefaultHasher,
        hash::{Hash as _, Hasher as _},
    };

    use proptest::prelude::*;

    use super::*;

    proptest! {
        #[test]
        fn property_test_identical_hash_impls(pod_uid in "[a-z0-9]{1,64}", container_name in "[a-z0-9]{1,64}", init_container in any::<bool>()) {
            let external_data = ExternalData::new(pod_uid.clone().into(), container_name.clone().into(), init_container);
            let external_data_ref = RawExternalData {
                pod_uid: &pod_uid,
                container_name: &container_name,
                init_container,
            };

            let mut hasher = DefaultHasher::new();
            external_data.hash(&mut hasher);
            let external_data_hash = hasher.finish();

            let mut hasher = DefaultHasher::new();
            external_data_ref.hash(&mut hasher);
            let external_data_ref_hash = hasher.finish();

            assert_eq!(external_data_hash, external_data_ref_hash);
        }
    }

    /// Parses `raw` and projects the result into a comparable tuple of its fields.
    fn parsed(raw: &str) -> Option<(&str, &str, bool)> {
        RawExternalData::try_from_str(raw).map(|data| (data.pod_uid, data.container_name, data.init_container))
    }

    #[test]
    fn try_from_str_parses_valid_and_malformed_external_data() {
        let cases: &[(&str, Option<(&str, &str, bool)>)] = &[
            // Empty input is the only case that yields no data at all.
            ("", None),
            // All three keys, in any order, populate their respective fields.
            ("pu-podid,cn-web,it-true", Some(("podid", "web", true))),
            ("it-false,cn-web,pu-podid", Some(("podid", "web", false))),
            // A single key populates just that field; the others keep their empty/false defaults.
            ("pu-podid", Some(("podid", "", false))),
            ("cn-web", Some(("", "web", false))),
            // The `it-` value is parsed as a bool; any non-bool value falls back to false.
            ("it-true", Some(("", "", true))),
            ("it-notabool", Some(("", "", false))),
            // Parts shorter than four characters (no room for an `xx-` prefix plus a value) are
            // skipped, but a non-empty input still parses to defaults rather than `None`.
            ("x", Some(("", "", false))),
            ("abc", Some(("", "", false))),
            // Unknown keys are ignored.
            ("zz-whatever", Some(("", "", false))),
            // Short/garbage parts are skipped while valid parts are still parsed.
            ("pu-podid,z,cn-web", Some(("podid", "web", false))),
            // Duplicate keys: the last occurrence wins.
            ("pu-first,pu-second", Some(("second", "", false))),
        ];

        for (raw, expected) in cases {
            assert_eq!(parsed(raw), *expected, "input: {raw:?}");
        }
    }

    #[test]
    fn origin_tag_cardinality_displays_lowercase_name() {
        let cases = [
            (OriginTagCardinality::None, "none"),
            (OriginTagCardinality::Low, "low"),
            (OriginTagCardinality::Orchestrator, "orchestrator"),
            (OriginTagCardinality::High, "high"),
        ];

        for (cardinality, expected) in cases {
            assert_eq!(cardinality.to_string(), expected);
        }
    }

    #[test]
    fn raw_external_data_display_lists_all_fields() {
        let data = RawExternalData::try_from_str("pu-podid,cn-web,it-true").unwrap();
        assert_eq!(data.to_string(), "pu-podid,cn-web,it-true");
    }

    #[test]
    fn raw_external_data_display_round_trips_through_parser() {
        // Display always emits the canonical `pu-,cn-,it-` form, which the parser accepts verbatim,
        // regardless of the field order in the original input.
        let original = RawExternalData::try_from_str("cn-web,pu-podid,it-true").unwrap();
        let rendered = original.to_string();
        let reparsed = RawExternalData::try_from_str(&rendered).unwrap();
        assert_eq!(reparsed, original);
    }

    #[test]
    fn raw_origin_display_empty() {
        assert_eq!(RawOrigin::default().to_string(), "RawOrigin()");
    }

    #[test]
    fn raw_origin_display_individual_fields() {
        let mut process_id = RawOrigin::default();
        process_id.set_process_id(1234);
        assert_eq!(process_id.to_string(), "RawOrigin(process_id=1234)");

        let mut local_data = RawOrigin::default();
        local_data.set_local_data("abc");
        assert_eq!(local_data.to_string(), "RawOrigin(local_data=abc)");

        let mut pod_uid = RawOrigin::default();
        pod_uid.set_pod_uid("pod-1");
        assert_eq!(pod_uid.to_string(), "RawOrigin(pod_uid=pod-1)");

        let mut cardinality = RawOrigin::default();
        cardinality.set_cardinality(OriginTagCardinality::Low);
        assert_eq!(cardinality.to_string(), "RawOrigin(cardinality=low)");

        let mut external_data = RawOrigin::default();
        external_data.set_external_data("pu-podid,cn-web,it-true");
        assert_eq!(
            external_data.to_string(),
            "RawOrigin(external_data=pu-podid,cn-web,it-true)"
        );
    }

    #[test]
    fn raw_origin_display_combines_present_fields() {
        let mut origin = RawOrigin::default();
        origin.set_process_id(1234);
        origin.set_local_data("abc");
        origin.set_pod_uid("pod-1");
        origin.set_cardinality(OriginTagCardinality::Low);

        // NOTE: the `process_id` branch does not set the internal `has_written` flag, so the first
        // separator space is dropped and `process_id` runs directly into `local_data`. This asserts
        // the current (buggy) spacing so a future fix has to update the expectation deliberately.
        assert_eq!(
            origin.to_string(),
            "RawOrigin(process_id=1234local_data=abc pod_uid=pod-1 cardinality=low)"
        );
    }
}