use std::{collections::HashMap, sync::Arc, 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::{
BackgroundTask, TaskContext, TaskEndpoint, TaskEndpointResponse, tasks::TaskEndpointHandler,
};
pub type BackupRunner = Arc<
dyn Fn(String, mpsc::UnboundedSender<Value>) -> 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>,
}
#[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) -> BoxStream<'static, Value> {
let mut running = self.running.lock().await;
if let Some(handle) = running.get(&backup_type) {
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 handle = Arc::new(RunHandle {
started_at: Timestamp::now(),
run_id: Mutex::new(None),
latest: Mutex::new(json!({ "event": "starting" })),
events: events.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);
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 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(),
};
};
TaskEndpointResponse::JsonLines(registry.ensure_run(backup_type).await)
})
})
};
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 {
name: "run",
handler: run_handler,
},
TaskEndpoint {
name: "running",
handler: running_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>| {
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()).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()).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>| {
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())))
.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"
);
}
}