wait for flush to terminate

DataDog · Feb 10, 2025 · fbd13cd · fbd13cd
1 parent 1470851
commit fbd13cd
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/dogstatsd/Cargo.toml b/dogstatsd/Cargo.toml
@@ -25,6 +25,7 @@ tokio = { version = "1.37.0", default-features = false, features = ["macros", "r
 tokio-util = { version = "0.7.11", default-features = false }
 tracing = { version = "0.1.40", default-features = false }
 regex = { version = "1.10.6", default-features = false }
+log = "0.4.22"
 
 [dev-dependencies]
 mockito = { version = "1.5.0", default-features = false }

diff --git a/dogstatsd/src/flusher.rs b/dogstatsd/src/flusher.rs
@@ -52,7 +52,7 @@ impl Flusher {
         debug!("Flushing {n_series} series and {n_distributions} distributions");
 
         let dd_api_clone = self.dd_api.clone();
-        tokio::spawn(async move {
+        let series_handle = tokio::spawn(async move {
             for a_batch in all_series {
                 let continue_shipping =
                     should_try_next_batch(dd_api_clone.ship_series(&a_batch).await).await;
@@ -62,7 +62,7 @@ impl Flusher {
             }
         });
         let dd_api_clone = self.dd_api.clone();
-        tokio::spawn(async move {
+        let distributions_handle = tokio::spawn(async move {
             for a_batch in all_distributions {
                 let continue_shipping =
                     should_try_next_batch(dd_api_clone.ship_distributions(&a_batch).await).await;
@@ -71,6 +71,15 @@ impl Flusher {
                 }
             }
         });
+
+        match tokio::try_join!(series_handle, distributions_handle) {
+            Ok(_) => {
+                debug!("Successfully flushed {n_series} series and {n_distributions} distributions")
+            }
+            Err(err) => {
+                error!("Failed to flush metrics{err}")
+            }
+        };
     }
 }