70 lines
2.3 KiB
Rust
70 lines
2.3 KiB
Rust
//! Tracks blocking operations until their system calls return, independent of cancellation.
|
|
use std::sync::{
|
|
Arc,
|
|
atomic::{AtomicUsize, Ordering},
|
|
};
|
|
use tokio::sync::Notify;
|
|
|
|
#[derive(Default)]
|
|
struct State {
|
|
count: AtomicUsize,
|
|
idle: Notify,
|
|
}
|
|
/// A per-process operation tracker. Stop the process before awaiting idle.
|
|
#[derive(Clone, Default)]
|
|
pub struct Process(Arc<State>);
|
|
struct Guard(Arc<State>);
|
|
impl Drop for Guard {
|
|
fn drop(&mut self) {
|
|
if self.0.count.fetch_sub(1, Ordering::SeqCst) == 1 {
|
|
self.0.idle.notify_waiters();
|
|
}
|
|
}
|
|
}
|
|
impl Process {
|
|
/// Run a system call on the blocking pool. Cancellation does not release its guard.
|
|
pub async fn blocking<T: Send + 'static>(
|
|
&self,
|
|
operation: impl FnOnce() -> std::io::Result<T> + Send + 'static,
|
|
) -> std::io::Result<T> {
|
|
self.0.count.fetch_add(1, Ordering::SeqCst);
|
|
let guard = Guard(self.0.clone());
|
|
tokio::task::spawn_blocking(move || {
|
|
let _guard = guard;
|
|
operation()
|
|
})
|
|
.await
|
|
.map_err(std::io::Error::other)?
|
|
}
|
|
/// Complete an asynchronous mutation and its observation even if the guest is cancelled.
|
|
/// Its guard participates in idle(), so no mutation outlives process cleanup.
|
|
pub async fn operation<T: Send + 'static>(
|
|
&self,
|
|
operation: impl std::future::Future<Output = std::io::Result<T>> + Send + 'static,
|
|
) -> std::io::Result<T> {
|
|
self.0.count.fetch_add(1, Ordering::SeqCst);
|
|
let guard = Guard(self.0.clone());
|
|
tokio::spawn(async move {
|
|
let _guard = guard;
|
|
operation.await
|
|
})
|
|
.await
|
|
.map_err(std::io::Error::other)?
|
|
}
|
|
/// Wait for all outstanding system calls and observed mutations.
|
|
/// First abort the guest task. Then await its JoinHandle, or drop its store.
|
|
/// Only then await idle before releasing locks. Abort alone does not stop a running poll.
|
|
/// This ordering also prevents later memory writes, which need no blocking guard.
|
|
pub async fn idle(&self) {
|
|
loop {
|
|
let notified = self.0.idle.notified();
|
|
tokio::pin!(notified);
|
|
notified.as_mut().enable();
|
|
if self.0.count.load(Ordering::SeqCst) == 0 {
|
|
return;
|
|
}
|
|
notified.await;
|
|
}
|
|
}
|
|
}
|