use std::{
collections::HashMap,
sync::{Arc, atomic::AtomicBool},
};
use tokio_util::sync::CancellationToken;
use super::completion::{OutcomeTx, RemovalCompletion};
use super::scheduler::{ActorHandle, AttemptReaper};
use crate::{core::deferred_drop::DropBundle, identity::TaskId};
pub(super) struct Handle {
join: ActorHandle,
pub(super) cancel: CancellationToken,
pub(super) done: Option<OutcomeTx>,
pub(super) completion: RemovalCompletion,
cleanup: HandleCleanup,
}
impl Handle {
pub(super) fn new(
join: ActorHandle,
cancel: CancellationToken,
done: Option<OutcomeTx>,
completion: RemovalCompletion,
cleanup: HandleCleanup,
) -> Self {
Self {
join,
cancel,
done,
completion,
cleanup,
}
}
pub(super) fn result_ready(&mut self) -> bool {
self.join.result_ready()
}
pub(super) fn join_mut(&mut self) -> &mut ActorHandle {
&mut self.join
}
pub(super) fn abort(&mut self) {
self.join.abort();
}
pub(super) fn into_report_parts(self) -> (Option<OutcomeTx>, DropBundle) {
let Self {
join,
done,
cleanup,
..
} = self;
drop(join);
(done, cleanup.into_bundle())
}
}
pub(super) struct HandleCleanup {
id: TaskId,
reaper: AttemptReaper,
completion: RemovalCompletion,
bundle: Option<DropBundle>,
}
impl HandleCleanup {
pub(super) fn new(
id: TaskId,
reaper: AttemptReaper,
completion: RemovalCompletion,
bundle: DropBundle,
) -> Self {
Self {
id,
reaper,
completion,
bundle: Some(bundle),
}
}
pub(super) fn into_bundle(mut self) -> DropBundle {
self.bundle
.take()
.expect("registry handle owns one terminal cleanup bundle")
}
}
impl Drop for HandleCleanup {
fn drop(&mut self) {
let Some(bundle) = self.bundle.take() else {
return;
};
self.reaper
.attach_terminal(self.id, bundle, None, self.completion.clone());
self.completion.complete_logical();
}
}
pub(super) enum EntryState {
Registered(Box<Handle>),
Removing {
completion: RemovalCompletion,
},
}
pub(super) struct Entry {
pub(super) label: Arc<str>,
pub(super) activity: Arc<AtomicBool>,
pub(super) state: EntryState,
}
#[derive(Default)]
pub(super) struct Inner {
pub(super) tasks: HashMap<TaskId, Entry>,
pub(super) by_label: HashMap<Arc<str>, TaskId>,
}