saluki_components/transforms/host_enrichment/
mod.rs

1use async_trait::async_trait;
2use saluki_core::accounting::{MemoryBounds, MemoryBoundsBuilder};
3use saluki_core::{components::transforms::*, topology::EventsBuffer};
4use saluki_core::{
5    components::BuildContext,
6    data_model::event::{eventd::EventD, service_check::ServiceCheck},
7};
8use saluki_env::{EnvironmentProvider, HostProvider as _};
9use saluki_error::GenericError;
10use stringtheory::MetaString;
11
12/// Host enrichment synchronous transform.
13///
14/// Enriches events and service checks with a hostname if one isn't already present. Metrics must carry their hostname
15/// in their context before this transform so metric identity is stable before fanout/encoding.
16pub struct HostEnrichmentConfiguration<E> {
17    env_provider: E,
18}
19
20impl<E> HostEnrichmentConfiguration<E> {
21    /// Creates a new `HostEnrichmentConfiguration` with the given environment provider.
22    pub fn from_environment_provider(env_provider: E) -> Self {
23        Self { env_provider }
24    }
25}
26
27#[async_trait]
28impl<E> SynchronousTransformBuilder for HostEnrichmentConfiguration<E>
29where
30    E: EnvironmentProvider + Send + Sync + 'static,
31{
32    async fn build(&self, _context: BuildContext) -> Result<Box<dyn SynchronousTransform + Send>, GenericError> {
33        Ok(Box::new(
34            HostEnrichment::from_environment_provider(&self.env_provider).await?,
35        ))
36    }
37}
38
39impl<E> MemoryBounds for HostEnrichmentConfiguration<E> {
40    fn specify_bounds(&self, builder: &mut MemoryBoundsBuilder) {
41        // TODO: We don't account for the size of the hostname since we only query it when we go to actually build the
42        // transform. We could move the querying to the point where we create `HostEnrichmentConfiguration` itself but
43        // that would mean it couldn't be updated dynamically.
44        //
45        // Not a relevant problem _right now_, but a _potential_ problem in the future. :shrug:
46
47        // Capture the size of the heap allocation when the component is built.
48        builder
49            .minimum()
50            .with_single_value::<HostEnrichment>("component struct");
51    }
52}
53
54pub struct HostEnrichment {
55    hostname: MetaString,
56}
57
58impl HostEnrichment {
59    pub async fn from_environment_provider<E>(env_provider: &E) -> Result<Self, GenericError>
60    where
61        E: EnvironmentProvider + Send + Sync + 'static,
62    {
63        Ok(Self {
64            hostname: env_provider
65                .host()
66                .get_hostname()
67                .await
68                .map(MetaString::from)
69                .map_err(Into::into)?,
70        })
71    }
72
73    fn enrich_eventd(&self, eventd: &mut EventD) {
74        // Only add the hostname if it's not already present.
75        if eventd.hostname().is_none() {
76            eventd.set_hostname(Some(self.hostname.clone()));
77        }
78    }
79
80    fn enrich_service_check(&self, service_check: &mut ServiceCheck) {
81        // Only add the hostname if it's not already present.
82        if service_check.hostname().is_none() {
83            service_check.set_hostname(Some(self.hostname.clone()));
84        }
85    }
86}
87
88impl SynchronousTransform for HostEnrichment {
89    fn transform_buffer(&mut self, event_buffer: &mut EventsBuffer) {
90        for event in event_buffer {
91            if let Some(eventd) = event.try_as_eventd_mut() {
92                self.enrich_eventd(eventd);
93            } else if let Some(service_check) = event.try_as_service_check_mut() {
94                self.enrich_service_check(service_check);
95            }
96        }
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use saluki_core::{
103        components::transforms::SynchronousTransform,
104        data_model::event::{
105            metric::{context::Context, Metric},
106            Event,
107        },
108        topology::EventsBuffer,
109    };
110    use stringtheory::MetaString;
111
112    use super::HostEnrichment;
113
114    fn host_enrichment() -> HostEnrichment {
115        HostEnrichment {
116            hostname: MetaString::from_static("default-host"),
117        }
118    }
119
120    #[test]
121    fn transform_leaves_metric_context_host_unchanged() {
122        let cases = [
123            None,
124            Some(MetaString::empty()),
125            Some(MetaString::from_static("custom-host")),
126        ];
127
128        for host in cases {
129            let context = Context::from_static_name("metric").with_host(host.clone());
130            let metric = Metric::gauge(context, 1.0);
131            let mut events = EventsBuffer::default();
132            assert!(events.try_push(Event::Metric(metric)).is_none());
133
134            host_enrichment().transform_buffer(&mut events);
135
136            let Event::Metric(metric) = events.into_iter().next().expect("metric event") else {
137                panic!("expected metric event");
138            };
139            assert_eq!(metric.context().host(), host.as_deref());
140        }
141    }
142}