itsi-server 0.1.1 → 0.1.11

data/ext/itsi_rb_helpers/src/lib.rs

@@ -1,23 +1,45 @@
 use std::{os::raw::c_void, ptr::null_mut};
 
+use magnus::{
+    RArray, Ruby, Thread, Value,
+    block::Proc,
+    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");
+static ID_BACKTRACE: LazyId = LazyId::new("backtrace");
+
+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 +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)).unwrap()
     }
 }
 
@@ -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,81 @@ 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, ())
+        .unwrap();
+    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();
+
+    eprintln!("Ruby exception {:?}", rb_err);
+    for line in backtrace {
+        eprintln!("{}", line);
+    }
+}

data/ext/itsi_scheduler/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "itsi-scheduler"
+version = "0.1.0"
+edition = "2021"
+authors = ["Wouter Coppieters <wc@pico.net.nz>"]
+license = "MIT"
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+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

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+require "mkmf"
+require "rb_sys/mkmf"
+
+create_rust_makefile("itsi/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())
+    }
+}

data/ext/itsi_scheduler/src/itsi_scheduler.rs

@@ -0,0 +1,308 @@
+mod io_helpers;
+mod io_waiter;
+mod timer;
+use io_helpers::{build_interest, poll_readiness, set_nonblocking};
+use io_waiter::IoWaiter;
+use itsi_error::ItsiError;
+use itsi_rb_helpers::{call_without_gvl, create_ruby_thread};
+use magnus::{
+    error::Result as MagnusResult,
+    value::{InnerValue, Opaque, ReprValue},
+    Module, RClass, Ruby, Value,
+};
+use mio::{Events, Poll, Token, Waker};
+use parking_lot::{Mutex, RwLock};
+use std::{
+    collections::{BinaryHeap, HashMap, VecDeque},
+    os::fd::RawFd,
+    sync::Arc,
+    time::Duration,
+};
+use timer::Timer;
+use tracing::{debug, info, warn};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub(crate) struct Readiness(i16);
+
+impl std::fmt::Debug for ItsiScheduler {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ItsiScheduler").finish()
+    }
+}
+
+const WAKE_TOKEN: Token = Token(0);
+
+#[magnus::wrap(class = "Itsi::Scheduler", free_immediately, size)]
+pub(crate) struct ItsiScheduler {
+    timers: Mutex<BinaryHeap<Timer>>,
+    io_waiters: Mutex<HashMap<Token, IoWaiter>>,
+    registry: Mutex<HashMap<RawFd, VecDeque<IoWaiter>>>,
+    poll: Mutex<Poll>,
+    events: Mutex<Events>,
+    waker: Mutex<Waker>,
+}
+
+impl Default for ItsiScheduler {
+    fn default() -> Self {
+        let poll = Poll::new().unwrap();
+        let waker = Waker::new(poll.registry(), WAKE_TOKEN).unwrap();
+        let events = Events::with_capacity(1024);
+
+        ItsiScheduler {
+            timers: Mutex::new(BinaryHeap::new()),
+            io_waiters: Mutex::new(HashMap::new()),
+            registry: Mutex::new(HashMap::new()),
+            poll: Mutex::new(poll),
+            events: Mutex::new(events),
+            waker: Mutex::new(waker),
+        }
+    }
+}
+
+impl ItsiScheduler {
+    pub fn initialize(&self) {}
+
+    pub fn wake(&self) -> MagnusResult<()> {
+        self.waker.lock().wake().map_err(|_| {
+            magnus::Error::new(
+                magnus::exception::exception(),
+                "Failed to wake the scheduler",
+            )
+        })?;
+        Ok(())
+    }
+    pub fn register_io_wait(
+        &self,
+        io_obj: i32,
+        events: i16,
+        timeout: Option<f64>,
+        token: usize,
+    ) -> MagnusResult<Option<i16>> {
+        debug!(
+            "Registering IO Wait for {:?}, {:?}, {:?}, {:?}",
+            io_obj, events, timeout, token
+        );
+        let fd: RawFd = io_obj;
+
+        let readiness = poll_readiness(fd, events).unwrap_or(Readiness(0));
+        if readiness == Readiness(events) {
+            return Ok(Some(readiness.0));
+        }
+
+        set_nonblocking(fd)?;
+        let interest = build_interest(events)?;
+        let token = Token(token);
+        let mut waiter = IoWaiter::new(fd, events, token);
+        self.io_waiters.lock().insert(token, waiter.clone());
+        let mut binding = self.registry.lock();
+        let queue = binding.entry(fd).or_default();
+
+        queue.push_back(waiter.clone());
+
+        if queue.len() == 1 {
+            self.poll
+                .lock()
+                .registry()
+                .register(&mut waiter, token, interest)
+                .map_err(|e| ItsiError::ArgumentError(format!("register error: {}", e)))?;
+        }
+        Ok(None)
+    }
+
+    pub fn start_timer(&self, timeout: Option<f64>, token: usize) {
+        if timeout.is_some_and(|t| t >= 0.0) {
+            let timer_entry = Timer::new(Duration::from_secs_f64(timeout.unwrap()), Token(token));
+            self.timers.lock().push(timer_entry);
+        }
+    }
+    pub fn has_pending_io(&self) -> bool {
+        !self.timers.lock().is_empty() || !self.io_waiters.lock().is_empty()
+    }
+
+    pub fn class_info(msg: String) {
+        info!(msg);
+    }
+
+    pub fn info(&self, msg: String) {
+        info!(msg);
+    }
+
+    pub fn warn(&self, msg: String) {
+        warn!(msg);
+    }
+
+    pub fn debug(&self, msg: String) {
+        debug!(msg);
+    }
+
+    pub fn fetch_due_events(&self) -> MagnusResult<Option<Vec<(usize, i16)>>> {
+        call_without_gvl(|| {
+            let timeout = if let Some(timer) = self.timers.lock().peek() {
+                timer.duration().or(Some(Duration::ZERO))
+            } else {
+                None
+            };
+            let mut due_fibers: Option<Vec<(usize, i16)>> = None;
+            let mut io_waiters = self.io_waiters.lock();
+            if !io_waiters.is_empty() || timeout.is_none() {
+                let mut events = self.events.lock();
+                {
+                    let mut poll = self.poll.lock();
+                    poll.poll(&mut events, timeout)
+                        .map_err(|e| ItsiError::ArgumentError(format!("poll error: {}", e)))?;
+                };
+
+                for event in events.iter() {
+                    let token = event.token();
+                    if token == WAKE_TOKEN {
+                        continue;
+                    }
+
+                    let waiter = io_waiters.remove(&token);
+                    if waiter.is_none() {
+                        continue;
+                    }
+                    let mut waiter = waiter.unwrap();
+                    let mut evt_readiness = 0;
+                    if event.is_readable() {
+                        evt_readiness |= 1;
+                    }
+                    if event.is_priority() {
+                        evt_readiness |= 2;
+                    }
+                    if event.is_writable() {
+                        evt_readiness |= 4
+                    }
+                    self.poll
+                        .lock()
+                        .registry()
+                        .deregister(&mut waiter)
+                        .map_err(|_| {
+                            ItsiError::ArgumentError("Failed to deregister".to_string())
+                        })?;
+
+                    due_fibers
+                        .get_or_insert_default()
+                        .push((waiter.token.0, evt_readiness));
+
+                    let mut binding = self.registry.lock();
+                    // Pop the current item for the current waiter off the queue
+                    let queue = binding.get_mut(&(waiter.fd)).unwrap();
+                    queue.pop_front();
+
+                    if let Some(head) = queue.get_mut(0) {
+                        // Register the next item in the queue if there is one.
+                        let interest = build_interest(head.readiness)?;
+                        self.poll
+                            .lock()
+                            .registry()
+                            .register(head, head.token, interest)
+                            .map_err(|_| {
+                                ItsiError::ArgumentError("Failed to deregister".to_string())
+                            })?;
+                    } else {
+                        // Otherwise we drop the queue altogether.
+                        binding.remove(&waiter.fd);
+                    }
+                }
+                return Ok(due_fibers);
+            }
+            Ok(None)
+        })
+    }
+
+    pub fn run_blocking_in_thread<T, F>(&self, ruby: &Ruby, work: F) -> MagnusResult<Option<T>>
+    where
+        T: Send + Sync + std::fmt::Debug + 'static,
+        F: FnOnce() -> Option<T> + Send + 'static,
+    {
+        let result: Arc<RwLock<Option<T>>> = Arc::new(RwLock::new(None));
+        let result_clone = Arc::clone(&result);
+
+        let current_fiber = Opaque::from(ruby.fiber_current());
+        let scheduler = Opaque::from(
+            ruby.module_kernel()
+                .const_get::<_, RClass>("Fiber")
+                .unwrap()
+                .funcall::<_, _, Value>("scheduler", ())
+                .unwrap(),
+        );
+
+        create_ruby_thread(move || {
+            call_without_gvl(|| {
+                let outcome = work();
+                *result_clone.write() = outcome;
+            });
+
+            let ruby = Ruby::get().unwrap();
+            scheduler
+                .get_inner_with(&ruby)
+                .funcall::<_, _, Value>("unblock", (None::<String>, current_fiber))
+                .unwrap();
+        });
+
+        scheduler
+            .get_inner_with(ruby)
+            .funcall::<_, _, Value>("block", (None::<Value>, None::<u64>))?;
+
+        let result_opt = Arc::try_unwrap(result).unwrap().write().take();
+        Ok(result_opt)
+    }
+
+    pub fn address_resolve(
+        ruby: &Ruby,
+        rself: &Self,
+        hostname: String,
+    ) -> MagnusResult<Option<Vec<String>>> {
+        let result: Option<Vec<String>> = rself.run_blocking_in_thread(ruby, move || {
+            use std::net::ToSocketAddrs;
+            let addrs_res = (hostname.as_str(), 0).to_socket_addrs();
+            match addrs_res {
+                Ok(addrs) => {
+                    let ips: Vec<String> = addrs.map(|s| s.ip().to_string()).collect();
+                    Some(ips)
+                }
+                Err(_) => None,
+            }
+        })?;
+        Ok(result)
+    }
+
+    pub fn fetch_due_timers(&self) -> MagnusResult<Option<Vec<usize>>> {
+        call_without_gvl(|| {
+            let mut timers = self.timers.lock();
+            let mut io_waiters = self.io_waiters.lock();
+            let mut due_fibers: Option<Vec<usize>> = None;
+            while let Some(timer) = timers.peek() {
+                if timer.is_due() {
+                    due_fibers.get_or_insert_default().push(timer.token.0);
+                    if let Some(waiter) = io_waiters.remove(&timer.token) {
+                        let mut binding = self.registry.lock();
+                        // Pop the current item for the current waiter off the queue
+                        let queue = binding.get_mut(&waiter.fd).unwrap();
+                        queue.pop_front();
+
+                        if let Some(head) = queue.get_mut(0) {
+                            // Register the next item in the queue if there is one.
+                            let interest = build_interest(head.readiness)?;
+                            self.poll
+                                .lock()
+                                .registry()
+                                .register(head, head.token, interest)
+                                .map_err(|_| {
+                                    ItsiError::ArgumentError("Failed to deregister".to_string())
+                                })?;
+                        } else {
+                            // Otherwise we drop the queue altogether.
+                            binding.remove(&waiter.fd);
+                        }
+                    }
+                    timers.pop();
+                } else {
+                    break;
+                }
+            }
+            Ok(due_fibers)
+        })
+    }
+}

data/ext/itsi_scheduler/src/lib.rs

@@ -0,0 +1,38 @@
+use itsi_scheduler::ItsiScheduler;
+use magnus::{function, method, Class, Error, Module, Object, Ruby};
+mod itsi_scheduler;
+
+#[magnus::init]
+fn init(ruby: &Ruby) -> Result<(), Error> {
+    itsi_tracing::init();
+    let module = ruby.define_module("Itsi")?;
+    let scheduler = module.define_class("Scheduler", ruby.class_object())?;
+    scheduler.define_singleton_method("info", function!(ItsiScheduler::class_info, 1))?;
+    scheduler.define_alloc_func::<ItsiScheduler>();
+    scheduler.define_method("initialize", method!(ItsiScheduler::initialize, 0))?;
+    scheduler.define_method("wake", method!(ItsiScheduler::wake, 0))?;
+    scheduler.define_method(
+        "register_io_wait",
+        method!(ItsiScheduler::register_io_wait, 4),
+    )?;
+    scheduler.define_method("info", method!(ItsiScheduler::info, 1))?;
+    scheduler.define_method("debug", method!(ItsiScheduler::debug, 1))?;
+    scheduler.define_method("warn", method!(ItsiScheduler::warn, 1))?;
+    scheduler.define_method("start_timer", method!(ItsiScheduler::start_timer, 2))?;
+    scheduler.define_method(
+        "address_resolve",
+        method!(ItsiScheduler::address_resolve, 1),
+    )?;
+    scheduler.define_method("has_pending_io?", method!(ItsiScheduler::has_pending_io, 0))?;
+
+    scheduler.define_method(
+        "fetch_due_timers",
+        method!(ItsiScheduler::fetch_due_timers, 0),
+    )?;
+    scheduler.define_method(
+        "fetch_due_events",
+        method!(ItsiScheduler::fetch_due_events, 0),
+    )?;
+
+    Ok(())
+}