saluki_components/sources/dogstatsd/
origin.rs

1use std::sync::Arc;
2
3use saluki_context::{
4    origin::{OriginTagCardinality, OriginTagsResolver, RawOrigin},
5    tags::SharedTagSet,
6};
7use saluki_env::{
8    workload::{origin::ResolvedOrigin, EntityId},
9    WorkloadProvider,
10};
11use saluki_io::deser::codec::dogstatsd::{EventPacket, MetricPacket, ServiceCheckPacket};
12use tracing::trace;
13
14use super::{replay::CapturedTaggerHandle, tags::WellKnownTags};
15
16const REPLAY_PROCESS_ID_MARKER: u32 = 1u32 << 31;
17
18#[derive(Clone, Debug, Eq, PartialEq)]
19pub(super) enum ProcessOrigin {
20    Pinned(Option<EntityId>),
21    Unpinned(u32),
22    Replay(u32),
23}
24
25impl ProcessOrigin {
26    pub(super) fn container_entity_id(&self) -> Option<&EntityId> {
27        match self {
28            Self::Pinned(entity_id) => entity_id.as_ref(),
29            Self::Unpinned(_) | Self::Replay(_) => None,
30        }
31    }
32}
33
34pub(super) fn mark_replay_process_id(process_id: u32) -> u32 {
35    process_id | REPLAY_PROCESS_ID_MARKER
36}
37
38fn captured_process_id_from_replay(process_id: u32) -> Option<u32> {
39    if process_id & REPLAY_PROCESS_ID_MARKER != 0 {
40        Some(process_id & !REPLAY_PROCESS_ID_MARKER)
41    } else {
42        None
43    }
44}
45
46/// Origin enrichment configuration.
47///
48/// Origin enrichment controls the when and how of enriching metrics ingested via DogStatsD based on various sources of
49/// "origin" information, such as specific metric tags or UDS socket credentials. Enrichment involves adding additional
50/// metric tags that describe the origin of the metric, such as the Kubernetes pod or container.
51#[derive(Clone)]
52#[cfg_attr(test, derive(Debug, PartialEq))]
53pub struct OriginEnrichmentConfiguration {
54    /// Whether or not to enable origin detection.
55    ///
56    /// If disabled, no origin tags will be added to events even if the origin information is detected.
57    pub enabled: bool,
58
59    /// Whether or not a client-provided entity ID should take precedence over automatically detected origin metadata.
60    ///
61    /// When a client-provided entity ID is specified, and an origin process ID has automatically been detected, setting
62    /// this to `true` will cause the origin process ID to be ignored. This only applies when
63    /// `origin_detection_unified` is `false`.
64    pub entity_id_precedence: bool,
65
66    /// The cardinality of tags to enrich metrics with, when the payload doesn't specify one itself.
67    pub tag_cardinality: OriginTagCardinality,
68
69    /// Whether or not to use the unified origin detection behavior.
70    ///
71    /// When set to `true`, all detected entity IDs -- UDS Origin Detection, `dd.internal.entity_id`, container ID from
72    /// DogStatsD payload -- will be used for querying tags to enrich with. When set to `false`, the original precedence
73    /// behavior will be used, which enriches with the entity ID detected via Origin Detection first, and then
74    /// potentially again with either the client-provided entity ID (`dd.internal.entity_id`) or the container ID from
75    /// the DogStatsD payload, with the client-provided entity ID taking precedence. An entity ID detected via Origin
76    /// Detection is only used when no client-provided entity ID was present, or when `entity_id_precedence` is
77    /// `false`.
78    pub origin_detection_unified: bool,
79
80    /// Whether or not to opt out of origin detection for DogStatsD metrics.
81    ///
82    /// When set to `true`, and the metric explicitly denotes a cardinality of `"none"`, origin enrichment will be
83    /// skipped. This is only applicable to DogStatsD metrics when unified origin detection behavior isn't enabled.
84    pub origin_detection_optout: bool,
85
86    /// Whether or not to parse client-provided origin fields from DogStatsD payloads.
87    ///
88    /// When enabled, the `c:` (Local Data), `e:` (External Data), and `card:` (Cardinality) protocol fields are
89    /// parsed and used for origin enrichment.
90    pub origin_detection_client: bool,
91}
92
93#[cfg(test)]
94impl OriginEnrichmentConfiguration {
95    /// Creates a fixture configuration with origin detection disabled, for tests that exercise enrichment behavior
96    /// rather than configuration.
97    pub(super) fn for_test() -> Self {
98        Self {
99            enabled: false,
100            entity_id_precedence: false,
101            tag_cardinality: OriginTagCardinality::Low,
102            origin_detection_unified: false,
103            origin_detection_optout: true,
104            origin_detection_client: false,
105        }
106    }
107}
108
109#[derive(Clone)]
110pub(super) struct DogStatsDOriginTagResolver {
111    config: OriginEnrichmentConfiguration,
112    workload_provider: Arc<dyn WorkloadProvider + Send + Sync>,
113    captured_tagger: CapturedTaggerHandle,
114}
115
116impl DogStatsDOriginTagResolver {
117    pub fn new(
118        config: OriginEnrichmentConfiguration, workload_provider: Arc<dyn WorkloadProvider + Send + Sync>,
119        captured_tagger: CapturedTaggerHandle,
120    ) -> Self {
121        Self {
122            config,
123            workload_provider,
124            captured_tagger,
125        }
126    }
127
128    fn collect_origin_tags_with_process_entity(
129        &self, origin: &ResolvedOrigin, process_entity_id: Option<&EntityId>,
130    ) -> SharedTagSet {
131        let mut collected_tags = SharedTagSet::default();
132
133        if !self.config.enabled {
134            return collected_tags;
135        }
136
137        // Examine the various possible entity ID values, and based on their state, use one or more of them to grab any
138        // enriched tags attached to the entities. We evalulate a number of possible entity IDs:
139        //
140        // - sender entity (resolved from UDS socket credentials when the packet is received)
141        // - Local Data-based container ID (extracted from special "container ID" extension in DogStatsD protocol; also known as Local Data)
142        // - Local Data-based pod UID (extracted from `dd.internal.entity_id` tag)
143        // - External Data-based container ID (derived from pod UID and container name in External Data)
144        // - External Data-based pod UID (raw pod UID from External Data)
145        let maybe_process_entity_id = process_entity_id;
146        let maybe_local_container_id = origin.local_data();
147        let maybe_local_pod_uid = origin.pod_uid();
148        let maybe_external_container_id = origin.resolved_external_data().map(|red| red.container_entity_id());
149        let maybe_external_pod_uid = origin.resolved_external_data().map(|red| red.pod_entity_id());
150
151        let tag_cardinality = origin.cardinality().unwrap_or(self.config.tag_cardinality);
152
153        if !self.config.origin_detection_unified {
154            if self.config.origin_detection_optout && tag_cardinality == OriginTagCardinality::None {
155                trace!("Skipping origin enrichment for DogStatsD metric with cardinality 'none'.");
156                return collected_tags;
157            }
158
159            // If we discovered an entity ID from socket credentials, and no client-provided entity ID was provided (or
160            // it was, but entity ID precedence is disabled), then try to get tags for the detected entity ID.
161            if let Some(entity_id) = maybe_process_entity_id {
162                if maybe_local_pod_uid.is_none() || !self.config.entity_id_precedence {
163                    if let Some(tags) = self.workload_provider.get_tags_for_entity(entity_id, tag_cardinality) {
164                        collected_tags.extend_from_shared(&tags);
165                    } else {
166                        trace!(
167                            ?entity_id,
168                            cardinality = tag_cardinality.as_str(),
169                            "No tags found for entity."
170                        );
171                    }
172                }
173            }
174
175            // If we have a client-provided entity ID, try to get tags for the entity based on those. A client-provided
176            // entity ID takes precedence over the container ID.
177            let maybe_entity_id = maybe_local_pod_uid.or(maybe_local_container_id);
178            if let Some(entity_id) = maybe_entity_id {
179                if let Some(tags) = self.workload_provider.get_tags_for_entity(entity_id, tag_cardinality) {
180                    collected_tags.extend_from_shared(&tags);
181                } else {
182                    trace!(
183                        ?entity_id,
184                        cardinality = tag_cardinality.as_str(),
185                        "No tags found for entity."
186                    );
187                }
188            }
189        } else {
190            if tag_cardinality == OriginTagCardinality::None {
191                trace!("Skipping origin enrichment for metric with cardinality 'none'.");
192                return collected_tags;
193            }
194
195            // Evaluate all available entity IDs in order of priority: Local Data-based container ID, sender entity,
196            // External Data-based container ID, Local Data-based pod UID, and External Data-based pod UID.
197            //
198            // As soon as the first set of tags for an entity ID is found, we skip the remaining entity IDs.
199            let maybe_entity_ids = &[
200                maybe_local_container_id,
201                maybe_process_entity_id,
202                maybe_external_container_id,
203                maybe_local_pod_uid,
204                maybe_external_pod_uid,
205            ];
206            for entity_id in maybe_entity_ids.iter().flatten() {
207                if let Some(tags) = self.workload_provider.get_tags_for_entity(entity_id, tag_cardinality) {
208                    if !tags.is_empty() {
209                        collected_tags.extend_from_shared(&tags);
210                        break;
211                    }
212                } else {
213                    trace!(
214                        ?entity_id,
215                        cardinality = tag_cardinality.as_str(),
216                        "No tags found for entity."
217                    );
218                }
219            }
220        }
221
222        collected_tags
223    }
224
225    fn collect_origin_tags(&self, origin: &ResolvedOrigin) -> SharedTagSet {
226        self.collect_origin_tags_with_process_entity(origin, origin.process_id())
227    }
228
229    pub(super) fn resolve_origin_tags_with_process_origin(
230        &self, mut origin: RawOrigin<'_>, process_origin: Option<&ProcessOrigin>,
231    ) -> SharedTagSet {
232        match process_origin {
233            Some(ProcessOrigin::Replay(process_id)) => {
234                origin.set_process_id(mark_replay_process_id(*process_id));
235                self.resolve_origin_tags(origin)
236            }
237            Some(ProcessOrigin::Unpinned(process_id)) => {
238                origin.set_process_id(*process_id);
239                self.resolve_origin_tags(origin)
240            }
241            Some(ProcessOrigin::Pinned(process_entity_id)) => {
242                let resolved_origin = self
243                    .workload_provider
244                    .get_resolved_origin(origin.clone())
245                    .unwrap_or_else(|| ResolvedOrigin::from_parts(origin.cardinality(), None, None, None, None));
246                self.collect_origin_tags_with_process_entity(&resolved_origin, process_entity_id.as_ref())
247            }
248            None => self.resolve_origin_tags(origin),
249        }
250    }
251}
252
253impl OriginTagsResolver for DogStatsDOriginTagResolver {
254    fn resolve_origin_tags(&self, origin: RawOrigin<'_>) -> SharedTagSet {
255        // Replay traffic is tagged by setting the marker bit on the origin process ID. It bypasses the live enrichment
256        // pipeline entirely: the captured `TaggerState` already contains fully-resolved tags per entity, so the
257        // resolver's entity-ID-walking logic has nothing to add.
258        if let Some(captured_process_id) = origin.process_id().and_then(captured_process_id_from_replay) {
259            if let Some(store) = self.captured_tagger.current() {
260                let cardinality = origin.cardinality().unwrap_or(self.config.tag_cardinality);
261                return store
262                    .lookup(captured_process_id as i32, cardinality)
263                    .unwrap_or_default();
264            }
265            trace!(?origin, "Replay-flagged origin but no captured tagger available.");
266            return SharedTagSet::default();
267        }
268
269        match self.workload_provider.get_resolved_origin(origin.clone()) {
270            Some(resolved_origin) => self.collect_origin_tags(&resolved_origin),
271            None => {
272                trace!(?origin, "No resolved origin found for origin.");
273                SharedTagSet::default()
274            }
275        }
276    }
277}
278
279/// Builds an `RawOrigin` object from the given metric packet.
280pub fn origin_from_metric_packet<'packet, 'tags>(
281    packet: &'tags MetricPacket<'packet>, well_known_tags: &'tags WellKnownTags<'tags>,
282) -> RawOrigin<'tags>
283where
284    'packet: 'tags,
285{
286    let cardinality = packet.cardinality.or(well_known_tags.cardinality);
287
288    let mut origin = RawOrigin::default();
289    origin.set_pod_uid(well_known_tags.pod_uid);
290    origin.set_local_data(packet.local_data);
291    origin.set_external_data(packet.external_data);
292    origin.set_cardinality(cardinality);
293    origin
294}
295
296/// Builds an `RawOrigin` object from the given event packet.
297pub fn origin_from_event_packet<'packet, 'tags>(
298    packet: &'tags EventPacket<'packet>, well_known_tags: &'tags WellKnownTags<'tags>,
299) -> RawOrigin<'tags>
300where
301    'packet: 'tags,
302{
303    let cardinality = packet.cardinality.or(well_known_tags.cardinality);
304
305    let mut origin = RawOrigin::default();
306    origin.set_pod_uid(well_known_tags.pod_uid);
307    origin.set_local_data(packet.local_data);
308    origin.set_external_data(packet.external_data);
309    origin.set_cardinality(cardinality);
310    origin
311}
312
313/// Builds an `RawOrigin` object from the given service check packet.
314pub fn origin_from_service_check_packet<'packet, 'tags>(
315    packet: &'tags ServiceCheckPacket<'packet>, well_known_tags: &'tags WellKnownTags<'tags>,
316) -> RawOrigin<'tags>
317where
318    'packet: 'tags,
319{
320    let cardinality = packet.cardinality.or(well_known_tags.cardinality);
321
322    let mut origin = RawOrigin::default();
323    origin.set_pod_uid(well_known_tags.pod_uid);
324    origin.set_local_data(packet.local_data);
325    origin.set_external_data(packet.external_data);
326    origin.set_cardinality(cardinality);
327    origin
328}
329
330#[cfg(test)]
331mod tests {
332    use std::collections::HashMap;
333
334    use saluki_context::tags::{RawTags, TagSet};
335    use saluki_core::data_model::event::{metric::MetricValues, service_check::CheckStatus};
336    use saluki_env::workload::{origin::ResolvedExternalData, providers::TestWorkloadProvider, EntityId};
337    use stringtheory::MetaString;
338
339    use super::*;
340
341    static EID_PID: EntityId = EntityId::ContainerPid(12345);
342    static EID_LOCAL_CID: EntityId = EntityId::Container(MetaString::from_static("local-cid"));
343    static EID_EXTERNAL_CID_VALID: EntityId = EntityId::Container(MetaString::from_static("external-cid"));
344    static EID_EXTERNAL_CID_INVALID: EntityId = EntityId::Container(MetaString::from_static("invalid-external-cid"));
345    static EID_LOCAL_POD: EntityId = EntityId::PodUid(MetaString::from_static("local-pod-uid"));
346    static EID_EXTERNAL_POD: EntityId = EntityId::PodUid(MetaString::from_static("external-pod-uid"));
347
348    fn single_tag(tag: &str) -> SharedTagSet {
349        let mut tag_set = TagSet::default();
350        tag_set.insert_tag(tag);
351        tag_set.into_shared()
352    }
353
354    fn tags_for_entity(entity_id: &EntityId) -> SharedTagSet {
355        if entity_id == &EID_PID {
356            single_tag("tag_source:pid")
357        } else if entity_id == &EID_LOCAL_CID {
358            single_tag("tag_source:local-cid")
359        } else if entity_id == &EID_EXTERNAL_CID_VALID {
360            single_tag("tag_source:external-cid")
361        } else if entity_id == &EID_LOCAL_POD {
362            single_tag("tag_source:local-pod")
363        } else if entity_id == &EID_EXTERNAL_POD {
364            single_tag("tag_source:external-pod")
365        } else {
366            SharedTagSet::default()
367        }
368    }
369
370    fn origin(
371        maybe_process_id: Option<&EntityId>, maybe_local_container_id: Option<&EntityId>,
372        maybe_local_pod_uid: Option<&EntityId>, maybe_external_data: Option<&ResolvedExternalData>,
373    ) -> ResolvedOrigin {
374        ResolvedOrigin::from_parts(
375            None,
376            maybe_process_id.cloned(),
377            maybe_local_container_id.cloned(),
378            maybe_local_pod_uid.cloned(),
379            maybe_external_data.cloned(),
380        )
381    }
382
383    fn build_tags_resolver_with_default_tags(config: OriginEnrichmentConfiguration) -> DogStatsDOriginTagResolver {
384        let mut workload_provider = TestWorkloadProvider::new();
385        workload_provider.add_entity_shared_tags(EID_PID.clone(), tags_for_entity(&EID_PID));
386        workload_provider.add_entity_shared_tags(EID_LOCAL_CID.clone(), tags_for_entity(&EID_LOCAL_CID));
387        workload_provider
388            .add_entity_shared_tags(EID_EXTERNAL_CID_VALID.clone(), tags_for_entity(&EID_EXTERNAL_CID_VALID));
389        workload_provider.add_entity_shared_tags(EID_LOCAL_POD.clone(), tags_for_entity(&EID_LOCAL_POD));
390        workload_provider.add_entity_shared_tags(EID_EXTERNAL_POD.clone(), tags_for_entity(&EID_EXTERNAL_POD));
391
392        let erased_workload_provider = Arc::new(workload_provider);
393
394        DogStatsDOriginTagResolver::new(config, erased_workload_provider, CapturedTaggerHandle::new())
395    }
396
397    #[test]
398    fn metric_cardinality_precedence() {
399        // Tests that the cardinality specified in a metric packet (`|card:high`, etc) takes precedence over the cardinality
400        // specified via the deprecated `dd.internal.card` tag.
401        let raw_tags_input = "dd.internal.card:high";
402        let raw_tags = RawTags::new(raw_tags_input, usize::MAX, usize::MAX);
403
404        let well_known_tags = WellKnownTags::from_raw_tags(&raw_tags);
405        assert_eq!(well_known_tags.cardinality, Some(OriginTagCardinality::High));
406
407        let packet_with_card = MetricPacket {
408            metric_name: "test_metric",
409            tags: raw_tags.clone(),
410            values: MetricValues::counter(1.0),
411            num_points: 1,
412            timestamp: None,
413            local_data: None,
414            external_data: None,
415            cardinality: Some(OriginTagCardinality::Low),
416            unit: None,
417        };
418
419        let packet_without_card = MetricPacket {
420            metric_name: "test_metric",
421            tags: raw_tags.clone(),
422            values: MetricValues::counter(1.0),
423            num_points: 1,
424            timestamp: None,
425            local_data: None,
426            external_data: None,
427            cardinality: None,
428            unit: None,
429        };
430
431        let with_card_origin = origin_from_metric_packet(&packet_with_card, &well_known_tags);
432        assert_ne!(packet_with_card.cardinality, well_known_tags.cardinality);
433        assert_eq!(with_card_origin.cardinality(), packet_with_card.cardinality);
434
435        let without_card_origin = origin_from_metric_packet(&packet_without_card, &well_known_tags);
436        assert_ne!(packet_without_card.cardinality, well_known_tags.cardinality);
437        assert_eq!(without_card_origin.cardinality(), well_known_tags.cardinality);
438    }
439
440    #[test]
441    fn event_cardinality_precedence() {
442        // Tests that the cardinality specified in an event packet (`|card:high`, etc) takes precedence over the cardinality
443        // specified via the deprecated `dd.internal.card` tag.
444        let raw_tags_input = "dd.internal.card:low";
445        let raw_tags = RawTags::new(raw_tags_input, usize::MAX, usize::MAX);
446
447        let well_known_tags = WellKnownTags::from_raw_tags(&raw_tags);
448        assert_eq!(well_known_tags.cardinality, Some(OriginTagCardinality::Low));
449
450        let packet_with_card = EventPacket {
451            title: MetaString::empty(),
452            text: MetaString::empty(),
453            timestamp: None,
454            hostname: None,
455            aggregation_key: None,
456            priority: None,
457            alert_type: None,
458            source_type_name: None,
459            tags: raw_tags.clone(),
460            local_data: None,
461            external_data: None,
462            cardinality: Some(OriginTagCardinality::Orchestrator),
463        };
464
465        let packet_without_card = EventPacket {
466            title: MetaString::empty(),
467            text: MetaString::empty(),
468            timestamp: None,
469            hostname: None,
470            aggregation_key: None,
471            priority: None,
472            alert_type: None,
473            source_type_name: None,
474            tags: raw_tags.clone(),
475            local_data: None,
476            external_data: None,
477            cardinality: None,
478        };
479
480        let with_card_origin = origin_from_event_packet(&packet_with_card, &well_known_tags);
481        assert_ne!(packet_with_card.cardinality, well_known_tags.cardinality);
482        assert_eq!(with_card_origin.cardinality(), packet_with_card.cardinality);
483
484        let without_card_origin = origin_from_event_packet(&packet_without_card, &well_known_tags);
485        assert_ne!(packet_without_card.cardinality, well_known_tags.cardinality);
486        assert_eq!(without_card_origin.cardinality(), well_known_tags.cardinality);
487    }
488
489    #[test]
490    fn service_check_cardinality_precedence() {
491        // Tests that the cardinality specified in an event packet (`|card:high`, etc) takes precedence over the cardinality
492        // specified via the deprecated `dd.internal.card` tag.
493        let raw_tags_input = "dd.internal.card:orchestrator";
494        let raw_tags = RawTags::new(raw_tags_input, usize::MAX, usize::MAX);
495
496        let well_known_tags = WellKnownTags::from_raw_tags(&raw_tags);
497        assert_eq!(well_known_tags.cardinality, Some(OriginTagCardinality::Orchestrator));
498
499        let packet_with_card = ServiceCheckPacket {
500            name: MetaString::empty(),
501            status: CheckStatus::Ok,
502            timestamp: None,
503            hostname: None,
504            message: None,
505            tags: raw_tags.clone(),
506            local_data: None,
507            external_data: None,
508            cardinality: Some(OriginTagCardinality::Low),
509        };
510
511        let packet_without_card = ServiceCheckPacket {
512            name: MetaString::empty(),
513            status: CheckStatus::Ok,
514            timestamp: None,
515            hostname: None,
516            message: None,
517            tags: raw_tags.clone(),
518            local_data: None,
519            external_data: None,
520            cardinality: None,
521        };
522
523        let with_card_origin = origin_from_service_check_packet(&packet_with_card, &well_known_tags);
524        assert_ne!(packet_with_card.cardinality, well_known_tags.cardinality);
525        assert_eq!(with_card_origin.cardinality(), packet_with_card.cardinality);
526
527        let without_card_origin = origin_from_service_check_packet(&packet_without_card, &well_known_tags);
528        assert_ne!(packet_without_card.cardinality, well_known_tags.cardinality);
529        assert_eq!(without_card_origin.cardinality(), well_known_tags.cardinality);
530    }
531
532    #[test]
533    fn origin_detection_legacy_precedence() {
534        let mut pid_plus_local_pod_tags = tags_for_entity(&EID_PID);
535        pid_plus_local_pod_tags.extend_from_shared(&tags_for_entity(&EID_LOCAL_POD));
536
537        let mut pid_plus_local_cid_tags = tags_for_entity(&EID_PID);
538        pid_plus_local_cid_tags.extend_from_shared(&tags_for_entity(&EID_LOCAL_CID));
539
540        let cases = [
541            // We only have the process ID, so entity ID precedence should be irrelevant.
542            (
543                false,
544                origin(Some(&EID_PID), None, None, None),
545                tags_for_entity(&EID_PID),
546            ),
547            (
548                true,
549                origin(Some(&EID_PID), None, None, None),
550                tags_for_entity(&EID_PID),
551            ),
552            // We have both the process ID and local pod UID, but entity ID precedence is disabled, so we get the
553            // process ID and local pod UID tags.
554            (
555                false,
556                origin(Some(&EID_PID), None, Some(&EID_LOCAL_POD), None),
557                pid_plus_local_pod_tags.clone(),
558            ),
559            // We have both the process ID and local pod UID, but entity ID precedence is enabled, so we should only get
560            // the local pod UID tags.
561            (
562                true,
563                origin(Some(&EID_PID), None, Some(&EID_LOCAL_POD), None),
564                tags_for_entity(&EID_LOCAL_POD),
565            ),
566            // We have the process ID, local container ID, and local pod UID, but entity ID precedence is disabled, so
567            // we should get the process ID and local pod UID tags.
568            (
569                false,
570                origin(Some(&EID_PID), Some(&EID_LOCAL_CID), Some(&EID_LOCAL_POD), None),
571                pid_plus_local_pod_tags,
572            ),
573            // We have the process ID, local container ID, and local pod UID, but entity ID precedence is enabled, so we
574            // should only get the local pod UID tags.
575            (
576                true,
577                origin(Some(&EID_PID), Some(&EID_LOCAL_CID), Some(&EID_LOCAL_POD), None),
578                tags_for_entity(&EID_LOCAL_POD),
579            ),
580            // We only have the process ID and local container ID, so entity ID precedence should be irrelevant, and so
581            // we should get the process ID and local container ID tags.
582            (
583                false,
584                origin(Some(&EID_PID), Some(&EID_LOCAL_CID), None, None),
585                pid_plus_local_cid_tags.clone(),
586            ),
587            (
588                true,
589                origin(Some(&EID_PID), Some(&EID_LOCAL_CID), None, None),
590                pid_plus_local_cid_tags,
591            ),
592        ];
593
594        for (entity_id_precedence, resolved_origin, expected_tags) in cases {
595            let tag_resolver_config = OriginEnrichmentConfiguration {
596                enabled: true,
597                entity_id_precedence,
598                tag_cardinality: OriginTagCardinality::High,
599                origin_detection_unified: false,
600                origin_detection_optout: false,
601                origin_detection_client: false,
602            };
603
604            let origin_tags_resolver = build_tags_resolver_with_default_tags(tag_resolver_config);
605
606            let actual_tags = origin_tags_resolver.collect_origin_tags(&resolved_origin);
607            assert_eq!(
608                actual_tags, expected_tags,
609                "failed to resolve the expected tags for origin {:?}",
610                resolved_origin
611            );
612        }
613    }
614
615    #[test]
616    fn origin_detection_unified_precedence() {
617        // We craft a "valid" and "invalid" variant for External Data, where the invalid one has a container ID with no tags
618        // assigned to it, which lets us exercise the tags resolver logic for when we have a pod UID through External Data,
619        // but not a container ID (or a container ID with no tags attached).
620        //
621        // We have to do it this way because we pass both External Data-based entity IDs through `ResolvedExternalData` when
622        // creating `ResolvedOrigin`, so we can't pass them separately.
623        let ext_data_valid = ResolvedExternalData::new(EID_EXTERNAL_POD.clone(), EID_EXTERNAL_CID_VALID.clone());
624        let ext_data_invalid = ResolvedExternalData::new(EID_EXTERNAL_POD.clone(), EID_EXTERNAL_CID_INVALID.clone());
625
626        let tag_resolver_config = OriginEnrichmentConfiguration {
627            enabled: true,
628            entity_id_precedence: false,
629            tag_cardinality: OriginTagCardinality::High,
630            origin_detection_unified: true,
631            origin_detection_optout: false,
632            origin_detection_client: false,
633        };
634
635        let origin_tags_resolver = build_tags_resolver_with_default_tags(tag_resolver_config);
636
637        // We craft our test cases to ensure that we always take the tags of the highest precedence entity ID available,
638        // and don't take any other tags.
639        let cases = [
640            // Cases where we're only setting a single entity ID. This is the happy path.
641            (origin(Some(&EID_PID), None, None, None), tags_for_entity(&EID_PID)),
642            (
643                origin(None, Some(&EID_LOCAL_CID), None, None),
644                tags_for_entity(&EID_LOCAL_CID),
645            ),
646            (
647                origin(None, None, Some(&EID_LOCAL_POD), None),
648                tags_for_entity(&EID_LOCAL_POD),
649            ),
650            (
651                origin(None, None, None, Some(&ext_data_valid)),
652                tags_for_entity(&EID_EXTERNAL_CID_VALID),
653            ),
654            (
655                origin(None, None, None, Some(&ext_data_invalid)),
656                tags_for_entity(&EID_EXTERNAL_POD),
657            ),
658            // Cases where we have multiple entity IDs to choose from. We work our way backwards here.
659            (
660                origin(
661                    Some(&EID_PID),
662                    Some(&EID_LOCAL_CID),
663                    Some(&EID_LOCAL_POD),
664                    Some(&ext_data_valid),
665                ),
666                tags_for_entity(&EID_LOCAL_CID),
667            ),
668            (
669                origin(Some(&EID_PID), None, Some(&EID_LOCAL_POD), Some(&ext_data_valid)),
670                tags_for_entity(&EID_PID),
671            ),
672            (
673                origin(None, None, Some(&EID_LOCAL_POD), Some(&ext_data_valid)),
674                tags_for_entity(&EID_EXTERNAL_CID_VALID),
675            ),
676            (
677                origin(None, None, Some(&EID_LOCAL_POD), Some(&ext_data_invalid)),
678                tags_for_entity(&EID_LOCAL_POD),
679            ),
680            (
681                origin(None, None, None, Some(&ext_data_invalid)),
682                tags_for_entity(&EID_EXTERNAL_POD),
683            ),
684        ];
685
686        for (resolved_origin, expected_tags) in cases {
687            let actual_tags = origin_tags_resolver.collect_origin_tags(&resolved_origin);
688            assert_eq!(
689                actual_tags, expected_tags,
690                "failed to resolve the expected tags for origin {:?}",
691                resolved_origin
692            );
693        }
694    }
695
696    #[test]
697    fn origin_detection_disabled() {
698        // When origin detection is disabled, no tags should be resolved even if we do have mapped tags for the given
699        // resolved origin.
700        let tag_resolver_config = OriginEnrichmentConfiguration::for_test();
701        assert!(!tag_resolver_config.enabled);
702
703        let origin_tags_resolver = build_tags_resolver_with_default_tags(tag_resolver_config);
704
705        let resolved_origin = origin(Some(&EID_PID), None, None, None);
706        let actual_tags = origin_tags_resolver.collect_origin_tags(&resolved_origin);
707        assert!(actual_tags.is_empty());
708    }
709
710    #[test]
711    fn resolve_origin_tags_dispatches_to_captured_store_when_replay_flag_set() {
712        use datadog_protos::agent::{Entity as ProtoEntity, TaggerState};
713
714        // Build a captured store with a single entity keyed by PID 7777.
715        let mut entities = HashMap::new();
716        entities.insert(
717            "container_id://captured-container".to_string(),
718            ProtoEntity {
719                low_cardinality_tags: vec!["env:captured".into(), "service:replayed".into()],
720                ..Default::default()
721            },
722        );
723        let mut pid_map = HashMap::new();
724        pid_map.insert(7777, "container_id://captured-container".to_string());
725        let state = TaggerState {
726            state: entities,
727            pid_map,
728            duration: 0,
729        };
730        let captured_tagger = CapturedTaggerHandle::new();
731        captured_tagger.set_current(Some(super::super::replay::CapturedTaggerStore::from_tagger_state(
732            state,
733        )));
734
735        // Build a resolver whose live workload provider is empty; if the replay path is wired wrong, we'd get
736        // no tags. If it's wired right, we get the captured tags.
737        let config = OriginEnrichmentConfiguration {
738            enabled: true,
739            tag_cardinality: OriginTagCardinality::Low,
740            ..OriginEnrichmentConfiguration::for_test()
741        };
742        let live = Arc::new(TestWorkloadProvider::new());
743        let resolver = DogStatsDOriginTagResolver::new(config, live, captured_tagger);
744
745        let mut origin = RawOrigin::default();
746        origin.set_process_id(mark_replay_process_id(7777));
747        origin.set_cardinality(OriginTagCardinality::Low);
748
749        let tags = resolver.resolve_origin_tags(origin);
750        let tag_strs: Vec<String> = tags.into_iter().map(|t| t.as_str().to_string()).collect();
751        assert!(tag_strs.contains(&"env:captured".to_string()));
752        assert!(tag_strs.contains(&"service:replayed".to_string()));
753    }
754
755    #[test]
756    fn resolve_origin_tags_returns_empty_when_replay_flag_set_but_no_captured_store() {
757        // If the replay flag is set but no captured store is current, the resolver returns an empty tag set
758        // (no fallback to the live tagger). This guards against accidentally serving live tags for replay packets.
759        let config = OriginEnrichmentConfiguration {
760            enabled: true,
761            tag_cardinality: OriginTagCardinality::Low,
762            ..OriginEnrichmentConfiguration::for_test()
763        };
764        let live = Arc::new(TestWorkloadProvider::new());
765        let resolver = DogStatsDOriginTagResolver::new(config, live, CapturedTaggerHandle::new());
766
767        let mut origin = RawOrigin::default();
768        origin.set_process_id(mark_replay_process_id(7777));
769
770        let tags = resolver.resolve_origin_tags(origin);
771        assert!(
772            tags.is_empty(),
773            "replay path with no captured store must return empty tags"
774        );
775    }
776
777    #[test]
778    fn resolve_origin_tags_live_path_resolves_tags_via_workload_provider() {
779        // Exercises the real, non-replay production entry point: `resolve_origin_tags` asks the workload provider to
780        // resolve the raw origin, then walks the resolved entity IDs to collect tags. Every other test either drives
781        // the internal `collect_origin_tags` directly or takes the replay bypass, so this is the only coverage of the
782        // live `get_resolved_origin` -> `collect_origin_tags` path through the trait method.
783        let config = OriginEnrichmentConfiguration {
784            enabled: true,
785            tag_cardinality: OriginTagCardinality::High,
786            ..OriginEnrichmentConfiguration::for_test()
787        };
788
789        let mut workload_provider = TestWorkloadProvider::new();
790        workload_provider.add_entity_shared_tags(EntityId::ContainerPid(4242), single_tag("tag_source:live-pid"));
791        let resolver =
792            DogStatsDOriginTagResolver::new(config, Arc::new(workload_provider), CapturedTaggerHandle::new());
793
794        // A plain process ID (no replay marker bit) resolves to `EntityId::ContainerPid(4242)` via the live path.
795        let mut origin = RawOrigin::default();
796        origin.set_process_id(4242);
797
798        let tags = resolver.resolve_origin_tags(origin);
799        assert_eq!(tags, single_tag("tag_source:live-pid"));
800    }
801
802    #[test]
803    fn resolve_origin_tags_uses_pinned_process_entity() {
804        let pinned_entity = EntityId::Container(MetaString::from_static("original-container"));
805        let config = OriginEnrichmentConfiguration {
806            enabled: true,
807            tag_cardinality: OriginTagCardinality::High,
808            ..OriginEnrichmentConfiguration::for_test()
809        };
810        let workload_provider =
811            TestWorkloadProvider::with_entity(pinned_entity.clone(), &["tag_source:pinned-container"]);
812        let resolver =
813            DogStatsDOriginTagResolver::new(config, Arc::new(workload_provider), CapturedTaggerHandle::new());
814        let process_origin = ProcessOrigin::Pinned(Some(pinned_entity));
815
816        let tags = resolver.resolve_origin_tags_with_process_origin(RawOrigin::default(), Some(&process_origin));
817
818        assert_eq!(tags, single_tag("tag_source:pinned-container"));
819    }
820
821    #[test]
822    fn resolve_origin_tags_live_path_returns_empty_when_origin_unresolved() {
823        // When the workload provider can't resolve the origin (here, an empty origin resolves to `None`), the live
824        // path returns an empty tag set rather than falling through to anything else.
825        let config = OriginEnrichmentConfiguration {
826            enabled: true,
827            tag_cardinality: OriginTagCardinality::High,
828            ..OriginEnrichmentConfiguration::for_test()
829        };
830        let resolver = DogStatsDOriginTagResolver::new(
831            config,
832            Arc::new(TestWorkloadProvider::new()),
833            CapturedTaggerHandle::new(),
834        );
835
836        let tags = resolver.resolve_origin_tags(RawOrigin::default());
837        assert!(tags.is_empty(), "unresolved origin must produce no tags");
838    }
839}
saluki_components/sources/dogstatsd/origin.rs

saluki_components/sources/dogstatsd/
origin.rs
