fix: more robust sampling thread

Author: ajbt200128

example/Example.ml

@@ -16,7 +16,6 @@
 (*****************************************************************************)
 (* Prelude *)
 (*****************************************************************************)
-
 let example_func2 () =
   My_module.do_short_thing () ;
   My_module.alloc_thing ()
@@ -32,18 +31,10 @@ let () =
   @@ fun () ->
   Printf.printf "Starting loop\n" ;
   flush_all () ;
-  let[@pyro_profile] do_main_thing () =
-    let i = 0 in
-    while i < 1 do
+  let do_main_thing () =
+    while true do
       example_func () ; example_func3 () ; example_func () ; example_func3 ()
     done
   in
-  let d1 = Domain.spawn (fun () -> do_main_thing ()) in
-  let d2 = Domain.spawn (fun () -> do_main_thing ()) in
-  let d3 = Domain.spawn (fun () -> do_main_thing ()) in
-  let d4 = Domain.spawn (fun () -> do_main_thing ()) in
-  do_main_thing () ;
-  Domain.join d1 ;
-  Domain.join d2 ;
-  Domain.join d3 ;
-  Domain.join d4
+  let domains = List.init 8 (fun _ -> Domain.spawn (fun () -> do_main_thing ())) in
+  List.iter Domain.join domains

example/My_module.ml

@@ -27,6 +27,7 @@ let f x y z =
   let c = b - y in
   c
 
+
 let alloc_thing () =
   let random_list =
     List.init 1000 (fun _ ->

example/dune

@@ -2,10 +2,10 @@
  (package pyro-caml-ppx)
  (name example)
  (public_name example)
- (libraries pyro-caml-instruments pyro-caml-ppx unix)
+ (libraries pyro-caml-instruments unix)
  (modes exe)
  (preprocess
-  (pps ppx_pyro_caml))
+  (pps ppx_pyro_caml --auto))
  (foreign_stubs
   (language c)
   (names example_stubs)

src/backend.rs

@@ -10,7 +10,7 @@ use std::{
     path::{Path, PathBuf},
     sync::{
         atomic::{AtomicBool, Ordering},
-        Arc, Mutex, OnceLock,
+        Arc, Mutex, MutexGuard, OnceLock,
     },
     thread::{self, JoinHandle},
 };
@@ -70,7 +70,7 @@ pub struct CamlSpy {
     ruleset: Arc<Mutex<Ruleset>>,
     backend_config: Arc<Mutex<BackendConfig>>,
     config: CamlSpyConfig,
-    sampler_thread: Option<JoinHandle<Result<()>>>,
+    sampler_thread: Option<JoinHandle<()>>,
 }
 
 impl CamlSpy {
@@ -126,15 +126,22 @@ impl Backend for CamlSpy {
             log::debug!(target:LOG_TAG, "starting sampler thread");
             while running.load(Ordering::Relaxed) {
                 log::trace!(target:LOG_TAG, "acquiring backend config and cursor");
-                let backend_config = backend_config.lock()?;
+                let backend_config = backend_config
+                    .lock()
+                    .expect("Could not take backend config lock"); // we only have one sampler thread and one reporting thread so this should never fail unless something is seriously wrong
 
                 let cursor = config.acquire_cursor();
                 log::trace!(target:LOG_TAG, "sampling...");
                 let mut stack_frames: Vec<StackTrace> = OCAML_GC
-                    .with_borrow(|gc| ocaml_intf::read_poll(gc, cursor, config.sample_interval()))?
+                    .with_borrow(|gc| ocaml_intf::read_poll(gc, cursor, config.sample_interval()))
+                    .unwrap_or_else(|e| {
+                        log::error!(target:LOG_TAG, "Error reading from OCaml runtime: {:?}", e);
+                        vec![]
+                    })
                     .into_iter()
                     .map(|st| st.into_stack_trace(backend_config.deref(), config.pid))
                     .collect();
+
                 log::trace!(target:LOG_TAG, "done sampling");
 
                 if stack_frames.is_empty() {
@@ -145,8 +152,12 @@ impl Backend for CamlSpy {
 
                 log::trace!(target:LOG_TAG, "recording stack frames");
                 for st in stack_frames.into_iter() {
-                    let stack_trace = st + &ruleset.lock()?.clone();
-                    buffer.lock()?.record(stack_trace).unwrap();
+                    let stack_trace = st + &ruleset.lock().expect("Failed to lock ruleset").clone();
+                    buffer
+                        .lock()
+                        .expect("Failed to lock buffer") // we only have one sampler thread and one sending reports so we should never fail to obtain the lock unless something is seriously wrong
+                        .record(stack_trace)
+                        .expect("Failed to record stack traces to buffer"); // So this seems to just be a result in the pyroscope sdk for no reason
                 }
                 log::trace!(target:LOG_TAG, "recorded stack trace");
                 log::trace!(target:LOG_TAG, "sleeping for {} ms", 1000 / config.sample_rate);
@@ -155,8 +166,7 @@ impl Backend for CamlSpy {
                     1000 / config.sample_rate as u64,
                 ));
             }
-            log::debug!(target:LOG_TAG, "sampler thread exiting");
-            Ok(())
+            log::debug!(target:LOG_TAG, "sampler thread exiting")
         });
         self.sampler_thread = Some(sampler);
         Ok(())
@@ -168,19 +178,26 @@ impl Backend for CamlSpy {
         self.sampler_thread
             .ok_or_else(|| PyroscopeError::new("CamlSpy: Failed to unwrap sampler thread"))?
             .join()
-            .unwrap_or_else(|_| {
-                Err(PyroscopeError::new(
-                    "CamlSpy: Failed to join sampler thread",
-                ))
+            .map_err(|e| {
+                PyroscopeError::new(
+                    format!("CamlSpy: Failed to join sampler thread: {:?}", e).as_str(),
+                )
             })?;
         Ok(())
     }
 
     fn report(&mut self) -> Result<Vec<Report>> {
-        let report: StackBuffer = self.buffer.lock()?.deref().to_owned();
-        let reports: Vec<Report> = report.into();
-
-        self.buffer.lock()?.clear();
+        // first check if the thread has finished
+        if let Some(handle) = &self.sampler_thread {
+            if handle.is_finished() {
+                return Err(PyroscopeError::new(
+                    "CamlSpy: Sampler thread has exited unexpectedly",
+                ));
+            }
+        }
+        let mut buffer: MutexGuard<'_, StackBuffer> = self.buffer.lock()?;
+        let reports: Vec<Report> = buffer.deref().to_owned().into();
+        buffer.clear();
         Ok(reports)
     }