use std::collections::HashMap;
use tokio::sync::{oneshot, RwLock};
use crate::common::task::{ValorMasterTask, ValorTaskId, ValorTaskStatus};
#[derive(Default)]
pub struct ValorResultAggregator {
waiters: RwLock<HashMap<ValorTaskId, Vec<oneshot::Sender<ValorMasterTask>>>>,
}
impl ValorResultAggregator {
pub fn new() -> Self {
Self {
waiters: RwLock::new(HashMap::new()),
}
}
pub async fn register_waiter(
&self,
task_id: &ValorTaskId,
) -> oneshot::Receiver<ValorMasterTask> {
let (tx, rx) = oneshot::channel();
let mut map = self.waiters.write().await;
map.entry(task_id.clone()).or_insert_with(Vec::new).push(tx);
rx
}
pub async fn notify_if_terminal(&self, task: &ValorMasterTask) {
if !is_terminal(task.status) {
return;
}
let mut map = self.waiters.write().await;
if let Some(waiters) = map.remove(&task.task_id) {
for tx in waiters {
let _ = tx.send(task.clone());
}
}
}
}
fn is_terminal(status: ValorTaskStatus) -> bool {
matches!(
status,
ValorTaskStatus::Completed | ValorTaskStatus::Failed | ValorTaskStatus::Cancelled
)
}