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_context::Context;
103    use saluki_core::components::transforms::SynchronousTransform;
104    use saluki_core::data_model::event::{metric::Metric, Event};
105    use saluki_core::topology::EventsBuffer;
106    use stringtheory::MetaString;
107
108    use super::HostEnrichment;
109
110    fn host_enrichment() -> HostEnrichment {
111        HostEnrichment {
112            hostname: MetaString::from_static("default-host"),
113        }
114    }
115
116    #[test]
117    fn transform_leaves_metric_context_host_unchanged() {
118        let cases = [
119            None,
120            Some(MetaString::empty()),
121            Some(MetaString::from_static("custom-host")),
122        ];
123
124        for host in cases {
125            let context = Context::from_static_name("metric").with_host(host.clone());
126            let metric = Metric::gauge(context, 1.0);
127            let mut events = EventsBuffer::default();
128            assert!(events.try_push(Event::Metric(metric)).is_none());
129
130            host_enrichment().transform_buffer(&mut events);
131
132            let Event::Metric(metric) = events.into_iter().next().expect("metric event") else {
133                panic!("expected metric event");
134            };
135            assert_eq!(metric.context().host(), host.as_deref());
136        }
137    }
138}