Reduce likelihood of timeout/signal related race condition

python/apsis/program/procstar/agent.py

@@ -477,7 +477,9 @@ async def timeout_handler():
                         await asyncio.sleep(sleep_duration)
 
                     if not self.stopping and self.proc is not None:
-                        elapsed = ora.now() - start
+                        # Introduce a small delay to allow updates to arrive in case Apsis and Procstar just reconnected.
+                        # If the process already terminated, the timeout task will be cancelled.
+                        await asyncio.sleep(0.01)
                         log.info(f"{self.run_id}: timeout")
                         self.timed_out = True
                         timeout_signal = Signals[self.program.timeout.signal]
@@ -523,6 +525,13 @@ def more_output():
                             yield base.ProgramUpdate(meta=meta)
                         else:
                             # Process terminated.
+                            if res.state == "terminated":
+                                try:
+                                    # cancel timeout task as soon as possible, if any
+                                    await tasks.cancel("timeout")
+                                    log.info(f"{self.run_id}: cancelled timeout task as process already terminated.")
+                                except KeyError:
+                                    pass
                             break
 
             else:

test/int/procstar/jobs/timeout-failure.yaml

@@ -0,0 +1,8 @@
+params: ["timeout", "sleep_duration"]
+
+program:
+  type: procstar-shell
+  command: "/usr/bin/sleep {{ sleep_duration }} && wrong_cmd_to_make_it_fail"
+  timeout:
+    duration: "{{ timeout }}"
+    signal: "SIGTERM"

test/int/procstar/jobs/timeout-handle-sigterm.yaml

@@ -0,0 +1,8 @@
+params: ["timeout", "sleep_duration"]
+
+program:
+  type: procstar-shell
+  command: "trap 'exit 0' SIGTERM; sleep {{ sleep_duration }}"
+  timeout:
+    duration: "{{ timeout }}"
+    signal: "SIGTERM"

test/int/procstar/test_timeout.py

@@ -7,7 +7,9 @@
 JOB_DIR = Path(__file__).parent / "jobs"
 
 
-@pytest.mark.parametrize("job_name", ["timeout", "timeout-shell"])
+@pytest.mark.parametrize(
+    "job_name", ["timeout", "timeout-shell", "timeout-handle-sigterm"]
+)
 def test_timeout(job_name):
     """
     Tests agent program timeout.
@@ -151,3 +153,45 @@ def test_timeout_with_delayed_reconnect(job_name):
         assert (
             abs(actual_elapsed - downtime) <= tolerance
         ), f"Elapsed time {actual_elapsed:.3f}s should be close to downtime {downtime}s (tolerance: {tolerance}s). Run should have been killed shortly after Apsis reconnected."
+
+
+@pytest.mark.parametrize(
+    "job_name,expected_state",
+    [
+        ("timeout", "success"),
+        ("timeout-shell", "success"),
+        ("timeout-failure", "failure"),
+    ],
+)
+def test_timeout_with_delayed_reconnect_process_terminated(job_name, expected_state):
+    """
+    Tests that the timeout signal is not sent if the process already terminated
+    while Apsis was restarting.
+    """
+    with ApsisService(job_dir=JOB_DIR) as svc, svc.agent(serve=True):
+        client = svc.client
+        timeout = 1
+        sleep_duration = 2
+        run_id = client.schedule(
+            job_name, {"timeout": timeout, "sleep_duration": sleep_duration}
+        )["run_id"]
+
+        # Start a run that will terminate while Apsis is restarting
+        res = svc.wait_run(run_id, wait_states=("starting",))
+        assert res["state"] == "running"
+
+        # Stop Apsis for longer than the run duration
+        svc.stop_serve()
+        downtime = 3
+        sleep(downtime)
+        svc.start_serve()
+        svc.wait_for_serve()
+
+        res = svc.wait_run(run_id, wait_states="running", timeout=sleep_duration + 1)
+        assert res["state"] == expected_state
+
+        # ensure no stop signal was sent
+        assert res["meta"]["program"]["stop"]["signals"] == []
+
+        run_log = svc.client.get_run_log(run_id)
+        assert "timeout" not in run_log[-1]["message"]