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