1use std::{collections::HashMap, sync::Arc, time::Duration};
22
23use futures::{StreamExt, future::BoxFuture, stream::BoxStream};
24use jiff::Timestamp;
25use miette::Result;
26use serde_json::{Value, json};
27use tokio::sync::{Mutex, Semaphore, broadcast, mpsc};
28use tokio_stream::wrappers::BroadcastStream;
29
30use crate::{
31 BackgroundTask, TaskContext, TaskEndpoint, TaskEndpointResponse, tasks::TaskEndpointHandler,
32};
33
34pub type BackupRunner = Arc<
39 dyn Fn(String, mpsc::UnboundedSender<Value>) -> BoxFuture<'static, ()> + Send + Sync + 'static,
40>;
41
42const HEARTBEAT: Duration = Duration::from_secs(5);
43const BROADCAST_CAPACITY: usize = 256;
44const QUIET_PERIOD: Duration = Duration::from_secs(30);
48
49struct RunHandle {
50 started_at: Timestamp,
51 run_id: Mutex<Option<String>>,
52 latest: Mutex<Value>,
54 events: broadcast::Sender<Value>,
55}
56
57#[derive(Clone, serde::Serialize, serde::Deserialize)]
59pub struct RunningBackup {
60 pub r#type: String,
61 pub run_id: Option<String>,
62 pub started_at: String,
63 pub latest: Value,
64}
65
66pub struct BackupRegistry {
68 runner: BackupRunner,
69 running: Mutex<HashMap<String, Arc<RunHandle>>>,
70 configured: Mutex<Vec<String>>,
74 run_slot: Semaphore,
77 quiet_period: Duration,
79}
80
81impl BackupRegistry {
82 pub fn new(runner: BackupRunner) -> Arc<Self> {
83 Self::with_quiet_period(runner, QUIET_PERIOD)
84 }
85
86 fn with_quiet_period(runner: BackupRunner, quiet_period: Duration) -> Arc<Self> {
87 Arc::new(Self {
88 runner,
89 running: Mutex::new(HashMap::new()),
90 configured: Mutex::new(Vec::new()),
91 run_slot: Semaphore::new(1),
92 quiet_period,
93 })
94 }
95
96 pub async fn set_configured(&self, mut types: Vec<String>) {
99 types.sort();
100 *self.configured.lock().await = types;
101 }
102
103 pub async fn configured(&self) -> Vec<String> {
105 self.configured.lock().await.clone()
106 }
107
108 pub async fn ensure_run(self: &Arc<Self>, backup_type: String) -> BoxStream<'static, Value> {
111 let mut running = self.running.lock().await;
112
113 if let Some(handle) = running.get(&backup_type) {
114 let attached = json!({
115 "event": "attached",
116 "runId": *handle.run_id.lock().await,
117 "startedAt": handle.started_at.to_string(),
118 "latest": *handle.latest.lock().await,
119 });
120 return subscription(Some(attached), handle.events.subscribe());
121 }
122
123 let (events, receiver) = broadcast::channel(BROADCAST_CAPACITY);
124 let handle = Arc::new(RunHandle {
125 started_at: Timestamp::now(),
126 run_id: Mutex::new(None),
127 latest: Mutex::new(json!({ "event": "starting" })),
128 events: events.clone(),
129 });
130 running.insert(backup_type.clone(), handle.clone());
131 drop(running);
132
133 let (sink, run_rx) = mpsc::unbounded_channel::<Value>();
134 let runner = (self.runner)(backup_type.clone(), sink);
135 let registry = self.clone();
136 tokio::spawn(async move {
137 let queued = json!({ "event": "queued" });
142 *handle.latest.lock().await = queued.clone();
143 let _ = events.send(queued);
144 let Ok(_permit) = registry.run_slot.acquire().await else {
145 return; };
147
148 tokio::spawn(runner);
149 registry
150 .clone()
151 .pump(backup_type, handle, run_rx, events)
152 .await;
153
154 tokio::time::sleep(registry.quiet_period).await;
157 });
158
159 subscription(None, receiver)
160 }
161
162 async fn pump(
166 self: Arc<Self>,
167 backup_type: String,
168 handle: Arc<RunHandle>,
169 mut run_rx: mpsc::UnboundedReceiver<Value>,
170 events: broadcast::Sender<Value>,
171 ) {
172 let mut heartbeat = tokio::time::interval(HEARTBEAT);
173 heartbeat.tick().await; loop {
175 tokio::select! {
176 message = run_rx.recv() => match message {
177 Some(event) => {
178 if let Some(id) = event.get("runId").and_then(Value::as_str) {
179 *handle.run_id.lock().await = Some(id.to_owned());
180 }
181 *handle.latest.lock().await = event.clone();
182 let _ = events.send(event);
183 }
184 None => break, },
186 _ = heartbeat.tick() => {
187 let _ = events.send(json!({ "event": "heartbeat" }));
188 }
189 }
190 }
191 self.running.lock().await.remove(&backup_type);
192 }
194
195 pub async fn running(&self) -> Vec<RunningBackup> {
197 let map = self.running.lock().await;
198 let mut out = Vec::with_capacity(map.len());
199 for (backup_type, handle) in map.iter() {
200 out.push(RunningBackup {
201 r#type: backup_type.clone(),
202 run_id: handle.run_id.lock().await.clone(),
203 started_at: handle.started_at.to_string(),
204 latest: handle.latest.lock().await.clone(),
205 });
206 }
207 out
208 }
209}
210
211fn subscription(
212 replay: Option<Value>,
213 receiver: broadcast::Receiver<Value>,
214) -> BoxStream<'static, Value> {
215 let live = BroadcastStream::new(receiver).filter_map(|item| async move { item.ok() });
217 match replay {
218 Some(value) => futures::stream::once(async move { value })
219 .chain(live)
220 .boxed(),
221 None => live.boxed(),
222 }
223}
224
225pub struct BackupTask {
229 registry: Arc<BackupRegistry>,
230}
231
232impl BackupTask {
233 pub fn new(registry: Arc<BackupRegistry>) -> Self {
234 Self { registry }
235 }
236}
237
238impl BackgroundTask for BackupTask {
239 fn name(&self) -> &'static str {
240 "backup"
241 }
242
243 fn interval(&self) -> Duration {
244 Duration::from_secs(3600)
247 }
248
249 fn run<'a>(&'a self, _ctx: &'a TaskContext) -> BoxFuture<'a, Result<()>> {
250 Box::pin(async { Ok(()) })
251 }
252
253 fn http_endpoints(&self) -> Vec<TaskEndpoint> {
254 let run_handler: TaskEndpointHandler = {
255 let registry = self.registry.clone();
256 Arc::new(move |ctx| {
257 let registry = registry.clone();
258 Box::pin(async move {
259 let Some(backup_type) = ctx.query.get("type").cloned() else {
260 return TaskEndpointResponse::Error {
261 status: 400,
262 message: "missing ?type= query parameter".into(),
263 };
264 };
265 TaskEndpointResponse::JsonLines(registry.ensure_run(backup_type).await)
266 })
267 })
268 };
269
270 let running_handler: TaskEndpointHandler = {
271 let registry = self.registry.clone();
272 Arc::new(move |_ctx| {
273 let registry = registry.clone();
274 Box::pin(async move {
275 TaskEndpointResponse::Json(json!({ "running": registry.running().await }))
276 })
277 })
278 };
279
280 vec![
281 TaskEndpoint {
282 name: "run",
283 handler: run_handler,
284 },
285 TaskEndpoint {
286 name: "running",
287 handler: running_handler,
288 },
289 ]
290 }
291}
292
293#[cfg(test)]
294mod tests {
295 use tokio::sync::Notify;
296
297 use super::*;
298
299 fn event(value: &Value) -> String {
300 value
301 .get("event")
302 .and_then(Value::as_str)
303 .unwrap_or_default()
304 .to_owned()
305 }
306
307 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
308 async fn start_then_attach_then_finish() {
309 let gate = Arc::new(Notify::new());
312 let runner: BackupRunner = {
313 let gate = gate.clone();
314 Arc::new(move |_type, sink: mpsc::UnboundedSender<Value>| {
315 let gate = gate.clone();
316 Box::pin(async move {
317 let _ = sink.send(json!({ "event": "started", "runId": "r1" }));
318 gate.notified().await;
319 let _ = sink.send(json!({ "event": "done", "success": true }));
320 })
321 })
322 };
323 let registry = BackupRegistry::new(runner);
324
325 let mut starter = registry.ensure_run("pg".into()).await;
326 assert_eq!(event(&starter.next().await.unwrap()), "queued");
329 assert_eq!(event(&starter.next().await.unwrap()), "started");
330
331 let mut attacher = registry.ensure_run("pg".into()).await;
334 assert_eq!(event(&attacher.next().await.unwrap()), "attached");
335 assert_eq!(registry.running().await.len(), 1);
336
337 gate.notify_one();
339 let starter_events: Vec<String> = starter.map(|v| event(&v)).collect().await;
340 assert!(starter_events.contains(&"done".to_owned()));
341 let attacher_events: Vec<String> = attacher.map(|v| event(&v)).collect().await;
342 assert!(attacher_events.contains(&"done".to_owned()));
343 }
344
345 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
346 async fn distinct_types_run_sequentially_not_concurrently() {
347 use std::sync::atomic::{AtomicUsize, Ordering};
348
349 let concurrent = Arc::new(AtomicUsize::new(0));
352 let max_seen = Arc::new(AtomicUsize::new(0));
353 let runner: BackupRunner = {
354 let concurrent = concurrent.clone();
355 let max_seen = max_seen.clone();
356 Arc::new(move |_type, sink: mpsc::UnboundedSender<Value>| {
357 let concurrent = concurrent.clone();
358 let max_seen = max_seen.clone();
359 Box::pin(async move {
360 let now = concurrent.fetch_add(1, Ordering::SeqCst) + 1;
361 max_seen.fetch_max(now, Ordering::SeqCst);
362 let _ = sink.send(json!({ "event": "started" }));
363 tokio::time::sleep(Duration::from_millis(50)).await;
364 concurrent.fetch_sub(1, Ordering::SeqCst);
365 let _ = sink.send(json!({ "event": "done", "success": true }));
366 })
367 })
368 };
369 let registry = BackupRegistry::with_quiet_period(runner, Duration::ZERO);
371
372 let streams =
373 futures::future::join_all(["a", "b", "c"].map(|t| registry.ensure_run(t.to_owned())))
374 .await;
375 for mut stream in streams {
376 while stream.next().await.is_some() {}
377 }
378
379 assert_eq!(
380 max_seen.load(Ordering::SeqCst),
381 1,
382 "backups of distinct types must not overlap"
383 );
384 }
385}