liz/veilid
1
0
Fork 0
Code Pull Requests Packages Releases Activity

refactor for cooperative cancellation

Browse Source
This commit is contained in:
John Smith
2022-06-12 20:58:02 -04:00
parent bcc1bfc1a3
commit 180628beef
19 changed files with 810 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
veilid-core/src/xx/mod.rs
View File

@@ -8,6 +8,8 @@ mod eventual_value_clone;
mod ip_addr_port;
mod ip_extra;
mod log_thru;
mod must_join_handle;
mod must_join_single_future;
mod mutable_future;
mod single_future;
mod single_shot_eventual;
@@ -25,6 +27,7 @@ pub use owo_colors::OwoColorize;
pub use parking_lot::*;
pub use split_url::*;
pub use static_assertions::*;
pub use stop_token::*;
pub use tracing::*;
pub type PinBox<T> = Pin<Box<T>>;
@@ -105,6 +108,8 @@ pub use eventual_value::*;
pub use eventual_value_clone::*;
pub use ip_addr_port::*;
pub use ip_extra::*;
pub use must_join_handle::*;
pub use must_join_single_future::*;
pub use mutable_future::*;
pub use single_future::*;
pub use single_shot_eventual::*;

43
veilid-core/src/xx/must_join_handle.rs Normal file
View File

@@ -0,0 +1,43 @@
use async_executors::JoinHandle;
use core::future::Future;
use core::pin::Pin;
use core::sync::atomic::{AtomicBool, Ordering};
use core::task::{Context, Poll};
#[derive(Debug)]
pub struct MustJoinHandle<T> {
join_handle: JoinHandle<T>,
completed: AtomicBool,
}
impl<T> MustJoinHandle<T> {
pub fn new(join_handle: JoinHandle<T>) -> Self {
Self {
join_handle,
completed: AtomicBool::new(false),
}
}
}
impl<T> Drop for MustJoinHandle<T> {
fn drop(&mut self) {
// panic if we haven't completed
if !self.completed.load(Ordering::Relaxed) {
panic!("MustJoinHandle was not completed upon drop. Add cooperative cancellation where appropriate to ensure this is completed before drop.")
}
}
}
impl<T: 'static> Future for MustJoinHandle<T> {
type Output = T;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match Pin::new(&mut self.join_handle).poll(cx) {
Poll::Ready(t) => {
self.completed.store(true, Ordering::Relaxed);
Poll::Ready(t)
}
Poll::Pending => Poll::Pending,
}
}
}

213
veilid-core/src/xx/must_join_single_future.rs Normal file
View File

@@ -0,0 +1,213 @@
use super::*;
use crate::intf::*;
use cfg_if::*;
use core::task::Poll;
use futures_util::poll;
#[derive(Debug)]
struct MustJoinSingleFutureInner<T>
where
T: 'static,
{
locked: bool,
join_handle: Option<MustJoinHandle<T>>,
}
/// Spawns a single background processing task idempotently, possibly returning the return value of the previously executed background task
/// This does not queue, just ensures that no more than a single copy of the task is running at a time, but allowing tasks to be retriggered
#[derive(Debug, Clone)]
pub struct MustJoinSingleFuture<T>
where
T: 'static,
{
inner: Arc<Mutex<MustJoinSingleFutureInner<T>>>,
}
impl<T> Default for MustJoinSingleFuture<T>
where
T: 'static,
{
fn default() -> Self {
Self::new()
}
}
impl<T> MustJoinSingleFuture<T>
where
T: 'static,
{
pub fn new() -> Self {
Self {
inner: Arc::new(Mutex::new(MustJoinSingleFutureInner {
locked: false,
join_handle: None,
})),
}
}
fn try_lock(&self) -> Result<Option<MustJoinHandle<T>>, ()> {
let mut inner = self.inner.lock();
if inner.locked {
// If already locked error out
return Err(());
}
inner.locked = true;
// If we got the lock, return what we have for a join handle if anything
Ok(inner.join_handle.take())
}
fn unlock(&self, jh: Option<MustJoinHandle<T>>) {
let mut inner = self.inner.lock();
assert!(inner.locked);
assert!(inner.join_handle.is_none());
inner.locked = false;
inner.join_handle = jh;
}
// Check the result
pub async fn check(&self) -> Result<Option<T>, ()> {
let mut out: Option<T> = None;
// See if we have a result we can return
let maybe_jh = match self.try_lock() {
Ok(v) => v,
Err(_) => {
// If we are already polling somewhere else, don't hand back a result
return Err(());
}
};
if maybe_jh.is_some() {
let mut jh = maybe_jh.unwrap();
// See if we finished, if so, return the value of the last execution
if let Poll::Ready(r) = poll!(&mut jh) {
out = Some(r);
// Task finished, unlock with nothing
self.unlock(None);
} else {
// Still running put the join handle back so we can check on it later
self.unlock(Some(jh));
}
} else {
// No task, unlock with nothing
self.unlock(None);
}
// Return the prior result if we have one
Ok(out)
}
// Wait for the result
pub async fn join(&self) -> Result<Option<T>, ()> {
let mut out: Option<T> = None;
// See if we have a result we can return
let maybe_jh = match self.try_lock() {
Ok(v) => v,
Err(_) => {
// If we are already polling somewhere else,
// that's an error because you can only join
// these things once
return Err(());
}
};
if maybe_jh.is_some() {
let jh = maybe_jh.unwrap();
// Wait for return value of the last execution
out = Some(jh.await);
// Task finished, unlock with nothing
} else {
// No task, unlock with nothing
}
self.unlock(None);
// Return the prior result if we have one
Ok(out)
}
// Possibly spawn the future possibly returning the value of the last execution
cfg_if! {
if #[cfg(target_arch = "wasm32")] {
pub async fn single_spawn(
&self,
future: impl Future<Output = T> + 'static,
) -> Result<Option<T>, ()> {
let mut out: Option<T> = None;
// See if we have a result we can return
let maybe_jh = match self.try_lock() {
Ok(v) => v,
Err(_) => {
// If we are already polling somewhere else, don't hand back a result
return Err(());
}
};
let mut run = true;
if maybe_jh.is_some() {
let mut jh = maybe_jh.unwrap();
// See if we finished, if so, return the value of the last execution
if let Poll::Ready(r) = poll!(&mut jh) {
out = Some(r);
// Task finished, unlock with a new task
} else {
// Still running, don't run again, unlock with the current join handle
run = false;
self.unlock(Some(jh));
}
}
// Run if we should do that
if run {
self.unlock(Some(MustJoinHandle::new(spawn_local(future))));
}
// Return the prior result if we have one
Ok(out)
}
}
}
}
cfg_if! {
if #[cfg(not(target_arch = "wasm32"))] {
impl<T> MustJoinSingleFuture<T>
where
T: 'static + Send,
{
pub async fn single_spawn(
&self,
future: impl Future<Output = T> + Send + 'static,
) -> Result<Option<T>, ()> {
let mut out: Option<T> = None;
// See if we have a result we can return
let maybe_jh = match self.try_lock() {
Ok(v) => v,
Err(_) => {
// If we are already polling somewhere else, don't hand back a result
return Err(());
}
};
let mut run = true;
if maybe_jh.is_some() {
let mut jh = maybe_jh.unwrap();
// See if we finished, if so, return the value of the last execution
if let Poll::Ready(r) = poll!(&mut jh) {
out = Some(r);
// Task finished, unlock with a new task
} else {
// Still running, don't run again, unlock with the current join handle
run = false;
self.unlock(Some(jh));
}
}
// Run if we should do that
if run {
self.unlock(Some(MustJoinHandle::new(spawn(future))));
}
// Return the prior result if we have one
Ok(out)
}
}
}
}

61
veilid-core/src/xx/tick_task.rs
View File

@@ -6,10 +6,10 @@ use once_cell::sync::OnceCell;
cfg_if! {
if #[cfg(target_arch = "wasm32")] {
type TickTaskRoutine =
dyn Fn(u64, u64) -> PinBoxFuture<Result<(), String>> + 'static;
dyn Fn(StopToken, u64, u64) -> PinBoxFuture<Result<(), String>> + 'static;
} else {
type TickTaskRoutine =
dyn Fn(u64, u64) -> SendPinBoxFuture<Result<(), String>> + Send + Sync + 'static;
dyn Fn(StopToken, u64, u64) -> SendPinBoxFuture<Result<(), String>> + Send + Sync + 'static;
}
}
@@ -20,7 +20,8 @@ pub struct TickTask {
last_timestamp_us: AtomicU64,
tick_period_us: u64,
routine: OnceCell<Box<TickTaskRoutine>>,
single_future: SingleFuture<Result<(), String>>,
stop_source: AsyncMutex<Option<StopSource>>,
single_future: MustJoinSingleFuture<Result<(), String>>,
}
impl TickTask {
@@ -29,7 +30,8 @@ impl TickTask {
last_timestamp_us: AtomicU64::new(0),
tick_period_us,
routine: OnceCell::new(),
single_future: SingleFuture::new(),
stop_source: AsyncMutex::new(None),
single_future: MustJoinSingleFuture::new(),
}
}
pub fn new_ms(tick_period_ms: u32) -> Self {
@@ -37,7 +39,8 @@ impl TickTask {
last_timestamp_us: AtomicU64::new(0),
tick_period_us: (tick_period_ms as u64) * 1000u64,
routine: OnceCell::new(),
single_future: SingleFuture::new(),
stop_source: AsyncMutex::new(None),
single_future: MustJoinSingleFuture::new(),
}
}
pub fn new(tick_period_sec: u32) -> Self {
@@ -45,7 +48,8 @@ impl TickTask {
last_timestamp_us: AtomicU64::new(0),
tick_period_us: (tick_period_sec as u64) * 1000000u64,
routine: OnceCell::new(),
single_future: SingleFuture::new(),
stop_source: AsyncMutex::new(None),
single_future: MustJoinSingleFuture::new(),
}
}
@@ -53,22 +57,31 @@ impl TickTask {
if #[cfg(target_arch = "wasm32")] {
pub fn set_routine(
&self,
routine: impl Fn(u64, u64) -> PinBoxFuture<Result<(), String>> + 'static,
routine: impl Fn(StopToken, u64, u64) -> PinBoxFuture<Result<(), String>> + 'static,
) {
self.routine.set(Box::new(routine)).map_err(drop).unwrap();
}
} else {
pub fn set_routine(
&self,
routine: impl Fn(u64, u64) -> SendPinBoxFuture<Result<(), String>> + Send + Sync + 'static,
routine: impl Fn(StopToken, u64, u64) -> SendPinBoxFuture<Result<(), String>> + Send + Sync + 'static,
) {
self.routine.set(Box::new(routine)).map_err(drop).unwrap();
}
}
}
pub async fn cancel(&self) -> Result<(), String> {
match self.single_future.cancel().await {
pub async fn stop(&self) -> Result<(), String> {
// drop the stop source if we have one
let opt_stop_source = &mut *self.stop_source.lock().await;
if opt_stop_source.is_none() {
// already stopped, just return
return Ok(());
}
*opt_stop_source = None;
// wait for completion of the tick task
match self.single_future.join().await {
Ok(Some(Err(err))) => Err(err),
_ => Ok(()),
}
@@ -80,27 +93,35 @@ impl TickTask {
if last_timestamp_us == 0u64 || (now - last_timestamp_us) >= self.tick_period_us {
// Run the singlefuture
let opt_stop_source = &mut *self.stop_source.lock().await;
let stop_source = StopSource::new();
match self
.single_future
.single_spawn(self.routine.get().unwrap()(last_timestamp_us, now))
.single_spawn(self.routine.get().unwrap()(
stop_source.token(),
last_timestamp_us,
now,
))
.await
{
Ok(Some(Err(err))) => {
// If the last execution errored out then we should pass that error up
// Single future ran this tick
Ok(Some(ret)) => {
// Set new timer
self.last_timestamp_us.store(now, Ordering::Release);
return Err(err);
// Save new stopper
*opt_stop_source = Some(stop_source);
ret
}
// Single future did not run this tick
Ok(None) | Err(()) => {
// If the execution didn't happen this time because it was already running
// then we should try again the next tick and not reset the timestamp so we try as soon as possible
}
_ => {
// Execution happened, next execution attempt should happen only after tick period
self.last_timestamp_us.store(now, Ordering::Release);
Ok(())
}
}
} else {
// It's not time yet
Ok(())
}
Ok(())
}
}