itsi-scheduler 0.2.22-aarch64-linux

data/ext/itsi_rb_helpers/src/lib.rs

@@ -0,0 +1,232 @@
+use std::{ffi::c_int, os::raw::c_void, ptr::null_mut};
+
+use magnus::{
+    block::Proc,
+    rb_sys::{protect, AsRawId, FromRawValue},
+    value::{IntoId, LazyId, ReprValue},
+    ArgList, RArray, Ruby, Thread, Value,
+};
+use rb_sys::{
+    rb_funcallv, rb_thread_call_with_gvl, rb_thread_call_without_gvl, rb_thread_create,
+    rb_thread_schedule, rb_thread_wakeup, VALUE,
+};
+
+mod heap_value;
+pub use heap_value::{HeapVal, HeapValue};
+static ID_FORK: LazyId = LazyId::new("fork");
+static ID_LIST: LazyId = LazyId::new("list");
+static ID_EQ: LazyId = LazyId::new("==");
+static ID_ALIVE: LazyId = LazyId::new("alive?");
+static ID_THREAD_VARIABLE_GET: LazyId = LazyId::new("thread_variable_get");
+static ID_BACKTRACE: LazyId = LazyId::new("backtrace");
+
+pub fn schedule_thread() {
+    unsafe {
+        rb_thread_schedule();
+    };
+}
+pub fn create_ruby_thread<F>(f: F) -> Option<Thread>
+where
+    F: FnOnce() + Send + 'static,
+{
+    extern "C" fn trampoline<F>(ptr: *mut c_void) -> u64
+    where
+        F: FnOnce(),
+    {
+        // Reconstruct the boxed Option<F> that holds our closure.
+        let boxed_closure: Box<Option<F>> = unsafe { Box::from_raw(ptr as *mut Option<F>) };
+        // Extract the closure. (The Option should be Some; panic otherwise.)
+        let closure = (*boxed_closure).expect("Closure already taken");
+        // Call the closure and return its result.
+        closure();
+        0
+    }
+
+    // Box the closure (wrapped in an Option) to create a stable pointer.
+    let boxed_closure = Box::new(Some(f));
+    let ptr = Box::into_raw(boxed_closure) as *mut c_void;
+
+    // Call rb_thread_create with our trampoline and boxed closure.
+    unsafe {
+        let thread = rb_thread_create(Some(trampoline::<F>), ptr);
+        rb_thread_wakeup(thread);
+        rb_thread_schedule();
+        Thread::from_value(Value::from_raw(thread))
+    }
+}
+
+pub fn call_without_gvl<F, R>(f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    // This is the function that Ruby calls “in the background” with the GVL released.
+    extern "C" fn trampoline<F, R>(arg: *mut c_void) -> *mut c_void
+    where
+        F: FnOnce() -> R,
+    {
+        // 1) Reconstruct the Box that holds our closure
+        let closure_ptr = arg as *mut Option<F>;
+        let closure = unsafe { (*closure_ptr).take().expect("Closure already taken") };
+
+        // 2) Call the user’s closure
+        let result = closure();
+
+        // 3) Box up the result so we can return a pointer to it
+        let boxed_result = Box::new(result);
+        Box::into_raw(boxed_result) as *mut c_void
+    }
+
+    // Box up the closure so we have a stable pointer
+    let mut closure_opt = Some(f);
+    let closure_ptr = &mut closure_opt as *mut Option<F> as *mut c_void;
+
+    // 4) Actually call `rb_thread_call_without_gvl`
+    let raw_result_ptr = unsafe {
+        rb_thread_call_without_gvl(Some(trampoline::<F, R>), closure_ptr, None, null_mut())
+    };
+
+    // 5) Convert the returned pointer back into R
+    let result_box = unsafe { Box::from_raw(raw_result_ptr as *mut R) };
+    *result_box
+}
+
+pub fn call_with_gvl<F, R>(f: F) -> R
+where
+    F: FnOnce(Ruby) -> R,
+{
+    extern "C" fn trampoline<F, R>(arg: *mut c_void) -> *mut c_void
+    where
+        F: FnOnce(Ruby) -> R,
+    {
+        // 1) Reconstruct the Box that holds our closure
+        let closure_ptr = arg as *mut Option<F>;
+        let closure = unsafe { (*closure_ptr).take().expect("Closure already taken") };
+
+        // 2) Call the user’s closure
+        let result = closure(Ruby::get().unwrap());
+
+        // 3) Box up the result so we can return a pointer to it
+        let boxed_result = Box::new(result);
+        Box::into_raw(boxed_result) as *mut c_void
+    }
+
+    // Box up the closure so we have a stable pointer
+    let mut closure_opt = Some(f);
+    let closure_ptr = &mut closure_opt as *mut Option<F> as *mut c_void;
+
+    // 4) Actually call `rb_thread_call_without_gvl`
+    let raw_result_ptr = unsafe { rb_thread_call_with_gvl(Some(trampoline::<F, R>), closure_ptr) };
+
+    // 5) Convert the returned pointer back into R
+    let result_box = unsafe { Box::from_raw(raw_result_ptr as *mut R) };
+    *result_box
+}
+
+pub fn fork(after_fork: Option<HeapValue<Proc>>) -> Option<i32> {
+    let ruby = Ruby::get().unwrap();
+    let fork_result = ruby
+        .module_kernel()
+        .funcall::<_, _, Option<i32>>(*ID_FORK, ())
+        .ok()
+        .flatten();
+    if fork_result.is_none() {
+        if let Some(proc) = after_fork {
+            call_proc_and_log_errors(proc)
+        }
+    }
+    fork_result
+}
+
+pub fn call_proc_and_log_errors(proc: HeapValue<Proc>) {
+    if let Err(e) = proc.call::<_, Value>(()) {
+        if let Some(value) = e.value() {
+            print_rb_backtrace(value);
+        } else {
+            eprintln!("Error occurred {:?}", e);
+        }
+    }
+}
+
+pub fn kill_threads<T>(threads: Vec<T>)
+where
+    T: ReprValue,
+{
+    for thr in &threads {
+        let alive: bool = thr
+            .funcall(*ID_ALIVE, ())
+            .expect("Failed to check if thread is alive");
+        if !alive {
+            eprintln!("Thread killed");
+            break;
+        }
+        thr.funcall::<_, _, Value>("terminate", ())
+            .expect("Failed to kill thread");
+    }
+}
+
+pub fn terminate_non_fork_safe_threads() {
+    let ruby = Ruby::get().unwrap();
+    let thread_class = ruby.class_thread();
+    let current: Thread = ruby.thread_current();
+    let threads: RArray = thread_class
+        .funcall(*ID_LIST, ())
+        .expect("Failed to list Ruby threads");
+
+    let non_fork_safe_threads = threads
+        .into_iter()
+        .filter_map(|v| {
+            let v_thread = Thread::from_value(v).unwrap();
+            let non_fork_safe = !v_thread
+                .funcall::<_, _, bool>(*ID_EQ, (current,))
+                .unwrap_or(false)
+                && !v_thread
+                    .funcall::<_, _, bool>(*ID_THREAD_VARIABLE_GET, (ruby.sym_new("fork_safe"),))
+                    .unwrap_or(false);
+            if non_fork_safe {
+                Some(v_thread)
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<_>>();
+
+    kill_threads(non_fork_safe_threads);
+}
+
+pub fn print_rb_backtrace(rb_err: Value) {
+    let backtrace = rb_err
+        .funcall::<_, _, Vec<String>>(*ID_BACKTRACE, ())
+        .unwrap_or_default();
+    let rust_backtrace = std::backtrace::Backtrace::capture().to_string();
+    eprintln!("Ruby exception {:?}", rb_err);
+    for line in backtrace {
+        eprintln!("{}", line);
+    }
+    for line in rust_backtrace.lines() {
+        eprintln!("{}", line);
+    }
+}
+
+pub fn funcall_no_ret<T, M, A>(target: T, method: M, args: A) -> magnus::error::Result<()>
+where
+    T: ReprValue,
+    M: IntoId,
+    A: ArgList,
+{
+    protect(|| {
+        let handle = Ruby::get().unwrap();
+        let method = method.into_id_with(&handle);
+        let args = args.into_arg_list_with(&handle);
+        let slice = args.as_ref();
+        unsafe {
+            rb_funcallv(
+                target.as_rb_value(),
+                method.as_raw(),
+                slice.len() as c_int,
+                slice.as_ptr() as *const VALUE,
+            );
+        }
+        0
+    })?;
+    Ok(())
+}

data/ext/itsi_scheduler/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "itsi-scheduler"
+version = "0.2.22"
+edition = "2021"
+authors = ["Wouter Coppieters <wc@pico.net.nz>"]
+license = "MIT"
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+magnus = { version = "0.8.2", features = ["rb-sys", "bytes"] }
+derive_more = { version = "2.0.1", features = ["debug"] }
+itsi_tracing = { path = "../itsi_tracing" }
+itsi_rb_helpers = { path = "../itsi_rb_helpers" }
+itsi_error = { path = "../itsi_error" }
+itsi_instrument_entry = { path = "../itsi_instrument_entry" }
+parking_lot = "0.12.3"
+mio = { version = "1.0.3", features = ["os-poll", "os-ext"] }
+rb-sys = "0.9.117"
+bytes = "1.10.1"
+nix = "0.29.0"
+tracing = "0.1.41"

data/ext/itsi_scheduler/extconf.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+require 'mkmf'
+require 'rb_sys/mkmf'
+
+create_rust_makefile('itsi/scheduler/itsi_scheduler') do |r|
+  r.extra_rustflags = []
+  r.env = {
+    'BINDGEN_EXTRA_CLANG_ARGS' => '-include stdbool.h -std=c99'
+  }
+end

data/ext/itsi_scheduler/src/itsi_scheduler/io_helpers.rs

@@ -0,0 +1,56 @@
+use std::os::fd::RawFd;
+
+use itsi_error::{ItsiError, Result};
+use mio::Interest;
+use nix::libc::{fcntl, poll, pollfd, F_GETFL, F_SETFL, O_NONBLOCK};
+
+use super::Readiness;
+
+pub fn set_nonblocking(fd: RawFd) -> itsi_error::Result<()> {
+    unsafe {
+        let flags = fcntl(fd, F_GETFL);
+        if flags < 0 {
+            return Err(ItsiError::ArgumentError(format!(
+                "fcntl(F_GETFL) error for fd {}: {}",
+                fd,
+                std::io::Error::last_os_error()
+            )));
+        }
+        let new_flags = flags | O_NONBLOCK;
+        if fcntl(fd, F_SETFL, new_flags) < 0 {
+            return Err(ItsiError::ArgumentError(format!(
+                "fcntl(F_SETFL) error for fd {}: {}",
+                fd,
+                std::io::Error::last_os_error()
+            )));
+        }
+    }
+    Ok(())
+}
+
+pub fn poll_readiness(fd: RawFd, events: i16) -> Option<Readiness> {
+    let mut pfd = pollfd {
+        fd,
+        events,
+        revents: 0,
+    };
+    let ret = unsafe { poll(&mut pfd as *mut pollfd, 1, 0) };
+    if ret > 0 {
+        return Some(Readiness(pfd.revents));
+    }
+    None
+}
+
+pub fn build_interest(events: i16) -> Result<Interest> {
+    let mut interest_opt = None;
+    if events & 1 != 0 {
+        interest_opt = Some(Interest::READABLE);
+    }
+    if events & 4 != 0 {
+        interest_opt = Some(match interest_opt {
+            Some(i) => i | Interest::WRITABLE,
+            None => Interest::WRITABLE,
+        });
+    }
+    interest_opt.ok_or_else(|| ItsiError::ArgumentError("No valid event specified".to_owned()))
+}

data/ext/itsi_scheduler/src/itsi_scheduler/io_waiter.rs

@@ -0,0 +1,44 @@
+use derive_more::Debug;
+use mio::{event::Source, unix::SourceFd, Interest, Token};
+use std::os::fd::RawFd;
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct IoWaiter {
+    pub fd: RawFd,
+    pub readiness: i16,
+    pub token: Token,
+}
+
+impl IoWaiter {
+    pub fn new(fd: RawFd, readiness: i16, token: Token) -> Self {
+        Self {
+            fd,
+            readiness,
+            token,
+        }
+    }
+}
+
+impl Source for IoWaiter {
+    fn register(
+        &mut self,
+        registry: &mio::Registry,
+        token: Token,
+        interests: Interest,
+    ) -> std::io::Result<()> {
+        SourceFd(&self.fd).register(registry, token, interests)
+    }
+
+    fn reregister(
+        &mut self,
+        registry: &mio::Registry,
+        token: Token,
+        interests: Interest,
+    ) -> std::io::Result<()> {
+        SourceFd(&self.fd).reregister(registry, token, interests)
+    }
+
+    fn deregister(&mut self, registry: &mio::Registry) -> std::io::Result<()> {
+        SourceFd(&self.fd).deregister(registry)
+    }
+}

data/ext/itsi_scheduler/src/itsi_scheduler/timer.rs

@@ -0,0 +1,44 @@
+use std::{
+    cmp::Ordering,
+    time::{Duration, Instant},
+};
+
+use mio::Token;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Timer {
+    pub wake_time: Instant,
+    pub token: Token,
+}
+impl PartialOrd for Timer {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for Timer {
+    fn cmp(&self, other: &Self) -> Ordering {
+        // Reverse the order: a timer with an earlier wake_time should be considered greater.
+        other
+            .wake_time
+            .cmp(&self.wake_time)
+            .then_with(|| other.token.cmp(&self.token))
+    }
+}
+
+impl Timer {
+    pub fn new(wake_in: Duration, token: Token) -> Self {
+        Self {
+            wake_time: Instant::now() + wake_in,
+            token,
+        }
+    }
+
+    pub fn is_due(&self) -> bool {
+        self.wake_time <= Instant::now()
+    }
+
+    pub(crate) fn duration(&self) -> Option<Duration> {
+        self.wake_time.checked_duration_since(Instant::now())
+    }
+}