fluxion_exec/

subscribe.rs

// Copyright 2025 Umberto Gotti <umberto.gotti@umbertogotti.dev>
// Licensed under the Apache License, Version 2.0
// http://www.apache.org/licenses/LICENSE-2.0

use alloc::boxed::Box;
use async_trait::async_trait;
use core::fmt::Debug;
use core::future::Future;
use fluxion_core::{CancellationToken, Result};
use futures::stream::Stream;
use futures::stream::StreamExt;

/// Extension trait providing async subscription capabilities for streams.
///
/// This trait enables processing stream items with async handlers in a sequential manner.
#[cfg(target_arch = "wasm32")]
#[async_trait(?Send)]
pub trait SubscribeExt<T>: Stream<Item = T> + Sized {
    /// Subscribes to the stream with an async handler, processing items sequentially.
    ///
    /// This method consumes the stream and processes each item with the provided handler.
    /// Items are processed in the order they arrive, with each item's handler completing
    /// before the next item is processed.
    ///
    /// # Behavior
    ///
    /// - Processes each stream item with the provided async handler sequentially
    /// - Waits for handler completion before processing next item
    /// - Continues until stream ends or cancellation token is triggered
    /// - Errors from handlers are passed to the error callback
    ///
    /// # Arguments
    ///
    /// * `on_next_func` - Async function called for each stream item. Receives the item
    ///                    and a cancellation token. Returns `Result<(), E>`.
    /// * `cancellation_token` - Optional token to stop processing. If `None`, a default
    ///                          token is created that never cancels.
    /// * `on_error_callback` - Error handler called when `on_next_func` returns an error.
    ///
    /// # Type Parameters
    ///
    /// * `F` - Function type for the item handler
    /// * `Fut` - Future type returned by the handler
    /// * `E` - Error type
    /// * `OnError` - Function type for error handling
    ///
    /// # Errors
    ///
    /// Returns `Ok(())` on stream completion. Errors from item handlers are passed
    /// to the error callback. The subscription continues processing subsequent items
    /// even if individual items fail, unless the cancellation token is triggered.
    ///
    /// # See Also
    ///
    /// - [`subscribe_latest`](crate::SubscribeLatestExt::subscribe_latest) - Cancels old work for new items
    ///
    /// # Examples
    ///
    /// ## Basic Usage
    ///
    /// Process all items sequentially:
    ///
    /// ```
    /// use fluxion_exec::SubscribeExt;
    /// use futures::channel::mpsc::unbounded;
    /// use futures::stream;
    /// use futures::StreamExt;
    /// use std::sync::Arc;
    /// use futures::lock::Mutex;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let results = Arc::new(Mutex::new(Vec::new()));
    /// let results_clone = results.clone();
    /// let (notify_tx, mut notify_rx) = unbounded();
    ///
    /// let stream = stream::iter(vec![1, 2, 3, 4, 5]);
    ///
    /// // Subscribe and process each item
    /// stream.subscribe(
    ///     move |item, _token| {
    ///         let results = results_clone.clone();
    ///         let notify_tx = notify_tx.clone();
    ///         async move {
    ///             results.lock().await.push(item * 2);
    ///             let _ = notify_tx.unbounded_send(());
    ///             Ok::<(), std::io::Error>(())
    ///         }
    ///     },
    ///     |_err| {}, // Ignore errors
    ///     None, // No cancellation
    /// ).await.unwrap();
    ///
    /// // Wait for all 5 items to be processed
    /// for _ in 0..5 {
    ///     notify_rx.next().await.unwrap();
    /// }
    ///
    /// let processed = results.lock().await;
    /// assert!(processed.contains(&2));
    /// assert!(processed.contains(&4));
    /// # }
    /// ```
    ///
    /// ## With Error Handling
    ///
    /// Use an error callback to handle errors without stopping the stream:
    ///
    /// ```
    /// use fluxion_exec::SubscribeExt;
    /// use futures::channel::mpsc::unbounded;
    /// use futures::stream;
    /// use futures::StreamExt;
    /// use std::sync::Arc;
    /// use futures::lock::Mutex;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// #[derive(Debug)]
    /// struct MyError(String);
    /// impl core::fmt::Display for MyError {
    ///     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
    ///         write!(f, "MyError: {}", self.0)
    ///     }
    /// }
    /// impl std::error::Error for MyError {}
    ///
    /// let error_count = Arc::new(Mutex::new(0));
    /// let error_count_clone = error_count.clone();
    /// let (notify_tx, mut notify_rx) = unbounded();
    ///
    /// let stream = stream::iter(vec![1, 2, 3, 4, 5]);
    ///
    /// stream.subscribe(
    ///     move |item, _token| {
    ///         let notify_tx = notify_tx.clone();
    ///         async move {
    ///             let res = if item % 2 == 0 {
    ///                 Err(MyError(format!("Even number: {}", item)))
    ///             } else {
    ///                 Ok(())
    ///             };
    ///             // Signal completion regardless of success/failure
    ///             // Note: In real code, you might signal in the error callback too
    ///             // but here we just want to know the handler finished.
    ///             // However, subscribe spawns the handler. If it errors,
    ///             // the error callback is called.
    ///             // We need to signal completion in both paths.
    ///             // Since the handler returns the error, we can't signal *after* returning Err.
    ///             // So we signal before returning.
    ///             let _ = notify_tx.unbounded_send(());
    ///             res
    ///         }
    ///     },
    ///     move |_err| {
    ///         let count = error_count_clone.clone();
    ///         tokio::spawn(async move {
    ///             *count.lock().await += 1;
    ///         });
    ///     },
    ///     None,
    /// ).await.unwrap();
    ///
    /// // Wait for 5 items
    /// for _ in 0..5 {
    ///     notify_rx.next().await.unwrap();
    /// }
    ///
    /// // Give a tiny bit of time for the error callback spawn to finish updating the count
    /// // (Since the callback spawns another task)
    /// // Alternatively, we could use a channel in the error callback too.
    /// tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
    ///
    /// assert_eq!(*error_count.lock().await, 2); // Items 2 and 4 errored
    /// # }
    /// ```
    ///
    /// ## With Cancellation
    ///
    /// Use a cancellation token to stop processing:
    ///
    /// ```
    /// use fluxion_exec::SubscribeExt;
    /// use futures::channel::mpsc::unbounded;
    /// use futures::StreamExt;
    /// use fluxion_core::CancellationToken;
    /// use std::sync::Arc;
    /// use futures::lock::Mutex;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let (tx, rx) = unbounded();
    /// let stream = rx;
    ///
    /// let cancel_token = CancellationToken::new();
    /// let cancel_clone = cancel_token.clone();
    ///
    /// let processed = Arc::new(Mutex::new(Vec::new()));
    /// let processed_clone = processed.clone();
    /// let (notify_tx, mut notify_rx) = unbounded();
    ///
    /// let handle = tokio::spawn(async move {
    ///     stream.subscribe(
    ///         move |item, token| {
    ///             let vec = processed_clone.clone();
    ///             let notify_tx = notify_tx.clone();
    ///             async move {
    ///                 if token.is_cancelled() {
    ///                     return Ok(());
    ///                 }
    ///                 vec.lock().await.push(item);
    ///                 let _ = notify_tx.unbounded_send(());
    ///                 Ok::<(), std::io::Error>(())
    ///             }
    ///         },
    ///         |_| {}, // Ignore errors
    ///         Some(cancel_token),
    ///     ).await
    /// });
    ///
    /// // Send items
    /// tx.unbounded_send(1).unwrap();
    /// tx.unbounded_send(2).unwrap();
    /// tx.unbounded_send(3).unwrap();
    ///
    /// // Wait for first item to be processed
    /// notify_rx.next().await.unwrap();
    ///
    /// // Cancel now
    /// cancel_clone.cancel();
    /// drop(tx);
    ///
    /// handle.await.unwrap().unwrap();
    ///
    /// // At least one item should be processed before cancellation
    /// assert!(!processed.lock().await.is_empty());
    /// # }
    /// ```
    ///
    /// ## Database Write Pattern
    ///
    /// Process events and persist to a database:
    ///
    /// ```
    /// use fluxion_exec::SubscribeExt;
    /// use futures::channel::mpsc::unbounded;
    /// use futures::stream;
    /// use futures::StreamExt;
    /// use std::sync::Arc;
    /// use futures::lock::Mutex;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// #[derive(Clone, Debug)]
    /// struct Event { id: u32, data: String }
    ///
    /// // Simulated database
    /// let db = Arc::new(Mutex::new(Vec::new()));
    /// let db_clone = db.clone();
    /// let (notify_tx, mut notify_rx) = unbounded();
    ///
    /// let events = vec![
    ///     Event { id: 1, data: "event1".to_string() },
    ///     Event { id: 2, data: "event2".to_string() },
    /// ];
    ///
    /// let stream = stream::iter(events);
    ///
    /// stream.subscribe(
    ///     move |event, _token| {
    ///         let db = db_clone.clone();
    ///         let notify_tx = notify_tx.clone();
    ///         async move {
    ///             // Simulate database write
    ///             db.lock().await.push(event);
    ///             let _ = notify_tx.unbounded_send(());
    ///             Ok::<(), std::io::Error>(())
    ///         }
    ///     },
    ///     |err| eprintln!("DB Error: {}", err),
    ///     None,
    /// ).await.unwrap();
    ///
    /// // Wait for 2 events
    /// notify_rx.next().await.unwrap();
    /// notify_rx.next().await.unwrap();
    ///
    /// assert_eq!(db.lock().await.len(), 2);
    /// # }
    /// ```
    ///
    /// # Thread Safety
    ///
    /// Handlers are executed sequentially on the calling task. The subscription
    /// completes when the stream ends or cancellation is triggered.
    async fn subscribe<F, Fut, E, OnError>(
        self,
        on_next_func: F,
        on_error_callback: OnError,
        cancellation_token: Option<CancellationToken>,
    ) -> Result<()>
    where
        F: Fn(T, CancellationToken) -> Fut + Clone + 'static,
        Fut: Future<Output = core::result::Result<(), E>> + 'static,
        OnError: Fn(E) + Clone + 'static,
        T: Debug + Clone + 'static,
        E: 'static;
}

// Non-WASM version with Send + Sync bounds for multi-threaded runtimes
#[cfg(not(target_arch = "wasm32"))]
#[async_trait]
pub trait SubscribeExt<T>: Stream<Item = T> + Sized {
    async fn subscribe<F, Fut, E, OnError>(
        self,
        on_next_func: F,
        on_error_callback: OnError,
        cancellation_token: Option<CancellationToken>,
    ) -> Result<()>
    where
        F: Fn(T, CancellationToken) -> Fut + Clone + Send + Sync + 'static,
        Fut: Future<Output = core::result::Result<(), E>> + Send + 'static,
        OnError: Fn(E) + Clone + Send + Sync + 'static,
        T: Debug + Send + Clone + 'static,
        E: Send + 'static;
}

// WASM implementation
#[cfg(target_arch = "wasm32")]
#[async_trait(?Send)]
impl<S, T> SubscribeExt<T> for S
where
    S: Stream<Item = T> + Unpin + 'static,
    T: 'static,
{
    async fn subscribe<F, Fut, E, OnError>(
        self,
        on_next_func: F,
        on_error_callback: OnError,
        cancellation_token: Option<CancellationToken>,
    ) -> Result<()>
    where
        F: Fn(T, CancellationToken) -> Fut + Clone + 'static,
        Fut: Future<Output = core::result::Result<(), E>> + 'static,
        OnError: Fn(E) + Clone + 'static,
        T: Debug + Clone + 'static,
        E: 'static,
    {
        subscribe_impl(self, on_next_func, on_error_callback, cancellation_token).await
    }
}

// Non-WASM implementation with Send + Sync bounds
#[cfg(not(target_arch = "wasm32"))]
#[async_trait]
impl<S, T> SubscribeExt<T> for S
where
    S: Stream<Item = T> + Send + Unpin + 'static,
    T: Send + 'static,
{
    async fn subscribe<F, Fut, E, OnError>(
        self,
        on_next_func: F,
        on_error_callback: OnError,
        cancellation_token: Option<CancellationToken>,
    ) -> Result<()>
    where
        F: Fn(T, CancellationToken) -> Fut + Clone + Send + Sync + 'static,
        Fut: Future<Output = core::result::Result<(), E>> + Send + 'static,
        OnError: Fn(E) + Clone + Send + Sync + 'static,
        T: Debug + Send + Clone + 'static,
        E: Send + 'static,
    {
        subscribe_impl(self, on_next_func, on_error_callback, cancellation_token).await
    }
}

// Shared implementation for both WASM and non-WASM
async fn subscribe_impl<S, T, F, Fut, E, OnError>(
    mut stream: S,
    on_next_func: F,
    on_error_callback: OnError,
    cancellation_token: Option<CancellationToken>,
) -> Result<()>
where
    S: Stream<Item = T> + Unpin,
    F: Fn(T, CancellationToken) -> Fut + Clone,
    Fut: Future<Output = core::result::Result<(), E>>,
    OnError: Fn(E) + Clone,
    T: Debug + Clone,
{
    let cancellation_token = cancellation_token.unwrap_or_default();

    while let Some(item) = stream.next().await {
        if cancellation_token.is_cancelled() {
            break;
        }

        // Call handler directly (sequential processing)
        let result = on_next_func(item.clone(), cancellation_token.clone()).await;

        if let Err(error) = result {
            on_error_callback(error);
        }
    }

    Ok(())
}