fix: more robust sampling thread (#4)

I was incorrectly just shutting the sample collection thread down whenever an exception from ocaml was raised. Really what we want to log it and try again, since occasionally we may get a "stream corrupted" exception which just means some things in the ring buffer were overwritten. This ultimately isn't a big deal unless it happens a lot, but that will be obvious by a lot of missing samples.

## Test plan
I modified Example.ml to generate enough samples that occasionally this exception will pop up, and then I ran pyro caml and confirmed it no longer takes down the collection 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)
     }