itsi-server 0.1.1 → 0.2.2

data/ext/itsi_instrument_entry/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "itsi_instrument_entry"
+version = "0.1.0"
+edition = "2021"
+authors = ["Wouter Coppieters <wc@pico.net.nz>"]
+license = "MIT"
+publish = false
+
+
+[lib]
+proc-macro = true
+[dependencies]
+proc-macro2 = "1.0"
+quote = "1.0"
+syn = { version = "1.0", features = ["full"] }

data/ext/itsi_instrument_entry/src/lib.rs

@@ -0,0 +1,31 @@
+use proc_macro::TokenStream;
+use proc_macro2::TokenStream as TokenStream2;
+use quote::quote;
+use syn::{parse_macro_input, ItemFn};
+
+#[proc_macro_attribute]
+pub fn instrument_with_entry(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let attr_tokens = TokenStream2::from(attr);
+    let input_fn = parse_macro_input!(item as ItemFn);
+    let attrs = input_fn.attrs;
+    let vis = input_fn.vis;
+    let sig = input_fn.sig;
+    let block = input_fn.block;
+    let output = quote! {
+        #[cfg(debug_assertions)]
+        #[tracing::instrument(#attr_tokens)]
+        #(#attrs)*
+        #vis #sig {
+            tracing::trace!("");
+            #block
+        }
+
+        #[cfg(not(debug_assertions))]
+        #(#attrs)*
+        #vis #sig {
+            #block
+        }
+    };
+
+    output.into()
+}

data/ext/itsi_rb_helpers/Cargo.toml

@@ -4,5 +4,8 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
+cfg-if = "1.0.0"
 magnus = { version = "0.7.1", features = ["rb-sys", "bytes"] }
+nix = "0.29.0"
 rb-sys = "0.9.105"
+serde = "1.0.219"

data/ext/itsi_rb_helpers/src/heap_value.rs

@@ -0,0 +1,139 @@
+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> HeapValue<T>
+where
+    T: ReprValue,
+{
+    pub fn into_inner(self) -> T {
+        *self.0
+    }
+}
+
+impl<T> HeapValue<T>
+where
+    T: ReprValue,
+{
+    pub fn cloned(&self) -> T {
+        *self.0
+    }
+}
+
+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,45 @@
-use std::{os::raw::c_void, ptr::null_mut};
+use std::{ffi::c_int, os::raw::c_void, ptr::null_mut};
 
+use magnus::{
+    ArgList, RArray, Ruby, Thread, Value,
+    block::Proc,
+    rb_sys::{AsRawId, FromRawValue, protect},
+    value::{IntoId, LazyId, ReprValue},
+};
 use rb_sys::{
-    rb_thread_call_with_gvl, rb_thread_call_without_gvl, rb_thread_create, rb_thread_wakeup,
+    VALUE, rb_funcallv, 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");
+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() -> 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 +48,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))
     }
 }
 
@@ -67,18 +92,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 +121,109 @@ where
     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;
+        }
+        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);
+}
+
+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(())
+}