prosody 0.3.0 → 0.5.0

data/ext/prosody/src/handler/context.rs

@@ -5,6 +5,10 @@
 //! information from Kafka messages and schedule timer events.
 
 use crate::bridge::Bridge;
+use crate::handler::state::{
+    NativeJsonDequeState, NativeJsonMapState, NativeJsonValueState, NativeMessageDequeState,
+    NativeMessageMapState, NativeMessageValueState, state_error,
+};
 use crate::tracing_util::extract_opentelemetry_context;
 use crate::{ROOT_MOD, id};
 use educe::Educe;
@@ -35,9 +39,8 @@ pub struct Context {
     ///
     /// This field is marked as hidden in debug output to prevent logging large
     /// data.
-    #[allow(dead_code)]
     #[educe(Debug(ignore))]
-    inner: BoxEventContext,
+    inner: BoxEventContext<serde_json::Value>,
 
     /// Bridge for handling async operations
     #[educe(Debug(ignore))]
@@ -57,7 +60,7 @@ impl Context {
     /// * `bridge` - The bridge for handling async operations
     /// * `propagator` - Shared OpenTelemetry propagator for distributed tracing
     pub fn new(
-        inner: BoxEventContext,
+        inner: BoxEventContext<serde_json::Value>,
         bridge: Bridge,
         propagator: Arc<TextMapCompositePropagator>,
     ) -> Self {
@@ -302,10 +305,10 @@ impl Context {
                 async move { inner.scheduled(TimerType::Application).await },
                 span,
             )?
-            .map_err(|e| {
+            .map_err(|error| {
                 Error::new(
                     ruby.exception_runtime_error(),
-                    format!("Failed to get scheduled times: {e}"),
+                    format!("Failed to get scheduled times: {error:#}"),
                 )
             })?;
 
@@ -319,6 +322,131 @@ impl Context {
 
         Ok(ruby_array.as_value())
     }
+
+    /// Vends the handle for the named JSON value collection.
+    ///
+    /// Vending verifies the collection's registration (core-side); no span is
+    /// opened here — vended handles outlive the call, and every operation opens
+    /// its own span.
+    ///
+    /// # Errors
+    ///
+    /// Returns a permanent state error if the name is unregistered or its
+    /// registered identity mismatches.
+    #[allow(clippy::needless_pass_by_value, reason = "Magnus method argument type")]
+    fn value_state(ruby: &Ruby, this: &Self, name: String) -> Result<NativeJsonValueState, Error> {
+        let handle = this
+            .inner
+            .value_state(&name)
+            .map_err(|error| state_error(ruby, &error))?;
+        Ok(NativeJsonValueState::new(
+            Arc::from(handle),
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        ))
+    }
+
+    /// Vends the handle for the named JSON map collection.
+    ///
+    /// # Errors
+    ///
+    /// See [`value_state`](Self::value_state).
+    #[allow(clippy::needless_pass_by_value, reason = "Magnus method argument type")]
+    fn map_state(ruby: &Ruby, this: &Self, name: String) -> Result<NativeJsonMapState, Error> {
+        let handle = this
+            .inner
+            .map_state(&name)
+            .map_err(|error| state_error(ruby, &error))?;
+        Ok(NativeJsonMapState::new(
+            Arc::from(handle),
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        ))
+    }
+
+    /// Vends the handle for the named JSON deque collection.
+    ///
+    /// # Errors
+    ///
+    /// See [`value_state`](Self::value_state).
+    #[allow(clippy::needless_pass_by_value, reason = "Magnus method argument type")]
+    fn deque_state(ruby: &Ruby, this: &Self, name: String) -> Result<NativeJsonDequeState, Error> {
+        let handle = this
+            .inner
+            .deque_state(&name)
+            .map_err(|error| state_error(ruby, &error))?;
+        Ok(NativeJsonDequeState::new(
+            Arc::from(handle),
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        ))
+    }
+
+    /// Vends the handle for the named Kafka-message value collection.
+    ///
+    /// # Errors
+    ///
+    /// See [`value_state`](Self::value_state).
+    #[allow(clippy::needless_pass_by_value, reason = "Magnus method argument type")]
+    fn message_value_state(
+        ruby: &Ruby,
+        this: &Self,
+        name: String,
+    ) -> Result<NativeMessageValueState, Error> {
+        let handle = this
+            .inner
+            .message_value_state(&name)
+            .map_err(|error| state_error(ruby, &error))?;
+        Ok(NativeMessageValueState::new(
+            Arc::from(handle),
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        ))
+    }
+
+    /// Vends the handle for the named Kafka-message map collection.
+    ///
+    /// # Errors
+    ///
+    /// See [`value_state`](Self::value_state).
+    #[allow(clippy::needless_pass_by_value, reason = "Magnus method argument type")]
+    fn message_map_state(
+        ruby: &Ruby,
+        this: &Self,
+        name: String,
+    ) -> Result<NativeMessageMapState, Error> {
+        let handle = this
+            .inner
+            .message_map_state(&name)
+            .map_err(|error| state_error(ruby, &error))?;
+        Ok(NativeMessageMapState::new(
+            Arc::from(handle),
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        ))
+    }
+
+    /// Vends the handle for the named Kafka-message deque collection.
+    ///
+    /// # Errors
+    ///
+    /// See [`value_state`](Self::value_state).
+    #[allow(clippy::needless_pass_by_value, reason = "Magnus method argument type")]
+    fn message_deque_state(
+        ruby: &Ruby,
+        this: &Self,
+        name: String,
+    ) -> Result<NativeMessageDequeState, Error> {
+        let handle = this
+            .inner
+            .message_deque_state(&name)
+            .map_err(|error| state_error(ruby, &error))?;
+        Ok(NativeMessageDequeState::new(
+            Arc::from(handle),
+            this.bridge.clone(),
+            Arc::clone(&this.propagator),
+        ))
+    }
 }
 
 /// Initializes the Context class in Ruby.
@@ -357,6 +485,23 @@ pub fn init(ruby: &Ruby) -> Result<(), Error> {
     )?;
     class.define_method(id!(ruby, "scheduled"), method!(Context::scheduled, 0))?;
 
+    // Keyed-state vend methods
+    class.define_method(id!(ruby, "value_state"), method!(Context::value_state, 1))?;
+    class.define_method(id!(ruby, "map_state"), method!(Context::map_state, 1))?;
+    class.define_method(id!(ruby, "deque_state"), method!(Context::deque_state, 1))?;
+    class.define_method(
+        id!(ruby, "message_value_state"),
+        method!(Context::message_value_state, 1),
+    )?;
+    class.define_method(
+        id!(ruby, "message_map_state"),
+        method!(Context::message_map_state, 1),
+    )?;
+    class.define_method(
+        id!(ruby, "message_deque_state"),
+        method!(Context::message_deque_state, 1),
+    )?;
+
     Ok(())
 }
 

data/ext/prosody/src/handler/message.rs

@@ -24,6 +24,15 @@ pub struct Message {
     inner: ConsumerMessage<serde_json::Value>,
 }
 
+/// Ruby wrapper for an excise record.
+#[derive(Educe, Clone)]
+#[educe(Debug)]
+#[magnus::wrap(class = "Prosody::ExciseMessage", frozen_shareable)]
+pub struct ExciseMessage {
+    #[educe(Debug(ignore))]
+    inner: ConsumerMessage<()>,
+}
+
 impl Message {
     /// Returns the topic name this message was published to.
     ///
@@ -100,9 +109,54 @@ impl Message {
     /// # Errors
     ///
     /// Returns an error if payload deserialization fails.
+    ///
+    /// The public RBS exposes this JSON value as `Message[Payload]#payload`.
+    /// That generic parameter is static-only: this native boundary continues
+    /// to deserialize JSON into ordinary Ruby objects without runtime schema
+    /// validation.
     fn payload(ruby: &Ruby, this: &Self) -> Result<Value, Error> {
         serialize(ruby, this.inner.payload())
     }
+
+    /// Clones the wrapped `ConsumerMessage` for a message-collection write.
+    ///
+    /// The wrapper holds the real `ConsumerMessage`, so a message write clones
+    /// the inner value directly — a cheap operation.
+    ///
+    /// # Returns
+    ///
+    /// An owned clone of the wrapped `ConsumerMessage`.
+    pub(crate) fn consumer_message(&self) -> ConsumerMessage<serde_json::Value> {
+        self.inner.clone()
+    }
+}
+
+impl ExciseMessage {
+    fn topic(&self) -> &'static str {
+        self.inner.topic().as_ref()
+    }
+
+    fn partition(&self) -> i32 {
+        self.inner.partition()
+    }
+
+    fn offset(&self) -> i64 {
+        self.inner.offset()
+    }
+
+    fn key(&self) -> &str {
+        self.inner.key()
+    }
+
+    fn timestamp(ruby: &Ruby, this: &Self) -> Result<Value, Error> {
+        let epoch_micros = this.inner.timestamp().timestamp_micros();
+        ruby.module_kernel()
+            .const_get::<_, RClass>(id!(ruby, "Time"))?
+            .funcall(
+                id!(ruby, "at"),
+                (epoch_micros, ruby.to_symbol("microsecond")),
+            )
+    }
 }
 
 impl From<ConsumerMessage<serde_json::Value>> for Message {
@@ -116,6 +170,12 @@ impl From<ConsumerMessage<serde_json::Value>> for Message {
     }
 }
 
+impl From<ConsumerMessage<()>> for ExciseMessage {
+    fn from(inner: ConsumerMessage<()>) -> Self {
+        Self { inner }
+    }
+}
+
 /// Initializes the `Prosody::Message` Ruby class and defines its methods.
 ///
 /// # Arguments
@@ -140,5 +200,12 @@ pub fn init(ruby: &Ruby) -> Result<(), Error> {
     class.define_method(id!(ruby, "timestamp"), method!(Message::timestamp, 0))?;
     class.define_method(id!(ruby, "payload"), method!(Message::payload, 0))?;
 
+    let class = module.define_class(id!(ruby, "ExciseMessage"), ruby.class_object())?;
+    class.define_method(id!(ruby, "topic"), method!(ExciseMessage::topic, 0))?;
+    class.define_method(id!(ruby, "partition"), method!(ExciseMessage::partition, 0))?;
+    class.define_method(id!(ruby, "offset"), method!(ExciseMessage::offset, 0))?;
+    class.define_method(id!(ruby, "key"), method!(ExciseMessage::key, 0))?;
+    class.define_method(id!(ruby, "timestamp"), method!(ExciseMessage::timestamp, 0))?;
+
     Ok(())
 }

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

@@ -11,7 +11,7 @@
 
 use crate::bridge::{Bridge, BridgeError};
 use crate::handler::context::Context;
-use crate::handler::message::Message;
+use crate::handler::message::{ExciseMessage, Message};
 use crate::handler::trigger::Timer;
 use crate::id;
 use crate::scheduler::result::ProcessingError;
@@ -19,7 +19,7 @@ use crate::scheduler::{Scheduler, SchedulerError};
 use crate::util::ThreadSafeValue;
 use futures::pin_mut;
 use magnus::value::ReprValue;
-use magnus::{Error, Ruby, Value};
+use magnus::{Error, IntoValue, Ruby, Value};
 use opentelemetry::propagation::TextMapCompositePropagator;
 use opentelemetry::trace::Status;
 use prosody::consumer::event_context::EventContext;
@@ -27,19 +27,26 @@ use prosody::consumer::message::ConsumerMessage;
 use prosody::consumer::middleware::FallibleHandler;
 use prosody::consumer::{DemandType, Keyed};
 use prosody::error::{ClassifyError, ErrorCategory};
+use prosody::high_level::{ClientHandler, JsonCodecs};
 use prosody::propagator::new_propagator;
 use prosody::timers::{TimerType, Trigger as ProsodyTrigger};
 use std::collections::HashMap;
 use std::sync::Arc;
 use thiserror::Error;
 use tokio::select;
-use tracing::{Instrument, info_span};
+use tracing::{Instrument, Span, info_span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 mod context;
 mod message;
+mod state;
 mod trigger;
 
+pub(crate) use state::{
+    NativeJsonDequeScan, NativeJsonMapScan, NativeMapKeyScan, parse_direction,
+    published_deque_scan, published_map_key_scan, published_map_scan,
+};
+
 /// A handler that bridges between Kafka messages and Ruby message processing
 /// code.
 ///
@@ -83,11 +90,78 @@ impl RubyHandler {
             propagator: Arc::new(new_propagator()),
         })
     }
+
+    async fn handle_record<C, P, M>(
+        &self,
+        context: C,
+        message: ConsumerMessage<P>,
+        method: &'static str,
+        event_type: &'static str,
+        span: Span,
+    ) -> Result<serde_json::Value, RubyHandlerError>
+    where
+        C: EventContext<Payload = serde_json::Value>,
+        M: From<ConsumerMessage<P>> + IntoValue + Send + 'static,
+        P: Send + Sync + 'static,
+    {
+        let cancel_future = context.clone().on_cancel();
+        let handler = self.handler.clone();
+        let task_id = format!(
+            "{}/{}:{}",
+            message.topic(),
+            message.partition(),
+            message.offset()
+        );
+        let event_context = HashMap::from([
+            ("event_type".into(), event_type.into()),
+            ("topic".into(), message.topic().to_string()),
+            ("partition".into(), message.partition().to_string()),
+            ("key".into(), message.key().to_string()),
+            ("offset".into(), message.offset().to_string()),
+        ]);
+        let context = Context::new(
+            context.boxed(),
+            self.bridge.clone(),
+            self.propagator.clone(),
+        );
+        let response_requested = message.response_requested();
+        let message = M::from(message);
+        let cloned_span = span.clone();
+
+        async move {
+            let task_handle = self
+                .scheduler
+                .schedule(task_id, &cloned_span, event_context, move |ruby| {
+                    let result = handler.get(ruby).funcall(method, (context, message))?;
+                    Ok(if response_requested {
+                        result
+                    } else {
+                        ruby.qnil().as_value()
+                    })
+                })
+                .await?;
+            let result_future = task_handle.result.receive();
+            pin_mut!(result_future);
+            let result = select! {
+                result = &mut result_future => {
+                    result.inspect_err(|error| cloned_span.set_status(Status::error(error.to_string())))?
+                }
+                () = cancel_future => {
+                    task_handle.cancellation_token.cancel(&self.bridge).await
+                        .inspect_err(|error| cloned_span.set_status(Status::error(error.to_string())))?;
+                    result_future.await?
+                }
+            };
+            Ok(result)
+        }
+        .instrument(span)
+        .await
+    }
 }
 
 impl FallibleHandler for RubyHandler {
     type Error = RubyHandlerError;
-    type Output = ();
+    type Output = serde_json::Value;
     type Payload = serde_json::Value;
 
     /// Processes a Kafka message by dispatching it to the Ruby handler.
@@ -115,12 +189,10 @@ impl FallibleHandler for RubyHandler {
         context: C,
         message: ConsumerMessage<Self::Payload>,
         _demand_type: DemandType,
-    ) -> Result<(), Self::Error>
+    ) -> Result<Self::Output, Self::Error>
     where
-        C: EventContext,
+        C: EventContext<Payload = Self::Payload>,
     {
-        // Create a new span for the on_message operation as a child of the message's
-        // span
         let span = info_span!(
             parent: message.span(),
             "on_message",
@@ -129,75 +201,29 @@ impl FallibleHandler for RubyHandler {
             offset = message.offset(),
             key = %message.key()
         );
+        self.handle_record::<_, _, Message>(context, message, "on_message", "message", span)
+            .await
+    }
 
-        // Get a future that completes when cancellation is signaled
-        let cloned_context = context.clone();
-        let cancel_future = cloned_context.on_cancel();
-
-        // Clone the handler reference for use in the closure
-        let handler = self.handler.clone();
-
-        // Create a unique task ID for this message
-        let task_id = format!(
-            "{}/{}:{}",
-            message.topic(),
-            message.partition(),
-            message.offset()
-        );
-
-        let event_context = HashMap::from([
-            ("event_type".into(), "message".into()),
-            ("topic".into(), message.topic().to_string()),
-            ("partition".into(), message.partition().to_string()),
-            ("key".into(), message.key().to_string()),
-            ("offset".into(), message.offset().to_string()),
-        ]);
-
-        // Convert the Kafka message and context to Ruby-compatible types
-        let context = Context::new(
-            context.boxed(),
-            self.bridge.clone(),
-            self.propagator.clone(),
+    async fn on_excise<C>(
+        &self,
+        context: C,
+        message: ConsumerMessage<()>,
+        _demand_type: DemandType,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        C: EventContext<Payload = Self::Payload>,
+    {
+        let span = info_span!(
+            parent: message.span(),
+            "on_excise",
+            topic = %message.topic(),
+            partition = message.partition(),
+            offset = message.offset(),
+            key = %message.key()
         );
-        let message: Message = message.into();
-
-        // Execute the entire message handling operation within the span
-        let cloned_span = span.clone();
-        async move {
-            // Schedule the task to run in Ruby
-            let task_handle = self
-                .scheduler
-                .schedule(task_id, &cloned_span, event_context, move |ruby| {
-                    let _: Value = handler
-                        .get(ruby)
-                        .funcall(id!(ruby, "on_message"), (context, message))?;
-
-                    Ok(())
-                })
-                .await?;
-
-            // Get the future that will complete when the task is done
-            let result_future = task_handle.result.receive();
-            pin_mut!(result_future);
-
-            // Wait for either task completion or shutdown signal
-            select! {
-                result = &mut result_future => {
-                    result.inspect_err(|e| cloned_span.set_status(Status::error(e.to_string())))?;
-                }
-                () = cancel_future => {
-                    // A cancel() failure is a genuine bridge error; mark the span.
-                    // The subsequent result_future error is expected cancellation, not a handler bug.
-                    task_handle.cancellation_token.cancel(&self.bridge).await
-                        .inspect_err(|e| cloned_span.set_status(Status::error(e.to_string())))?;
-                    result_future.await?;
-                }
-            }
-
-            Ok(())
-        }
-        .instrument(span)
-        .await
+        self.handle_record::<_, _, ExciseMessage>(context, message, "on_excise", "excise", span)
+            .await
     }
 
     async fn on_timer<C>(
@@ -205,13 +231,13 @@ impl FallibleHandler for RubyHandler {
         context: C,
         trigger: ProsodyTrigger,
         _demand_type: DemandType,
-    ) -> Result<(), Self::Error>
+    ) -> Result<Self::Output, Self::Error>
     where
-        C: EventContext,
+        C: EventContext<Payload = Self::Payload>,
     {
         // Only process application timers; internal timers are handled by middleware
         if trigger.timer_type != TimerType::Application {
-            return Ok(());
+            return Ok(serde_json::Value::Null);
         }
 
         // Create a new span for the on_timer operation as a child of the trigger's span
@@ -256,8 +282,7 @@ impl FallibleHandler for RubyHandler {
                     let _: Value = handler
                         .get(ruby)
                         .funcall(id!(ruby, "on_timer"), (context, timer))?;
-
-                    Ok(())
+                    Ok(ruby.qnil().as_value())
                 })
                 .await?;
 
@@ -266,20 +291,20 @@ impl FallibleHandler for RubyHandler {
             pin_mut!(result_future);
 
             // Wait for either task completion or shutdown signal
-            select! {
+            let result = select! {
                 result = &mut result_future => {
-                    result.inspect_err(|e| cloned_span.set_status(Status::error(e.to_string())))?;
+                    result.inspect_err(|e| cloned_span.set_status(Status::error(e.to_string())))?
                 }
                 () = cancel_future => {
                     // A cancel() failure is a genuine bridge error; mark the span.
                     // The subsequent result_future error is expected cancellation, not a handler bug.
                     task_handle.cancellation_token.cancel(&self.bridge).await
                         .inspect_err(|e| cloned_span.set_status(Status::error(e.to_string())))?;
-                    result_future.await?;
+                    result_future.await?
                 }
-            }
+            };
 
-            Ok(())
+            Ok(result)
         }
         .instrument(span)
         .await
@@ -294,6 +319,10 @@ impl FallibleHandler for RubyHandler {
     }
 }
 
+impl ClientHandler for RubyHandler {
+    type Codecs = JsonCodecs;
+}
+
 impl ClassifyError for RubyHandlerError {
     /// Categorizes errors for proper retry handling in the Kafka consumer.
     ///
@@ -340,6 +369,7 @@ pub enum RubyHandlerError {
 pub fn init(ruby: &Ruby) -> Result<(), Error> {
     context::init(ruby)?;
     message::init(ruby)?;
+    state::register(ruby)?;
     trigger::init(ruby)?;
 
     Ok(())