Implement stream listener (#20)

* First attempt for stream_listener

* add test to stream listener

* add util fn to convert receiver to stream

* add bounded channel

* add broadcast listener

* fix `spawn_broadcast_listener`

* unify spawn_listener & remove oneline functions

* doc update

* add impl of sync spawn listener

* rename spawn_listener to spawn_listener_from_iter, and port spawn_listener

* add bound channel to threads concurrency

* merge duplicated code

* add cancel token with 'flaky' test

* unflaky the test

* add cancellation to task impl of spawn_listener

* docs & clippy

* use futures select inside spawn listener

* remove bounded channels from tasks impl

* remove sync channels from threads impl

* deprecate spawn_listener for threads impl

* fix imports

* reword example for clarity

* add comment for clarity

---------

Co-authored-by: Esteban Dimitroff Hodi <[email protected]>

Author: juan518munoz

.tool-versions

@@ -1 +1 @@
-rust 1.85.1
+rust 1.88.0

Cargo.lock

@@ -8,5 +8,10 @@ spawned-rt = { workspace = true }
 tracing = { workspace = true }
 futures = "0.3.1"
 
+[dev-dependencies]
+# This tokio imports are only used in tests, we should not use them in the library code.
+tokio-stream = { version = "0.1.17" }
+tokio = { version = "1", features = ["full"] }
+
 [lib]
-path = "./src/lib.rs"
+path = "./src/lib.rs"

concurrency/Cargo.toml

@@ -3,11 +3,15 @@
 
 mod gen_server;
 mod process;
+mod stream;
 mod time;
 
+#[cfg(test)]
+mod stream_tests;
 #[cfg(test)]
 mod timer_tests;
 
 pub use gen_server::{CallResponse, CastResponse, GenServer, GenServerHandle, GenServerInMsg};
 pub use process::{send, Process, ProcessInfo};
+pub use stream::spawn_listener;
 pub use time::{send_after, send_interval};

concurrency/src/tasks/mod.rs

@@ -0,0 +1,57 @@
+use crate::tasks::{GenServer, GenServerHandle};
+use futures::{future::select, Stream, StreamExt};
+use spawned_rt::tasks::{CancellationToken, JoinHandle};
+
+/// Spawns a listener that listens to a stream and sends messages to a GenServer.
+///
+/// Items sent through the stream are required to be wrapped in a Result type.
+///
+/// This function returns a handle to the spawned task and a cancellation token
+/// to stop it.
+pub fn spawn_listener<T, F, S, I, E>(
+    mut handle: GenServerHandle<T>,
+    message_builder: F,
+    mut stream: S,
+) -> (JoinHandle<()>, CancellationToken)
+where
+    T: GenServer + 'static,
+    F: Fn(I) -> T::CastMsg + Send + 'static + std::marker::Sync,
+    I: Send,
+    E: std::fmt::Debug + Send,
+    S: Unpin + Send + Stream<Item = Result<I, E>> + 'static,
+{
+    let cancelation_token = CancellationToken::new();
+    let cloned_token = cancelation_token.clone();
+    let join_handle = spawned_rt::tasks::spawn(async move {
+        let result = select(
+            Box::pin(cloned_token.cancelled()),
+            Box::pin(async {
+                loop {
+                    match stream.next().await {
+                        Some(Ok(i)) => match handle.cast(message_builder(i)).await {
+                            Ok(_) => tracing::trace!("Message sent successfully"),
+                            Err(e) => {
+                                tracing::error!("Failed to send message: {e:?}");
+                                break;
+                            }
+                        },
+                        Some(Err(e)) => {
+                            tracing::trace!("Received Error in msg {e:?}");
+                            break;
+                        }
+                        None => {
+                            tracing::trace!("Stream finished");
+                            break;
+                        }
+                    }
+                }
+            }),
+        )
+        .await;
+        match result {
+            futures::future::Either::Left(_) => tracing::trace!("Listener cancelled"),
+            futures::future::Either::Right(_) => (), // Stream finished or errored out
+        }
+    });
+    (join_handle, cancelation_token)
+}

concurrency/src/tasks/stream.rs

@@ -0,0 +1,174 @@
+use std::time::Duration;
+
+use spawned_rt::tasks::{self as rt, BroadcastStream, ReceiverStream};
+
+use crate::tasks::{
+    stream::spawn_listener, CallResponse, CastResponse, GenServer, GenServerHandle,
+};
+
+type SummatoryHandle = GenServerHandle<Summatory>;
+
+struct Summatory;
+
+type SummatoryState = u16;
+
+#[derive(Clone)]
+struct UpdateSumatory {
+    added_value: u16,
+}
+
+impl Summatory {
+    pub async fn get_value(server: &mut SummatoryHandle) -> Result<SummatoryState, ()> {
+        server.call(()).await.map_err(|_| ())
+    }
+}
+
+impl GenServer for Summatory {
+    type CallMsg = (); // We only handle one type of call, so there is no need for a specific message type.
+    type CastMsg = UpdateSumatory;
+    type OutMsg = SummatoryState;
+    type State = SummatoryState;
+    type Error = ();
+
+    fn new() -> Self {
+        Self
+    }
+
+    async fn handle_cast(
+        &mut self,
+        message: Self::CastMsg,
+        _handle: &GenServerHandle<Self>,
+        state: Self::State,
+    ) -> CastResponse<Self> {
+        let new_state = state + message.added_value;
+        CastResponse::NoReply(new_state)
+    }
+
+    async fn handle_call(
+        &mut self,
+        _message: Self::CallMsg,
+        _handle: &SummatoryHandle,
+        state: Self::State,
+    ) -> CallResponse<Self> {
+        let current_value = state;
+        CallResponse::Reply(state, current_value)
+    }
+}
+
+// In this example, the stream sends u8 values, which are converted to the type
+// supported by the GenServer (UpdateSumatory / u16).
+fn message_builder(value: u8) -> UpdateSumatory {
+    UpdateSumatory {
+        added_value: value as u16,
+    }
+}
+
+#[test]
+pub fn test_sum_numbers_from_stream() {
+    let runtime = rt::Runtime::new().unwrap();
+    runtime.block_on(async move {
+        let mut summatory_handle = Summatory::start(0);
+        let stream = tokio_stream::iter(vec![1u8, 2, 3, 4, 5].into_iter().map(Ok::<u8, ()>));
+
+        spawn_listener(summatory_handle.clone(), message_builder, stream);
+
+        // Wait for 1 second so the whole stream is processed
+        rt::sleep(Duration::from_secs(1)).await;
+
+        let val = Summatory::get_value(&mut summatory_handle).await.unwrap();
+        assert_eq!(val, 15);
+    })
+}
+
+#[test]
+pub fn test_sum_numbers_from_channel() {
+    let runtime = rt::Runtime::new().unwrap();
+    runtime.block_on(async move {
+        let mut summatory_handle = Summatory::start(0);
+        let (tx, rx) = spawned_rt::tasks::mpsc::channel::<Result<u8, ()>>();
+
+        // Spawn a task to send numbers to the channel
+        spawned_rt::tasks::spawn(async move {
+            for i in 1..=5 {
+                tx.send(Ok(i)).unwrap();
+            }
+        });
+
+        spawn_listener(
+            summatory_handle.clone(),
+            message_builder,
+            ReceiverStream::new(rx),
+        );
+
+        // Wait for 1 second so the whole stream is processed
+        rt::sleep(Duration::from_secs(1)).await;
+
+        let val = Summatory::get_value(&mut summatory_handle).await.unwrap();
+        assert_eq!(val, 15);
+    })
+}
+
+#[test]
+pub fn test_sum_numbers_from_broadcast_channel() {
+    let runtime = rt::Runtime::new().unwrap();
+    runtime.block_on(async move {
+        let mut summatory_handle = Summatory::start(0);
+        let (tx, rx) = tokio::sync::broadcast::channel::<u8>(5);
+
+        // Spawn a task to send numbers to the channel
+        spawned_rt::tasks::spawn(async move {
+            for i in 1u8..=5 {
+                tx.send(i).unwrap();
+            }
+        });
+
+        spawn_listener(
+            summatory_handle.clone(),
+            message_builder,
+            BroadcastStream::new(rx),
+        );
+
+        // Wait for 1 second so the whole stream is processed
+        rt::sleep(Duration::from_secs(1)).await;
+
+        let val = Summatory::get_value(&mut summatory_handle).await.unwrap();
+        assert_eq!(val, 15);
+    })
+}
+
+#[test]
+pub fn test_stream_cancellation() {
+    const RUNNING_TIME: u64 = 1000;
+
+    let runtime = rt::Runtime::new().unwrap();
+    runtime.block_on(async move {
+        let mut summatory_handle = Summatory::start(0);
+        let (tx, rx) = spawned_rt::tasks::mpsc::channel::<Result<u8, ()>>();
+
+        // Spawn a task to send numbers to the channel
+        spawned_rt::tasks::spawn(async move {
+            for i in 1..=5 {
+                tx.send(Ok(i)).unwrap();
+                rt::sleep(Duration::from_millis(RUNNING_TIME / 4)).await;
+            }
+        });
+
+        let (_handle, cancellation_token) = spawn_listener(
+            summatory_handle.clone(),
+            message_builder,
+            ReceiverStream::new(rx),
+        );
+
+        // Wait for 1 second so the whole stream is processed
+        rt::sleep(Duration::from_millis(RUNNING_TIME)).await;
+
+        cancellation_token.cancel();
+
+        // The reasoning for this assertion is that each message takes a quarter of the total time
+        // to be processed, so having a stream of 5 messages, the last one won't be processed.
+        // We could safely assume that it will get to process 4 messages, but in case of any extenal
+        // slowdown, it could process less.
+        let val = Summatory::get_value(&mut summatory_handle).await.unwrap();
+        assert!(val > 0 && val < 15);
+    })
+}

concurrency/src/tasks/stream_tests.rs

@@ -3,11 +3,13 @@
 
 mod gen_server;
 mod process;
+mod stream;
 mod time;
 
 #[cfg(test)]
 mod timer_tests;
 
 pub use gen_server::{CallResponse, CastResponse, GenServer, GenServerHandle, GenServerInMsg};
 pub use process::{send, Process, ProcessInfo};
+pub use stream::spawn_listener;
 pub use time::{send_after, send_interval};

concurrency/src/threads/mod.rs

@@ -0,0 +1,17 @@
+use crate::threads::{GenServer, GenServerHandle};
+
+use futures::Stream;
+
+/// Spawns a listener that listens to a stream and sends messages to a GenServer.
+///
+/// Items sent through the stream are required to be wrapped in a Result type.
+pub fn spawn_listener<T, F, S, I, E>(_handle: GenServerHandle<T>, _message_builder: F, _stream: S)
+where
+    T: GenServer + 'static,
+    F: Fn(I) -> T::CastMsg + Send + 'static,
+    I: Send + 'static,
+    E: std::fmt::Debug + Send + 'static,
+    S: Unpin + Send + Stream<Item = Result<I, E>> + 'static,
+{
+    unimplemented!("Unsupported function in threads mode")
+}

concurrency/src/threads/stream.rs

@@ -6,9 +6,10 @@ edition = "2021"
 [dependencies]
 tokio = { version = "1", features = ["full"] }
 tokio-util = { version = "0.7.15" }
+tokio-stream = { version = "0.1.17", features = ["sync"] }
 crossbeam = { version = "0.7.3" }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 
 [lib]
-path = "./src/lib.rs"
+path = "./src/lib.rs"

rt/Cargo.toml

@@ -18,6 +18,7 @@ pub use crate::tasks::tokio::oneshot;
 pub use crate::tasks::tokio::sleep;
 pub use crate::tasks::tokio::CancellationToken;
 pub use crate::tasks::tokio::{spawn, spawn_blocking, JoinHandle, Runtime};
+pub use crate::tasks::tokio::{BroadcastStream, ReceiverStream};
 use std::future::Future;
 
 pub fn run<F: Future>(future: F) -> F::Output {