Wait for telemetry client to finish so the metrics are flushed. (#862)

* Wait for telemetry client to finish so the metrics are flushed.
DataDog · Feb 10, 2025 · 1a7e673 · 1a7e673
1 parent b3b2852
commit 1a7e673
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 12 deletions.
diff --git a/data-pipeline-ffi/src/trace_exporter.rs b/data-pipeline-ffi/src/trace_exporter.rs
@@ -901,7 +901,8 @@ mod tests {
             assert_eq!(response.assume_init().body.to_string_lossy(), response_body);
 
             ddog_trace_exporter_free(exporter);
-            mock_metrics.assert();
+            // It should receive 3 payloads: app-started, metrics and app-closing.
+            mock_metrics.assert_hits(3);
         }
     }
 }
diff --git a/data-pipeline/src/telemetry/mod.rs b/data-pipeline/src/telemetry/mod.rs
@@ -169,15 +169,18 @@ impl TelemetryClient {
             self.handle.abort();
         }
     }
+
     /// Shutdowns the telemetry client.
-    pub async fn shutdown(&self) {
+    pub async fn shutdown(self) {
         if let Err(_e) = self
             .worker
             .send_msg(TelemetryActions::Lifecycle(LifecycleAction::Stop))
             .await
         {
             self.handle.abort();
         }
+
+        let _ = self.handle.await;
     }
 }
 
@@ -264,7 +267,6 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
-        let _ = client.handle.await;
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -304,7 +306,6 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
-        let _ = client.handle.await;
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -344,7 +345,6 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
-        let _ = client.handle.await;
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -384,7 +384,6 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
-        let _ = client.handle.await;
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -424,7 +423,6 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
-        let _ = client.handle.await;
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -464,7 +462,6 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
-        let _ = client.handle.await;
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -504,7 +501,6 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
-        let _ = client.handle.await;
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -544,7 +540,6 @@ mod tests {
         client.start().await;
         let _ = client.send(&data);
         client.shutdown().await;
-        let _ = client.handle.await;
         telemetry_srv.assert_hits_async(1).await;
     }
 
@@ -577,8 +572,6 @@ mod tests {
 
         client.start().await;
         client.shutdown().await;
-        let _ = client.handle.await;
-
         // Check for 2 hits: app-started and app-closing.
         telemetry_srv.assert_hits_async(2).await;
     }