use std::sync::mpsc;
use std::time::Duration;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum JobResult<T> {
Done(T),
TimedOut,
Failed,
}
pub async fn run<T, F>(work: F, timeout: Duration) -> JobResult<T>
where
T: Send + 'static,
F: FnOnce() -> T + Send + 'static,
{
let (tx, rx) = mpsc::channel();
std::thread::spawn(move || {
let _ = tx.send(work());
});
match rx.recv_timeout(timeout) {
Ok(value) => JobResult::Done(value),
Err(mpsc::RecvTimeoutError::Timeout) => JobResult::TimedOut,
Err(mpsc::RecvTimeoutError::Disconnected) => JobResult::Failed,
}
}
pub fn stream<T, F>(work: F, timeout: Duration) -> impl futures::Stream<Item = JobResult<T>>
where
T: Send + 'static,
F: FnOnce(&mut dyn FnMut(T)) + Send + 'static,
{
let (tx, rx) = mpsc::channel();
std::thread::spawn(move || {
let emitter = tx.clone();
work(&mut move |item| {
let _ = emitter.send(Some(item));
});
let _ = tx.send(None);
});
futures::stream::unfold(Some(rx), move |state| async move {
let rx = state?;
match rx.recv_timeout(timeout) {
Ok(Some(item)) => Some((JobResult::Done(item), Some(rx))),
Ok(None) => None,
Err(mpsc::RecvTimeoutError::Timeout) => Some((JobResult::TimedOut, None)),
Err(mpsc::RecvTimeoutError::Disconnected) => Some((JobResult::Failed, None)),
}
})
}
#[cfg(test)]
mod tests {
use super::*;
use futures::StreamExt;
use std::time::Duration;
#[test]
fn work_that_finishes_returns_its_value() {
let got = pollster_block(run(|| 21 * 2, Duration::from_secs(5)));
assert_eq!(got, JobResult::Done(42));
}
#[test]
fn work_that_hangs_times_out_rather_than_blocking_forever() {
let got: JobResult<()> = pollster_block(run(
|| std::thread::sleep(Duration::from_secs(30)),
Duration::from_millis(50),
));
assert_eq!(got, JobResult::TimedOut);
}
#[test]
fn work_that_panics_is_reported_rather_than_killing_the_app() {
let got: JobResult<()> = pollster_block(run(|| panic!("boom"), Duration::from_secs(5)));
assert_eq!(got, JobResult::Failed);
}
#[test]
fn a_timed_out_job_does_not_block_the_caller_for_the_works_duration() {
let started = std::time::Instant::now();
let _: JobResult<()> = pollster_block(run(
|| std::thread::sleep(Duration::from_secs(30)),
Duration::from_millis(50),
));
assert!(
started.elapsed() < Duration::from_secs(5),
"returned after {:?}; the timeout did not bound the wait",
started.elapsed()
);
}
#[test]
fn every_item_arrives_in_order_and_the_stream_then_ends() {
let got: Vec<_> = futures::executor::block_on(
stream(
|emit| {
for i in 0..3 {
emit(i);
}
},
Duration::from_secs(5),
)
.collect(),
);
assert_eq!(
got,
vec![JobResult::Done(0), JobResult::Done(1), JobResult::Done(2)],
"the stream must end when the work returns, with nothing after it"
);
}
#[test]
fn an_item_reaches_the_consumer_before_the_work_has_finished() {
let started = std::time::Instant::now();
let mut arrivals = Vec::new();
futures::executor::block_on(
stream(
|emit| {
emit(());
std::thread::sleep(Duration::from_millis(300));
emit(());
},
Duration::from_secs(5),
)
.for_each(|_| {
arrivals.push(started.elapsed());
std::future::ready(())
}),
);
assert_eq!(arrivals.len(), 2);
assert!(
arrivals[0] + Duration::from_millis(100) < arrivals[1],
"both items arrived at once ({arrivals:?}); that is the buffered \
behaviour `run` already had"
);
}
#[test]
fn work_that_goes_silent_times_out_rather_than_hanging_the_stream() {
let got: Vec<_> = futures::executor::block_on(
stream(
|emit| {
emit(1);
std::thread::sleep(Duration::from_secs(30));
},
Duration::from_millis(50),
)
.collect(),
);
assert_eq!(got, vec![JobResult::Done(1), JobResult::TimedOut]);
}
#[test]
fn work_that_keeps_producing_is_not_timed_out() {
let got: Vec<_> = futures::executor::block_on(
stream(
|emit| {
for i in 0..5 {
std::thread::sleep(Duration::from_millis(30));
emit(i);
}
},
Duration::from_millis(200),
)
.collect(),
);
assert_eq!(got.len(), 5, "got {got:?}");
assert!(
got.iter().all(|item| matches!(item, JobResult::Done(_))),
"a job that kept producing was cut off: {got:?}"
);
}
#[test]
fn work_that_panics_ends_the_stream_as_a_failure() {
let got: Vec<_> = futures::executor::block_on(
stream(
|emit| {
emit(1);
panic!("boom");
},
Duration::from_secs(5),
)
.collect(),
);
assert_eq!(got, vec![JobResult::Done(1), JobResult::Failed]);
}
fn pollster_block<F: std::future::Future>(mut f: F) -> F::Output {
use std::task::{Context, Poll, Waker};
let mut f = unsafe { std::pin::Pin::new_unchecked(&mut f) };
let waker = Waker::noop();
let mut cx = Context::from_waker(waker);
loop {
if let Poll::Ready(v) = f.as_mut().poll(&mut cx) {
return v;
}
std::thread::yield_now();
}
}
}