pf2 0.1.0 → 0.2.0

data/ext/pf2/src/profile.rs

@@ -0,0 +1,50 @@
+use std::collections::HashSet;
+use std::time::Instant;
+
+use rb_sys::*;
+
+use super::ringbuffer::Ringbuffer;
+use super::sample::Sample;
+
+// Capacity large enough to hold 1 second worth of samples for 16 threads
+// 16 threads * 20 samples per second * 1 second = 320
+const DEFAULT_RINGBUFFER_CAPACITY: usize = 320;
+
+#[derive(Debug)]
+pub struct Profile {
+    pub start_timestamp: Instant,
+    pub samples: Vec<Sample>,
+    pub temporary_sample_buffer: Ringbuffer,
+    known_values: HashSet<VALUE>,
+}
+
+impl Profile {
+    pub fn new() -> Self {
+        Self {
+            start_timestamp: Instant::now(),
+            samples: vec![],
+            temporary_sample_buffer: Ringbuffer::new(DEFAULT_RINGBUFFER_CAPACITY),
+            known_values: HashSet::new(),
+        }
+    }
+
+    pub fn flush_temporary_sample_buffer(&mut self) {
+        while let Some(sample) = self.temporary_sample_buffer.pop() {
+            self.known_values.insert(sample.ruby_thread);
+            for frame in sample.frames.iter() {
+                if frame == &0 {
+                    break;
+                }
+                self.known_values.insert(*frame);
+            }
+            self.samples.push(sample);
+        }
+    }
+
+    pub unsafe fn dmark(&self) {
+        for value in self.known_values.iter() {
+            rb_gc_mark(*value);
+        }
+        self.temporary_sample_buffer.dmark();
+    }
+}

data/ext/pf2/src/profile_serializer.rs

@@ -0,0 +1,130 @@
+use std::{collections::HashMap, ffi::CStr};
+
+use rb_sys::*;
+
+use crate::profile::Profile;
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct ProfileSerializer {
+    threads: HashMap<ThreadId, ThreadProfile>,
+}
+
+type ThreadId = VALUE;
+
+#[derive(Debug, Deserialize, Serialize)]
+struct ThreadProfile {
+    thread_id: ThreadId,
+    stack_tree: StackTreeNode,
+    #[serde(rename = "frames")]
+    frame_table: HashMap<FrameTableId, FrameTableEntry>,
+    samples: Vec<ProfileSample>,
+}
+
+impl ThreadProfile {
+    fn new(thread_id: ThreadId) -> ThreadProfile {
+        ThreadProfile {
+            thread_id,
+            // The root node
+            stack_tree: StackTreeNode {
+                children: HashMap::new(),
+                node_id: 0,
+                frame_id: 0,
+            },
+            frame_table: HashMap::new(),
+            samples: vec![],
+        }
+    }
+}
+
+type StackTreeNodeId = i32;
+
+// Arbitary value which is used inside StackTreeNode.
+// This VALUE should not be dereferenced as a pointer; we're merely using its pointer as a unique value.
+// (Probably should be reconsidered)
+type FrameTableId = VALUE;
+
+#[derive(Debug, Deserialize, Serialize)]
+struct StackTreeNode {
+    // TODO: Maybe a Vec<StackTreeNode> is enough?
+    // There's no particular meaning in using FrameTableId as key
+    children: HashMap<FrameTableId, StackTreeNode>,
+    // An arbitary ID (no particular meaning)
+    node_id: StackTreeNodeId,
+    // ?
+    frame_id: FrameTableId,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+struct FrameTableEntry {
+    full_label: String,
+}
+
+// Represents leaf (末端)
+#[derive(Debug, Deserialize, Serialize)]
+struct ProfileSample {
+    elapsed_ns: u128,
+    stack_tree_id: StackTreeNodeId,
+}
+
+impl ProfileSerializer {
+    pub fn serialize(profile: &Profile) -> String {
+        let mut sequence = 1;
+
+        let mut serializer = ProfileSerializer {
+            threads: HashMap::new(),
+        };
+
+        unsafe {
+            // Process each sample
+            for sample in profile.samples.iter() {
+                // Find the Thread profile for this sample
+                let thread_serializer = serializer
+                    .threads
+                    .entry(sample.ruby_thread)
+                    .or_insert(ThreadProfile::new(sample.ruby_thread));
+
+                // Stack frames, shallow to deep
+                let mut stack_tree = &mut thread_serializer.stack_tree;
+
+                for i in (0..(sample.line_count - 1)).rev() {
+                    let frame = sample.frames[i as usize];
+
+                    // Register frame metadata to frame table, if not registered yet
+                    let frame_table_id: FrameTableId = frame;
+                    thread_serializer
+                        .frame_table
+                        .entry(frame_table_id)
+                        .or_insert(FrameTableEntry {
+                            full_label: CStr::from_ptr(rb_string_value_cstr(
+                                &mut rb_profile_frame_full_label(frame),
+                            ))
+                            .to_str()
+                            .unwrap()
+                            .to_string(),
+                        });
+
+                    stack_tree = stack_tree.children.entry(frame_table_id).or_insert({
+                        let node = StackTreeNode {
+                            children: HashMap::new(),
+                            node_id: sequence,
+                            frame_id: frame_table_id,
+                        };
+                        sequence += 1;
+                        node
+                    });
+
+                    if i == 0 {
+                        // This is the leaf node, record a Sample
+                        let elapsed_ns = (sample.timestamp - profile.start_timestamp).as_nanos();
+                        thread_serializer.samples.push(ProfileSample {
+                            elapsed_ns,
+                            stack_tree_id: stack_tree.node_id,
+                        });
+                    }
+                }
+            }
+        }
+
+        serde_json::to_string(&serializer).unwrap()
+    }
+}

data/ext/pf2/src/ringbuffer.rs

@@ -0,0 +1,145 @@
+use crate::sample::Sample;
+
+#[derive(Debug)]
+pub struct Ringbuffer {
+    capacity: usize,
+    buffer: Vec<Option<Sample>>,
+    read_index: usize,
+    write_index: usize,
+}
+
+#[derive(Debug, PartialEq)]
+pub enum RingbufferError {
+    Full,
+}
+
+impl Ringbuffer {
+    pub fn new(capacity: usize) -> Self {
+        Self {
+            capacity,
+            buffer: std::iter::repeat_with(|| None)
+                .take(capacity + 1)
+                .collect::<Vec<_>>(),
+            read_index: 0,
+            write_index: 0,
+        }
+    }
+
+    // async-signal-safe
+    pub fn push(&mut self, sample: Sample) -> Result<(), RingbufferError> {
+        let next = (self.write_index + 1) % (self.capacity + 1);
+        if next == self.read_index {
+            return Err(RingbufferError::Full);
+        }
+        self.buffer[self.write_index] = Some(sample);
+        self.write_index = next;
+        Ok(())
+    }
+
+    pub fn pop(&mut self) -> Option<Sample> {
+        if self.read_index == self.write_index {
+            return None;
+        }
+        let sample = self.buffer[self.read_index].take();
+        self.read_index = (self.read_index + 1) % (self.capacity + 1);
+        sample
+    }
+
+    // This will call rb_gc_mark() for capacity * Sample::MAX_STACK_DEPTH * 2 times, which is a lot!
+    pub fn dmark(&self) {
+        for sample in self.buffer.iter().flatten() {
+            unsafe {
+                sample.dmark();
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::Instant;
+
+    #[test]
+    fn test_ringbuffer() {
+        let mut ringbuffer = Ringbuffer::new(2);
+        assert_eq!(ringbuffer.pop(), None);
+
+        let sample1 = Sample {
+            ruby_thread: 1,
+            timestamp: Instant::now(),
+            line_count: 0,
+            frames: [0; 500],
+            linenos: [0; 500],
+        };
+        let sample2 = Sample {
+            ruby_thread: 2,
+            timestamp: Instant::now(),
+            line_count: 0,
+            frames: [0; 500],
+            linenos: [0; 500],
+        };
+
+        ringbuffer.push(sample1).unwrap();
+        ringbuffer.push(sample2).unwrap();
+
+        assert_eq!(ringbuffer.pop().unwrap().ruby_thread, 1);
+        assert_eq!(ringbuffer.pop().unwrap().ruby_thread, 2);
+        assert_eq!(ringbuffer.pop(), None);
+    }
+
+    #[test]
+    fn test_ringbuffer_full() {
+        let mut ringbuffer = Ringbuffer::new(1);
+        let sample1 = Sample {
+            ruby_thread: 1,
+            timestamp: Instant::now(),
+            line_count: 0,
+            frames: [0; 500],
+            linenos: [0; 500],
+        };
+        let sample2 = Sample {
+            ruby_thread: 2,
+            timestamp: Instant::now(),
+            line_count: 0,
+            frames: [0; 500],
+            linenos: [0; 500],
+        };
+
+        ringbuffer.push(sample1).unwrap();
+        assert_eq!(ringbuffer.push(sample2), Err(RingbufferError::Full));
+    }
+
+    #[test]
+    fn test_ringbuffer_write_a_lot() {
+        let mut ringbuffer = Ringbuffer::new(2);
+        let sample1 = Sample {
+            ruby_thread: 1,
+            timestamp: Instant::now(),
+            line_count: 0,
+            frames: [0; 500],
+            linenos: [0; 500],
+        };
+        let sample2 = Sample {
+            ruby_thread: 2,
+            timestamp: Instant::now(),
+            line_count: 0,
+            frames: [0; 500],
+            linenos: [0; 500],
+        };
+        let sample3 = Sample {
+            ruby_thread: 3,
+            timestamp: Instant::now(),
+            line_count: 0,
+            frames: [0; 500],
+            linenos: [0; 500],
+        };
+
+        ringbuffer.push(sample1).unwrap();
+        ringbuffer.pop().unwrap();
+        ringbuffer.push(sample2).unwrap();
+        ringbuffer.pop().unwrap();
+        ringbuffer.push(sample3).unwrap();
+        assert_eq!(ringbuffer.pop().unwrap().ruby_thread, 3);
+    }
+}

data/ext/pf2/src/ruby_init.rs

@@ -0,0 +1,62 @@
+#![deny(unsafe_op_in_unsafe_fn)]
+
+use rb_sys::*;
+
+#[cfg(target_os = "linux")]
+use crate::signal_scheduler::SignalScheduler;
+use crate::timer_thread_scheduler::TimerThreadScheduler;
+use crate::util::*;
+
+#[allow(non_snake_case)]
+#[no_mangle]
+extern "C" fn Init_pf2() {
+    #[cfg(feature = "debug")]
+    {
+        env_logger::builder()
+            .format_timestamp(None)
+            .format_module_path(false)
+            .init();
+    }
+
+    unsafe {
+        let rb_mPf2: VALUE = rb_define_module(cstr!("Pf2"));
+
+        #[cfg(target_os = "linux")]
+        {
+            let rb_mPf2_SignalScheduler =
+                rb_define_class_under(rb_mPf2, cstr!("SignalScheduler"), rb_cObject);
+            rb_define_alloc_func(rb_mPf2_SignalScheduler, Some(SignalScheduler::rb_alloc));
+            rb_define_method(
+                rb_mPf2_SignalScheduler,
+                cstr!("start"),
+                Some(to_ruby_cfunc3(SignalScheduler::rb_start)),
+                2,
+            );
+            rb_define_method(
+                rb_mPf2_SignalScheduler,
+                cstr!("stop"),
+                Some(to_ruby_cfunc1(SignalScheduler::rb_stop)),
+                0,
+            );
+        }
+
+        let rb_mPf2_TimerThreadScheduler =
+            rb_define_class_under(rb_mPf2, cstr!("TimerThreadScheduler"), rb_cObject);
+        rb_define_alloc_func(
+            rb_mPf2_TimerThreadScheduler,
+            Some(TimerThreadScheduler::rb_alloc),
+        );
+        rb_define_method(
+            rb_mPf2_TimerThreadScheduler,
+            cstr!("start"),
+            Some(to_ruby_cfunc3(TimerThreadScheduler::rb_start)),
+            2,
+        );
+        rb_define_method(
+            rb_mPf2_TimerThreadScheduler,
+            cstr!("stop"),
+            Some(to_ruby_cfunc1(TimerThreadScheduler::rb_stop)),
+            0,
+        );
+    }
+}

data/ext/pf2/src/sample.rs

@@ -0,0 +1,45 @@
+use std::time::Instant;
+
+use rb_sys::*;
+
+const MAX_STACK_DEPTH: usize = 500;
+
+#[derive(Debug, PartialEq)]
+pub struct Sample {
+    pub ruby_thread: VALUE,
+    pub timestamp: Instant,
+    pub line_count: i32,
+    pub frames: [VALUE; MAX_STACK_DEPTH],
+    pub linenos: [i32; MAX_STACK_DEPTH],
+}
+
+impl Sample {
+    // Nearly async-signal-safe
+    // (rb_profile_thread_frames isn't defined as a-s-s)
+    pub fn capture(ruby_thread: VALUE) -> Self {
+        let mut sample = Sample {
+            ruby_thread,
+            timestamp: Instant::now(),
+            line_count: 0,
+            frames: [0; MAX_STACK_DEPTH],
+            linenos: [0; MAX_STACK_DEPTH],
+        };
+        unsafe {
+            sample.line_count = rb_profile_thread_frames(
+                ruby_thread,
+                0,
+                2000,
+                sample.frames.as_mut_ptr(),
+                sample.linenos.as_mut_ptr(),
+            );
+        };
+        sample
+    }
+
+    pub unsafe fn dmark(&self) {
+        rb_gc_mark(self.ruby_thread);
+        for frame in self.frames.iter() {
+            rb_gc_mark(*frame);
+        }
+    }
+}

data/ext/pf2/src/siginfo_t.c

@@ -0,0 +1,5 @@
+#include <signal.h>
+
+void *extract_si_value_sival_ptr(siginfo_t *siginfo) {
+  return siginfo->si_value.sival_ptr;
+}

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

@@ -0,0 +1,24 @@
+use std::str::FromStr;
+
+#[derive(Clone, Debug)]
+pub struct Configuration {
+    pub time_mode: TimeMode,
+}
+
+#[derive(Clone, Debug)]
+pub enum TimeMode {
+    CpuTime,
+    WallTime,
+}
+
+impl FromStr for TimeMode {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "cpu" => Ok(Self::CpuTime),
+            "wall" => Ok(Self::WallTime),
+            _ => Err(()),
+        }
+    }
+}

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

@@ -0,0 +1,192 @@
+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,
+    target_ruby_threads: HashSet<VALUE>,
+    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,
+        ruby_threads: &HashSet<VALUE>,
+        profile: Arc<RwLock<Profile>>,
+        track_new_threads: bool,
+    ) {
+        let registrar = Self {
+            internal: Box::new(Mutex::new(Internal {
+                configuration,
+                target_ruby_threads: ruby_threads.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 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.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.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 10 ms of CPU time
+        let mut timer: libc::timer_t = unsafe { mem::zeroed() };
+        match configuration.time_mode {
+            crate::signal_scheduler::TimeMode::CpuTime => {
+                let err = unsafe {
+                    libc::timer_create(libc::CLOCK_THREAD_CPUTIME_ID, &mut sigevent, &mut timer)
+                };
+                if err != 0 {
+                    panic!("timer_create failed: {}", err);
+                }
+            }
+            crate::signal_scheduler::TimeMode::WallTime => {
+                todo!("WallTime is not supported yet");
+            }
+        };
+        let mut its: libc::itimerspec = unsafe { mem::zeroed() };
+        its.it_interval.tv_sec = 0;
+        its.it_interval.tv_nsec = 10_000_000; // 10 ms
+        its.it_value.tv_sec = 0;
+        its.it_value.tv_nsec = 10_000_000;
+        let err = unsafe { libc::timer_settime(timer, 0, &its, null_mut()) };
+        if err != 0 {
+            panic!("timer_settime failed: {}", err);
+        }
+
+        log::debug!("timer registered for thread {}", current_pthread_id);
+    }
+}