refactor: decouple ClaimCheckPublisher from JetStreamPublisher trait

Split serve() to accept JetStreamContext and ClaimCheckPublisher as
separate parameters instead of requiring a single type that implements
both traits. ClaimCheckPublisher is now a plain struct with a
publish_event method — no trait implementations, no type erasure.

Signed-off-by: Yordis Prieto <yordis.prieto@gmail.com>

Author: yordis

rsworkspace/crates/trogon-nats/src/jetstream/claim_check.rs

@@ -1,16 +1,13 @@
 use std::fmt;
-use std::future::IntoFuture;
-use std::pin::Pin;
+use std::time::Duration;
 
 use async_nats::HeaderMap;
-use async_nats::jetstream::publish::PublishAck;
-use async_nats::subject::ToSubject;
 use bytes::Bytes;
-
-use async_nats::jetstream::stream;
+use tracing::error;
 
 use super::object_store::{ObjectStoreGet, ObjectStorePut};
-use super::traits::{JetStreamContext, JetStreamPublisher};
+use super::publish::PublishOutcome;
+use super::traits::JetStreamPublisher;
 
 pub const HEADER_CLAIM_CHECK: &str = "Trogon-Claim-Check";
 pub const HEADER_CLAIM_BUCKET: &str = "Trogon-Claim-Bucket";
@@ -89,42 +86,18 @@ impl<E: fmt::Display + fmt::Debug + Send + Sync + 'static> std::error::Error
 {
 }
 
-#[derive(Debug)]
-pub enum ClaimCheckPublishError<PE: fmt::Display, SE: fmt::Display> {
-    PublishFailed(PE),
-    StoreFailed(SE),
-}
-
-impl<PE: fmt::Display, SE: fmt::Display> fmt::Display for ClaimCheckPublishError<PE, SE> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            Self::PublishFailed(e) => write!(f, "publish failed: {e}"),
-            Self::StoreFailed(e) => {
-                write!(f, "claim check object store put failed: {e}")
-            }
-        }
-    }
-}
-
-impl<
-    PE: fmt::Display + fmt::Debug + Send + Sync + 'static,
-    SE: fmt::Display + fmt::Debug + Send + Sync + 'static,
-> std::error::Error for ClaimCheckPublishError<PE, SE>
-{
-}
-
 #[derive(Clone, Debug)]
 pub struct ClaimCheckPublisher<P, S> {
-    inner: P,
+    publisher: P,
     store: S,
     bucket_name: String,
     max_payload: MaxPayload,
 }
 
 impl<P, S> ClaimCheckPublisher<P, S> {
-    pub fn new(inner: P, store: S, bucket_name: String, max_payload: MaxPayload) -> Self {
+    pub fn new(publisher: P, store: S, bucket_name: String, max_payload: MaxPayload) -> Self {
         Self {
-            inner,
+            publisher,
             store,
             bucket_name,
             max_payload,
@@ -137,105 +110,48 @@ fn claim_object_key(subject: &str) -> String {
     format!("{subject}/{id}")
 }
 
-pub struct ClaimCheckAckFuture<PE, SE>
-where
-    PE: fmt::Display + Send,
-    SE: fmt::Display + Send,
-{
-    inner: Pin<Box<dyn std::future::Future<Output = Result<PublishAck, PE>> + Send>>,
-    _se: std::marker::PhantomData<SE>,
-}
-
-impl<PE, SE> IntoFuture for ClaimCheckAckFuture<PE, SE>
-where
-    PE: fmt::Display + Send + 'static,
-    SE: fmt::Display + Send + 'static,
-{
-    type Output = Result<PublishAck, ClaimCheckPublishError<PE, SE>>;
-    type IntoFuture = Pin<
-        Box<
-            dyn std::future::Future<Output = Result<PublishAck, ClaimCheckPublishError<PE, SE>>>
-                + Send,
-        >,
-    >;
-
-    fn into_future(self) -> Self::IntoFuture {
-        Box::pin(async move {
-            self.inner
-                .await
-                .map_err(ClaimCheckPublishError::PublishFailed)
-        })
-    }
-}
-
-impl<P, S> JetStreamPublisher for ClaimCheckPublisher<P, S>
-where
-    P: JetStreamPublisher,
-    S: ObjectStorePut,
-{
-    type PublishError = ClaimCheckPublishError<P::PublishError, S::Error>;
-    type AckFuture = ClaimCheckAckFuture<P::PublishError, S::Error>;
-
-    async fn publish_with_headers<Sub: ToSubject + Send>(
+impl<P: JetStreamPublisher, S: ObjectStorePut> ClaimCheckPublisher<P, S> {
+    pub async fn publish_event(
         &self,
-        subject: Sub,
+        subject: String,
         headers: HeaderMap,
         payload: Bytes,
-    ) -> Result<Self::AckFuture, Self::PublishError> {
+        ack_timeout: Duration,
+    ) -> PublishOutcome<P::PublishError> {
         if payload.len() <= self.max_payload.threshold() {
-            let ack = self
-                .inner
-                .publish_with_headers(subject, headers, payload)
-                .await
-                .map_err(ClaimCheckPublishError::PublishFailed)?;
-            return Ok(ClaimCheckAckFuture {
-                inner: Box::pin(ack.into_future()),
-                _se: std::marker::PhantomData,
-            });
+            return super::publish::publish_event(
+                &self.publisher,
+                subject,
+                headers,
+                payload,
+                ack_timeout,
+            )
+            .await;
         }
 
-        let subject_str = subject.to_subject();
-        let key = claim_object_key(&subject_str);
+        let key = claim_object_key(&subject);
 
         // Store-then-publish: if publish fails, the object becomes orphaned.
         // Cleanup relies on the object store bucket's TTL — see #101.
         let mut cursor = std::io::Cursor::new(payload);
-        self.store
-            .put(&key, &mut cursor)
-            .await
-            .map_err(ClaimCheckPublishError::StoreFailed)?;
+        if let Err(e) = self.store.put(&key, &mut cursor).await {
+            error!(error = %e, "claim check: failed to store payload in object store");
+            return PublishOutcome::StoreFailed;
+        }
 
         let mut claim_headers = headers;
         claim_headers.insert(HEADER_CLAIM_CHECK, CLAIM_CHECK_VERSION);
         claim_headers.insert(HEADER_CLAIM_BUCKET, self.bucket_name.as_str());
         claim_headers.insert(HEADER_CLAIM_KEY, key.as_str());
 
-        let ack = self
-            .inner
-            .publish_with_headers(subject_str, claim_headers, Bytes::new())
-            .await
-            .map_err(ClaimCheckPublishError::PublishFailed)?;
-
-        Ok(ClaimCheckAckFuture {
-            inner: Box::pin(ack.into_future()),
-            _se: std::marker::PhantomData,
-        })
-    }
-}
-
-impl<P, S> JetStreamContext for ClaimCheckPublisher<P, S>
-where
-    P: JetStreamContext,
-    S: Send + Sync + Clone + 'static,
-{
-    type Error = P::Error;
-    type Stream = P::Stream;
-
-    fn get_or_create_stream<C: Into<stream::Config> + Send>(
-        &self,
-        config: C,
-    ) -> impl std::future::Future<Output = Result<Self::Stream, Self::Error>> + Send {
-        self.inner.get_or_create_stream(config)
+        super::publish::publish_event(
+            &self.publisher,
+            subject,
+            claim_headers,
+            Bytes::new(),
+            ack_timeout,
+        )
+        .await
     }
 }
 
@@ -289,58 +205,57 @@ mod integration_tests {
     use super::*;
     use crate::jetstream::mocks::MockJetStreamPublisher;
     use crate::jetstream::mocks::MockObjectStore;
-    use crate::jetstream::publish::publish_event;
 
     #[tokio::test]
     async fn small_payload_publishes_directly() {
-        let inner = MockJetStreamPublisher::new();
+        let publisher = MockJetStreamPublisher::new();
         let store = MockObjectStore::new();
-        let publisher = ClaimCheckPublisher::new(
-            inner.clone(),
+        let cc = ClaimCheckPublisher::new(
+            publisher.clone(),
             store.clone(),
             "test-bucket".to_string(),
             MaxPayload::from_server_limit(1024 + PROTOCOL_OVERHEAD),
         );
 
-        let outcome = publish_event(
-            &publisher,
-            "test.subject".to_string(),
-            HeaderMap::new(),
-            Bytes::from(vec![0u8; 512]),
-            Duration::from_secs(5),
-        )
-        .await;
+        let outcome = cc
+            .publish_event(
+                "test.subject".to_string(),
+                HeaderMap::new(),
+                Bytes::from(vec![0u8; 512]),
+                Duration::from_secs(5),
+            )
+            .await;
 
         assert!(outcome.is_ok());
-        let messages = inner.published_messages();
+        let messages = publisher.published_messages();
         assert_eq!(messages.len(), 1);
         assert_eq!(messages[0].payload.len(), 512);
         assert!(store.stored_objects().is_empty());
     }
 
     #[tokio::test]
     async fn large_payload_stores_in_object_store_and_publishes_claim() {
-        let inner = MockJetStreamPublisher::new();
+        let publisher = MockJetStreamPublisher::new();
         let store = MockObjectStore::new();
-        let publisher = ClaimCheckPublisher::new(
-            inner.clone(),
+        let cc = ClaimCheckPublisher::new(
+            publisher.clone(),
             store.clone(),
             "test-bucket".to_string(),
             MaxPayload::from_server_limit(1024 + PROTOCOL_OVERHEAD),
         );
 
-        let outcome = publish_event(
-            &publisher,
-            "test.subject".to_string(),
-            HeaderMap::new(),
-            Bytes::from(vec![0u8; 2048]),
-            Duration::from_secs(5),
-        )
-        .await;
+        let outcome = cc
+            .publish_event(
+                "test.subject".to_string(),
+                HeaderMap::new(),
+                Bytes::from(vec![0u8; 2048]),
+                Duration::from_secs(5),
+            )
+            .await;
 
         assert!(outcome.is_ok());
 
-        let messages = inner.published_messages();
+        let messages = publisher.published_messages();
         assert_eq!(messages.len(), 1);
         assert!(messages[0].payload.is_empty());
         assert_eq!(
@@ -375,60 +290,60 @@ mod integration_tests {
 
     #[tokio::test]
     async fn payload_at_exact_threshold_publishes_directly() {
-        let inner = MockJetStreamPublisher::new();
+        let publisher = MockJetStreamPublisher::new();
         let store = MockObjectStore::new();
-        let publisher = ClaimCheckPublisher::new(
-            inner.clone(),
+        let cc = ClaimCheckPublisher::new(
+            publisher.clone(),
             store.clone(),
             "test-bucket".to_string(),
             MaxPayload::from_server_limit(1024 + PROTOCOL_OVERHEAD),
         );
 
-        let outcome = publish_event(
-            &publisher,
-            "test.subject".to_string(),
-            HeaderMap::new(),
-            Bytes::from(vec![0u8; 1024]),
-            Duration::from_secs(5),
-        )
-        .await;
+        let outcome = cc
+            .publish_event(
+                "test.subject".to_string(),
+                HeaderMap::new(),
+                Bytes::from(vec![0u8; 1024]),
+                Duration::from_secs(5),
+            )
+            .await;
 
         assert!(outcome.is_ok());
-        assert_eq!(inner.published_messages()[0].payload.len(), 1024);
+        assert_eq!(publisher.published_messages()[0].payload.len(), 1024);
         assert!(store.stored_objects().is_empty());
     }
 
     #[tokio::test]
-    async fn object_store_failure_returns_publish_failed() {
-        let inner = MockJetStreamPublisher::new();
+    async fn object_store_failure_returns_store_failed() {
+        let publisher = MockJetStreamPublisher::new();
         let store = MockObjectStore::new();
         store.fail_next_put();
-        let publisher = ClaimCheckPublisher::new(
-            inner.clone(),
+        let cc = ClaimCheckPublisher::new(
+            publisher.clone(),
             store,
             "test-bucket".to_string(),
             MaxPayload::from_server_limit(1024 + PROTOCOL_OVERHEAD),
         );
 
-        let outcome = publish_event(
-            &publisher,
-            "test.subject".to_string(),
-            HeaderMap::new(),
-            Bytes::from(vec![0u8; 2048]),
-            Duration::from_secs(5),
-        )
-        .await;
+        let outcome = cc
+            .publish_event(
+                "test.subject".to_string(),
+                HeaderMap::new(),
+                Bytes::from(vec![0u8; 2048]),
+                Duration::from_secs(5),
+            )
+            .await;
 
         assert!(!outcome.is_ok());
-        assert!(inner.published_messages().is_empty());
+        assert!(publisher.published_messages().is_empty());
     }
 
     #[tokio::test]
     async fn large_payload_preserves_original_headers() {
-        let inner = MockJetStreamPublisher::new();
+        let publisher = MockJetStreamPublisher::new();
         let store = MockObjectStore::new();
-        let publisher = ClaimCheckPublisher::new(
-            inner.clone(),
+        let cc = ClaimCheckPublisher::new(
+            publisher.clone(),
             store,
             "test-bucket".to_string(),
             MaxPayload::from_server_limit(1024 + PROTOCOL_OVERHEAD),
@@ -437,17 +352,17 @@ mod integration_tests {
         let mut headers = HeaderMap::new();
         headers.insert("X-Custom", "value");
 
-        let outcome = publish_event(
-            &publisher,
-            "test.subject".to_string(),
-            headers,
-            Bytes::from(vec![0u8; 2048]),
-            Duration::from_secs(5),
-        )
-        .await;
+        let outcome = cc
+            .publish_event(
+                "test.subject".to_string(),
+                headers,
+                Bytes::from(vec![0u8; 2048]),
+                Duration::from_secs(5),
+            )
+            .await;
 
         assert!(outcome.is_ok());
-        let msg = &inner.published_messages()[0];
+        let msg = &publisher.published_messages()[0];
         assert_eq!(msg.headers.get("X-Custom").unwrap().as_str(), "value");
         assert!(msg.headers.get(HEADER_CLAIM_CHECK).is_some());
     }