ext/pf2/src/signal_scheduler.rs in pf2-0.2.0 vs ext/pf2/src/signal_scheduler.rs in pf2-0.3.0

@@ -9,12 +9,13 @@
 use crate::profile_serializer::ProfileSerializer;
 use crate::sample::Sample;
 
 use core::panic;
 use std::collections::HashSet;
-use std::ffi::{c_int, c_void, CString};
+use std::ffi::{c_int, c_void, CStr, CString};
 use std::mem::ManuallyDrop;
+use std::str::FromStr;
 use std::sync::{Arc, RwLock};
 use std::thread;
 use std::time::Duration;
 use std::{mem, ptr::null_mut};
 
@@ -22,11 +23,11 @@
 
 use crate::util::*;
 
 #[derive(Debug)]
 pub struct SignalScheduler {
-    configuration: configuration::Configuration,
+    configuration: Option<configuration::Configuration>,
     profile: Option<Arc<RwLock<Profile>>>,
 }
 
 pub struct SignalHandlerArgs {
     profile: Arc<RwLock<Profile>>,
@@ -34,41 +35,104 @@
 }
 
 impl SignalScheduler {
     fn new() -> Self {
         Self {
-            configuration: Configuration {
-                time_mode: TimeMode::CpuTime,
-            },
+            configuration: None,
             profile: None,
         }
     }
 
-    fn start(
-        &mut self,
-        _rbself: VALUE,
-        ruby_threads_rary: VALUE,
-        track_new_threads: VALUE,
-    ) -> VALUE {
-        let track_new_threads = RTEST(track_new_threads);
+    fn initialize(&mut self, argc: c_int, argv: *const VALUE, _rbself: VALUE) -> VALUE {
+        // Parse arguments
+        let kwargs: VALUE = Qnil.into();
+        unsafe {
+            rb_scan_args(argc, argv, cstr!(":"), &kwargs);
+        };
+        let mut kwargs_values: [VALUE; 4] = [Qnil.into(); 4];
+        unsafe {
+            rb_get_kwargs(
+                kwargs,
+                [
+                    rb_intern(cstr!("interval_ms")),
+                    rb_intern(cstr!("threads")),
+                    rb_intern(cstr!("time_mode")),
+                    rb_intern(cstr!("track_new_threads")),
+                ]
+                .as_mut_ptr(),
+                0,
+                4,
+                kwargs_values.as_mut_ptr(),
+            );
+        };
+        let interval: Duration = if kwargs_values[0] != Qundef as VALUE {
+            let interval_ms = unsafe { rb_num2long(kwargs_values[0]) };
+            Duration::from_millis(interval_ms.try_into().unwrap_or_else(|_| {
+                eprintln!(
+                    "[Pf2] Warning: Specified interval ({}) is not valid. Using default value (49ms).",
+                    interval_ms
+                );
+                49
+            }))
+        } else {
+            Duration::from_millis(49)
+        };
+        let threads: VALUE = if kwargs_values[1] != Qundef as VALUE {
+            kwargs_values[1]
+        } else {
+            unsafe { rb_funcall(rb_cThread, rb_intern(cstr!("list")), 0) }
+        };
+        let time_mode: configuration::TimeMode = if kwargs_values[2] != Qundef as VALUE {
+            let specified_mode = unsafe {
+                let mut str = rb_funcall(kwargs_values[2], rb_intern(cstr!("to_s")), 0);
+                let ptr = rb_string_value_ptr(&mut str);
+                CStr::from_ptr(ptr).to_str().unwrap()
+            };
+            configuration::TimeMode::from_str(specified_mode).unwrap_or_else(|_| {
+                // Raise an ArgumentError
+                unsafe {
+                    rb_raise(
+                        rb_eArgError,
+                        cstr!("Invalid time mode. Valid values are 'cpu' and 'wall'."),
+                    )
+                }
+            })
+        } else {
+            configuration::TimeMode::CpuTime
+        };
+        let track_new_threads: bool = if kwargs_values[3] != Qundef as VALUE {
+            RTEST(kwargs_values[3])
+        } else {
+            false
+        };
 
-        let profile = Arc::new(RwLock::new(Profile::new()));
-        self.start_profile_buffer_flusher_thread(&profile);
-        self.install_signal_handler();
-
         let mut target_ruby_threads = HashSet::new();
         unsafe {
-            for i in 0..RARRAY_LEN(ruby_threads_rary) {
-                let ruby_thread: VALUE = rb_ary_entry(ruby_threads_rary, i);
+            for i in 0..RARRAY_LEN(threads) {
+                let ruby_thread: VALUE = rb_ary_entry(threads, i);
                 target_ruby_threads.insert(ruby_thread);
             }
         }
+
+        self.configuration = Some(Configuration {
+            interval,
+            target_ruby_threads,
+            time_mode,
+            track_new_threads,
+        });
+
+        Qnil.into()
+    }
+
+    fn start(&mut self, _rbself: VALUE) -> VALUE {
+        let profile = Arc::new(RwLock::new(Profile::new()));
+        self.start_profile_buffer_flusher_thread(&profile);
+        self.install_signal_handler();
+
         TimerInstaller::install_timer_to_ruby_threads(
-            self.configuration.clone(),
-            &target_ruby_threads,
+            self.configuration.as_ref().unwrap().clone(), // FIXME: don't clone
             Arc::clone(&profile),
-            track_new_threads,
         );
 
         self.profile = Some(profile);
 
         Qtrue.into()
@@ -130,11 +194,11 @@
                 log::trace!("Failed to acquire profile lock (garbage collection possibly in progress). Dropping sample.");
                 return;
             }
         };
 
-        let sample = Sample::capture(args.context_ruby_thread); // NOT async-signal-safe
+        let sample = Sample::capture(args.context_ruby_thread, &profile.backtrace_state); // NOT async-signal-safe
         if profile.temporary_sample_buffer.push(sample).is_err() {
             log::debug!("Temporary sample buffer full. Dropping sample.");
         }
     }
 
@@ -154,16 +218,21 @@
         });
     }
 
     // Ruby Methods
 
-    pub unsafe extern "C" fn rb_start(
+    pub unsafe extern "C" fn rb_initialize(
+        argc: c_int,
+        argv: *const VALUE,
         rbself: VALUE,
-        ruby_threads: VALUE,
-        track_new_threads: VALUE,
     ) -> VALUE {
         let mut collector = unsafe { Self::get_struct_from(rbself) };
-        collector.start(rbself, ruby_threads, track_new_threads)
+        collector.initialize(argc, argv, rbself)
+    }
+
+    pub unsafe extern "C" fn rb_start(rbself: VALUE) -> VALUE {
+        let mut collector = unsafe { Self::get_struct_from(rbself) };
+        collector.start(rbself)
     }
 
     pub unsafe extern "C" fn rb_stop(rbself: VALUE) -> VALUE {
         let mut collector = unsafe { Self::get_struct_from(rbself) };
         collector.stop(rbself)