pf2 0.1.0 → 0.3.0

data/ext/pf2/src/signal_scheduler/timer_installer.rs

@@ -0,0 +1,199 @@
+use std::collections::HashMap;
+use std::ffi::c_void;
+use std::mem;
+use std::mem::ManuallyDrop;
+use std::ptr::null_mut;
+use std::sync::{Mutex, RwLock};
+use std::{collections::HashSet, sync::Arc};
+
+use rb_sys::*;
+
+use crate::signal_scheduler::SignalHandlerArgs;
+
+use super::configuration::Configuration;
+use crate::profile::Profile;
+
+// We could avoid deferring the timer creation by combining pthread_getcpuclockid(3) and timer_create(2) here,
+// but we're not doing so since (1) Ruby does not expose the pthread_self() of a Ruby Thread
+// (which is actually stored in th->nt->thread_id), and (2) pthread_getcpuclockid(3) is not portable
+// in the first place (e.g. not available on macOS).
+pub struct TimerInstaller {
+    internal: Box<Mutex<Internal>>,
+}
+
+struct Internal {
+    configuration: Configuration,
+    registered_pthread_ids: HashSet<libc::pthread_t>,
+    kernel_thread_id_to_ruby_thread_map: HashMap<libc::pid_t, VALUE>,
+    profile: Arc<RwLock<Profile>>,
+}
+
+impl TimerInstaller {
+    // Register a callback that gets called when a Ruby Thread is resumed.
+    // The callback should create a timer for the thread.
+    pub fn install_timer_to_ruby_threads(
+        configuration: Configuration,
+        profile: Arc<RwLock<Profile>>,
+    ) {
+        let registrar = Self {
+            internal: Box::new(Mutex::new(Internal {
+                configuration: configuration.clone(),
+                registered_pthread_ids: HashSet::new(),
+                kernel_thread_id_to_ruby_thread_map: HashMap::new(),
+                profile,
+            })),
+        };
+
+        let ptr = Box::into_raw(registrar.internal);
+        unsafe {
+            rb_internal_thread_add_event_hook(
+                Some(Self::on_thread_resume),
+                RUBY_INTERNAL_THREAD_EVENT_RESUMED,
+                ptr as *mut c_void,
+            );
+            // Spawn a no-op Thread to fire the event hook
+            // (at least 2 Ruby Threads must be active for the RESUMED hook to be fired)
+            rb_thread_create(Some(Self::do_nothing), null_mut());
+        };
+
+        if configuration.track_new_threads {
+            unsafe {
+                rb_internal_thread_add_event_hook(
+                    Some(Self::on_thread_start),
+                    RUBY_INTERNAL_THREAD_EVENT_STARTED,
+                    ptr as *mut c_void,
+                );
+            };
+        }
+    }
+
+    unsafe extern "C" fn do_nothing(_: *mut c_void) -> VALUE {
+        Qnil.into()
+    }
+
+    // Thread resume callback
+    unsafe extern "C" fn on_thread_resume(
+        _flag: rb_event_flag_t,
+        data: *const rb_internal_thread_event_data,
+        custom_data: *mut c_void,
+    ) {
+        // The SignalScheduler (as a Ruby obj) should be passed as custom_data
+        let internal =
+            unsafe { ManuallyDrop::new(Box::from_raw(custom_data as *mut Mutex<Internal>)) };
+        let mut internal = internal.lock().unwrap();
+
+        // Check if the current thread is a target Ruby Thread
+        let current_ruby_thread: VALUE = unsafe { (*data).thread };
+        if !internal
+            .configuration
+            .target_ruby_threads
+            .contains(&current_ruby_thread)
+        {
+            return;
+        }
+
+        // Check if the current thread is already registered
+        let current_pthread_id = unsafe { libc::pthread_self() };
+        if internal
+            .registered_pthread_ids
+            .contains(&current_pthread_id)
+        {
+            return;
+        }
+
+        // Record the pthread ID of the current thread
+        internal.registered_pthread_ids.insert(current_pthread_id);
+        // Keep a mapping from kernel thread ID to Ruby Thread
+        internal
+            .kernel_thread_id_to_ruby_thread_map
+            .insert(unsafe { libc::gettid() }, current_ruby_thread);
+
+        Self::register_timer_to_current_thread(
+            &internal.configuration,
+            &internal.profile,
+            &internal.kernel_thread_id_to_ruby_thread_map,
+        );
+
+        // TODO: Remove the hook when all threads have been registered
+    }
+
+    // Thread resume callback
+    unsafe extern "C" fn on_thread_start(
+        _flag: rb_event_flag_t,
+        data: *const rb_internal_thread_event_data,
+        custom_data: *mut c_void,
+    ) {
+        // The SignalScheduler (as a Ruby obj) should be passed as custom_data
+        let internal =
+            unsafe { ManuallyDrop::new(Box::from_raw(custom_data as *mut Mutex<Internal>)) };
+        let mut internal = internal.lock().unwrap();
+
+        let current_ruby_thread: VALUE = unsafe { (*data).thread };
+        internal
+            .configuration
+            .target_ruby_threads
+            .insert(current_ruby_thread);
+    }
+
+    // Creates a new POSIX timer which invocates sampling for the thread that called this function.
+    fn register_timer_to_current_thread(
+        configuration: &Configuration,
+        profile: &Arc<RwLock<Profile>>,
+        kernel_thread_id_to_ruby_thread_map: &HashMap<libc::pid_t, VALUE>,
+    ) {
+        let current_pthread_id = unsafe { libc::pthread_self() };
+        let context_ruby_thread: VALUE = unsafe {
+            *(kernel_thread_id_to_ruby_thread_map
+                .get(&(libc::gettid()))
+                .unwrap())
+        };
+
+        // NOTE: This Box is never dropped
+        let signal_handler_args = Box::new(SignalHandlerArgs {
+            profile: Arc::clone(profile),
+            context_ruby_thread,
+        });
+
+        // Create a signal event
+        let mut sigevent: libc::sigevent = unsafe { mem::zeroed() };
+        // Note: SIGEV_THREAD_ID is Linux-specific. In other platforms, we would need to
+        // "tranpoline" the signal as any pthread can receive the signal.
+        sigevent.sigev_notify = libc::SIGEV_THREAD_ID;
+        sigevent.sigev_notify_thread_id =
+            unsafe { libc::syscall(libc::SYS_gettid).try_into().unwrap() }; // The kernel thread ID
+        sigevent.sigev_signo = libc::SIGALRM;
+        // Pass required args to the signal handler
+        sigevent.sigev_value.sival_ptr = Box::into_raw(signal_handler_args) as *mut c_void;
+
+        // Create and configure timer to fire every _interval_ ms of CPU time
+        let mut timer: libc::timer_t = unsafe { mem::zeroed() };
+        let clockid = match configuration.time_mode {
+            crate::signal_scheduler::TimeMode::CpuTime => libc::CLOCK_THREAD_CPUTIME_ID,
+            crate::signal_scheduler::TimeMode::WallTime => libc::CLOCK_MONOTONIC,
+        };
+        let err = unsafe { libc::timer_create(clockid, &mut sigevent, &mut timer) };
+        if err != 0 {
+            panic!("timer_create failed: {}", err);
+        }
+        let itimerspec = Self::duration_to_itimerspec(&configuration.interval);
+        let err = unsafe { libc::timer_settime(timer, 0, &itimerspec, null_mut()) };
+        if err != 0 {
+            panic!("timer_settime failed: {}", err);
+        }
+
+        log::debug!("timer registered for thread {}", current_pthread_id);
+    }
+
+    fn duration_to_itimerspec(duration: &std::time::Duration) -> libc::itimerspec {
+        let nanos = duration.as_nanos();
+        let seconds_part: i64 = (nanos / 1_000_000_000).try_into().unwrap();
+        let nanos_part: i64 = (nanos % 1_000_000_000).try_into().unwrap();
+
+        let mut its: libc::itimerspec = unsafe { mem::zeroed() };
+        its.it_interval.tv_sec = seconds_part;
+        its.it_interval.tv_nsec = nanos_part;
+        its.it_value.tv_sec = seconds_part;
+        its.it_value.tv_nsec = nanos_part;
+        its
+    }
+}

data/ext/pf2/src/signal_scheduler.rs

@@ -0,0 +1,311 @@
+#![deny(unsafe_op_in_unsafe_fn)]
+
+mod configuration;
+mod timer_installer;
+
+use self::configuration::{Configuration, TimeMode};
+use self::timer_installer::TimerInstaller;
+use crate::profile::Profile;
+use crate::profile_serializer::ProfileSerializer;
+use crate::sample::Sample;
+
+use core::panic;
+use std::collections::HashSet;
+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};
+
+use rb_sys::*;
+
+use crate::util::*;
+
+#[derive(Debug)]
+pub struct SignalScheduler {
+    configuration: Option<configuration::Configuration>,
+    profile: Option<Arc<RwLock<Profile>>>,
+}
+
+pub struct SignalHandlerArgs {
+    profile: Arc<RwLock<Profile>>,
+    context_ruby_thread: VALUE,
+}
+
+impl SignalScheduler {
+    fn new() -> Self {
+        Self {
+            configuration: None,
+            profile: None,
+        }
+    }
+
+    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 mut target_ruby_threads = HashSet::new();
+        unsafe {
+            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.as_ref().unwrap().clone(), // FIXME: don't clone
+            Arc::clone(&profile),
+        );
+
+        self.profile = Some(profile);
+
+        Qtrue.into()
+    }
+
+    fn stop(&mut self, _rbself: VALUE) -> VALUE {
+        if let Some(profile) = &self.profile {
+            // Finalize
+            match profile.try_write() {
+                Ok(mut profile) => {
+                    profile.flush_temporary_sample_buffer();
+                }
+                Err(_) => {
+                    println!("[pf2 ERROR] stop: Failed to acquire profile lock.");
+                    return Qfalse.into();
+                }
+            }
+
+            let profile = profile.try_read().unwrap();
+            log::debug!("Number of samples: {}", profile.samples.len());
+
+            let serialized = ProfileSerializer::serialize(&profile);
+            let serialized = CString::new(serialized).unwrap();
+            unsafe { rb_str_new_cstr(serialized.as_ptr()) }
+        } else {
+            panic!("stop() called before start()");
+        }
+    }
+
+    // Install signal handler for profiling events to the current process.
+    fn install_signal_handler(&self) {
+        let mut sa: libc::sigaction = unsafe { mem::zeroed() };
+        sa.sa_sigaction = Self::signal_handler as usize;
+        sa.sa_flags = libc::SA_SIGINFO;
+        let err = unsafe { libc::sigaction(libc::SIGALRM, &sa, null_mut()) };
+        if err != 0 {
+            panic!("sigaction failed: {}", err);
+        }
+    }
+
+    // Respond to the signal and collect a sample.
+    // This function is called when a timer fires.
+    //
+    // Expected to be async-signal-safe, but the current implementation is not.
+    extern "C" fn signal_handler(
+        _sig: c_int,
+        info: *mut libc::siginfo_t,
+        _ucontext: *mut libc::ucontext_t,
+    ) {
+        let args = unsafe {
+            let ptr = extract_si_value_sival_ptr(info) as *mut SignalHandlerArgs;
+            ManuallyDrop::new(Box::from_raw(ptr))
+        };
+
+        let mut profile = match args.profile.try_write() {
+            Ok(profile) => profile,
+            Err(_) => {
+                // FIXME: Do we want to properly collect GC samples? I don't know yet.
+                log::trace!("Failed to acquire profile lock (garbage collection possibly in progress). Dropping sample.");
+                return;
+            }
+        };
+
+        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.");
+        }
+    }
+
+    fn start_profile_buffer_flusher_thread(&self, profile: &Arc<RwLock<Profile>>) {
+        let profile = Arc::clone(profile);
+        thread::spawn(move || loop {
+            log::trace!("Flushing temporary sample buffer");
+            match profile.try_write() {
+                Ok(mut profile) => {
+                    profile.flush_temporary_sample_buffer();
+                }
+                Err(_) => {
+                    log::debug!("flusher: Failed to acquire profile lock");
+                }
+            }
+            thread::sleep(Duration::from_millis(500));
+        });
+    }
+
+    // Ruby Methods
+
+    pub unsafe extern "C" fn rb_initialize(
+        argc: c_int,
+        argv: *const VALUE,
+        rbself: VALUE,
+    ) -> VALUE {
+        let mut collector = unsafe { Self::get_struct_from(rbself) };
+        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)
+    }
+
+    // Functions for TypedData
+
+    // Extract the SignalScheduler struct from a Ruby object
+    unsafe fn get_struct_from(obj: VALUE) -> ManuallyDrop<Box<Self>> {
+        unsafe {
+            let ptr = rb_check_typeddata(obj, &RBDATA);
+            ManuallyDrop::new(Box::from_raw(ptr as *mut SignalScheduler))
+        }
+    }
+
+    #[allow(non_snake_case)]
+    pub unsafe extern "C" fn rb_alloc(_rbself: VALUE) -> VALUE {
+        let collector = Box::new(SignalScheduler::new());
+        unsafe { Arc::increment_strong_count(&collector) };
+
+        unsafe {
+            let rb_mPf2: VALUE = rb_define_module(cstr!("Pf2"));
+            let rb_cSignalScheduler =
+                rb_define_class_under(rb_mPf2, cstr!("SignalScheduler"), rb_cObject);
+
+            // "Wrap" the SignalScheduler struct into a Ruby object
+            rb_data_typed_object_wrap(
+                rb_cSignalScheduler,
+                Box::into_raw(collector) as *mut c_void,
+                &RBDATA,
+            )
+        }
+    }
+
+    unsafe extern "C" fn dmark(ptr: *mut c_void) {
+        unsafe {
+            let collector = ManuallyDrop::new(Box::from_raw(ptr as *mut SignalScheduler));
+            if let Some(profile) = &collector.profile {
+                match profile.read() {
+                    Ok(profile) => {
+                        profile.dmark();
+                    }
+                    Err(_) => {
+                        panic!("[pf2 FATAL] dmark: Failed to acquire profile lock.");
+                    }
+                }
+            }
+        }
+    }
+
+    unsafe extern "C" fn dfree(ptr: *mut c_void) {
+        unsafe {
+            drop(Box::from_raw(ptr as *mut SignalScheduler));
+        }
+    }
+
+    unsafe extern "C" fn dsize(_: *const c_void) -> size_t {
+        // FIXME: Report something better
+        mem::size_of::<SignalScheduler>() as size_t
+    }
+}
+
+static mut RBDATA: rb_data_type_t = rb_data_type_t {
+    wrap_struct_name: cstr!("SignalScheduler"),
+    function: rb_data_type_struct__bindgen_ty_1 {
+        dmark: Some(SignalScheduler::dmark),
+        dfree: Some(SignalScheduler::dfree),
+        dsize: Some(SignalScheduler::dsize),
+        dcompact: None,
+        reserved: [null_mut(); 1],
+    },
+    parent: null_mut(),
+    data: null_mut(),
+    flags: 0,
+};