67 lines
1.8 KiB
Rust
67 lines
1.8 KiB
Rust
|
use crate::util::BoxFuture;
|
||
|
|
use std::error;
|
||
|
|
use tokio_util::sync::CancellationToken;
|
||
|
|
|
||
|
|
pub struct TaskStack<E>
|
||
|
|
where
|
||
|
|
E: error::Error + Send + 'static,
|
||
|
|
{
|
||
|
|
wait_group: Vec<BoxFuture<'static, Result<(), E>>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl<E> TaskStack<E>
|
||
|
|
where
|
||
|
|
E: error::Error + Send + 'static,
|
||
|
|
{
|
||
|
|
pub fn new() -> TaskStack<E> {
|
||
|
|
TaskStack {
|
||
|
|
wait_group: Vec::new(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// push adds the given Future to the stack, to be executed once stop is called.
|
||
|
|
pub fn push<Fut>(&mut self, f: Fut)
|
||
|
|
where
|
||
|
|
Fut: futures::Future<Output = Result<(), E>> + Send + 'static,
|
||
|
|
{
|
||
|
|
self.wait_group.push(Box::pin(f))
|
||
|
|
}
|
||
|
|
|
||
|
|
/// push_spawn will spawn the given closure in a tokio task. Once the CancellationToken is
|
||
|
|
/// cancelled the closure is expected to return.
|
||
|
|
pub fn push_spawn<F, Fut>(&mut self, mut f: F)
|
||
|
|
where
|
||
|
|
Fut: futures::Future<Output = Result<(), E>> + Send + 'static,
|
||
|
|
F: FnMut(CancellationToken) -> Fut,
|
||
|
|
{
|
||
|
|
let canceller = CancellationToken::new();
|
||
|
|
let handle = tokio::spawn(f(canceller.clone()));
|
||
|
|
self.push(async move {
|
||
|
|
canceller.cancel();
|
||
|
|
handle.await.expect("failed to join task")
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
/// stop will process all operations which have been pushed onto the stack in the reverse order
|
||
|
|
/// they were pushed.
|
||
|
|
pub async fn stop(mut self) -> Result<(), E> {
|
||
|
|
// reverse wait_group in place, so we stop the most recently added first. Since this method
|
||
|
|
// consumes self this is fine.
|
||
|
|
self.wait_group.reverse();
|
||
|
|
|
||
|
|
for fut in self.wait_group {
|
||
|
|
fut.await?;
|
||
|
|
}
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl<E> Default for TaskStack<E>
|
||
|
|
where
|
||
|
|
E: error::Error + Send + 'static,
|
||
|
|
{
|
||
|
|
fn default() -> Self {
|
||
|
|
Self::new()
|
||
|
|
}
|
||
|
|
}
|