use crate::plugin::PLUGIN_NAME;
use ahash::RandomState;
use parking_lot::Mutex;
use std::borrow::BorrowMut;
use std::collections::HashMap;
use std::future::Future;
use std::sync::{Arc, LazyLock};
use std::time::Duration;
use tokio::task::{AbortHandle, JoinHandle};
use tokio::time::interval;
pub(crate) static TASK_MANAGER: LazyLock<TaskManager> = LazyLock::new(TaskManager::init);
pub(crate) struct TaskManager {
pub(crate) handles: Arc<Mutex<TaskAbortHandles>>,
}
impl TaskManager {
pub(crate) fn init() -> Self {
let handles = Arc::new(Mutex::new(TaskAbortHandles::default()));
let handles_clone = handles.clone();
tokio::spawn(async move {
let mut interval = interval(Duration::from_secs(20)); loop {
interval.tick().await;
log::debug!("Kovi task thread is cleaning up task handles");
let mut handles_lock = handles_clone.lock();
handles_lock.clear();
}
});
Self { handles }
}
pub(crate) fn disable_plugin(&self, plugin_name: &str) {
let mut task_manager = self.handles.lock();
let map = task_manager.map.borrow_mut();
let vec = match map.get(plugin_name) {
Some(v) => v,
None => return,
};
for abort in vec {
abort.abort();
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct TaskAbortHandles {
map: HashMap<String, Vec<AbortHandle>, RandomState>,
}
impl Default for TaskAbortHandles {
fn default() -> Self {
Self {
map: HashMap::with_hasher(RandomState::new()),
}
}
}
impl TaskAbortHandles {
pub(crate) fn clear(&mut self) {
for vec in self.map.values_mut() {
vec.retain(|abort| !abort.is_finished());
vec.shrink_to_fit();
}
}
}
pub fn spawn<F>(future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
PLUGIN_NAME.with(|name| {
let join = {
let name = name.clone();
tokio::spawn(PLUGIN_NAME.scope(name, future))
};
let about_join = join.abort_handle();
task_manager_handler(name, about_join);
join
})
}
pub(crate) fn task_manager_handler(name: &str, about_join: AbortHandle) {
let mut task_abort_handles = TASK_MANAGER.handles.lock();
let aborts = task_abort_handles.map.entry(name.to_string()).or_default();
aborts.push(about_join);
}