use std::{
collections::HashMap,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
use futures::{StreamExt, future::BoxFuture, stream::BoxStream};
use jiff::Timestamp;
use miette::Result;
use serde_json::{Value, json};
use tokio::sync::{Mutex, Semaphore, broadcast, mpsc};
use tokio_stream::wrappers::BroadcastStream;
use crate::alertd::{
BackgroundTask, TaskContext, TaskEndpoint, TaskEndpointResponse, tasks::TaskEndpointHandler,
};
pub type BackupRunner = Arc<
dyn Fn(String, mpsc::UnboundedSender<Value>, Arc<AtomicBool>) -> BoxFuture<'static, ()>
+ Send
+ Sync
+ 'static,
>;
const HEARTBEAT: Duration = Duration::from_secs(5);
const BROADCAST_CAPACITY: usize = 256;
const QUIET_PERIOD: Duration = Duration::from_secs(30);
struct RunHandle {
started_at: Timestamp,
run_id: Mutex<Option<String>>,
latest: Mutex<Value>,
events: broadcast::Sender<Value>,
hold: Arc<AtomicBool>,
}
#[derive(Clone, serde::Serialize, serde::Deserialize)]
pub struct RunningBackup {
pub r#type: String,
pub run_id: Option<String>,
pub started_at: String,
pub latest: Value,
}
pub struct BackupRegistry {
runner: BackupRunner,
running: Mutex<HashMap<String, Arc<RunHandle>>>,
configured: Mutex<Vec<String>>,
run_slot: Semaphore,
quiet_period: Duration,
}
impl BackupRegistry {
pub fn new(runner: BackupRunner) -> Arc<Self> {
Self::with_quiet_period(runner, QUIET_PERIOD)
}
fn with_quiet_period(runner: BackupRunner, quiet_period: Duration) -> Arc<Self> {
Arc::new(Self {
runner,
running: Mutex::new(HashMap::new()),
configured: Mutex::new(Vec::new()),
run_slot: Semaphore::new(1),
quiet_period,
})
}
pub async fn set_configured(&self, mut types: Vec<String>) {
types.sort();
*self.configured.lock().await = types;
}
pub async fn configured(&self) -> Vec<String> {
self.configured.lock().await.clone()
}
pub async fn ensure_run(
self: &Arc<Self>,
backup_type: String,
hold: bool,
) -> BoxStream<'static, Value> {
let mut running = self.running.lock().await;
if let Some(handle) = running.get(&backup_type) {
if hold {
handle.hold.store(true, Ordering::Relaxed);
}
let attached = json!({
"event": "attached",
"runId": *handle.run_id.lock().await,
"startedAt": handle.started_at.to_string(),
"latest": *handle.latest.lock().await,
});
return subscription(Some(attached), handle.events.subscribe());
}
let (events, receiver) = broadcast::channel(BROADCAST_CAPACITY);
let hold = Arc::new(AtomicBool::new(hold));
let handle = Arc::new(RunHandle {
started_at: Timestamp::now(),
run_id: Mutex::new(None),
latest: Mutex::new(json!({ "event": "starting" })),
events: events.clone(),
hold: hold.clone(),
});
running.insert(backup_type.clone(), handle.clone());
drop(running);
let (sink, run_rx) = mpsc::unbounded_channel::<Value>();
let runner = (self.runner)(backup_type.clone(), sink, hold);
let registry = self.clone();
tokio::spawn(async move {
let queued = json!({ "event": "queued" });
*handle.latest.lock().await = queued.clone();
let _ = events.send(queued);
let Ok(_permit) = registry.run_slot.acquire().await else {
return; };
tokio::spawn(runner);
registry
.clone()
.pump(backup_type, handle, run_rx, events)
.await;
tokio::time::sleep(registry.quiet_period).await;
});
subscription(None, receiver)
}
async fn pump(
self: Arc<Self>,
backup_type: String,
handle: Arc<RunHandle>,
mut run_rx: mpsc::UnboundedReceiver<Value>,
events: broadcast::Sender<Value>,
) {
let mut heartbeat = tokio::time::interval(HEARTBEAT);
heartbeat.tick().await; loop {
tokio::select! {
message = run_rx.recv() => match message {
Some(event) => {
if let Some(id) = event.get("runId").and_then(Value::as_str) {
*handle.run_id.lock().await = Some(id.to_owned());
}
*handle.latest.lock().await = event.clone();
let _ = events.send(event);
}
None => break, },
_ = heartbeat.tick() => {
let _ = events.send(json!({ "event": "heartbeat" }));
}
}
}
self.running.lock().await.remove(&backup_type);
}
pub async fn request_hold(&self, backup_type: &str) -> bool {
match self.running.lock().await.get(backup_type) {
Some(handle) => {
handle.hold.store(true, Ordering::Relaxed);
true
}
None => false,
}
}
pub async fn running(&self) -> Vec<RunningBackup> {
let map = self.running.lock().await;
let mut out = Vec::with_capacity(map.len());
for (backup_type, handle) in map.iter() {
out.push(RunningBackup {
r#type: backup_type.clone(),
run_id: handle.run_id.lock().await.clone(),
started_at: handle.started_at.to_string(),
latest: handle.latest.lock().await.clone(),
});
}
out
}
}
fn subscription(
replay: Option<Value>,
receiver: broadcast::Receiver<Value>,
) -> BoxStream<'static, Value> {
let live = BroadcastStream::new(receiver).filter_map(|item| async move { item.ok() });
match replay {
Some(value) => futures::stream::once(async move { value })
.chain(live)
.boxed(),
None => live.boxed(),
}
}
pub struct BackupTask {
registry: Arc<BackupRegistry>,
}
impl BackupTask {
pub fn new(registry: Arc<BackupRegistry>) -> Self {
Self { registry }
}
}
impl BackgroundTask for BackupTask {
fn name(&self) -> &'static str {
"backup"
}
fn interval(&self) -> Duration {
Duration::from_secs(3600)
}
fn run<'a>(&'a self, _ctx: &'a TaskContext) -> BoxFuture<'a, Result<()>> {
Box::pin(async { Ok(()) })
}
fn http_endpoints(&self) -> Vec<TaskEndpoint> {
let run_handler: TaskEndpointHandler = {
let registry = self.registry.clone();
Arc::new(move |ctx| {
let registry = registry.clone();
Box::pin(async move {
let Some(backup_type) = ctx.query.get("type").cloned() else {
return TaskEndpointResponse::Error {
status: 400,
message: "missing ?type= query parameter".into(),
};
};
let hold = ctx
.query
.get("hold")
.is_some_and(|v| v == "true" || v == "1");
TaskEndpointResponse::JsonLines(registry.ensure_run(backup_type, hold).await)
})
})
};
let hold_handler: TaskEndpointHandler = {
let registry = self.registry.clone();
Arc::new(move |ctx| {
let registry = registry.clone();
Box::pin(async move {
let Some(backup_type) = ctx.query.get("type").cloned() else {
return TaskEndpointResponse::Error {
status: 400,
message: "missing ?type= query parameter".into(),
};
};
let held = registry.request_hold(&backup_type).await;
TaskEndpointResponse::Json(json!({
"type": backup_type,
"held": held,
}))
})
})
};
let running_handler: TaskEndpointHandler = {
let registry = self.registry.clone();
Arc::new(move |_ctx| {
let registry = registry.clone();
Box::pin(async move {
TaskEndpointResponse::Json(json!({ "running": registry.running().await }))
})
})
};
vec![
TaskEndpoint::open("run", run_handler),
TaskEndpoint::open("running", running_handler),
TaskEndpoint::open("hold", hold_handler),
]
}
}
#[cfg(test)]
mod tests {
use tokio::sync::Notify;
use super::*;
fn event(value: &Value) -> String {
value
.get("event")
.and_then(Value::as_str)
.unwrap_or_default()
.to_owned()
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn start_then_attach_then_finish() {
let gate = Arc::new(Notify::new());
let runner: BackupRunner = {
let gate = gate.clone();
Arc::new(move |_type, sink: mpsc::UnboundedSender<Value>, _hold| {
let gate = gate.clone();
Box::pin(async move {
let _ = sink.send(json!({ "event": "started", "runId": "r1" }));
gate.notified().await;
let _ = sink.send(json!({ "event": "done", "success": true }));
})
})
};
let registry = BackupRegistry::new(runner);
let mut starter = registry.ensure_run("pg".into(), false).await;
assert_eq!(event(&starter.next().await.unwrap()), "queued");
assert_eq!(event(&starter.next().await.unwrap()), "started");
let mut attacher = registry.ensure_run("pg".into(), false).await;
assert_eq!(event(&attacher.next().await.unwrap()), "attached");
assert_eq!(registry.running().await.len(), 1);
gate.notify_one();
let starter_events: Vec<String> = starter.map(|v| event(&v)).collect().await;
assert!(starter_events.contains(&"done".to_owned()));
let attacher_events: Vec<String> = attacher.map(|v| event(&v)).collect().await;
assert!(attacher_events.contains(&"done".to_owned()));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn distinct_types_run_sequentially_not_concurrently() {
use std::sync::atomic::{AtomicUsize, Ordering};
let concurrent = Arc::new(AtomicUsize::new(0));
let max_seen = Arc::new(AtomicUsize::new(0));
let runner: BackupRunner = {
let concurrent = concurrent.clone();
let max_seen = max_seen.clone();
Arc::new(move |_type, sink: mpsc::UnboundedSender<Value>, _hold| {
let concurrent = concurrent.clone();
let max_seen = max_seen.clone();
Box::pin(async move {
let now = concurrent.fetch_add(1, Ordering::SeqCst) + 1;
max_seen.fetch_max(now, Ordering::SeqCst);
let _ = sink.send(json!({ "event": "started" }));
tokio::time::sleep(Duration::from_millis(50)).await;
concurrent.fetch_sub(1, Ordering::SeqCst);
let _ = sink.send(json!({ "event": "done", "success": true }));
})
})
};
let registry = BackupRegistry::with_quiet_period(runner, Duration::ZERO);
let streams = futures::future::join_all(
["a", "b", "c"].map(|t| registry.ensure_run(t.to_owned(), false)),
)
.await;
for mut stream in streams {
while stream.next().await.is_some() {}
}
assert_eq!(
max_seen.load(Ordering::SeqCst),
1,
"backups of distinct types must not overlap"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn holding_reaches_a_run_already_in_flight() {
let gate = Arc::new(Notify::new());
let observed = Arc::new(Mutex::new(None::<bool>));
let runner: BackupRunner = {
let gate = gate.clone();
let observed = observed.clone();
Arc::new(move |_type, sink: mpsc::UnboundedSender<Value>, hold| {
let gate = gate.clone();
let observed = observed.clone();
Box::pin(async move {
let _ = sink.send(json!({ "event": "started", "runId": "r1" }));
gate.notified().await;
*observed.lock().await = Some(hold.load(Ordering::Relaxed));
let _ = sink.send(json!({ "event": "done", "success": true }));
})
})
};
let registry = BackupRegistry::with_quiet_period(runner, Duration::ZERO);
let mut stream = registry.ensure_run("pg".into(), false).await;
while let Some(value) = stream.next().await {
if event(&value) == "started" {
break;
}
}
assert!(
registry.request_hold("pg").await,
"the run should be reachable"
);
gate.notify_waiters();
while stream.next().await.is_some() {}
assert_eq!(
*observed.lock().await,
Some(true),
"the run should have seen the hold requested mid-flight"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn holding_an_idle_type_reports_no_run() {
let runner: BackupRunner =
Arc::new(move |_type, _sink: mpsc::UnboundedSender<Value>, _hold| Box::pin(async {}));
let registry = BackupRegistry::with_quiet_period(runner, Duration::ZERO);
assert!(!registry.request_hold("pg").await);
}
}