drop-guard.rs/src/lib.rs

97 lines
2.0 KiB
Rust

pub struct DropGuard<T: DropGuarded> {
inner: Option<T>,
}
impl<T: DropGuarded> DropGuard<T> {
#[inline]
pub fn unguard(mut self) -> T {
self.inner.take().expect("only None in drop")
}
}
pub trait DropGuarded {
fn cancel(self);
}
#[cfg(feature = "tokio1-task")]
impl<T> DropGuarded for tokio1::task::JoinHandle<T> {
#[inline]
fn cancel(self) {
self.abort();
}
}
#[cfg(feature = "tokio1-sync")]
impl DropGuarded for tokio1::sync::oneshot::Sender<()> {
#[inline]
fn cancel(self) {
let _ = self.send(());
}
}
#[cfg(feature = "tokio1-sync")]
impl DropGuarded for std::sync::Arc<tokio1::sync::Semaphore> {
#[inline]
fn cancel(self) {
self.close();
}
}
#[cfg(feature = "triggered")]
impl DropGuarded for triggered::Trigger {
#[inline]
fn cancel(self) {
self.trigger();
}
}
#[repr(transparent)]
pub struct ArbitraryDropGuard<F: FnOnce() -> ()>(Option<F>);
impl<F: FnOnce() -> ()> ArbitraryDropGuard<F> {
#[inline]
pub fn new(f: F) -> Self {
Self(Some(f))
}
}
impl<F: FnOnce() -> ()> DropGuarded for ArbitraryDropGuard<F> {
fn cancel(mut self) {
if let Some(f) = self.0.take() {
f();
}
}
}
impl<T: DropGuarded> DropGuard<T> {
#[inline]
pub fn new(guarded: T) -> Self {
Self {
inner: Some(guarded),
}
}
}
impl<T: DropGuarded> Drop for DropGuard<T> {
#[inline]
fn drop(&mut self) {
if let Some(inner) = self.inner.take() {
inner.cancel();
}
}
}
#[cfg(feature = "tokio1-task")]
impl<T> std::future::Future for DropGuard<tokio1::task::JoinHandle<T>> {
type Output = Result<T, tokio1::task::JoinError>;
fn poll(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
use futures_lite::future::FutureExt;
let handle = (*self).inner.as_mut().expect("can only be None in drop");
handle.poll(cx)
}
}