prosody 0.3.0 → 0.5.0

data/ext/prosody/src/handler/state/scan.rs

@@ -0,0 +1,218 @@
+//! Concrete typed cursors for native state scans.
+//!
+//! Cancellation can discard an orphaned chunk. The attempt then closes the
+//! cursor through the Ruby `ensure`, so no later operation observes that chunk.
+
+use super::state_error;
+use crate::bridge::{Bridge, QUEUE_CLASS};
+use crate::handler::message::Message;
+use crate::id;
+use crate::tracing_util::extract_opentelemetry_context;
+use crate::util::ThreadSafeValue;
+use magnus::value::ReprValue;
+use magnus::{Error, IntoValue, Ruby, Value};
+use opentelemetry::propagation::TextMapCompositePropagator;
+use opentelemetry::trace::FutureExt;
+use prosody::consumer::event_context::StateCursor;
+use prosody::consumer::message::ConsumerMessage;
+use serde_json::Value as JsonValue;
+use serde_magnus::serialize;
+use std::cell::RefCell;
+use std::collections::VecDeque;
+use std::num::NonZeroUsize;
+use std::sync::Arc;
+use tracing::Span;
+
+#[allow(clippy::unwrap_used, reason = "256 is a nonzero literal")]
+const SCAN_READY_CHUNK_SIZE: NonZeroUsize = NonZeroUsize::new(256).unwrap();
+
+struct ScanInner<T> {
+    cursor: Arc<StateCursor<T>>,
+    buffer: VecDeque<T>,
+    done: bool,
+}
+
+macro_rules! drive_scan {
+    ($ruby:expr, $this:expr, $inner:expr, |$item:ident| $convert:block) => {{
+        loop {
+            if let Some($item) = $inner.buffer.pop_front() {
+                return $convert;
+            }
+            if $inner.done {
+                return Ok($ruby.qnil().as_value());
+            }
+            let cursor = Arc::clone(&$inner.cursor);
+            let context = extract_opentelemetry_context($ruby, &$this.propagator)?;
+            let chunk = $this
+                .bridge
+                .wait_for(
+                    $ruby,
+                    async move {
+                        cursor
+                            .next_ready_chunk(SCAN_READY_CHUNK_SIZE)
+                            .with_context(context)
+                            .await
+                    },
+                    Span::current(),
+                )?
+                .map_err(|error| state_error($ruby, &error))?;
+            match chunk {
+                Some(items) => $inner.buffer.extend(items),
+                None => {
+                    $inner.done = true;
+                    let cursor = Arc::clone(&$inner.cursor);
+                    $this.bridge.wait_for(
+                        $ruby,
+                        async move { cursor.close().await },
+                        Span::current(),
+                    )?;
+                    return Ok($ruby.qnil().as_value());
+                }
+            }
+        }
+    }};
+}
+
+macro_rules! native_scan {
+    ($name:ident, $class:literal, $item:ty, |$ruby:ident, $item_name:ident| $convert:block) => {
+        /// Native cursor with one item type.
+        #[magnus::wrap(class = $class)]
+        pub struct $name {
+            inner: RefCell<ScanInner<$item>>,
+            lock: ThreadSafeValue,
+            bridge: Bridge,
+            propagator: Arc<TextMapCompositePropagator>,
+        }
+
+        impl $name {
+            pub(super) fn new(
+                ruby: &Ruby,
+                cursor: Box<StateCursor<$item>>,
+                bridge: Bridge,
+                propagator: Arc<TextMapCompositePropagator>,
+            ) -> Result<Self, Error> {
+                let queue: Value = ruby.get_inner(&QUEUE_CLASS).funcall(id!(ruby, "new"), ())?;
+                let _: Value = queue.funcall(id!(ruby, "push"), (ruby.qnil(),))?;
+                Ok(Self {
+                    inner: RefCell::new(ScanInner {
+                        cursor: Arc::from(cursor),
+                        buffer: VecDeque::new(),
+                        done: false,
+                    }),
+                    lock: ThreadSafeValue::new(queue, bridge.clone()),
+                    bridge,
+                    propagator,
+                })
+            }
+
+            fn acquire(&self, ruby: &Ruby) -> Result<(), Error> {
+                let _: Value = self.lock.get(ruby).funcall(id!(ruby, "pop"), ())?;
+                Ok(())
+            }
+
+            fn release(&self, ruby: &Ruby) {
+                let _: Result<Value, Error> = self
+                    .lock
+                    .get(ruby)
+                    .funcall(id!(ruby, "push"), (ruby.qnil(),));
+            }
+
+            pub(super) fn next(ruby: &Ruby, this: &Self) -> Result<Value, Error> {
+                this.acquire(ruby)?;
+                let out = Self::next_locked(ruby, this);
+                this.release(ruby);
+                out
+            }
+
+            fn next_locked(ruby: &Ruby, this: &Self) -> Result<Value, Error> {
+                let inner = &mut *this.inner.borrow_mut();
+                drive_scan!(ruby, this, inner, |$item_name| {
+                    let $ruby = ruby;
+                    $convert
+                })
+            }
+
+            pub(super) fn close(ruby: &Ruby, this: &Self) -> Result<Value, Error> {
+                this.acquire(ruby)?;
+                let out = Self::close_locked(ruby, this);
+                this.release(ruby);
+                out
+            }
+
+            fn close_locked(ruby: &Ruby, this: &Self) -> Result<Value, Error> {
+                let inner = &mut *this.inner.borrow_mut();
+                inner.done = true;
+                inner.buffer.clear();
+                let cursor = Arc::clone(&inner.cursor);
+                this.bridge
+                    .wait_for(ruby, async move { cursor.close().await }, Span::current())?;
+                Ok(ruby.qnil().as_value())
+            }
+        }
+    };
+}
+
+native_scan!(
+    NativeJsonDequeScan,
+    "Prosody::NativeJsonDequeScan",
+    JsonValue,
+    |ruby, item| { serialize(ruby, &item) }
+);
+native_scan!(
+    NativeJsonMapScan,
+    "Prosody::NativeJsonMapScan",
+    (String, JsonValue),
+    |ruby, item| {
+        let (key, value) = item;
+        let value: Value = serialize(ruby, &value)?;
+        Ok((key, value).into_value_with(ruby))
+    }
+);
+native_scan!(
+    NativeMessageDequeScan,
+    "Prosody::NativeMessageDequeScan",
+    ConsumerMessage<JsonValue>,
+    |ruby, item| { Ok(Message::from(item).into_value_with(ruby)) }
+);
+native_scan!(
+    NativeMessageMapScan,
+    "Prosody::NativeMessageMapScan",
+    (String, ConsumerMessage<JsonValue>),
+    |ruby, item| {
+        let (key, message) = item;
+        Ok((key, Message::from(message)).into_value_with(ruby))
+    }
+);
+native_scan!(
+    NativeMapKeyScan,
+    "Prosody::NativeMapKeyScan",
+    String,
+    |ruby, item| { Ok(item.into_value_with(ruby)) }
+);
+
+pub(crate) fn published_map_scan(
+    ruby: &Ruby,
+    cursor: Box<StateCursor<(String, JsonValue)>>,
+    bridge: Bridge,
+    propagator: Arc<TextMapCompositePropagator>,
+) -> Result<NativeJsonMapScan, Error> {
+    NativeJsonMapScan::new(ruby, cursor, bridge, propagator)
+}
+
+pub(crate) fn published_map_key_scan(
+    ruby: &Ruby,
+    cursor: Box<StateCursor<String>>,
+    bridge: Bridge,
+    propagator: Arc<TextMapCompositePropagator>,
+) -> Result<NativeMapKeyScan, Error> {
+    NativeMapKeyScan::new(ruby, cursor, bridge, propagator)
+}
+
+pub(crate) fn published_deque_scan(
+    ruby: &Ruby,
+    cursor: Box<StateCursor<JsonValue>>,
+    bridge: Bridge,
+    propagator: Arc<TextMapCompositePropagator>,
+) -> Result<NativeJsonDequeScan, Error> {
+    NativeJsonDequeScan::new(ruby, cursor, bridge, propagator)
+}

data/ext/prosody/src/lib.rs

@@ -17,6 +17,8 @@ use crate::bridge::Bridge;
 use magnus::value::Lazy;
 use magnus::{Error, RModule, Ruby};
 use mimalloc::MiMalloc;
+use std::io::{self, Write};
+use std::process;
 use std::sync::{LazyLock, OnceLock};
 use tokio::runtime::Runtime;
 
@@ -26,6 +28,7 @@ mod client;
 mod gvl;
 mod handler;
 mod logging;
+mod published;
 mod scheduler;
 mod tracing_util;
 mod util;
@@ -44,9 +47,16 @@ pub static TRACING_INIT: OnceLock<()> = OnceLock::new();
 ///
 /// This runtime powers all async operations in the extension, including
 /// message processing, scheduling, and communication with Ruby.
-#[allow(clippy::expect_used)]
-static RUNTIME: LazyLock<Runtime> =
-    LazyLock::new(|| Runtime::new().expect("Failed to create Tokio runtime"));
+static RUNTIME: LazyLock<Runtime> = LazyLock::new(|| match Runtime::new() {
+    Ok(runtime) => runtime,
+    Err(error) => {
+        drop(writeln!(
+            io::stderr().lock(),
+            "failed to create Tokio runtime: {error:#}"
+        ));
+        process::abort();
+    }
+});
 
 /// Reference to the root Ruby module for this extension.
 ///
@@ -75,7 +85,9 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     admin::init(ruby)?;
     bridge::init(ruby)?;
     handler::init(ruby)?;
+    published::init(ruby)?;
     client::init(ruby)?;
+    util::init(ruby)?;
 
     Ok(())
 }

data/ext/prosody/src/published.rs

@@ -0,0 +1,273 @@
+//! Read-only published-state handles for Ruby.
+
+use crate::bridge::Bridge;
+use crate::handler::{
+    NativeJsonDequeScan, NativeJsonMapScan, NativeMapKeyScan, parse_direction,
+    published_deque_scan, published_map_key_scan, published_map_scan,
+};
+use crate::{ROOT_MOD, id};
+use magnus::value::ReprValue;
+use magnus::{Error, Module, Ruby, StaticSymbol, Value, method};
+use opentelemetry::propagation::TextMapCompositePropagator;
+use prosody::JsonCodec;
+use prosody::high_level::erased::{
+    ErasedDirection, SharedDequeReader, SharedMapReader, SharedValueReader,
+};
+use prosody::state::Direction;
+use serde_magnus::serialize;
+use std::sync::Arc;
+use tracing::Span;
+
+fn read_error(ruby: &Ruby, error: &impl ToString) -> Error {
+    Error::new(ruby.exception_runtime_error(), error.to_string())
+}
+
+fn erased_direction(direction: Direction) -> ErasedDirection {
+    match direction {
+        Direction::Forward => ErasedDirection::Forward,
+        Direction::Backward => ErasedDirection::Backward,
+    }
+}
+
+#[magnus::wrap(class = "Prosody::NativePublishedValue")]
+pub(crate) struct NativePublishedValue {
+    pub(crate) inner: SharedValueReader<JsonCodec>,
+    pub(crate) bridge: Bridge,
+}
+
+impl NativePublishedValue {
+    fn get(ruby: &Ruby, this: &Self, key: String) -> Result<Value, Error> {
+        let inner = Arc::clone(&this.inner);
+        let value = this
+            .bridge
+            .wait_for(ruby, async move { inner.get(key).await }, Span::current())?
+            .map_err(|error| read_error(ruby, &error))?;
+        match value {
+            Some(value) => serialize(ruby, &value),
+            None => Ok(ruby.qnil().as_value()),
+        }
+    }
+}
+
+#[magnus::wrap(class = "Prosody::NativePublishedMap")]
+pub(crate) struct NativePublishedMap {
+    pub(crate) inner: SharedMapReader<JsonCodec>,
+    pub(crate) bridge: Bridge,
+    pub(crate) propagator: Arc<TextMapCompositePropagator>,
+}
+
+impl NativePublishedMap {
+    fn get(ruby: &Ruby, this: &Self, key: String, map_key: String) -> Result<Value, Error> {
+        let inner = Arc::clone(&this.inner);
+        let value = this
+            .bridge
+            .wait_for(
+                ruby,
+                async move { inner.get(key, map_key).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))?;
+        match value {
+            Some(value) => serialize(ruby, &value),
+            None => Ok(ruby.qnil().as_value()),
+        }
+    }
+
+    fn get_many(
+        ruby: &Ruby,
+        this: &Self,
+        key: String,
+        map_keys: Vec<String>,
+    ) -> Result<Value, Error> {
+        let inner = Arc::clone(&this.inner);
+        let values = this
+            .bridge
+            .wait_for(
+                ruby,
+                async move { inner.get_many(key, map_keys).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))?;
+        serialize(ruby, &values)
+    }
+
+    fn contains_key(ruby: &Ruby, this: &Self, key: String, map_key: String) -> Result<bool, Error> {
+        let inner = Arc::clone(&this.inner);
+        this.bridge
+            .wait_for(
+                ruby,
+                async move { inner.contains_key(key, map_key).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))
+    }
+
+    fn scan(
+        ruby: &Ruby,
+        this: &Self,
+        key: String,
+        direction: StaticSymbol,
+    ) -> Result<NativeJsonMapScan, Error> {
+        let direction = erased_direction(parse_direction(ruby, direction)?);
+        let inner = Arc::clone(&this.inner);
+        let cursor = this
+            .bridge
+            .wait_for(
+                ruby,
+                async move { inner.stream(key, direction).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))?;
+        published_map_scan(
+            ruby,
+            cursor,
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        )
+    }
+
+    fn keys(
+        ruby: &Ruby,
+        this: &Self,
+        key: String,
+        direction: StaticSymbol,
+    ) -> Result<NativeMapKeyScan, Error> {
+        let direction = erased_direction(parse_direction(ruby, direction)?);
+        let inner = Arc::clone(&this.inner);
+        let cursor = this
+            .bridge
+            .wait_for(
+                ruby,
+                async move { inner.keys(key, direction).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))?;
+        published_map_key_scan(
+            ruby,
+            cursor,
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        )
+    }
+}
+
+#[magnus::wrap(class = "Prosody::NativePublishedDeque")]
+pub(crate) struct NativePublishedDeque {
+    pub(crate) inner: SharedDequeReader<JsonCodec>,
+    pub(crate) bridge: Bridge,
+    pub(crate) propagator: Arc<TextMapCompositePropagator>,
+}
+
+impl NativePublishedDeque {
+    fn get(ruby: &Ruby, this: &Self, key: String, index: usize) -> Result<Value, Error> {
+        let inner = Arc::clone(&this.inner);
+        let value = this
+            .bridge
+            .wait_for(
+                ruby,
+                async move { inner.get(key, index).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))?;
+        match value {
+            Some(value) => serialize(ruby, &value),
+            None => Ok(ruby.qnil().as_value()),
+        }
+    }
+
+    fn length(ruby: &Ruby, this: &Self, key: String) -> Result<usize, Error> {
+        let inner = Arc::clone(&this.inner);
+        this.bridge
+            .wait_for(ruby, async move { inner.len(key).await }, Span::current())?
+            .map_err(|error| read_error(ruby, &error))
+    }
+
+    fn is_empty(ruby: &Ruby, this: &Self, key: String) -> Result<bool, Error> {
+        let inner = Arc::clone(&this.inner);
+        this.bridge
+            .wait_for(
+                ruby,
+                async move { inner.is_empty(key).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))
+    }
+
+    fn peek_front(ruby: &Ruby, this: &Self, key: String) -> Result<Value, Error> {
+        let inner = Arc::clone(&this.inner);
+        let value = this
+            .bridge
+            .wait_for(
+                ruby,
+                async move { inner.peek_front(key).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))?;
+        match value {
+            Some(value) => serialize(ruby, &value),
+            None => Ok(ruby.qnil().as_value()),
+        }
+    }
+
+    fn peek_back(ruby: &Ruby, this: &Self, key: String) -> Result<Value, Error> {
+        let inner = Arc::clone(&this.inner);
+        let value = this
+            .bridge
+            .wait_for(
+                ruby,
+                async move { inner.peek_back(key).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))?;
+        match value {
+            Some(value) => serialize(ruby, &value),
+            None => Ok(ruby.qnil().as_value()),
+        }
+    }
+
+    fn scan(
+        ruby: &Ruby,
+        this: &Self,
+        key: String,
+        direction: StaticSymbol,
+    ) -> Result<NativeJsonDequeScan, Error> {
+        let direction = erased_direction(parse_direction(ruby, direction)?);
+        let inner = Arc::clone(&this.inner);
+        let cursor = this
+            .bridge
+            .wait_for(
+                ruby,
+                async move { inner.stream(key, direction).await },
+                Span::current(),
+            )?
+            .map_err(|error| read_error(ruby, &error))?;
+        published_deque_scan(
+            ruby,
+            cursor,
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        )
+    }
+}
+
+pub(crate) fn init(ruby: &Ruby) -> Result<(), Error> {
+    let module = ruby.get_inner(&ROOT_MOD);
+    let value = module.define_class(id!(ruby, "NativePublishedValue"), ruby.class_object())?;
+    value.define_method("get", method!(NativePublishedValue::get, 1))?;
+
+    let map = module.define_class(id!(ruby, "NativePublishedMap"), ruby.class_object())?;
+    map.define_method("get", method!(NativePublishedMap::get, 2))?;
+    map.define_method("get_many", method!(NativePublishedMap::get_many, 2))?;
+    map.define_method("contains_key", method!(NativePublishedMap::contains_key, 2))?;
+    map.define_method("scan", method!(NativePublishedMap::scan, 2))?;
+    map.define_method("keys", method!(NativePublishedMap::keys, 2))?;
+
+    let deque = module.define_class(id!(ruby, "NativePublishedDeque"), ruby.class_object())?;
+    deque.define_method("get", method!(NativePublishedDeque::get, 2))?;
+    deque.define_method("length", method!(NativePublishedDeque::length, 1))?;
+    deque.define_method("is_empty", method!(NativePublishedDeque::is_empty, 1))?;
+    deque.define_method("peek_front", method!(NativePublishedDeque::peek_front, 1))?;
+    deque.define_method("peek_back", method!(NativePublishedDeque::peek_back, 1))?;
+    deque.define_method("scan", method!(NativePublishedDeque::scan, 2))?;
+    Ok(())
+}

data/ext/prosody/src/scheduler/mod.rs

@@ -14,7 +14,7 @@ use crate::bridge::{Bridge, BridgeError};
 use crate::scheduler::handle::TaskHandle;
 use crate::scheduler::processor::RubyProcessor;
 use crate::scheduler::result::result_channel;
-use magnus::{Error, Ruby};
+use magnus::{Error, Ruby, Value};
 use opentelemetry::propagation::{TextMapCompositePropagator, TextMapPropagator};
 use prosody::propagator::new_propagator;
 use std::collections::HashMap;
@@ -99,7 +99,7 @@ impl Scheduler {
         function: F,
     ) -> Result<TaskHandle, SchedulerError>
     where
-        F: FnOnce(&Ruby) -> Result<(), Error> + Send + 'static,
+        F: FnOnce(&Ruby) -> Result<Value, Error> + Send + 'static,
     {
         let mut carrier: HashMap<String, String> = HashMap::with_capacity(2);
         self.propagator

data/ext/prosody/src/scheduler/processor.rs

@@ -104,7 +104,7 @@ impl RubyProcessor {
         function: F,
     ) -> Result<CancellationToken, Error>
     where
-        F: FnOnce(&Ruby) -> Result<(), Error> + Send + 'static,
+        F: FnOnce(&Ruby) -> Result<Value, Error> + Send + 'static,
     {
         if self.is_shutdown.load(Relaxed) {
             return Err(Error::new(
@@ -122,7 +122,7 @@ impl RubyProcessor {
             if let Some(function) = maybe_function.take() {
                 function(ruby)
             } else {
-                Ok(())
+                Ok(ruby.qnil().as_value())
             }
         });
 

data/ext/prosody/src/scheduler/result.rs

@@ -14,6 +14,7 @@ use magnus::block::Proc;
 use magnus::value::ReprValue;
 use magnus::{Error, Ruby, TryConvert, Value, kwargs};
 use prosody::error::{ClassifyError, ErrorCategory};
+use serde_magnus::deserialize;
 use thiserror::Error;
 use tokio::sync::oneshot;
 use tracing::debug;
@@ -40,7 +41,7 @@ pub fn result_channel() -> (ResultSender, ResultReceiver) {
 #[educe(Debug)]
 pub struct ResultSender {
     #[educe(Debug(ignore))]
-    result_tx: AtomicTake<oneshot::Sender<Result<(), ProcessingError>>>,
+    result_tx: AtomicTake<oneshot::Sender<Result<serde_json::Value, ProcessingError>>>,
 }
 
 /// Receives task results in Rust from Ruby.
@@ -51,7 +52,7 @@ pub struct ResultSender {
 #[educe(Debug)]
 pub struct ResultReceiver {
     #[educe(Debug(ignore))]
-    result_rx: oneshot::Receiver<Result<(), ProcessingError>>,
+    result_rx: oneshot::Receiver<Result<serde_json::Value, ProcessingError>>,
 }
 
 impl ResultSender {
@@ -80,7 +81,9 @@ impl ResultSender {
         };
 
         if is_success {
-            if result_tx.send(Ok(())).is_err() {
+            let result = deserialize(ruby, result)
+                .map_err(|error| ProcessingError::Permanent(error.to_string()));
+            if result_tx.send(result).is_err() {
                 debug!("discarding result; receiver went away");
             }
 
@@ -153,7 +156,7 @@ impl ResultReceiver {
     /// Returns a `ProcessingError` if:
     /// - The task failed (with either a permanent or transient error)
     /// - The channel was closed unexpectedly (e.g., the Ruby VM terminated)
-    pub async fn receive(self) -> Result<(), ProcessingError> {
+    pub async fn receive(self) -> Result<serde_json::Value, ProcessingError> {
         self.result_rx.await.map_err(|_| ProcessingError::Closed)?
     }
 }