fix(microvm-worker): keep a suspend_after request for the run hook

packages/aws-durable-execution-sdk-python-microvm-worker/README.md

@@ -131,15 +131,15 @@ A session job can carry `"suspendAfter": true`. The durable function sets it for
 1. Only a JSON `true` turns it on. The worker ignores any other value, and does not reject the job. The `run` hook payload never carries the field.
 2. The worker answers the job with HTTP 202 and the header `X-Microvm-Suspend-After: accepted`. A second delivery of the job that gets the duplicate HTTP 202 carries the header when the first delivery carried the field. A second delivery refused with HTTP 503 carries no header. An older worker ignores the field and sends no header. The durable function then logs a warning, and the MicroVM keeps running.
 3. When the job has ended, and its outcome reports have ended, the worker sets a pending suspend request. This happens whether the handler succeeded or failed. A job that ends while the worker refuses jobs for a suspend sets no request. That refusal comes from the worker's own call, a `suspend` hook, or the 30 seconds that follow either. The MicroVM is suspending already. So the request is done.
-4. The request fires when no job runs, no callback failure report runs, no `run` hook request is being handled, no `SuspendMicrovm` call runs, and the worker is not terminating, not closed, and not refusing jobs for a suspend. The worker checks these conditions again when a job ends, a report ends, a `run` hook request ends, and a call ends. So a request from a job that ends while another job runs fires when the last job ends.
+4. The request fires when no job runs, no callback failure report runs, no `run` hook request is being handled, and no `SuspendMicrovm` call runs. The worker checks these conditions again when a job ends, a report ends, a `run` hook request ends, and a call ends. So a request from a job that ends while another job runs fires when the last job ends.
 5. Firing clears the request before the call starts. A `resume` hook never sets it again. A call after the idle time clears it too, because that call suspends the MicroVM as well. A `suspend` hook clears it too, because Lambda suspends the MicroVM after that hook.
 6. Firing stops the idle time, when the session also suspends when idle.
-7. The `terminate` hook and `worker.close()` drop the request.
-8. The worker calls `SuspendMicrovm` with the MicroVM ID of the first accepted `run` hook with a payload. It never uses the MicroVM ID of a job request. A forged request could otherwise suspend another MicroVM of the same image. Without such a `run` hook, the worker logs the warning `cannot suspend the MicroVM after a job: the run hook named no MicroVM`, with `callbackId`, and does not suspend.
+7. The `terminate` hook and `worker.close()` drop the request, and a job end after either records none.
+8. The worker calls `SuspendMicrovm` with the MicroVM ID of the first accepted `run` hook with a payload. It never uses the MicroVM ID of a job request. A forged request could otherwise suspend another MicroVM of the same image. A request that comes before the first accepted run hook waits for it. The request fires when that run hook names a MicroVM, and is dropped with the warning when it names none. The warning is `cannot suspend the MicroVM after a job: the run hook named no MicroVM`, with `callbackId`. The worker then does not suspend.
 9. The call follows items 3 to 9 of [Suspending when idle](#suspending-when-idle): the refusal of jobs from the moment the request fires, one attempt and none after a `resume` hook, the outcome classification, and the 30-second refusal after the call.
 10. A call that the service refused for now is not retried without `autoSuspendIdleSeconds`. With it, the worker tries again after the idle time, at least 30 seconds later, as in item 7 of [Suspending when idle](#suspending-when-idle). A rejected call is not retried.
 
-Before the call, the worker logs `suspending the MicroVM after a job` at info level, with `microvmId` and `callbackId`. A failed call logs the warning of item 9 of [Suspending when idle](#suspending-when-idle), with `trigger` set to `job`.
+Before the call, the worker logs `suspending the MicroVM after a job` at info level, with `microvmId` and `callbackId`. When a request starts to wait because no `run` hook has been accepted yet, the worker logs the warning `a job asked to suspend the MicroVM before the run hook arrived. The worker waits for the run hook to name the MicroVM.`, with `callbackId`. It logs this warning once per worker, not once per job. An image without the `run` hook never ends the wait, and this warning is then the only sign of it. A failed call logs the warning of item 9 of [Suspending when idle](#suspending-when-idle), with `trigger` set to `job`.
 
 Enable the `resume` hook in the image. Without it, the worker refuses the next job with HTTP 503 until 30 seconds after the call. The next job's heartbeat timeout counts that time. The MicroVM's execution role needs `lambda:SuspendMicrovm` on the image. Without it, each call is rejected and logged, and the MicroVM keeps running.
 

packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py

@@ -424,14 +424,16 @@ class MicrovmWorkerListener:
        pending request. A job that ends while the listener refuses jobs for
        a suspend sets none, because the MicroVM is suspending already. The
        request fires when no job, failure report, ``run`` hook request, or
-       SuspendMicrovm call runs, and the listener is not terminating, not
-       closed, and not refusing jobs. Each of those ends checks it again.
+       SuspendMicrovm call runs. Each of those ends checks it again.
     3. Firing clears the request, stops the idle time, and makes the same
        call as items 3 to 9, with the MicroVM ID of the first accepted
-       ``run`` hook with a payload. Without that hook, the listener logs a
-       warning and does not suspend. A ``resume`` hook never repeats the
-       request. A ``suspend`` hook, the ``terminate`` hook, and ``close()``
-       drop it.
+       ``run`` hook with a payload. A request that comes before the first
+       accepted ``run`` hook waits for it. The request fires when that hook
+       names a MicroVM. When it names none, the listener logs a warning,
+       drops the request, and does not suspend. A ``resume`` hook never
+       repeats the request. A ``suspend`` hook drops it. The ``terminate``
+       hook and ``close()`` drop it, and a job end after either records
+       none.
     4. A call refused for now is retried only when the session suspends
        when idle.
 
@@ -530,9 +532,15 @@ def __init__(
         self._suspend_target: _SuspendTarget | None = None
         # The callback ID of the suspendAfter job whose end asked for a
         # suspend that has not started yet. The request fires when no work
-        # runs. A call that starts clears it, so a resume hook never repeats
-        # it. A suspend hook, the terminate hook, and close() clear it too.
+        # runs, and waits for the first accepted run hook, which sets the
+        # target. A call that starts clears it, so a resume hook never
+        # repeats it. A suspend hook, the terminate hook, and close() clear
+        # it too.
         self._suspend_requested: str | None = None
+        # Set when the worker has logged that a request waits for the run
+        # hook. An image without the run hook would make every such job log
+        # it. So the worker logs it once.
+        self._warned_run_hook_wait = False
         # Each timer action carries the generation it was started with, and
         # acts only when it is still current. Cancelling a timer increments
         # the generation.
@@ -1249,7 +1257,11 @@ def _run_job(
             # now. The suspend runs after this block, so that no lock is
             # held while it logs or starts its thread. During a suspend
             # refusal, the MicroVM is suspending already. So a job that ends
-            # then asks for nothing.
+            # then asks for nothing. A request that comes before the first
+            # accepted run hook waits for it. An image without the run hook
+            # never ends that wait. So the worker warns the first time a
+            # request starts to wait.
+            warn_run_hook_wait = False
             with self._locked():
                 callback_id = job.callback_id
                 suspend_after = callback_id in self._suspend_after_jobs
@@ -1264,8 +1276,18 @@ def _run_job(
                     self._terminating or self._closed or self._suspending
                 ):
                     self._suspend_requested = callback_id
+                    if not (self._run_hook_accepted or self._warned_run_hook_wait):
+                        self._warned_run_hook_wait = True
+                        warn_run_hook_wait = True
                 self._changed.notify_all()
                 decision = self._take_job_suspend()
+            if warn_run_hook_wait:
+                self._logger.warning(
+                    "a job asked to suspend the MicroVM before the run hook "
+                    "arrived. The worker waits for the run hook to name the "
+                    "MicroVM.",
+                    extra={"callbackId": callback_id},
+                )
             self._act_on_job_suspend(decision)
 
     def _run_job_and_report(
@@ -1685,11 +1707,21 @@ def _take_job_suspend(self) -> _JobSuspend | None:
         fires now.
 
         The request fires when no job, failure report, run hook, or
-        SuspendMicrovm call runs, and the listener is not terminating, not
-        closed, and not refusing jobs for a suspend. Each event that can
-        change one of these calls this method: a job end, a report end, a
-        run hook end, and a call end. Firing clears the request, and stops
-        an idle timer that auto-suspend armed.
+        SuspendMicrovm call runs. Each event that can change one of these
+        calls this method: a job end, a report end, a run hook end, and a
+        call end. Firing clears the request, and stops an idle timer that
+        auto-suspend armed.
+
+        The terminate hook and ``close()`` drop the request, and a job end
+        after either records none. A suspend refusal drops the request too,
+        and a job end during a refusal records none. So these states need
+        no check here.
+
+        HTTP jobs can arrive before the run hook. A request that is due
+        before the first accepted run hook stays pending, because that hook
+        sets the target. The end of that hook calls this method again. The
+        request then fires when the hook named a MicroVM, and is dropped
+        with a warning when it named none.
 
         The caller holds the lock through :meth:`_locked`. It passes the
         result to :meth:`_act_on_job_suspend` after it releases the lock. So
@@ -1702,14 +1734,13 @@ def _take_job_suspend(self) -> _JobSuspend | None:
             or self._jobs
             or self._reports
             or self._run_hooks
-            or self._terminating
-            or self._closed
-            or self._suspending
             or self._suspend_calls
         ):
             return None
-        self._suspend_requested = None
         target = self._suspend_target
+        if target is None and not self._run_hook_accepted:
+            return None
+        self._suspend_requested = None
         if target is None:
             return _JobSuspend(callback_id=callback_id, call=None)
         self._stop_idle_time()

packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py

@@ -3506,6 +3506,10 @@ def create(region: str) -> Any:
 NO_TARGET_WARNING = (
     "cannot suspend the MicroVM after a job: the run hook named no MicroVM"
 )
+RUN_HOOK_WAIT_WARNING = (
+    "a job asked to suspend the MicroVM before the run hook arrived. "
+    "The worker waits for the run hook to name the MicroVM."
+)
 
 
 def post_job(h: Harness, callback_id: str, **fields: Any) -> tuple[int, str | None]:
@@ -3825,13 +3829,11 @@ def test_suspends_the_run_hook_microvm_not_the_one_of_the_job_request(
     assert h.microvms.suspended == ["mvm-1"]
 
 
-@pytest.mark.parametrize("run_hook", ["none", "without-payload"])
 def test_warns_and_does_not_suspend_when_no_run_hook_named_a_microvm(
-    start: Any, logger: Any, run_hook: str
+    start: Any, logger: Any
 ) -> None:
     h = start()
-    if run_hook == "without-payload":
-        assert h.post(f"{HOOK_PATH_PREFIX}run", {"microvmId": "mvm-1"}) == 200
+    assert h.post(f"{HOOK_PATH_PREFIX}run", {"microvmId": "mvm-1"}) == 200
     post_job(h, "cb-1", suspendAfter=True, microvmId="mvm-other")
 
     end_job(h, "cb-1")
@@ -3847,6 +3849,101 @@ def test_warns_and_does_not_suspend_when_no_run_hook_named_a_microvm(
     assert h.job("cb-3") == 202
 
 
+def test_a_request_before_the_run_hook_fires_when_the_run_hook_names_a_microvm(
+    start: Any, logger: Any
+) -> None:
+    # HTTP jobs can arrive before the run hook. The job ends before the run
+    # hook, so no target exists yet. The request waits for the run hook,
+    # and fires with the run hook's MicroVM ID.
+    h = start()
+    assert post_job(h, "cb-1", suspendAfter=True, microvmId="mvm-other") == (
+        202,
+        "accepted",
+    )
+    end_job(h, "cb-1")
+    assert h.microvms.suspended == []
+    assert h.listener._suspend_requested == "cb-1"  # noqa: SLF001
+    assert NO_TARGET_WARNING not in logger.messages("warning")
+
+    assert h.run_hook() == 200
+
+    assert eventually(lambda: h.microvms.suspended == ["mvm-1"])
+    h.settle()
+    assert info_lines(logger, JOB_SUSPEND_INFO) == [
+        {"microvmId": "mvm-1", "callbackId": "cb-1"}
+    ]
+    assert NO_TARGET_WARNING not in logger.messages("warning")
+
+
+def test_a_request_before_the_run_hook_is_dropped_when_the_run_hook_names_none(
+    start: Any, logger: Any
+) -> None:
+    # The run hook without a payload names no MicroVM to suspend. So the
+    # request that waited for it is dropped with the warning.
+    h = start()
+    post_job(h, "cb-1", suspendAfter=True, microvmId="mvm-other")
+    end_job(h, "cb-1")
+    assert NO_TARGET_WARNING not in logger.messages("warning")
+
+    assert h.post(f"{HOOK_PATH_PREFIX}run", {"microvmId": "mvm-1"}) == 200
+
+    assert eventually(lambda: NO_TARGET_WARNING in logger.messages("warning"))
+    assert [line for line in logger.lines if line[1] == NO_TARGET_WARNING] == [
+        ("warning", NO_TARGET_WARNING, {"callbackId": "cb-1"})
+    ]
+    assert h.listener._suspend_requested is None  # noqa: SLF001
+    post_job(h, "cb-2")
+    end_job(h, "cb-2")
+    assert h.microvms.suspended == []
+    assert logger.messages("warning").count(NO_TARGET_WARNING) == 1
+
+
+def test_warns_once_when_requests_wait_for_the_run_hook(
+    start: Any, logger: Any
+) -> None:
+    # An image without the run hook never ends the wait. The warning is the
+    # only sign of that. Two jobs that wait log it once.
+    h = start()
+    post_job(h, "cb-1", suspendAfter=True)
+    end_job(h, "cb-1")
+    post_job(h, "cb-2", suspendAfter=True)
+    end_job(h, "cb-2")
+
+    assert [line for line in logger.lines if line[1] == RUN_HOOK_WAIT_WARNING] == [
+        ("warning", RUN_HOOK_WAIT_WARNING, {"callbackId": "cb-1"})
+    ]
+    assert h.microvms.suspended == []
+
+
+def test_does_not_warn_of_the_wait_when_the_run_hook_came_first(
+    start: Any, logger: Any
+) -> None:
+    h = start()
+    # The job request arrives before the run hook, and ends after it.
+    post_job(h, "cb-1", suspendAfter=True)
+    h.run_hook()
+
+    end_job(h, "cb-1")
+
+    assert h.microvms.suspended == ["mvm-1"]
+    assert RUN_HOOK_WAIT_WARNING not in logger.messages("warning")
+
+
+def test_close_drops_a_request_that_waits_for_the_run_hook(start: Any) -> None:
+    h = start()
+    post_job(h, "cb-1", suspendAfter=True)
+    end_job(h, "cb-1")
+    assert h.listener._suspend_requested == "cb-1"  # noqa: SLF001
+
+    h.listener.close()
+    assert h.listener._suspend_requested is None  # noqa: SLF001
+    assert h.run_hook() == 200
+    h.settle()
+
+    assert h.microvms.suspended == []
+    assert h.listener._suspend_requested is None  # noqa: SLF001
+
+
 def test_a_resume_hook_does_not_repeat_the_suspend(start: Any) -> None:
     h = start()
     h.run_hook()
@@ -3941,6 +4038,24 @@ def test_does_not_suspend_after_the_terminate_hook(start: Any) -> None:
     assert h.listener._suspend_requested is None  # noqa: SLF001
 
 
+def test_records_no_request_for_a_job_that_ends_after_the_terminate_hook(
+    start: Any,
+) -> None:
+    # The terminate hook clears no request here, because none is pending
+    # yet. Only the job end's check of the terminate hook stops the
+    # suspend.
+    h = start()
+    h.run_hook()
+    post_job(h, "cb-1", suspendAfter=True)
+
+    assert h.hook("terminate", LAMBDA_HEADERS) == 200
+    end_job(h, "cb-1")
+    assert h.worker.idle(WAIT)
+
+    assert h.listener._suspend_requested is None  # noqa: SLF001
+    assert h.microvms.suspended == []
+
+
 def test_does_not_suspend_after_close(start: Any) -> None:
     h = start()
     h.run_hook()