A pit full of rusty nails
Something went wrong. Try again.
12 kB · 347 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348mod tree_node;mod util;
use crate::util::MaybeDangling;use core::future::Future;use core::pin::Pin;use core::task::{Context, Poll};use std::sync::Arc;
use event_listener::EventListener;use pin_project_lite::pin_project;
/// A token which can be used to signal a cancellation request to one or more/// tasks.////// Tasks can call [`CancelToken::cancelled()`] in order to/// obtain a Future which will be resolved when cancellation is requested.////// Cancellation can be requested through the [`CancelToken::cancel`] method.pub struct CancelToken { inner: Arc<tree_node::TreeNode>,}
impl std::panic::UnwindSafe for CancelToken {}impl std::panic::RefUnwindSafe for CancelToken {}
pin_project! { /// A Future that is resolved once the corresponding [`CancelToken`] /// is cancelled. #[must_use = "futures do nothing unless polled"] pub struct WaitForCancellationFuture<'a> { cancellation_token: &'a CancelToken, #[pin] future: EventListener<()>, }}
pin_project! { /// A Future that is resolved once the corresponding [`CancelToken`] /// is cancelled. /// /// This is the counterpart to [`WaitForCancellationFuture`] that takes /// [`CancelToken`] by value instead of using a reference. #[must_use = "futures do nothing unless polled"] pub struct WaitForCancellationFutureOwned { // This field internally has a reference to the cancellation token, but camouflages // the relationship with `'static`. To avoid Undefined Behavior, we must ensure // that the reference is only used while the cancellation token is still alive. To // do that, we ensure that the future is the first field, so that it is dropped // before the cancellation token. // // We use `MaybeDanglingFuture` here because without it, the compiler could assert // the reference inside `future` to be valid even after the destructor of that // field runs. (Specifically, when the `WaitForCancellationFutureOwned` is passed // as an argument to a function, the reference can be asserted to be valid for the // rest of that function.) To avoid that, we use `MaybeDangling` which tells the // compiler that the reference stored inside it might not be valid. // // See <https://users.rust-lang.org/t/unsafe-code-review-semi-owning-weak-rwlock-t-guard/95706> // for more info. #[pin] future: MaybeDangling<EventListener<()>>, cancellation_token: CancelToken, }}
impl core::fmt::Debug for CancelToken { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("CancelToken") .field("cancelled", &self.is_cancelled()) .finish() }}
impl Clone for CancelToken { /// Creates a clone of the `CancelToken` which will get cancelled /// whenever the current token gets cancelled, and vice versa. fn clone(&self) -> Self { self.inner.increase_handle_refcount(); CancelToken { inner: self.inner.clone(), } }}
impl Drop for CancelToken { fn drop(&mut self) { self.inner.decrease_handle_refcount(); }}
impl Default for CancelToken { fn default() -> CancelToken { CancelToken::new() }}
impl CancelToken { /// Creates a new `CancelToken` in the non-cancelled state. pub fn new() -> CancelToken { CancelToken { inner: Arc::new(tree_node::TreeNode::new()), } }
/// Creates a `CancelToken` which will get cancelled whenever the /// current token gets cancelled. Unlike a cloned `CancelToken`, /// cancelling a child token does not cancel the parent token. /// /// If the current token is already cancelled, the child token will get /// returned in cancelled state. pub fn child_token(&self) -> CancelToken { CancelToken { inner: self.inner.child_node(), } }
/// Cancel the [`CancelToken`] and all child tokens which had been /// derived from it. /// /// This will wake up all tasks which are waiting for cancellation. /// /// Be aware that cancellation is not an atomic operation. It is possible /// for another thread running in parallel with a call to `cancel` to first /// receive `true` from `is_cancelled` on one child node, and then receive /// `false` from `is_cancelled` on another child node. However, once the /// call to `cancel` returns, all child nodes have been fully cancelled. pub fn cancel(&self) { self.inner.cancel(); }
/// Returns `true` if the `CancelToken` is cancelled. pub fn is_cancelled(&self) -> bool { self.inner.is_cancelled() }
/// Returns a `Future` that gets fulfilled when cancellation is requested. /// /// The future will complete immediately if the token is already cancelled /// when this method is called. /// /// # Cancel safety /// /// This method is cancel safe. pub fn cancelled(&self) -> WaitForCancellationFuture<'_> { WaitForCancellationFuture { cancellation_token: self, future: self.inner.notified(), } }
/// Returns a `Future` that gets fulfilled when cancellation is requested. /// /// The future will complete immediately if the token is already cancelled /// when this method is called. /// /// The function takes self by value and returns a future that owns the /// token. /// /// # Cancel safety /// /// This method is cancel safe. pub fn cancelled_owned(self) -> WaitForCancellationFutureOwned { WaitForCancellationFutureOwned::new(self) }
/// Creates a `DropGuard` for this token. /// /// Returned guard will cancel this token (and all its children) on drop /// unless disarmed. pub fn drop_guard(self) -> DropGuard { DropGuard::new(self) }
/// Runs a future to completion and returns its result wrapped inside of an `Option` /// unless the `CancelToken` is cancelled. In that case the function returns /// `None` and the future gets dropped. /// /// # Cancel safety /// /// This method is only cancel safe if `fut` is cancel safe. pub async fn run_until_cancelled<F>(&self, fut: F) -> Option<F::Output> where F: Future, { pin_project! { /// A Future that is resolved once the corresponding [`CancelToken`] /// is cancelled or a given Future gets resolved. It is biased towards the /// Future completion. #[must_use = "futures do nothing unless polled"] struct RunUntilCancelledFuture<'a, F: Future> { #[pin] cancellation: WaitForCancellationFuture<'a>, #[pin] future: F, } }
impl<'a, F: Future> Future for RunUntilCancelledFuture<'a, F> { type Output = Option<F::Output>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { let this = self.project(); if let Poll::Ready(res) = this.future.poll(cx) { Poll::Ready(Some(res)) } else if this.cancellation.poll(cx).is_ready() { Poll::Ready(None) } else { Poll::Pending } } }
RunUntilCancelledFuture { cancellation: self.cancelled(), future: fut, } .await }
/// Runs a future to completion and returns its result wrapped inside of an `Option` /// unless the `CancelToken` is cancelled. In that case the function returns /// `None` and the future gets dropped. /// /// The function takes self by value and returns a future that owns the token. /// /// # Cancel safety /// /// This method is only cancel safe if `fut` is cancel safe. pub async fn run_until_cancelled_owned<F>(self, fut: F) -> Option<F::Output> where F: Future, { self.run_until_cancelled(fut).await }}
impl<'a> core::fmt::Debug for WaitForCancellationFuture<'a> { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("WaitForCancellationFuture") .finish_non_exhaustive() }}
impl<'a> Future for WaitForCancellationFuture<'a> { type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { let mut this = self.project(); loop { if this.cancellation_token.is_cancelled() { return Poll::Ready(()); }
// No wakeups can be lost here because there is always a call to // `is_cancelled` between the creation of the future and the call to // `poll`, and the code that sets the cancelled flag does so before // waking the `Notified`. if this.future.as_mut().poll(cx).is_pending() { return Poll::Pending; }
this.future.set(this.cancellation_token.inner.notified()); } }}
impl core::fmt::Debug for WaitForCancellationFutureOwned { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("WaitForCancellationFutureOwned") .finish_non_exhaustive() }}
impl WaitForCancellationFutureOwned { fn new(cancellation_token: CancelToken) -> Self { WaitForCancellationFutureOwned { // cancellation_token holds a heap allocation and is guaranteed to have a // stable deref, thus it would be ok to move the cancellation_token while // the future holds a reference to it.
// Safety: Cancellation_token is dropped after future due to the field ordering. future: MaybeDangling::new(unsafe { Self::new_future(&cancellation_token) }), cancellation_token, } }
/// # Safety /// The returned future must be destroyed before the cancellation token is /// destroyed. unsafe fn new_future(cancellation_token: &CancelToken) -> EventListener<()> { let inner_ptr = Arc::as_ptr(&cancellation_token.inner); // SAFETY: The `Arc::as_ptr` method guarantees that `inner_ptr` remains // valid until the strong count of the Arc drops to zero, and the caller // guarantees that they will drop the future before that happens. unsafe { (*inner_ptr).notified() } }}
impl Future for WaitForCancellationFutureOwned { type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> { let mut this = self.project();
loop { if this.cancellation_token.is_cancelled() { return Poll::Ready(()); }
// No wakeups can be lost here because there is always a call to // `is_cancelled` between the creation of the future and the call to // `poll`, and the code that sets the cancelled flag does so before // waking the `Notified`. if this.future.as_mut().poll(cx).is_pending() { return Poll::Pending; }
// SAFETY: Cancellation_token is dropped after future due to the field ordering. this.future.set(MaybeDangling::new(unsafe { Self::new_future(this.cancellation_token) })); } }}
pub struct DropGuard { token: Option<CancelToken>,}
impl DropGuard { const fn new(token: CancelToken) -> Self { Self { token: Some(token) } }
pub fn disarm(&mut self) { self.token.take(); }}
impl Drop for DropGuard { fn drop(&mut self) { if let Some(token) = self.token.take() { token.cancel(); } }}