use std::collections::HashMap;
use std::sync::Mutex;
use tokio::sync::watch;
use crate::core::RunId;
#[derive(Debug, Default)]
pub struct AppendSignals {
runs: Mutex<HashMap<(String, RunId), watch::Sender<u64>>>,
}
impl AppendSignals {
pub fn subscribe(&self, tenant: &str, run: RunId) -> AppendWatch {
let mut runs = self
.runs
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let rx = runs
.entry((tenant.to_owned(), run))
.or_insert_with(|| watch::channel(0).0)
.subscribe();
AppendWatch { rx }
}
pub fn raise(&self, tenant: &str, run: RunId, head: u64) {
let mut runs = self
.runs
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let key = (tenant.to_owned(), run);
let Some(tx) = runs.get(&key) else { return };
if tx.receiver_count() == 0 {
runs.remove(&key);
return;
}
tx.send_if_modified(|seen| {
if head > *seen {
*seen = head;
true
} else {
false
}
});
}
}
#[derive(Debug, Clone)]
pub struct AppendWatch {
rx: watch::Receiver<u64>,
}
impl AppendWatch {
pub async fn changed(&mut self) -> Option<u64> {
self.rx.changed().await.ok()?;
Some(*self.rx.borrow_and_update())
}
#[must_use]
pub fn head(&self) -> u64 {
*self.rx.borrow()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_run_nobody_follows_any_more_is_forgotten() {
let signals = AppendSignals::default();
let run = RunId::generate();
drop(signals.subscribe("t", run));
signals.raise("t", run, 1);
assert!(signals.runs.lock().expect("lock").is_empty());
}
}