Fix replaying journal in tasks with multiple submits

Author: spirali

CHANGELOG.md

@@ -13,6 +13,7 @@
 
 * Fixed crash in auto alloc query in the case of multi-node tasks
 * Job and worker IDs are no longer reused after the server is restarted from a pruned journal
+* Fixes invalid counters when journal restores job with multiple submits
 
 
 ## v0.27.0

crates/hyperqueue/src/server/restore.rs

@@ -52,7 +52,7 @@ fn is_task_completed(tasks: &Map<JobTaskId, RestorerTaskInfo>, task_id: JobTaskI
 
 impl RestorerJob {
     pub fn restore_job(
-        mut self,
+        self,
         job_id: JobId,
         state: &mut State,
         server_ref: &ServerRef,
@@ -76,18 +76,22 @@ impl RestorerJob {
                 submit.submitted_at(),
             );
             let job = state.get_job_mut(job_id).unwrap();
+            let submitted_ids: Vec<JobTaskId> =
+                new_tasks.tasks.iter().map(|t| t.id.job_task_id()).collect();
 
             new_tasks.tasks.retain_mut(|t| {
                 t.task_deps
                     .retain(|d| !is_task_completed(&self.tasks, d.job_task_id()));
                 !is_task_completed(&self.tasks, t.id.job_task_id())
             });
 
-            for (task_id, job_task) in job.tasks.iter_mut() {
-                if let Some(task) = self.tasks.get_mut(task_id) {
+            for task_id in submitted_ids {
+                if let (Some(task), Some(job_task)) =
+                    (self.tasks.get(&task_id), job.tasks.get_mut(&task_id))
+                {
                     if task.crash_counter > 0 || task.instance_id.is_some() {
                         new_tasks.adjust_instance_id_and_crash_counters.insert(
-                            TaskId::new(job_id, *task_id),
+                            TaskId::new(job_id, task_id),
                             (
                                 task.instance_id.map(|x| x.as_num() + 1).unwrap_or(0).into(),
                                 task.crash_counter,

tests/test_journal.py

@@ -305,6 +305,35 @@ def test_restore_open_job_more_submits(hq_env: HqEnv, tmp_path):
     wait_for_job_state(hq_env, 1, "OPENED")
 
 
+def test_restore_open_job_more_submits_counters(hq_env: HqEnv, tmp_path):
+    journal_path = os.path.join(tmp_path, "my.journal")
+
+    hq_env.start_server(args=["--journal", journal_path])
+    w = hq_env.start_worker()
+    hq_env.command(["job", "open"])
+    hq_env.command(["job", "submit", "--job=1", "--array=0-1", "--", "uname"])
+    wait_for_task_state(hq_env, 1, [0, 1], ["finished", "finished"])
+    hq_env.command(["worker", "stop", "1"])
+    wait_for_worker_state(hq_env, 1, "STOPPED", check_running_processes=False)
+    hq_env.check_process_exited(w)
+    for _ in range(2):
+        hq_env.command(["job", "submit", "--job=1", "--", "uname"])
+
+    expected = {"aborted": 0, "canceled": 0, "failed": 0, "finished": 2, "running": 0, "waiting": 2}
+    out = hq_env.command(["--output-mode=json", "job", "info", "1"], as_json=True)
+    assert out[0]["info"]["task_stats"] == expected
+    hq_env.stop_server()
+
+    hq_env.start_server(args=["--journal", journal_path])
+    out = hq_env.command(["--output-mode=json", "job", "info", "1"], as_json=True)
+    assert out[0]["info"]["task_stats"] == expected
+
+    hq_env.start_worker()
+    wait_for_task_state(hq_env, 1, [2, 3], ["finished", "finished"])
+    out = hq_env.command(["--output-mode=json", "job", "info", "1"], as_json=True)
+    assert out[0]["info"]["task_stats"]["finished"] == 4
+
+
 def test_restore_streaming(hq_env: HqEnv, tmp_path):
     journal_path = os.path.join(tmp_path, "my.journal")
     stream_path = os.path.join(tmp_path, "stream")