fix join_all

This commit is contained in:
Rob Ede 2021-11-03 14:55:24 +00:00
parent 010b557492
commit c33fb4ef9a
3 changed files with 77 additions and 10 deletions

View File

@ -4,13 +4,15 @@ use std::{
task::{Context, Poll},
};
use futures_core::future::{BoxFuture, LocalBoxFuture};
// a poor man's join future. joined future is only used when starting/stopping the server.
// pin_project and pinned futures are overkill for this task.
pub(crate) struct JoinAll<T> {
fut: Vec<JoinFuture<T>>,
}
pub(crate) fn join_all<T>(fut: Vec<impl Future<Output = T> + 'static>) -> JoinAll<T> {
pub(crate) fn join_all<T>(fut: Vec<impl Future<Output = T> + Send + 'static>) -> JoinAll<T> {
let fut = fut
.into_iter()
.map(|f| JoinFuture::Future(Box::pin(f)))
@ -20,7 +22,7 @@ pub(crate) fn join_all<T>(fut: Vec<impl Future<Output = T> + 'static>) -> JoinAl
}
enum JoinFuture<T> {
Future(Pin<Box<dyn Future<Output = T>>>),
Future(BoxFuture<'static, T>),
Result(Option<T>),
}
@ -59,6 +61,63 @@ impl<T> Future for JoinAll<T> {
}
}
pub(crate) fn join_all_local<T>(
fut: Vec<impl Future<Output = T> + 'static>,
) -> JoinAllLocal<T> {
let fut = fut
.into_iter()
.map(|f| JoinLocalFuture::LocalFuture(Box::pin(f)))
.collect();
JoinAllLocal { fut }
}
// a poor man's join future. joined future is only used when starting/stopping the server.
// pin_project and pinned futures are overkill for this task.
pub(crate) struct JoinAllLocal<T> {
fut: Vec<JoinLocalFuture<T>>,
}
enum JoinLocalFuture<T> {
LocalFuture(LocalBoxFuture<'static, T>),
Result(Option<T>),
}
impl<T> Unpin for JoinAllLocal<T> {}
impl<T> Future for JoinAllLocal<T> {
type Output = Vec<T>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut ready = true;
let this = self.get_mut();
for fut in this.fut.iter_mut() {
if let JoinLocalFuture::LocalFuture(f) = fut {
match f.as_mut().poll(cx) {
Poll::Ready(t) => {
*fut = JoinLocalFuture::Result(Some(t));
}
Poll::Pending => ready = false,
}
}
}
if ready {
let mut res = Vec::new();
for fut in this.fut.iter_mut() {
if let JoinLocalFuture::Result(f) = fut {
res.push(f.take().unwrap());
}
}
Poll::Ready(res)
} else {
Poll::Pending
}
}
}
#[cfg(test)]
mod test {
use super::*;
@ -73,4 +132,13 @@ mod test {
assert_eq!(Err(3), res.next().unwrap());
assert_eq!(Ok(9), res.next().unwrap());
}
#[actix_rt::test]
async fn test_join_all_local() {
let futs = vec![ready(Ok(1)), ready(Err(3)), ready(Ok(9))];
let mut res = join_all_local(futs).await.into_iter();
assert_eq!(Ok(1), res.next().unwrap());
assert_eq!(Err(3), res.next().unwrap());
assert_eq!(Ok(9), res.next().unwrap());
}
}

View File

@ -7,7 +7,7 @@ use std::{
};
use actix_rt::{time::sleep, System};
use futures_core::future::LocalBoxFuture;
use futures_core::future::BoxFuture;
use log::{error, info, trace};
use tokio::sync::{
mpsc::{UnboundedReceiver, UnboundedSender},
@ -45,8 +45,6 @@ pub(crate) enum ServerCommand {
},
}
// TODO: docs + must use
/// Server
///
/// # Shutdown Signals
@ -232,11 +230,11 @@ pub struct ServerInner {
cmd_rx: UnboundedReceiver<ServerCommand>,
signals: Option<Signals>,
waker_queue: WakerQueue,
stop_task: Option<LocalBoxFuture<'static, ()>>,
stop_task: Option<BoxFuture<'static, ()>>,
}
impl ServerInner {
fn handle_cmd(&mut self, item: ServerCommand) -> Option<LocalBoxFuture<'static, ()>> {
fn handle_cmd(&mut self, item: ServerCommand) -> Option<BoxFuture<'static, ()>> {
match item {
ServerCommand::Pause(tx) => {
self.waker_queue.wake(WakerInterest::Pause);
@ -319,7 +317,7 @@ impl ServerInner {
}
}
fn handle_signal(&mut self, signal: Signal) -> Option<LocalBoxFuture<'static, ()>> {
fn handle_signal(&mut self, signal: Signal) -> Option<BoxFuture<'static, ()>> {
match signal {
Signal::Int => {
info!("SIGINT received; starting forced shutdown");

View File

@ -24,7 +24,7 @@ use tokio::sync::{
};
use crate::{
join_all::join_all,
join_all::join_all_local,
service::{BoxedServerService, InternalServiceFactory},
socket::MioStream,
waker_queue::{WakerInterest, WakerQueue},
@ -324,10 +324,11 @@ impl ServerWorker {
// a second spawn to run !Send future tasks.
spawn(async move {
let res = join_all(fut)
let res = join_all_local(fut)
.await
.into_iter()
.collect::<Result<Vec<_>, _>>();
let services = match res {
Ok(res) => res
.into_iter()