Adding a time out when closing the app insight channels

Author: chkeita

src/agent/onefuzz-telemetry/src/lib.rs

@@ -3,8 +3,9 @@
 
 use chrono::DateTime;
 use serde::{Deserialize, Serialize};
-use std::fmt;
 use std::sync::{LockResult, RwLockReadGuard, RwLockWriteGuard};
+use std::time::Duration;
+use std::fmt;
 use uuid::Uuid;
 
 pub use chrono::Utc;
@@ -15,6 +16,8 @@ use tokio::sync::broadcast::{self, Receiver};
 #[macro_use]
 extern crate lazy_static;
 
+const DEAFAULT_CHANNEL_CLOSING_TIMEOUT: Duration = Duration::from_secs(30);
+
 #[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
 #[serde(transparent)]
 pub struct MicrosoftTelemetryKey(Uuid);
@@ -355,18 +358,22 @@ pub async fn set_appinsights_clients(
     global::set_clients(instance_client, microsoft_client);
 }
 
+pub async fn try_flush_and_close() {
+    _try_flush_and_close(DEAFAULT_CHANNEL_CLOSING_TIMEOUT).await
+}
+
 /// Try to submit any pending telemetry with a blocking call.
 ///
 /// Meant for a final attempt at flushing pending items before an abnormal exit.
 /// After calling this function, any existing telemetry client will be dropped,
 /// and subsequent telemetry submission will be a silent no-op.
-pub async fn try_flush_and_close() {
+pub async fn _try_flush_and_close(timeout: Duration) {
     let clients = global::take_clients();
-
     for client in clients {
-        client.close_channel().await;
+        if let Err(e) = tokio::time::timeout(timeout, client.close_channel()).await {
+            log::warn!("Failed to close telemetry client: {}", e);
+        }
     }
-
     // dropping the broadcast sender to make sure all pending events are sent
     let _global_event_source = global::EVENT_SOURCE.write().unwrap().take();
 }
@@ -468,11 +475,15 @@ pub fn try_broadcast_trace(timestamp: DateTime<Utc>, msg: String, level: log::Le
 }
 
 pub fn subscribe_to_events() -> Result<Receiver<LoggingEvent>> {
-    let global_event_source = global::EVENT_SOURCE.read().unwrap();
-    if let Some(evs) = global_event_source.clone() {
-        Ok(evs.subscribe())
-    } else {
-        bail!("Event source not initialized");
+    match global::EVENT_SOURCE.read() {
+        Ok(global_event_source) => {
+            if let Some(evs) = global_event_source.clone() {
+                Ok(evs.subscribe())
+            } else {
+                bail!("Event source not initialized");
+            }
+        }
+        Err(e) => bail!("failed to acquire event source lock: {}", e),
     }
 }