omq-backend-rust 0.1.6 → 0.2.0

data/ext/omq_backend_rust/src/rb.rs

@@ -0,0 +1,433 @@
+use std::ffi::{CStr, CString, c_char, c_long, c_void};
+use std::panic::{AssertUnwindSafe, catch_unwind};
+use std::ptr;
+
+use rb_sys::{VALUE, rb_data_type_t};
+
+pub type RbResult<T = VALUE> = Result<T, RubyErr>;
+
+#[derive(Debug)]
+pub enum RubyErr {
+    Exception(VALUE),
+    Error { class: VALUE, message: String },
+}
+
+impl RubyErr {
+    pub fn new(class: VALUE, message: impl Into<String>) -> Self {
+        Self::Error {
+            class,
+            message: message.into(),
+        }
+    }
+
+    pub fn arg(message: impl Into<String>) -> Self {
+        Self::new(unsafe { rb_sys::rb_eArgError }, message)
+    }
+
+    pub fn io(message: impl Into<String>) -> Self {
+        Self::new(unsafe { rb_sys::rb_eIOError }, message)
+    }
+
+    pub fn runtime(message: impl Into<String>) -> Self {
+        Self::new(unsafe { rb_sys::rb_eRuntimeError }, message)
+    }
+
+    pub fn type_error(message: impl Into<String>) -> Self {
+        Self::new(unsafe { rb_sys::rb_eTypeError }, message)
+    }
+
+    fn current_exception() -> Self {
+        let err = unsafe { rb_sys::rb_errinfo() };
+        if err == qnil() {
+            Self::runtime("Ruby exception")
+        } else {
+            Self::Exception(err)
+        }
+    }
+}
+
+pub fn wrap<F>(f: F) -> VALUE
+where
+    F: FnOnce() -> RbResult<VALUE>,
+{
+    match catch_unwind(AssertUnwindSafe(f)) {
+        Ok(Ok(value)) => value,
+        Ok(Err(err)) => raise(err),
+        Err(_) => raise(RubyErr::runtime("native Rust panic")),
+    }
+}
+
+pub fn wrap_init<F>(f: F)
+where
+    F: FnOnce() -> RbResult<()>,
+{
+    match catch_unwind(AssertUnwindSafe(f)) {
+        Ok(Ok(())) => {}
+        Ok(Err(err)) => raise(err),
+        Err(_) => raise(RubyErr::runtime("native Rust panic")),
+    }
+}
+
+pub fn raise(err: RubyErr) -> ! {
+    match err {
+        RubyErr::Exception(exc) => unsafe { rb_sys::rb_exc_raise(exc) },
+        RubyErr::Error { class, message } => {
+            let message = message.replace('\0', "\\0");
+            let c_message =
+                CString::new(message).unwrap_or_else(|_| CString::new("Ruby error").unwrap());
+            let exc = unsafe { rb_sys::rb_exc_new_cstr(class, c_message.as_ptr()) };
+            unsafe { rb_sys::rb_exc_raise(exc) }
+        }
+    }
+}
+
+struct ProtectData<F> {
+    func: Option<F>,
+    panicked: bool,
+}
+
+pub fn protect_value<F>(func: F) -> RbResult<VALUE>
+where
+    F: FnOnce() -> VALUE,
+{
+    unsafe extern "C" fn call<F>(arg: VALUE) -> VALUE
+    where
+        F: FnOnce() -> VALUE,
+    {
+        let data = unsafe { &mut *(arg as *mut ProtectData<F>) };
+        let Some(func) = data.func.take() else {
+            data.panicked = true;
+            return qnil();
+        };
+
+        match catch_unwind(AssertUnwindSafe(func)) {
+            Ok(value) => value,
+            Err(_) => {
+                data.panicked = true;
+                qnil()
+            }
+        }
+    }
+
+    let mut data = ProtectData {
+        func: Some(func),
+        panicked: false,
+    };
+    let mut state = 0;
+    let value = unsafe {
+        rb_sys::rb_protect(
+            Some(call::<F>),
+            &mut data as *mut ProtectData<F> as VALUE,
+            &mut state,
+        )
+    };
+
+    if state != 0 {
+        Err(RubyErr::current_exception())
+    } else if data.panicked {
+        Err(RubyErr::runtime("native Rust panic"))
+    } else {
+        Ok(value)
+    }
+}
+
+pub fn protect_unit<F>(func: F) -> RbResult<()>
+where
+    F: FnOnce(),
+{
+    protect_value(|| {
+        func();
+        qnil()
+    })?;
+    Ok(())
+}
+
+pub const fn qnil() -> VALUE {
+    rb_sys::ruby_special_consts::RUBY_Qnil as VALUE
+}
+
+pub const fn qtrue() -> VALUE {
+    rb_sys::ruby_special_consts::RUBY_Qtrue as VALUE
+}
+
+pub const fn qfalse() -> VALUE {
+    rb_sys::ruby_special_consts::RUBY_Qfalse as VALUE
+}
+
+pub const fn qundef() -> VALUE {
+    rb_sys::ruby_special_consts::RUBY_Qundef as VALUE
+}
+
+pub fn bool_value(value: bool) -> VALUE {
+    if value { qtrue() } else { qfalse() }
+}
+
+pub fn check_hash(value: VALUE) -> RbResult<()> {
+    let is_hash = unsafe { rb_sys::rb_obj_is_kind_of(value, rb_sys::rb_cHash) };
+    if is_hash == qtrue() {
+        Ok(())
+    } else {
+        Err(RubyErr::type_error("expected Hash"))
+    }
+}
+
+pub fn check_array(value: VALUE) -> RbResult<()> {
+    let is_array = unsafe { rb_sys::rb_obj_is_kind_of(value, rb_sys::rb_cArray) };
+    if is_array == qtrue() {
+        Ok(())
+    } else {
+        Err(RubyErr::type_error("expected Array"))
+    }
+}
+
+pub fn hash_get(hash: VALUE, key: &str) -> RbResult<Option<VALUE>> {
+    check_hash(hash)?;
+    let key = new_utf8_string(key)?;
+    let value = protect_value(|| unsafe { rb_sys::rb_hash_lookup2(hash, key, qundef()) })?;
+    if value == qundef() {
+        Ok(None)
+    } else {
+        Ok(Some(value))
+    }
+}
+
+pub fn hash_new() -> RbResult<VALUE> {
+    protect_value(|| unsafe { rb_sys::rb_hash_new() })
+}
+
+pub fn hash_aset(hash: VALUE, key: VALUE, value: VALUE) -> RbResult<()> {
+    protect_value(|| unsafe { rb_sys::rb_hash_aset(hash, key, value) })?;
+    Ok(())
+}
+
+pub fn array_new() -> RbResult<VALUE> {
+    protect_value(|| unsafe { rb_sys::rb_ary_new() })
+}
+
+pub fn array_new_capa(capacity: usize) -> RbResult<VALUE> {
+    let capacity = c_long_len(capacity)?;
+    protect_value(|| unsafe { rb_sys::rb_ary_new_capa(capacity) })
+}
+
+pub fn array_len(array: VALUE) -> RbResult<usize> {
+    check_array(array)?;
+    let len = unsafe { rb_sys::RARRAY_LEN(array) };
+    if len < 0 {
+        Err(RubyErr::runtime("negative Array length"))
+    } else {
+        Ok(len as usize)
+    }
+}
+
+pub fn array_entry(array: VALUE, index: usize) -> RbResult<VALUE> {
+    check_array(array)?;
+    let index = c_long_len(index)?;
+    protect_value(|| unsafe { rb_sys::rb_ary_entry(array, index) })
+}
+
+pub fn array_push(array: VALUE, value: VALUE) -> RbResult<()> {
+    protect_value(|| unsafe { rb_sys::rb_ary_push(array, value) })?;
+    Ok(())
+}
+
+pub fn symbol(name: &str) -> RbResult<VALUE> {
+    let len = c_long_len(name.len())?;
+    protect_value(|| unsafe {
+        let id = rb_sys::rb_intern2(name.as_ptr() as *const c_char, len);
+        rb_sys::rb_id2sym(id)
+    })
+}
+
+pub fn new_binary_string(bytes: &[u8]) -> RbResult<VALUE> {
+    let len = c_long_len(bytes.len())?;
+    let ptr = if bytes.is_empty() {
+        ptr::null()
+    } else {
+        bytes.as_ptr() as *const c_char
+    };
+    let value = protect_value(|| unsafe { rb_sys::rb_str_new(ptr, len) })?;
+    protect_unit(|| unsafe { rb_sys::RB_OBJ_FREEZE(value) })?;
+    Ok(value)
+}
+
+pub fn new_utf8_string(text: &str) -> RbResult<VALUE> {
+    let len = c_long_len(text.len())?;
+    protect_value(|| unsafe { rb_sys::rb_utf8_str_new(text.as_ptr() as *const c_char, len) })
+}
+
+pub fn string_value(value: VALUE) -> RbResult<VALUE> {
+    protect_value(|| unsafe { rb_sys::rb_str_to_str(value) })
+}
+
+pub fn value_to_bytes(value: VALUE) -> RbResult<Vec<u8>> {
+    let string = string_value(value)?;
+    let len = unsafe { rb_sys::RSTRING_LEN(string) };
+    if len < 0 {
+        return Err(RubyErr::runtime("negative String length"));
+    }
+    if len == 0 {
+        return Ok(Vec::new());
+    }
+
+    let ptr = unsafe { rb_sys::RSTRING_PTR(string) };
+    if ptr.is_null() {
+        return Err(RubyErr::runtime("null String pointer"));
+    }
+
+    let bytes = unsafe { std::slice::from_raw_parts(ptr as *const u8, len as usize) };
+    Ok(bytes.to_vec())
+}
+
+pub fn value_to_string(value: VALUE) -> RbResult<String> {
+    let bytes = value_to_bytes(value)?;
+    String::from_utf8(bytes).map_err(|_| RubyErr::type_error("expected UTF-8 String"))
+}
+
+pub fn value_to_i64(value: VALUE) -> RbResult<i64> {
+    let mut out = 0i64;
+    protect_value(|| {
+        out = unsafe { rb_sys::rb_num2long(value) as i64 };
+        qnil()
+    })?;
+    Ok(out)
+}
+
+pub fn value_to_f64(value: VALUE) -> RbResult<f64> {
+    let mut out = 0.0f64;
+    protect_value(|| {
+        out = unsafe { rb_sys::rb_num2dbl(value) };
+        qnil()
+    })?;
+    Ok(out)
+}
+
+pub fn value_to_bool(value: VALUE) -> RbResult<bool> {
+    if value == qtrue() {
+        Ok(true)
+    } else if value == qfalse() {
+        Ok(false)
+    } else {
+        Err(RubyErr::type_error("expected true or false"))
+    }
+}
+
+pub fn int_value(value: i32) -> VALUE {
+    unsafe { rb_sys::rb_int2inum(value as isize) }
+}
+
+pub unsafe fn wrap_typed_data<T>(
+    class: VALUE,
+    value: Box<T>,
+    data_type: *const rb_data_type_t,
+) -> RbResult<VALUE> {
+    let raw = Box::into_raw(value);
+    match protect_value(|| unsafe {
+        rb_sys::rb_data_typed_object_wrap(class, raw as *mut c_void, data_type)
+    }) {
+        Ok(value) => Ok(value),
+        Err(err) => {
+            unsafe { drop(Box::from_raw(raw)) };
+            Err(err)
+        }
+    }
+}
+
+pub unsafe fn typed_data_ref<T>(
+    value: VALUE,
+    data_type: *const rb_data_type_t,
+    type_name: &str,
+) -> RbResult<&'static T> {
+    let mut ptr = std::ptr::null_mut();
+    protect_unit(|| unsafe {
+        ptr = rb_sys::rb_check_typeddata(value, data_type);
+    })
+    .map_err(|_| RubyErr::type_error(format!("expected {type_name}")))?;
+    if ptr.is_null() {
+        return Err(RubyErr::runtime(format!(
+            "{type_name} data pointer is null"
+        )));
+    }
+
+    Ok(unsafe { &*(ptr as *const T) })
+}
+
+pub unsafe fn define_module(name: &CStr) -> RbResult<VALUE> {
+    protect_value(|| unsafe { rb_sys::rb_define_module(name.as_ptr()) })
+}
+
+pub unsafe fn define_module_under(outer: VALUE, name: &CStr) -> RbResult<VALUE> {
+    protect_value(|| unsafe { rb_sys::rb_define_module_under(outer, name.as_ptr()) })
+}
+
+pub unsafe fn define_class_under(outer: VALUE, name: &CStr, superclass: VALUE) -> RbResult<VALUE> {
+    protect_value(|| unsafe { rb_sys::rb_define_class_under(outer, name.as_ptr(), superclass) })
+}
+
+pub unsafe fn undef_alloc_func(class: VALUE) -> RbResult<()> {
+    protect_unit(|| unsafe { rb_sys::rb_undef_alloc_func(class) })
+}
+
+pub unsafe fn define_module_function_1(
+    module: VALUE,
+    name: &CStr,
+    func: unsafe extern "C" fn(VALUE, VALUE) -> VALUE,
+) -> RbResult<()> {
+    protect_unit(|| unsafe {
+        rb_sys::rb_define_module_function(module, name.as_ptr(), Some(transmute_1(func)), 1)
+    })
+}
+
+pub unsafe fn define_singleton_method_1(
+    object: VALUE,
+    name: &CStr,
+    func: unsafe extern "C" fn(VALUE, VALUE) -> VALUE,
+) -> RbResult<()> {
+    protect_unit(|| unsafe {
+        rb_sys::rb_define_singleton_method(object, name.as_ptr(), Some(transmute_1(func)), 1)
+    })
+}
+
+pub unsafe fn define_method_0(
+    class: VALUE,
+    name: &CStr,
+    func: unsafe extern "C" fn(VALUE) -> VALUE,
+) -> RbResult<()> {
+    protect_unit(|| unsafe {
+        rb_sys::rb_define_method(class, name.as_ptr(), Some(transmute_0(func)), 0)
+    })
+}
+
+pub unsafe fn define_method_1(
+    class: VALUE,
+    name: &CStr,
+    func: unsafe extern "C" fn(VALUE, VALUE) -> VALUE,
+) -> RbResult<()> {
+    protect_unit(|| unsafe {
+        rb_sys::rb_define_method(class, name.as_ptr(), Some(transmute_1(func)), 1)
+    })
+}
+
+unsafe fn transmute_0(
+    func: unsafe extern "C" fn(VALUE) -> VALUE,
+) -> unsafe extern "C" fn() -> VALUE {
+    unsafe {
+        std::mem::transmute::<unsafe extern "C" fn(VALUE) -> VALUE, unsafe extern "C" fn() -> VALUE>(
+            func,
+        )
+    }
+}
+
+unsafe fn transmute_1(
+    func: unsafe extern "C" fn(VALUE, VALUE) -> VALUE,
+) -> unsafe extern "C" fn() -> VALUE {
+    unsafe {
+        std::mem::transmute::<
+            unsafe extern "C" fn(VALUE, VALUE) -> VALUE,
+            unsafe extern "C" fn() -> VALUE,
+        >(func)
+    }
+}
+
+fn c_long_len(len: usize) -> RbResult<c_long> {
+    c_long::try_from(len).map_err(|_| RubyErr::arg("length too large"))
+}

data/ext/omq_backend_rust/src/runtime.rs

@@ -27,10 +27,10 @@ pub fn ensure_runtime(io_threads: usize) -> Handle {
     }
     let mut guard = RUNTIME.lock().unwrap();
     let pid = std::process::id();
-    if let Some(ref rt) = *guard {
-        if rt.pid == pid {
-            return rt.handle.clone();
-        }
+    if let Some(ref rt) = *guard
+        && rt.pid == pid
+    {
+        return rt.handle.clone();
     }
     let (tx, rx) = flume::unbounded::<Job>();
     let (handle_tx, handle_rx) = flume::bounded::<Handle>(1);
@@ -69,28 +69,18 @@ pub fn ensure_runtime(io_threads: usize) -> Handle {
 
 fn submit_job(io_threads: usize) -> flume::Sender<Job> {
     let guard = RUNTIME.lock().unwrap();
-    if let Some(ref rt) = *guard {
-        if rt.pid == std::process::id() {
-            return rt.submit.clone();
-        }
+    if let Some(ref rt) = *guard
+        && rt.pid == std::process::id()
+    {
+        return rt.submit.clone();
     }
     drop(guard);
     ensure_runtime(io_threads);
     RUNTIME.lock().unwrap().as_ref().unwrap().submit.clone()
 }
 
-pub fn spawn_blocking<F, T>(io_threads: usize, fut: F) -> T
-where
-    F: Future<Output = T> + Send + 'static,
-    T: Send + 'static,
-{
-    let handle = ensure_runtime(io_threads);
-    let (otx, orx) = flume::bounded::<T>(1);
-    handle.spawn(async move {
-        let out = fut.await;
-        let _ = otx.send(out);
-    });
-
+#[cfg(ruby_engine = "mri")]
+fn recv_blocking<T>(rx: flume::Receiver<T>, missing: &'static str) -> T {
     struct RecvBox<U> {
         rx: flume::Receiver<U>,
         result: Option<U>,
@@ -102,10 +92,7 @@ where
         std::ptr::null_mut()
     }
 
-    let mut rd = RecvBox {
-        rx: orx,
-        result: None,
-    };
+    let mut rd = RecvBox { rx, result: None };
     unsafe {
         rb_sys::rb_thread_call_without_gvl(
             Some(blocking_recv::<T>),
@@ -114,7 +101,27 @@ where
             std::ptr::null_mut(),
         );
     }
-    rd.result.expect("omq-backend-rust: runtime dropped result")
+    rd.result.expect(missing)
+}
+
+#[cfg(not(ruby_engine = "mri"))]
+fn recv_blocking<T>(rx: flume::Receiver<T>, missing: &'static str) -> T {
+    rx.recv().expect(missing)
+}
+
+pub fn spawn_blocking<F, T>(io_threads: usize, fut: F) -> T
+where
+    F: Future<Output = T> + Send + 'static,
+    T: Send + 'static,
+{
+    let handle = ensure_runtime(io_threads);
+    let (otx, orx) = flume::bounded::<T>(1);
+    handle.spawn(async move {
+        let out = fut.await;
+        let _ = otx.send(out);
+    });
+
+    recv_blocking(orx, "omq-backend-rust: runtime dropped result")
 }
 
 pub struct Materialized {
@@ -171,10 +178,10 @@ fn convert_monitor_event(event: &omq_tokio::MonitorEvent) -> MonitorEventData {
         },
         HandshakeSucceeded { endpoint, peer } => {
             let mut detail = vec![("connection_id", peer.connection_id.to_string())];
-            if let Some(ref ident) = peer.peer_identity {
-                if !ident.is_empty() {
-                    detail.push(("identity", format!("{:?}", ident)));
-                }
+            if let Some(ref ident) = peer.peer_identity
+                && !ident.is_empty()
+            {
+                detail.push(("identity", format!("{:?}", ident)));
             }
             MonitorEventData {
                 event_type: "handshake_succeeded",
@@ -299,20 +306,15 @@ pub fn materialize(
     let job: Job = Box::new(move || {
         let sock = Arc::new(InnerSocket::new(socket_type, options));
 
-        const SEND_YIELD_INTERVAL: u32 = 256;
         let s = sock.clone();
         let sn = send_notify.clone();
         let send_pump = tokio::spawn(async move {
-            futures::pin_mut!(send_cons);
-            let mut batch = 0u32;
+            let mut send_cons = send_cons;
             while let Some(msg) = futures::StreamExt::next(&mut send_cons).await {
-                let _ = s.send(msg).await;
+                send_cons.release();
                 sn.notify();
-                batch += 1;
-                if batch >= SEND_YIELD_INTERVAL {
-                    batch = 0;
-                    tokio::task::yield_now().await;
-                }
+                let _ = s.send(msg).await;
+                tokio::task::yield_now().await;
             }
             sn.notify();
         });
@@ -343,6 +345,7 @@ pub fn materialize(
         });
 
         let monitor_sock = sock.clone();
+        let peer_ready_sock = sock.clone();
         let monitor_pump = tokio::spawn(async move {
             let mut stream = monitor_sock.monitor();
             let mut peer_count: u32 = 0;
@@ -359,6 +362,7 @@ pub fn materialize(
                                 had_peers = true;
                                 if !peer_connected_fired {
                                     peer_connected_fired = true;
+                                    let _ = peer_ready_sock.connections().await;
                                     peer_connected_notify.force_wake();
                                 }
                             }
@@ -368,11 +372,12 @@ pub fn materialize(
                                     all_peers_gone_notify.force_wake();
                                 }
                             }
-                            omq_tokio::MonitorEvent::SubscribeReceived { .. } => {
-                                if !subscriber_joined_fired {
-                                    subscriber_joined_fired = true;
-                                    subscriber_joined_notify.force_wake();
-                                }
+                            omq_tokio::MonitorEvent::SubscribeReceived { .. }
+                            | omq_tokio::MonitorEvent::JoinReceived { .. }
+                                if !subscriber_joined_fired =>
+                            {
+                                subscriber_joined_fired = true;
+                                subscriber_joined_notify.force_wake();
                             }
                             _ => {}
                         }
@@ -392,64 +397,43 @@ pub fn materialize(
     });
     tx.send(job).expect("omq-backend-rust: tokio runtime gone");
 
-    struct RecvBox {
-        rx: flume::Receiver<(
-            Arc<InnerSocket>,
-            JoinHandle<()>,
-            JoinHandle<()>,
-            JoinHandle<()>,
-        )>,
-        result: Option<(
-            Arc<InnerSocket>,
-            JoinHandle<()>,
-            JoinHandle<()>,
-            JoinHandle<()>,
-        )>,
-    }
-
-    extern "C" fn blocking_recv(data: *mut libc::c_void) -> *mut libc::c_void {
-        let rd = unsafe { &mut *(data as *mut RecvBox) };
-        rd.result = rd.rx.recv().ok();
-        std::ptr::null_mut()
-    }
-
-    let mut rd = RecvBox {
-        rx: orx,
-        result: None,
-    };
-    unsafe {
-        rb_sys::rb_thread_call_without_gvl(
-            Some(blocking_recv),
-            &mut rd as *mut RecvBox as *mut libc::c_void,
-            None,
-            std::ptr::null_mut(),
-        );
-    }
-    rd.result.expect("omq-backend-rust: materialize failed")
+    recv_blocking(orx, "omq-backend-rust: materialize failed")
 }
 
 pub fn destroy_socket(
     io_threads: usize,
     sock: Arc<InnerSocket>,
     send_prod: Mutex<yring::AsyncProducer<omq_tokio::Message>>,
-    send_pump: JoinHandle<()>,
+    mut send_pump: JoinHandle<()>,
     recv_pump: JoinHandle<()>,
     monitor_pump: JoinHandle<()>,
     linger: Option<Duration>,
 ) {
     recv_pump.abort();
     monitor_pump.abort();
-    send_pump.abort();
-    drop(send_prod);
-    let Ok(handle) = (|| -> std::result::Result<Handle, ()> { Ok(ensure_runtime(io_threads)) })()
-    else {
-        return;
-    };
+    let handle = ensure_runtime(io_threads);
     let close_timeout = linger
         .unwrap_or(Duration::from_secs(30))
         .max(Duration::from_millis(10));
-    handle.spawn(async move {
+    let fut = async move {
+        drop(send_prod);
+        if tokio::time::timeout(close_timeout, &mut send_pump)
+            .await
+            .is_err()
+        {
+            send_pump.abort();
+            let _ = send_pump.await;
+        }
+
         let s = Arc::try_unwrap(sock).unwrap_or_else(|arc| (*arc).clone());
         let _ = tokio::time::timeout(close_timeout, s.close()).await;
+    };
+
+    let (otx, orx) = flume::bounded::<()>(1);
+    handle.spawn(async move {
+        fut.await;
+        let _ = otx.send(());
     });
+
+    recv_blocking(orx, "omq-backend-rust: close failed")
 }