Skip to content

Instantly share code, notes, and snippets.

@rklaehn
Created April 17, 2024 06:13
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save rklaehn/da278012973f4dfd1ff0c4bddc698966 to your computer and use it in GitHub Desktop.
Save rklaehn/da278012973f4dfd1ff0c4bddc698966 to your computer and use it in GitHub Desktop.
use futures::StreamExt;
struct Response {
inner: flume::Receiver<u64>,
total: u64,
sum: u64,
}
impl Response {
#[cfg(feature = "async")]
pub fn recv_all_async(&mut self) -> impl futures::Stream<Item = u64> + '_ {
futures::stream::unfold(self, |res| async move {
match res.inner.recv_async().await {
Ok(v) => {
res.total += 1;
res.sum += v;
Some((v, res))
},
Err(_) => None,
}
})
}
#[cfg(feature = "sync")]
pub fn recv_all(&mut self) -> impl Iterator<Item = u64> + '_ {
std::iter::from_fn(move || {
match self.inner.recv() {
Ok(v) => {
self.total += 1;
self.sum += v;
Some(v)
},
Err(_) => None,
}
})
}
}
#[cfg(feature = "async")]
#[tokio::main]
async fn main() {
let (s, r) = flume::unbounded();
let mut res = Response {
inner: r,
total: 0,
sum: 0,
};
s.send(1).unwrap();
s.send(2).unwrap();
s.send(3).unwrap();
drop(s);
res.recv_all_async().for_each(|v| async move {
println!("Received: {}", v);
}).await;
println!("Total: {}", res.total);
println!("Sum: {}", res.sum);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment