itsi-scheduler 0.1.0 → 0.1.2

data/ext/itsi_rb_helpers/src/heap_value.rs

@@ -0,0 +1,121 @@
+use magnus::IntoValue;
+use magnus::rb_sys::AsRawValue;
+use magnus::value::BoxValue;
+use magnus::{Ruby, Value, value::ReprValue};
+use std::fmt::{self, Debug, Formatter};
+use std::ops::Deref;
+
+/// HeapVal is a wrapper for heap-allocated magnus ReprVa;ies
+/// that is marked as thread-safe(Send and Sync)
+/// It's up to the user to actually ensure this though,
+/// typically by only interacting with the value from a thread which
+/// holds the GVL.
+pub struct HeapValue<T>(pub BoxValue<T>)
+where
+    T: ReprValue;
+
+impl<T> PartialEq for HeapValue<T>
+where
+    T: ReprValue,
+{
+    fn eq(&self, other: &Self) -> bool {
+        self.0.as_raw() == other.0.as_raw()
+    }
+}
+
+impl<T> Deref for HeapValue<T>
+where
+    T: ReprValue,
+{
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T> HeapValue<T>
+where
+    T: ReprValue,
+{
+    pub fn inner(self) -> T {
+        *self.0
+    }
+}
+
+impl<T> IntoValue for HeapValue<T>
+where
+    T: ReprValue,
+{
+    fn into_value_with(self, _: &Ruby) -> Value {
+        self.0.into_value()
+    }
+}
+
+impl<T> From<T> for HeapValue<T>
+where
+    T: ReprValue,
+{
+    fn from(value: T) -> Self {
+        HeapValue(BoxValue::new(value))
+    }
+}
+
+impl<T> Clone for HeapValue<T>
+where
+    T: ReprValue + Clone,
+{
+    fn clone(&self) -> Self {
+        HeapValue(BoxValue::new(*self.0.deref()))
+    }
+}
+
+impl<T> Debug for HeapValue<T>
+where
+    T: ReprValue + Debug,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{:?}", self.0)
+    }
+}
+
+unsafe impl<T> Send for HeapValue<T> where T: ReprValue {}
+unsafe impl<T> Sync for HeapValue<T> where T: ReprValue {}
+
+/// HeapVal is a wrapper for heap-allocated magnus Values
+/// that is marked as thread-safe(Send and Sync)
+/// It's up to the user to actually ensure this though,
+/// typically by only interacting with the value from a thread which
+/// holds the GVL.
+pub struct HeapVal(HeapValue<Value>);
+impl Deref for HeapVal {
+    type Target = Value;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl IntoValue for HeapVal {
+    fn into_value_with(self, _: &Ruby) -> Value {
+        self.0.into_value()
+    }
+}
+
+impl From<Value> for HeapVal {
+    fn from(value: Value) -> Self {
+        HeapVal(HeapValue(BoxValue::new(value)))
+    }
+}
+
+impl Clone for HeapVal {
+    fn clone(&self) -> Self {
+        HeapVal(HeapValue(BoxValue::new(*self.0.deref())))
+    }
+}
+
+impl Debug for HeapVal {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{:?}", self.0)
+    }
+}

data/ext/itsi_rb_helpers/src/lib.rs

@@ -1,23 +1,43 @@
-use std::{os::raw::c_void, ptr::null_mut};
+use std::{os::raw::c_void, ptr::null_mut, sync::Arc};
 
+use magnus::{
+    RArray, Ruby, Thread, Value,
+    rb_sys::FromRawValue,
+    value::{LazyId, ReprValue},
+};
 use rb_sys::{
-    rb_thread_call_with_gvl, rb_thread_call_without_gvl, rb_thread_create, rb_thread_wakeup,
+    rb_thread_call_with_gvl, rb_thread_call_without_gvl, rb_thread_create, rb_thread_schedule,
+    rb_thread_wakeup,
 };
 
-pub fn create_ruby_thread<F>(f: F)
+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");
+
+pub fn schedule_thread() {
+    unsafe {
+        rb_thread_schedule();
+    };
+}
+pub fn create_ruby_thread<F>(f: F) -> Thread
 where
-    F: FnOnce() -> u64 + Send + 'static,
+    F: FnOnce() + Send + 'static,
 {
     extern "C" fn trampoline<F>(ptr: *mut c_void) -> u64
     where
-        F: FnOnce() -> u64,
+        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()
+        closure();
+        0
     }
 
     // Box the closure (wrapped in an Option) to create a stable pointer.
@@ -26,7 +46,10 @@ where
 
     // Call rb_thread_create with our trampoline and boxed closure.
     unsafe {
-        rb_thread_wakeup(rb_thread_create(Some(trampoline::<F>), ptr));
+        let thread = rb_thread_create(Some(trampoline::<F>), ptr);
+        rb_thread_wakeup(thread);
+        rb_thread_schedule();
+        Thread::from_value(Value::from_raw(thread)).unwrap()
     }
 }
 
@@ -67,18 +90,18 @@ where
 
 pub fn call_with_gvl<F, R>(f: F) -> R
 where
-    F: FnOnce() -> R,
+    F: FnOnce(Ruby) -> R,
 {
     extern "C" fn trampoline<F, R>(arg: *mut c_void) -> *mut c_void
     where
-        F: FnOnce() -> R,
+        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();
+        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);
@@ -96,3 +119,60 @@ where
     let result_box = unsafe { Box::from_raw(raw_result_ptr as *mut R) };
     *result_box
 }
+
+pub fn fork(after_fork: Arc<Option<impl Fn()>>) -> Option<i32> {
+    let ruby = Ruby::get().unwrap();
+    let fork_result = ruby
+        .module_kernel()
+        .funcall::<_, _, Option<i32>>(*ID_FORK, ())
+        .unwrap();
+    if fork_result.is_none() {
+        if let Some(f) = &*after_fork {
+            f()
+        }
+    }
+    fork_result
+}
+
+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;
+        }
+        eprintln!("Killing thread {:?}", thr.as_value());
+        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);
+}

data/ext/itsi_scheduler/Cargo.toml

@@ -10,7 +10,15 @@ publish = false
 crate-type = ["cdylib"]
 
 [dependencies]
-magnus = { version = "0.6.2" }
+magnus = { version = "0.7.1", 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.105"
+bytes = "1.10.1"
+nix = "0.29.0"
+tracing = "0.1.41"

data/ext/itsi_scheduler/extconf.rb

@@ -3,4 +3,4 @@
 require "mkmf"
 require "rb_sys/mkmf"
 
-create_rust_makefile("itsi_scheduler/itsi_scheduler")
+create_rust_makefile("scheduler/itsi_scheduler")

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())
+    }
+}