imp(session): kill SessionManager on Agent drop

Author: omarabid

src/pyroscope.rs

@@ -16,6 +16,7 @@ use crate::backends::Backend;
 use crate::error::Result;
 use crate::session::Session;
 use crate::session::SessionManager;
+use crate::session::SessionSignal;
 use crate::timer::Timer;
 
 /// Represent PyroscopeAgent Configuration
@@ -165,15 +166,18 @@ impl Drop for PyroscopeAgent {
         self.timer.drop_listeners().unwrap(); // Drop listeners
         self.timer.handle.take().unwrap().join().unwrap().unwrap(); // Wait for the Timer thread to finish
 
-        // Wait for main thread to finish
-        self.handle.take().unwrap().join().unwrap().unwrap();
+        // Stop the SessionManager
+        self.session_manager.push(SessionSignal::Kill).unwrap();
         self.session_manager
             .handle
             .take()
             .unwrap()
             .join()
             .unwrap()
             .unwrap();
+
+        // Wait for main thread to finish
+        self.handle.take().unwrap().join().unwrap().unwrap();
     }
 }
 
@@ -212,7 +216,11 @@ impl PyroscopeAgent {
             while let Ok(time) = rx.recv() {
                 let report = backend.lock()?.report()?;
                 // Send new Session to SessionManager
-                stx.send(Session::new(time, config.clone(), report)?)?;
+                stx.send(SessionSignal::Session(Session::new(
+                    time,
+                    config.clone(),
+                    report,
+                )?))?;
 
                 if time == 0 {
                     let (lock, cvar) = &*pair;

src/session.rs

@@ -14,23 +14,39 @@ use crate::pyroscope::PyroscopeConfig;
 use crate::utils::merge_tags_with_app_name;
 use crate::Result;
 
+/// Session Signal
+#[derive(Debug)]
+pub enum SessionSignal {
+    Session(Session),
+    Kill,
+}
+
 /// SessionManager
 #[derive(Debug)]
 pub struct SessionManager {
     pub handle: Option<JoinHandle<Result<()>>>,
-    pub tx: SyncSender<Session>,
+    pub tx: SyncSender<SessionSignal>,
 }
 
 impl SessionManager {
     /// Create a new SessionManager
     pub fn new() -> Result<Self> {
         // Create a channel for sending and receiving sessions
-        let (tx, rx): (SyncSender<Session>, Receiver<Session>) = sync_channel(10);
+        let (tx, rx): (SyncSender<SessionSignal>, Receiver<SessionSignal>) = sync_channel(10);
 
         // Create a thread for the SessionManager
         let handle = Some(thread::spawn(move || {
-            while let Ok(session) = rx.recv() {
-                session.send()?;
+            while let Ok(signal) = rx.recv() {
+                match signal {
+                    SessionSignal::Session(session) => {
+                        // Send the session
+                        session.send()?;
+                    }
+                    SessionSignal::Kill => {
+                        // Kill the session manager
+                        return Ok(());
+                    }
+                }
             }
             Ok(())
         }));
@@ -39,7 +55,7 @@ impl SessionManager {
     }
 
     /// Push a new session into the SessionManager
-    pub fn push(&self, session: Session) -> Result<()> {
+    pub fn push(&self, session: SessionSignal) -> Result<()> {
         self.tx.send(session)?;
         Ok(())
     }