prosody 0.4.0 → 0.5.0

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

@@ -6,8 +6,8 @@
 
 use crate::bridge::Bridge;
 use crate::handler::state::{
-    DequeStateVariant, MapStateVariant, NativeDequeState, NativeMapState, NativeValueState,
-    ValueStateVariant, state_error,
+    NativeJsonDequeState, NativeJsonMapState, NativeJsonValueState, NativeMessageDequeState,
+    NativeMessageMapState, NativeMessageValueState, state_error,
 };
 use crate::tracing_util::extract_opentelemetry_context;
 use crate::{ROOT_MOD, id};
@@ -334,13 +334,13 @@ impl Context {
     /// 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<NativeValueState, Error> {
+    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(NativeValueState::new(
-            ValueStateVariant::Json(Arc::from(handle)),
+        Ok(NativeJsonValueState::new(
+            Arc::from(handle),
             this.bridge.clone(),
             Arc::clone(&this.propagator),
         ))
@@ -352,13 +352,13 @@ impl Context {
     ///
     /// 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<NativeMapState, Error> {
+    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(NativeMapState::new(
-            MapStateVariant::Json(Arc::from(handle)),
+        Ok(NativeJsonMapState::new(
+            Arc::from(handle),
             this.bridge.clone(),
             Arc::clone(&this.propagator),
         ))
@@ -370,13 +370,13 @@ impl Context {
     ///
     /// 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<NativeDequeState, Error> {
+    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(NativeDequeState::new(
-            DequeStateVariant::Json(Arc::from(handle)),
+        Ok(NativeJsonDequeState::new(
+            Arc::from(handle),
             this.bridge.clone(),
             Arc::clone(&this.propagator),
         ))
@@ -392,13 +392,13 @@ impl Context {
         ruby: &Ruby,
         this: &Self,
         name: String,
-    ) -> Result<NativeValueState, Error> {
+    ) -> Result<NativeMessageValueState, Error> {
         let handle = this
             .inner
             .message_value_state(&name)
             .map_err(|error| state_error(ruby, &error))?;
-        Ok(NativeValueState::new(
-            ValueStateVariant::Message(Arc::from(handle)),
+        Ok(NativeMessageValueState::new(
+            Arc::from(handle),
             this.bridge.clone(),
             Arc::clone(&this.propagator),
         ))
@@ -410,13 +410,17 @@ impl Context {
     ///
     /// 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<NativeMapState, Error> {
+    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(NativeMapState::new(
-            MapStateVariant::Message(Arc::from(handle)),
+        Ok(NativeMessageMapState::new(
+            Arc::from(handle),
             this.bridge.clone(),
             Arc::clone(&this.propagator),
         ))
@@ -432,13 +436,13 @@ impl Context {
         ruby: &Ruby,
         this: &Self,
         name: String,
-    ) -> Result<NativeDequeState, Error> {
+    ) -> Result<NativeMessageDequeState, Error> {
         let handle = this
             .inner
             .message_deque_state(&name)
             .map_err(|error| state_error(ruby, &error))?;
-        Ok(NativeDequeState::new(
-            DequeStateVariant::Message(Arc::from(handle)),
+        Ok(NativeMessageDequeState::new(
+            Arc::from(handle),
             this.bridge.clone(),
             Arc::clone(&this.propagator),
         ))

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.
     ///
@@ -122,6 +131,34 @@ impl Message {
     }
 }
 
+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 {
     /// Creates a new Message wrapper from a Prosody `ConsumerMessage`.
     ///
@@ -133,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
@@ -157,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,13 +27,14 @@ 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;
@@ -41,6 +42,11 @@ 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.
 ///
@@ -84,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.
@@ -116,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<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",
@@ -130,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>(
@@ -206,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<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
@@ -257,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?;
 
@@ -267,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
@@ -295,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.
     ///
@@ -341,7 +369,7 @@ pub enum RubyHandlerError {
 pub fn init(ruby: &Ruby) -> Result<(), Error> {
     context::init(ruby)?;
     message::init(ruby)?;
-    state::init(ruby)?;
+    state::register(ruby)?;
     trigger::init(ruby)?;
 
     Ok(())