Event for batch canceling tasks

Author: spirali

crates/hyperqueue/src/client/commands/journal/output.rs

@@ -87,10 +87,9 @@ fn format_payload(event: EventPayload) -> serde_json::Value {
             "job": task_id.job_id(),
             "task": task_id.job_task_id(),
         }),
-        EventPayload::TaskCanceled { task_id } => json!({
+        EventPayload::TaskCanceled { task_ids } => json!({
             "type": "task-canceled",
-            "job": task_id.job_id(),
-            "task": task_id.job_task_id(),
+            "tasks": task_ids,
         }),
         EventPayload::TaskFailed { task_id, error } => json!({
             "type": "task-failed",

crates/hyperqueue/src/server/client/mod.rs

@@ -353,12 +353,22 @@ fn reconstruct_historical_events(
                     ));
                 }
                 JobTaskState::Canceled { cancelled_date, .. } => {
-                    events.push(Event::at(
-                        *cancelled_date,
-                        EventPayload::TaskCanceled {
-                            task_id: TaskId::new(job.job_id, *id),
-                        },
-                    ));
+                    if let Some(task_ids) = events.last_mut().and_then(|e| {
+                        if let EventPayload::TaskCanceled { task_ids, .. } = &mut e.payload {
+                            (e.time == *cancelled_date).then_some(task_ids)
+                        } else {
+                            None
+                        }
+                    }) {
+                        task_ids.push(TaskId::new(job.job_id, *id))
+                    } else {
+                        events.push(Event::at(
+                            *cancelled_date,
+                            EventPayload::TaskCanceled {
+                                task_ids: vec![TaskId::new(job.job_id, *id)],
+                            },
+                        ));
+                    }
                 }
                 JobTaskState::Waiting | JobTaskState::Running { .. } => {}
             };

crates/hyperqueue/src/server/event/journal/prune.rs

@@ -10,32 +10,37 @@ pub(crate) fn prune_journal(
     live_worker_ids: &Set<WorkerId>,
 ) -> crate::Result<()> {
     for event in reader {
-        let event = event?;
-        let retain = match &event.payload {
+        let mut event = event?;
+        let event = match &mut event.payload {
             EventPayload::WorkerConnected(worker_id, _)
-            | EventPayload::WorkerLost(worker_id, _) => live_worker_ids.contains(worker_id),
+            | EventPayload::WorkerLost(worker_id, _) => {
+                live_worker_ids.contains(worker_id).then_some(event)
+            }
             EventPayload::WorkerOverviewReceived(overview) => {
-                live_worker_ids.contains(&overview.id)
+                live_worker_ids.contains(&overview.id).then_some(event)
             }
             EventPayload::Submit { job_id, .. }
             | EventPayload::JobCompleted(job_id)
             | EventPayload::JobOpen(job_id, _)
-            | EventPayload::JobClose(job_id) => live_job_ids.contains(job_id),
+            | EventPayload::JobClose(job_id) => live_job_ids.contains(job_id).then_some(event),
             EventPayload::TaskStarted { task_id, .. }
             | EventPayload::TaskFinished { task_id, .. }
-            | EventPayload::TaskFailed { task_id, .. }
-            | EventPayload::TaskCanceled { task_id, .. } => {
-                live_job_ids.contains(&task_id.job_id())
+            | EventPayload::TaskFailed { task_id, .. } => {
+                live_job_ids.contains(&task_id.job_id()).then_some(event)
+            }
+            EventPayload::TaskCanceled { task_ids, .. } => {
+                task_ids.retain(|id| live_job_ids.contains(&id.job_id()));
+                (!task_ids.is_empty()).then_some(event)
             }
             EventPayload::AllocationQueueCreated(_, _)
             | EventPayload::AllocationQueueRemoved(_)
             | EventPayload::AllocationQueued { .. }
             | EventPayload::AllocationStarted(_, _)
             | EventPayload::AllocationFinished(_, _)
             | EventPayload::ServerStart { .. }
-            | EventPayload::ServerStop => true,
+            | EventPayload::ServerStop => Some(event),
         };
-        if retain {
+        if let Some(event) = event {
             writer.store(event)?;
         }
     }

crates/hyperqueue/src/server/event/payload.rs

@@ -48,7 +48,7 @@ pub enum EventPayload {
     },
     /// Task has been canceled
     TaskCanceled {
-        task_id: TaskId,
+        task_ids: Vec<TaskId>,
     },
     /// New allocation queue has been created
     AllocationQueueCreated(QueueId, Box<AllocationQueueParams>),

crates/hyperqueue/src/server/event/streamer.rs

@@ -136,9 +136,9 @@ impl EventStreamer {
         );
     }
 
-    pub fn on_task_canceled(&self, task_id: TaskId, now: DateTime<Utc>) {
+    pub fn on_task_canceled(&self, task_ids: Vec<TaskId>, now: DateTime<Utc>) {
         self.send_event(
-            EventPayload::TaskCanceled { task_id },
+            EventPayload::TaskCanceled { task_ids },
             Some(now),
             ForwardMode::StreamAndPersist,
         );

crates/hyperqueue/src/server/job.rs

@@ -433,10 +433,7 @@ impl Job {
         }
 
         self.counters.n_canceled_tasks += task_ids.len() as JobTaskCount;
-        // TODO: on_task_canceled should take vec
-        for task_id in task_ids {
-            senders.events.on_task_canceled(task_id, now);
-        }
+        senders.events.on_task_canceled(task_ids, now);
         self.check_termination(senders, now);
     }
 

crates/hyperqueue/src/server/restore.rs

@@ -290,36 +290,39 @@ impl StateRestorer {
                         }
                     }
                 }
-                EventPayload::TaskCanceled { task_id } => {
-                    log::debug!("Replaying: TaskCanceled {task_id}");
-                    if let Some(job) = self.jobs.get_mut(&task_id.job_id()) {
-                        let task = job.tasks.get_mut(&task_id.job_task_id());
-                        if let Some(task) = task {
-                            task.state =
-                                match std::mem::replace(&mut task.state, JobTaskState::Waiting) {
-                                    JobTaskState::Running { started_data } => {
-                                        JobTaskState::Canceled {
-                                            started_data: Some(started_data),
-                                            cancelled_date: event.time,
+                EventPayload::TaskCanceled { task_ids } => {
+                    log::debug!("Replaying: TaskCanceled {task_ids:?}");
+                    for task_id in task_ids {
+                        if let Some(job) = self.jobs.get_mut(&task_id.job_id()) {
+                            let task = job.tasks.get_mut(&task_id.job_task_id());
+                            if let Some(task) = task {
+                                task.state =
+                                    match std::mem::replace(&mut task.state, JobTaskState::Waiting)
+                                    {
+                                        JobTaskState::Running { started_data } => {
+                                            JobTaskState::Canceled {
+                                                started_data: Some(started_data),
+                                                cancelled_date: event.time,
+                                            }
                                         }
+                                        _ => JobTaskState::Canceled {
+                                            started_data: None,
+                                            cancelled_date: event.time,
+                                        },
                                     }
-                                    _ => JobTaskState::Canceled {
-                                        started_data: None,
-                                        cancelled_date: event.time,
-                                    },
-                                }
-                        } else {
-                            job.tasks.insert(
-                                task_id.job_task_id(),
-                                RestorerTaskInfo {
-                                    state: JobTaskState::Canceled {
-                                        started_data: None,
-                                        cancelled_date: event.time,
+                            } else {
+                                job.tasks.insert(
+                                    task_id.job_task_id(),
+                                    RestorerTaskInfo {
+                                        state: JobTaskState::Canceled {
+                                            started_data: None,
+                                            cancelled_date: event.time,
+                                        },
+                                        instance_id: None,
+                                        crash_counter: 0,
                                     },
-                                    instance_id: None,
-                                    crash_counter: 0,
-                                },
-                            );
+                                );
+                            }
                         }
                     }
                 }