Fixing lifetimes for dogstatsd flusher.

I think Rust is too dumb :-) Signed-off-by: Bob Weinand <bob.weinand@datadoghq.com>
DataDog · Sep 18, 2024 · ef7adb3 · ef7adb3
1 parent 8fecb9a
commit ef7adb3
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 29 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/data-pipeline/src/trace_exporter.rs b/data-pipeline/src/trace_exporter.rs
@@ -13,6 +13,9 @@ use hyper::http::uri::PathAndQuery;
 use hyper::{Body, Client, Method, Uri};
 use log::error;
 use std::{borrow::Borrow, collections::HashMap, str::FromStr};
+use std::iter::Chain;
+use std::marker::PhantomData;
+use std::sync::Mutex;
 use tokio::runtime::Runtime;
 
 /// TraceExporterInputFormat represents the format of the input traces.
@@ -209,7 +212,7 @@ impl TraceExporter {
                     Err(err) => {
                         self.emit_metric(HealthMetric::Count(STAT_SEND_ERRORS, 1), None);
                         anyhow::bail!("Failed to send traces: {err}")
-                    },
+                    }
                 }
             })
             .or_else(|err| {
@@ -220,18 +223,38 @@ impl TraceExporter {
 
     /// Emit a health metric to dogstatsd
     fn emit_metric(&self, metric: HealthMetric, custom_tags: Option<Vec<Tag>>) {
+        struct WrapIt<'a, 'b: 'a, V: 'b, I: IntoIterator<Item = &'a V>> {
+            it: Mutex<Option<I>>,
+            _p: PhantomData<&'b I>,
+        }
+        impl<'a, 'b, V, I: IntoIterator<Item = &'a V>> WrapIt<'a, 'b, V, I> {
+            pub fn wrap(iter: I) -> Self {
+                WrapIt {
+                    it: Mutex::new(Some(iter)),
+                    _p: PhantomData,
+                }
+            }
+        }
+        impl<'a, 'b, V, I: IntoIterator<Item = &'a V>> IntoIterator for &'b WrapIt<'a, 'b, V, I> {
+            type Item = I::Item;
+            type IntoIter = I::IntoIter;
+
+            fn into_iter(self) -> Self::IntoIter {
+                self.it.lock().unwrap().take().unwrap().into_iter()
+            }
+        }
         if let Some(flusher) = &self.dogstatsd {
             if custom_tags.is_some() {
                 match metric {
-                    HealthMetric::Count(name, c) => flusher.send(vec![DogStatsDAction::Count(
+                    HealthMetric::Count(name, c) => flusher.send(&vec![DogStatsDAction::Count(
                         name,
                         c,
-                        &self.common_stats_tags.iter().chain(&custom_tags.unwrap()),
+                        &WrapIt::wrap(self.common_stats_tags.iter().chain(&custom_tags.unwrap())),
                     )]),
                 }
             } else {
                 match metric {
-                    HealthMetric::Count(name, c) => flusher.send(vec![DogStatsDAction::Count(
+                    HealthMetric::Count(name, c) => flusher.send(&vec![DogStatsDAction::Count(
                         name,
                         c,
                         &self.common_stats_tags,
@@ -258,7 +281,10 @@ impl TraceExporter {
             return Ok(String::from("{}"));
         }
 
-        self.emit_metric(HealthMetric::Count(STAT_DESER_TRACES, traces.len() as i64), None);
+        self.emit_metric(
+            HealthMetric::Count(STAT_DESER_TRACES, traces.len() as i64),
+            None,
+        );
 
         let header_tags: TracerHeaderTags<'_> = (&self.tags).into();
 

diff --git a/dogstatsd-client/src/lib.rs b/dogstatsd-client/src/lib.rs
@@ -24,8 +24,6 @@ const QUEUE_SIZE: usize = 32 * 1024;
 /// The `DogStatsDActionRef` enum gathers the metric types that can be sent to the DogStatsD server.
 #[derive(Debug, Serialize, Deserialize)]
 pub enum DogStatsDAction<T: AsRef<str>, V: std::ops::Deref>
-where
-    for<'a> &'a <V as std::ops::Deref>::Target: IntoIterator<Item = &'a Tag>,
 {
     // TODO: instead of AsRef<str> we can accept a marker Trait that users of this crate implement
     Count(T, i64, V),
@@ -70,34 +68,34 @@ impl Flusher {
         Ok(())
     }
 
-    pub fn send<T: AsRef<str>, V: std::ops::Deref>(&self, actions: Vec<DogStatsDAction<T, V>>)
+    pub fn send<'a, T: AsRef<str> + 'a, V: std::ops::Deref + 'a>(&self, actions: &'a Vec<DogStatsDAction<T, V>>)
     where
-        for<'a> &'a <V as std::ops::Deref>::Target: IntoIterator<Item = &'a Tag>,
+        &'a <V as std::ops::Deref>::Target: IntoIterator<Item = &'a Tag> + 'a,
     {
         if self.client.is_none() {
             return;
         }
         let client = self.client.as_ref().unwrap();
 
-        for action in actions {
+        for action in actions.iter() {
             if let Err(err) = match action {
-                DogStatsDAction::Count(metric, value, ref tags) => {
-                    let metric_builder = client.count_with_tags(metric.as_ref(), value);
-                    do_send(metric_builder, tags.deref())
+                DogStatsDAction::Count(metric, value, tags) => {
+                    let metric_builder = client.count_with_tags(metric.as_ref(), *value);
+                    do_send(metric_builder, &mut tags.deref().into_iter())
                 }
-                DogStatsDAction::Distribution(metric, value, ref tags) => do_send(
-                    client.distribution_with_tags(metric.as_ref(), value),
-                    tags.deref(),
+                DogStatsDAction::Distribution(metric, value, tags) => do_send(
+                    client.distribution_with_tags(metric.as_ref(), *value),
+                    &mut tags.deref().into_iter(),
                 ),
-                DogStatsDAction::Gauge(metric, value, ref tags) => {
-                    do_send(client.gauge_with_tags(metric.as_ref(), value), tags.deref())
+                DogStatsDAction::Gauge(metric, value, tags) => {
+                    do_send(client.gauge_with_tags(metric.as_ref(), *value), &mut tags.deref().into_iter())
                 }
-                DogStatsDAction::Histogram(metric, value, ref tags) => do_send(
-                    client.histogram_with_tags(metric.as_ref(), value),
-                    tags.deref(),
+                DogStatsDAction::Histogram(metric, value, tags) => do_send(
+                    client.histogram_with_tags(metric.as_ref(), *value),
+                    &mut tags.deref().into_iter(),
                 ),
-                DogStatsDAction::Set(metric, value, ref tags) => {
-                    do_send(client.set_with_tags(metric.as_ref(), value), tags.deref())
+                DogStatsDAction::Set(metric, value, tags) => {
+                    do_send(client.set_with_tags(metric.as_ref(), *value), &mut tags.deref().into_iter())
                 }
             } {
                 error!("Error while sending metric: {}", err);
@@ -106,15 +104,14 @@ impl Flusher {
     }
 }
 
-fn do_send<'m, 't, T, V: IntoIterator<Item = &'t Tag>>(
+fn do_send<'m, 't, T, V: Iterator<Item = &'t Tag> + 't>(
     mut builder: MetricBuilder<'m, '_, T>,
-    tags: V,
+    tags_iter: &mut V,
 ) -> anyhow::Result<()>
 where
     T: Metric + From<String>,
     't: 'm,
 {
-    let mut tags_iter = tags.into_iter();
     let mut tag_opt = tags_iter.next();
     while tag_opt.is_some() {
         builder = builder.with_tag_value(tag_opt.unwrap().as_ref());
@@ -185,7 +182,7 @@ mod test {
         _ = flusher.set_endpoint(Endpoint::from_slice(
             socket.local_addr().unwrap().to_string().as_str(),
         ));
-        flusher.send(vec![
+        flusher.send(&vec![
             Count("test_count", 3, vec![tag!("foo", "bar")]),
             Count("test_neg_count", -2, vec![]),
             Distribution("test_distribution", 4.2, vec![]),

diff --git a/sidecar/src/service/sidecar_server.rs b/sidecar/src/service/sidecar_server.rs
@@ -850,7 +850,7 @@ impl SidecarInterface for SidecarServer {
         tokio::spawn(async move {
             self.get_session(&instance_id.session_id)
                 .get_dogstatsd()
-                .send(actions);
+                .send(&actions);
         });
 
         no_response()