1use crate::api::history::JobHistory;
28use crate::websocket::{JobUpdate, QueueStats, ServerMessage, WebSocketState};
29use chrono::{DateTime, Utc};
30use hammerwork::events::{EventSubscription, JobLifecycleEvent, JobLifecycleEventType};
31use std::collections::{HashMap, HashSet};
32use std::sync::Arc;
33use std::time::Duration;
34use tokio::sync::RwLock;
35use tokio::task::JoinHandle;
36use tracing::{debug, warn};
37
38pub const POLL_OVERLAP: chrono::Duration = chrono::Duration::seconds(5);
41
42type SentState = (String, i32, DateTime<Utc>);
44
45type QueueCounts = [u64; 5];
48
49pub struct LiveUpdates<Q> {
51 queue: Arc<Q>,
52 websocket: Arc<RwLock<WebSocketState>>,
53 max_jobs: u32,
54 last_poll: Option<DateTime<Utc>>,
57 sent: HashMap<String, SentState>,
59 queues: HashMap<String, QueueCounts>,
61 running: HashSet<String>,
65}
66
67#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
69pub struct PollOutcome {
70 pub jobs: usize,
72 pub queues: usize,
74}
75
76impl<Q: JobHistory + 'static> LiveUpdates<Q> {
77 pub fn new(queue: Arc<Q>, websocket: Arc<RwLock<WebSocketState>>, max_jobs: u32) -> Self {
80 Self {
81 queue,
82 websocket,
83 max_jobs: max_jobs.max(1),
84 last_poll: None,
85 sent: HashMap::new(),
86 queues: HashMap::new(),
87 running: HashSet::new(),
88 }
89 }
90
91 pub fn spawn(mut self, interval: Duration) -> JoinHandle<()> {
94 tokio::spawn(async move {
95 let mut ticker = tokio::time::interval(interval);
96 ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
97 loop {
98 ticker.tick().await;
99 if let Err(e) = self.poll_once().await {
100 warn!(error = %e, "Dashboard live update poll failed");
101 }
102 }
103 })
104 }
105
106 pub async fn poll_once(&mut self) -> hammerwork::Result<PollOutcome> {
109 if self.websocket.read().await.connection_count() == 0 {
110 self.last_poll = None;
111 self.sent.clear();
112 self.queues.clear();
113 self.running.clear();
114 return Ok(PollOutcome::default());
115 }
116
117 let now = self.queue.database_now().await?;
118 let since = self.last_poll.unwrap_or(now) - POLL_OVERLAP;
119 let mut changes = self.queue.recent_job_changes(since, self.max_jobs).await?;
120 let running = self.queue.running_jobs(self.max_jobs).await?;
121 let mut left = self.left_running(&changes, &running);
124 let unchecked = left.split_off(left.len().min(self.max_jobs as usize));
125 let states = if left.is_empty() {
126 Vec::new()
127 } else {
128 self.queue.job_states(&left).await?
129 };
130 let (still_running, stopped) = Self::stopped_running(&left, states, now);
131 self.running = running
132 .iter()
133 .chain(&changes)
134 .filter(|job| job.status == "Running")
135 .map(|job| job.id.clone())
136 .chain(still_running)
137 .chain(unchecked)
138 .collect();
139 changes.extend(stopped);
140 let stats = self.queue.get_all_queue_stats().await?;
141 self.last_poll = Some(now);
142
143 let jobs = self.fresh_changes(changes, since);
144 let queues = self.changed_queues(stats);
145 let outcome = PollOutcome {
146 jobs: jobs.len(),
147 queues: queues.len(),
148 };
149
150 let websocket = self.websocket.read().await;
151 for job in jobs {
152 websocket
153 .broadcast_to_subscribed(ServerMessage::JobUpdate { job }, "job_updates")
154 .await
155 .map_err(|e| hammerwork::HammerworkError::Processing(e.to_string()))?;
156 }
157 for (queue_name, stats) in queues {
158 websocket
159 .broadcast_to_subscribed(
160 ServerMessage::QueueUpdate { queue_name, stats },
161 "queue_updates",
162 )
163 .await
164 .map_err(|e| hammerwork::HammerworkError::Processing(e.to_string()))?;
165 }
166 if outcome != PollOutcome::default() {
167 debug!(?outcome, "Pushed dashboard live updates");
168 }
169 Ok(outcome)
170 }
171
172 fn left_running(&self, changes: &[JobUpdate], running: &[JobUpdate]) -> Vec<String> {
177 let seen: HashSet<&str> = changes
178 .iter()
179 .chain(running)
180 .map(|job| job.id.as_str())
181 .collect();
182 let mut left: Vec<String> = self
183 .running
184 .iter()
185 .filter(|id| !seen.contains(id.as_str()))
186 .cloned()
187 .collect();
188 left.sort_unstable();
189 left
190 }
191
192 fn stopped_running(
198 left: &[String],
199 states: Vec<JobUpdate>,
200 now: DateTime<Utc>,
201 ) -> (Vec<String>, Vec<JobUpdate>) {
202 let left: HashSet<&str> = left.iter().map(String::as_str).collect();
203 let (running, stopped): (Vec<JobUpdate>, Vec<JobUpdate>) = states
204 .into_iter()
205 .filter(|job| left.contains(job.id.as_str()))
206 .partition(|job| job.status == "Running");
207 let stopped = stopped
208 .into_iter()
209 .map(|mut job| {
210 job.updated_at = now;
211 job
212 })
213 .collect();
214 (running.into_iter().map(|job| job.id).collect(), stopped)
215 }
216
217 fn fresh_changes(&mut self, changes: Vec<JobUpdate>, since: DateTime<Utc>) -> Vec<JobUpdate> {
220 self.sent
221 .retain(|_, (_, _, changed_at)| *changed_at > since);
222 let mut fresh: Vec<JobUpdate> = changes
223 .into_iter()
224 .filter(|job| {
225 let state = (job.status.clone(), job.attempts, job.updated_at);
226 self.sent.insert(job.id.clone(), state.clone()) != Some(state)
227 })
228 .collect();
229 fresh.sort_by_key(|job| job.updated_at);
230 fresh
231 }
232
233 fn changed_queues(
235 &mut self,
236 stats: Vec<hammerwork::stats::QueueStats>,
237 ) -> Vec<(String, QueueStats)> {
238 let present: HashSet<&str> = stats.iter().map(|s| s.queue_name.as_str()).collect();
239 self.queues
240 .retain(|name, _| present.contains(name.as_str()));
241 let mut changed = Vec::new();
242 for queue in stats {
243 let counts = [
244 queue.pending_count,
245 queue.running_count,
246 queue.completed_count,
247 queue.dead_count,
248 queue.timed_out_count,
249 ];
250 if self.queues.insert(queue.queue_name.clone(), counts) == Some(counts) {
251 continue;
252 }
253 changed.push((
254 queue.queue_name,
255 QueueStats {
256 pending_count: queue.pending_count,
257 running_count: queue.running_count,
258 completed_count: queue.completed_count,
259 failed_count: queue.statistics.failed,
260 dead_count: queue.dead_count,
261 throughput_per_minute: queue.statistics.throughput_per_minute,
262 avg_processing_time_ms: queue.statistics.avg_processing_time_ms,
263 error_rate: queue.statistics.error_rate,
264 updated_at: Utc::now(),
265 },
266 ));
267 }
268 changed
269 }
270}
271
272pub fn forward_job_events(
281 websocket: Arc<RwLock<WebSocketState>>,
282 mut subscription: EventSubscription,
283) -> JoinHandle<()> {
284 use tokio::sync::broadcast::error::RecvError;
285 tokio::spawn(async move {
286 loop {
287 match subscription.receiver.recv().await {
288 Ok(event) if subscription.filter.matches(&event) => {
289 let message = ServerMessage::JobUpdate {
290 job: job_update_from_event(&event),
291 };
292 let state = websocket.read().await;
293 if let Err(e) = state.broadcast_to_subscribed(message, "job_updates").await {
294 warn!(error = %e, "Failed to forward a job event to dashboard clients");
295 }
296 }
297 Ok(_) => {}
298 Err(RecvError::Lagged(missed)) => {
299 debug!(missed, "Dashboard event forwarding fell behind");
300 }
301 Err(RecvError::Closed) => break,
302 }
303 }
304 })
305}
306
307pub fn job_update_from_event(event: &JobLifecycleEvent) -> JobUpdate {
309 use JobLifecycleEventType as E;
310 let status = match event.event_type {
311 E::Enqueued | E::Retried | E::Restored => "Pending",
312 E::Started => "Running",
313 E::Completed => "Completed",
314 E::Failed => "Failed",
315 E::Dead => "Dead",
316 E::TimedOut => "TimedOut",
317 E::Cancelled => "Cancelled",
318 E::Archived => "Archived",
319 };
320 let attempts = event
321 .error
322 .as_ref()
323 .and_then(|error| error.retry_attempt)
324 .and_then(|attempt| i32::try_from(attempt).ok())
325 .unwrap_or_default();
326 JobUpdate {
327 id: event.job_id.to_string(),
328 queue_name: event.queue_name.clone(),
329 status: status.to_string(),
330 priority: event.priority.to_string(),
331 attempts,
332 updated_at: event.timestamp,
333 }
334}
335
336#[cfg(test)]
337mod tests {
338 use super::*;
339 use crate::config::WebSocketConfig;
340 use hammerwork::events::{EventFilter, EventManager, JobLifecycleEventBuilder};
341 use hammerwork::priority::JobPriority;
342 use uuid::Uuid;
343 use warp::Filter;
344
345 type Shared = Arc<RwLock<WebSocketState>>;
346
347 fn new_state() -> Shared {
348 Arc::new(RwLock::new(WebSocketState::new(WebSocketConfig::default())))
349 }
350
351 async fn connect(state: &Shared) -> warp::test::WsClient {
352 let state = state.clone();
353 let route = warp::path("ws")
354 .and(warp::ws())
355 .map(move |ws: warp::ws::Ws| {
356 let state = state.clone();
357 ws.on_upgrade(move |socket| async move {
358 let _ = WebSocketState::serve_connection(state, socket).await;
359 })
360 });
361 warp::test::ws()
362 .path("/ws")
363 .handshake(route)
364 .await
365 .expect("handshake")
366 }
367
368 async fn wait_for_connections(state: &Shared, expected: usize) {
369 for _ in 0..200 {
370 if state.read().await.connection_count() == expected {
371 return;
372 }
373 tokio::time::sleep(Duration::from_millis(10)).await;
374 }
375 panic!("expected {expected} connections");
376 }
377
378 async fn next_json(client: &mut warp::test::WsClient) -> Option<serde_json::Value> {
380 match tokio::time::timeout(Duration::from_millis(500), client.recv()).await {
381 Ok(Ok(message)) if message.is_text() => {
382 Some(serde_json::from_str(message.to_str().unwrap()).unwrap())
383 }
384 _ => None,
385 }
386 }
387
388 async fn subscribe(client: &mut warp::test::WsClient, event_types: &[&str]) {
390 let message = serde_json::json!({"type": "Subscribe", "event_types": event_types});
391 client.send_text(message.to_string()).await;
392 tokio::time::sleep(Duration::from_millis(100)).await;
393 }
394
395 #[test]
396 fn job_events_map_to_dashboard_statuses() {
397 let id = Uuid::new_v4();
398 let event = JobLifecycleEvent::failed(
399 id,
400 "emails".to_string(),
401 JobPriority::High,
402 hammerwork::events::JobError {
403 message: "boom".to_string(),
404 error_type: None,
405 details: None,
406 retry_attempt: Some(2),
407 },
408 );
409 let update = job_update_from_event(&event);
410 assert_eq!(update.id, id.to_string());
411 assert_eq!(update.queue_name, "emails");
412 assert_eq!(update.status, "Failed");
413 assert_eq!(update.priority, "high");
414 assert_eq!(update.attempts, 2);
415 assert_eq!(update.updated_at, event.timestamp);
416
417 let completed =
418 JobLifecycleEvent::completed(id, "emails".to_string(), JobPriority::Normal, 12);
419 assert_eq!(job_update_from_event(&completed).status, "Completed");
420 let started = JobLifecycleEvent::started(id, "emails".to_string(), JobPriority::Normal);
421 assert_eq!(job_update_from_event(&started).status, "Running");
422 }
423
424 #[tokio::test]
427 async fn in_process_job_events_reach_subscribed_clients() {
428 let state = new_state();
429 let events = EventManager::new_default();
430 let subscription = events.subscribe(EventFilter::new()).await.unwrap();
431 let forwarder = forward_job_events(state.clone(), subscription);
432
433 let mut jobs = connect(&state).await;
434 let mut archive_only = connect(&state).await;
435 wait_for_connections(&state, 2).await;
436 subscribe(&mut jobs, &["job_updates"]).await;
437 subscribe(&mut archive_only, &["archive_events"]).await;
438
439 let id = Uuid::new_v4();
440 events
441 .publish_event(JobLifecycleEvent::completed(
442 id,
443 "emails".to_string(),
444 JobPriority::Normal,
445 5,
446 ))
447 .await
448 .unwrap();
449
450 let message = next_json(&mut jobs).await.expect("job update delivered");
451 assert_eq!(message["type"], "JobUpdate");
452 assert_eq!(message["job"]["id"], id.to_string());
453 assert_eq!(message["job"]["status"], "Completed");
454 assert!(next_json(&mut archive_only).await.is_none());
455
456 drop(events);
458 tokio::time::timeout(Duration::from_secs(2), forwarder)
459 .await
460 .expect("forwarder stopped")
461 .unwrap();
462 }
463
464 fn change(id: &str, status: &str, at: DateTime<Utc>) -> JobUpdate {
465 JobUpdate {
466 id: id.to_string(),
467 queue_name: "q".to_string(),
468 status: status.to_string(),
469 priority: "normal".to_string(),
470 attempts: 0,
471 updated_at: at,
472 }
473 }
474
475 #[tokio::test]
479 async fn overlapping_polls_push_each_change_once() {
480 let mut live = LiveUpdates::new(
481 crate::api::test_support::unreachable_queue(),
482 new_state(),
483 10,
484 );
485 let t0 = Utc::now();
486 let first = live.fresh_changes(
487 vec![
488 change("b", "Running", t0),
489 change("a", "Pending", t0 - POLL_OVERLAP / 2),
490 ],
491 t0 - POLL_OVERLAP,
492 );
493 let ids: Vec<&str> = first.iter().map(|j| j.id.as_str()).collect();
494 assert_eq!(ids, ["a", "b"], "oldest first");
495
496 let t1 = t0 + chrono::Duration::seconds(1);
497 let second = live.fresh_changes(
498 vec![
499 change("a", "Running", t1),
500 change("b", "Running", t0),
501 change("c", "Pending", t1),
502 ],
503 t1 - POLL_OVERLAP,
504 );
505 let pushed: Vec<(&str, &str)> = second
506 .iter()
507 .map(|j| (j.id.as_str(), j.status.as_str()))
508 .collect();
509 assert_eq!(pushed, [("a", "Running"), ("c", "Pending")]);
510
511 live.fresh_changes(Vec::new(), t1 + chrono::Duration::hours(1));
513 assert!(live.sent.is_empty());
514 }
515
516 #[tokio::test]
517 async fn queue_statistics_are_pushed_when_they_change() {
518 let mut live = LiveUpdates::new(
519 crate::api::test_support::unreachable_queue(),
520 new_state(),
521 10,
522 );
523 let stats = |name: &str, pending: u64| hammerwork::stats::QueueStats {
524 queue_name: name.to_string(),
525 pending_count: pending,
526 running_count: 0,
527 dead_count: 0,
528 timed_out_count: 0,
529 completed_count: 0,
530 statistics: Default::default(),
531 };
532 let first = live.changed_queues(vec![stats("a", 1), stats("b", 0)]);
533 assert_eq!(first.len(), 2);
534 let second = live.changed_queues(vec![stats("a", 1), stats("b", 3)]);
535 assert_eq!(second.len(), 1);
536 assert_eq!(second[0].0, "b");
537 assert_eq!(second[0].1.pending_count, 3);
538 live.changed_queues(vec![stats("a", 1)]);
540 assert_eq!(live.queues.len(), 1);
541 }
542
543 #[tokio::test]
546 async fn without_clients_nothing_is_polled() {
547 let mut live = LiveUpdates::new(
548 crate::api::test_support::unreachable_queue(),
549 new_state(),
550 10,
551 );
552 assert_eq!(live.poll_once().await.unwrap(), PollOutcome::default());
553 }
554
555 async fn clients_receive_job_state_changes<Q: JobHistory + 'static>(queue: Arc<Q>) {
558 use hammerwork::Job;
559 let state = new_state();
560 let mut live = LiveUpdates::new(queue.clone(), state.clone(), 100);
561 let mut jobs = connect(&state).await;
562 let mut alerts_only = connect(&state).await;
563 wait_for_connections(&state, 2).await;
564 subscribe(&mut jobs, &["job_updates"]).await;
565 subscribe(&mut alerts_only, &["system_alerts"]).await;
566 live.poll_once().await.unwrap();
567 while next_json(&mut jobs).await.is_some() {}
569
570 let queue_name = format!("live_{}", Uuid::new_v4().simple());
571 let id = queue
572 .enqueue(Job::new(queue_name.clone(), serde_json::json!({})))
573 .await
574 .unwrap();
575 let pending = wait_for(&mut live, &mut jobs, id, "Pending").await;
576 assert_eq!(pending["job"]["queue_name"], queue_name.as_str());
577
578 let run = queue.dequeue(&queue_name).await.unwrap().expect("job");
579 queue
580 .finish_job_run(&run, hammerwork::queue::JobOutcome::Completed)
581 .await
582 .unwrap();
583 wait_for(&mut live, &mut jobs, id, "Completed").await;
584 assert!(
585 next_json(&mut alerts_only).await.is_none(),
586 "not subscribed to job updates"
587 );
588
589 queue.delete_job(id).await.unwrap();
590 }
591
592 async fn clients_see_running_jobs_go_back_to_pending<Q: JobHistory + 'static>(queue: Arc<Q>) {
596 use hammerwork::Job;
597 let state = new_state();
598 let mut live = LiveUpdates::new(queue.clone(), state.clone(), 100);
599 let mut jobs = connect(&state).await;
600 wait_for_connections(&state, 1).await;
601 subscribe(&mut jobs, &["job_updates"]).await;
602 live.poll_once().await.unwrap();
603 while next_json(&mut jobs).await.is_some() {}
604
605 let queue_name = format!("live_back_{}", Uuid::new_v4().simple());
606 let mut job = Job::new(queue_name.clone(), serde_json::json!({}));
607 job.created_at = Utc::now() - chrono::Duration::hours(1);
608 job.scheduled_at = job.created_at;
609 let id = queue.enqueue(job).await.unwrap();
610 let cleanup = CleanupJob {
611 queue: queue.clone(),
612 id,
613 };
614
615 let run = queue.dequeue(&queue_name).await.unwrap().expect("job");
616 wait_for(&mut live, &mut jobs, id, "Running").await;
617 assert!(queue.release_job_run(&run).await.unwrap());
618 let released = wait_for(&mut live, &mut jobs, id, "Pending").await;
619 assert_eq!(released["job"]["attempts"], 0);
620
621 let run = queue.dequeue(&queue_name).await.unwrap().expect("job");
623 wait_for(&mut live, &mut jobs, id, "Running").await;
624 queue
625 .finish_job_run(
626 &run,
627 hammerwork::queue::JobOutcome::Retry {
628 retry_at: Utc::now() + chrono::Duration::hours(1),
629 error: "again".to_string(),
630 timed_out: false,
631 },
632 )
633 .await
634 .unwrap();
635 let retried = wait_for(&mut live, &mut jobs, id, "Pending").await;
636 assert_eq!(retried["job"]["attempts"], 1);
637 drop(cleanup);
638 }
639
640 struct CleanupJob<Q: JobHistory + 'static> {
642 queue: Arc<Q>,
643 id: Uuid,
644 }
645
646 impl<Q: JobHistory + 'static> Drop for CleanupJob<Q> {
647 fn drop(&mut self) {
648 let (queue, id) = (self.queue.clone(), self.id);
649 let _ = tokio::task::block_in_place(|| {
651 tokio::runtime::Handle::current().block_on(queue.delete_job(id))
652 });
653 }
654 }
655
656 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
657 #[ignore = "requires DATABASE_URL (PostgreSQL)"]
658 async fn postgres_clients_see_running_jobs_go_back_to_pending() {
659 let url = std::env::var("DATABASE_URL").expect("DATABASE_URL");
660 let pool = sqlx::PgPool::connect(&url).await.unwrap();
661 clients_see_running_jobs_go_back_to_pending(Arc::new(hammerwork::JobQueue::new(pool)))
662 .await;
663 }
664
665 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
666 #[ignore = "requires MYSQL_DATABASE_URL"]
667 async fn mysql_clients_see_running_jobs_go_back_to_pending() {
668 let url = std::env::var("MYSQL_DATABASE_URL").expect("MYSQL_DATABASE_URL");
669 let pool = sqlx::MySqlPool::connect(&url).await.unwrap();
670 clients_see_running_jobs_go_back_to_pending(Arc::new(hammerwork::JobQueue::new(pool)))
671 .await;
672 }
673
674 #[tokio::test]
678 async fn jobs_that_stop_running_without_a_timestamp_are_reported() {
679 let mut live = LiveUpdates::new(
680 crate::api::test_support::unreachable_queue(),
681 new_state(),
682 10,
683 );
684 let t0 = Utc::now();
685 live.running = ["a", "b", "c", "d"].map(String::from).into_iter().collect();
686 let left = live.left_running(
687 &[change("d", "Completed", t0)],
688 &[change("c", "Running", t0)],
689 );
690 assert_eq!(left, ["a", "b"]);
691
692 let now = t0 + chrono::Duration::seconds(5);
693 let (still_running, stopped) =
694 LiveUpdates::<hammerwork::JobQueue<sqlx::Postgres>>::stopped_running(
695 &left,
696 vec![
697 change("a", "Pending", t0 - chrono::Duration::hours(1)),
698 change("b", "Running", t0),
699 change("x", "Pending", t0),
700 ],
701 now,
702 );
703 assert_eq!(still_running, ["b"]);
704 assert_eq!(stopped.len(), 1);
705 assert_eq!(stopped[0].id, "a");
706 assert_eq!(stopped[0].status, "Pending");
707 assert_eq!(
708 stopped[0].updated_at, now,
709 "reported when the change was seen"
710 );
711 }
712
713 async fn wait_for<Q: JobHistory + 'static>(
716 live: &mut LiveUpdates<Q>,
717 client: &mut warp::test::WsClient,
718 id: Uuid,
719 status: &str,
720 ) -> serde_json::Value {
721 for _ in 0..20 {
722 live.poll_once().await.unwrap();
723 while let Some(message) = next_json(client).await {
724 assert_eq!(message["type"], "JobUpdate", "{message}");
725 if message["job"]["id"] == id.to_string().as_str()
726 && message["job"]["status"] == status
727 {
728 return message;
729 }
730 }
731 }
732 panic!("no {status} update for job {id}");
733 }
734
735 #[tokio::test]
736 #[ignore = "requires DATABASE_URL (PostgreSQL)"]
737 async fn postgres_clients_receive_job_state_changes() {
738 let url = std::env::var("DATABASE_URL").expect("DATABASE_URL");
739 let pool = sqlx::PgPool::connect(&url).await.unwrap();
740 clients_receive_job_state_changes(Arc::new(hammerwork::JobQueue::new(pool))).await;
741 }
742
743 #[tokio::test]
744 #[ignore = "requires MYSQL_DATABASE_URL"]
745 async fn mysql_clients_receive_job_state_changes() {
746 let url = std::env::var("MYSQL_DATABASE_URL").expect("MYSQL_DATABASE_URL");
747 let pool = sqlx::MySqlPool::connect(&url).await.unwrap();
748 clients_receive_job_state_changes(Arc::new(hammerwork::JobQueue::new(pool))).await;
749 }
750}