feat(microvm-worker): act on lifecycle hooks

packages/aws-durable-execution-sdk-python-extras/README.md

@@ -249,7 +249,7 @@ Suspending is off by default, because it needs more permissions and changes what
 - With `auto_suspend_on_idle`, the MicroVM's role needs `lambda:SuspendMicrovm` on the image. That permission lets code in any MicroVM of the image suspend any other. So use one image per trust boundary.
 - Background processes that an earlier job started stop while the MicroVM is suspended. The worker counts only its running jobs.
 
-The worker must support `auto_suspend_on_idle`. The JavaScript worker [`@aws/durable-execution-sdk-js-microvm-worker`](https://github.com/aws/aws-durable-execution-sdk-js/tree/main/packages/aws-durable-execution-sdk-js-microvm-worker) does. The Python worker in this repository does not yet. It accepts the setting, and its MicroVM keeps running. `idle_policy` works with both workers.
+The worker must support `auto_suspend_on_idle`. The JavaScript worker [`@aws/durable-execution-sdk-js-microvm-worker`](https://github.com/aws/aws-durable-execution-sdk-js/tree/main/packages/aws-durable-execution-sdk-js-microvm-worker) and the Python worker [`aws-durable-execution-sdk-python-microvm-worker`](../aws-durable-execution-sdk-python-microvm-worker) both do. `idle_policy` works with both workers. Enable the `suspend` and `resume` hooks in the image. The worker then refuses jobs while the MicroVM freezes, and accepts them again as soon as it resumes. Without the `resume` hook, the first job after a resume can wait a little over 30 seconds, until the worker's refusal ends and the request step delivers again.
 
 ### Session permissions
 

packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py

@@ -289,9 +289,9 @@ class MicrovmSessionConfig(MicrovmBaseConfig):
     not extend the MicroVM's lifetime.
 
     The worker must be a version that supports the setting. The JavaScript
-    worker ``@aws/durable-execution-sdk-js-microvm-worker`` does. The Python
-    worker in this repository does not yet: it accepts the setting, and its
-    MicroVM keeps running.
+    worker ``@aws/durable-execution-sdk-js-microvm-worker`` and the Python
+    worker ``aws-durable-execution-sdk-python-microvm-worker`` in this
+    repository both do.
     """
 
     auto_suspend_idle_time: Duration | None = None

packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py

@@ -17,8 +17,12 @@
 
 from __future__ import annotations
 
+import functools
 import http.client
 import json
+import socket
+import threading
+import time
 from collections.abc import Callable, Iterator
 from typing import Any
 
@@ -36,6 +40,7 @@
     MicrovmWorker,
     start_microvm_worker,
 )
+from aws_durable_execution_sdk_python_microvm_worker import worker as worker_module
 from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner
 
 from aws_durable_execution_sdk_python_extras.microvm import (
@@ -54,6 +59,49 @@
 WAIT = 30
 
 
+@functools.cache
+def _has_ipv6_loopback() -> bool:
+    """Whether this host can use ``::1``.
+
+    A host can support IPv6 sockets and still have no ``::1`` address. A
+    Docker container with IPv6 turned off is one example. A connection to
+    ``::1`` then fails with ``EADDRNOTAVAIL``.
+    """
+    if not socket.has_dualstack_ipv6():
+        return False
+    try:
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as probe:
+            probe.bind(("::1", 0))
+    except OSError:
+        return False
+    return True
+
+
+def host_of(worker: MicrovmWorker) -> str:
+    """The loopback address on which a test reaches ``worker``.
+
+    A dual-stack worker binds ``[::]`` with ``SO_REUSEADDR``. On macOS, that
+    bind succeeds on a port that another process holds on ``127.0.0.1``. A
+    connection to ``127.0.0.1`` then reaches the other process, and the test
+    fails with its answer. So the tests connect to ``::1`` when the worker
+    is dual-stack. Linux refuses such a bind, so a MicroVM is not affected.
+    """
+    server = worker._server  # noqa: SLF001
+    if server.address_family == socket.AF_INET6 and _has_ipv6_loopback():
+        return "::1"
+    return "127.0.0.1"
+
+
+def url_of(worker: MicrovmWorker) -> str:
+    """The worker's endpoint URL for the durable function."""
+    host = host_of(worker)
+    return (
+        f"http://[{host}]:{worker.port}"
+        if ":" in host
+        else f"http://{host}:{worker.port}"
+    )
+
+
 class RunnerCallbacks:
     """A Lambda client for the worker that sends each callback call to the
     local test runner."""
@@ -94,7 +142,9 @@ def run_hook_to(worker: MicrovmWorker) -> Callable[[dict[str, Any]], None]:
     """A RunMicrovm answer that sends the run hook to ``worker``."""
 
     def deliver(request: dict[str, Any]) -> None:
-        connection = http.client.HTTPConnection("127.0.0.1", worker.port, timeout=WAIT)
+        connection = http.client.HTTPConnection(
+            host_of(worker), worker.port, timeout=WAIT
+        )
         try:
             body = {"microvmId": "mvm-1", "runHookPayload": request["runHookPayload"]}
             connection.request("POST", f"{HOOK_PATH_PREFIX}run", body=json.dumps(body))
@@ -142,7 +192,7 @@ def handler(event: Any, context: DurableContext) -> Any:
         closers.append(worker.close)
         client = fake_client(
             launches=[run_hook_to(worker)],
-            endpoint=f"http://127.0.0.1:{worker.port}",
+            endpoint=url_of(worker),
         )
         holder["config"] = MicrovmConfig(
             image_identifier="img",
@@ -264,7 +314,7 @@ def handler(event: Any, context: DurableContext) -> Any:
     try:
         client = fake_client(
             launches=[run_hook_to(worker)],
-            endpoint=f"http://127.0.0.1:{worker.port}",
+            endpoint=url_of(worker),
         )
         holder["config"] = MicrovmSessionConfig(
             image_identifier="img",
@@ -284,3 +334,170 @@ def handler(event: Any, context: DurableContext) -> Any:
     assert len(client.requests("run_microvm")) == 1
     assert len(client.requests("create_microvm_auth_token")) == 2
     assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}]
+
+
+def post_hook(worker: MicrovmWorker, hook: str) -> None:
+    """Send a lifecycle hook as Lambda does, with a local Host."""
+    connection = http.client.HTTPConnection(host_of(worker), worker.port, timeout=WAIT)
+    try:
+        connection.request(
+            "POST",
+            f"{HOOK_PATH_PREFIX}{hook}",
+            body=b"{}",
+            headers={"Host": "localhost:8080"},
+        )
+        assert connection.getresponse().status == 200
+    finally:
+        connection.close()
+
+
+class RecordingLogger:
+    """Records the worker's log messages."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self.messages: list[str] = []
+
+    def _add(self, msg: object, *args: object, extra: Any = None) -> None:
+        with self._lock:
+            self.messages.append(str(msg))
+
+    info = warning = error = _add
+
+    def count(self, message: str) -> int:
+        with self._lock:
+            return self.messages.count(message)
+
+
+@pytest.mark.parametrize("resume_hook", [True, False], ids=["resume-hook", "none"])
+def test_resumes_a_worker_that_suspended_itself_between_two_session_jobs(
+    fake_client: Any, monkeypatch: pytest.MonkeyPatch, resume_hook: bool
+) -> None:
+    """The worker suspends its idle MicroVM, and the next job resumes it.
+
+    A stand-in for Lambda keeps the MicroVM's state. SuspendMicrovm from the
+    worker sets it to ``SUSPENDED`` and sends the ``suspend`` hook.
+    ResumeMicrovm from the session sets it to ``RUNNING``, and sends the
+    ``resume`` hook when the image enables it. The worker's timers run 20
+    times faster, so the session's 60-second idle time takes 3 seconds, and
+    the 30-second refusal takes 1.5 seconds.
+
+    Without the ``resume`` hook, the worker answers the second job with 503
+    until its refusal ends. The request step delivers it again after each
+    503, after 0.25 seconds at first, doubling up to 4 seconds. The scaled
+    idle time must stay longer than the gap between two deliveries.
+    Otherwise the worker suspends again before each delivery. A 10-second
+    idle time scaled to 1 second did that. Unscaled, the shortest idle time,
+    10 seconds, is longer than the longest gap.
+    """
+    original_timer = worker_module._start_timer  # noqa: SLF001
+    monkeypatch.setattr(
+        worker_module,
+        "_start_timer",
+        lambda delay, action: original_timer(delay / 20, action),
+    )
+    state = {"value": "RUNNING"}
+    lock = threading.Lock()
+    holder: dict[Any, Any] = {}
+
+    class StatefulMicrovms(fake_client):  # type: ignore[misc, valid-type]
+        def get_microvm(self, **request: Any) -> dict[str, Any]:
+            with lock:
+                current = state["value"]
+            response = super().get_microvm(**request)
+            return {**response, "state": current}
+
+        def resume_microvm(self, **request: Any) -> dict[str, Any]:
+            super().resume_microvm(**request)
+            with lock:
+                state["value"] = "RUNNING"
+            if resume_hook:
+                post_hook(holder["worker"], "resume")
+            return {}
+
+    class WorkerMicrovms:
+        """The worker's lambda-microvms client."""
+
+        def __init__(self) -> None:
+            self.suspended: list[str] = []
+
+        def suspend_microvm(self, *, microvmIdentifier: str) -> dict[str, Any]:  # noqa: N803
+            self.suspended.append(microvmIdentifier)
+            with lock:
+                state["value"] = "SUSPENDED"
+            # Lambda calls the suspend hook before it freezes the MicroVM.
+            post_hook(holder["worker"], "suspend")
+            return {}
+
+    worker_microvms = WorkerMicrovms()
+
+    def job(job_input: Any, context: MicrovmJobContext) -> Any:
+        return f"{job_input} done"
+
+    def wait_until_suspended(_: Any) -> str:
+        deadline = time.monotonic() + WAIT
+        while time.monotonic() < deadline:
+            with lock:
+                if state["value"] == "SUSPENDED":
+                    return "suspended"
+            time.sleep(0.01)
+        return "still running"
+
+    def pipeline(vm: MicrovmSession, ctx: DurableContext) -> Any:
+        options = MicrovmInvokeOptions(timeout=Duration.from_minutes(1))
+        first = vm.invoke(ctx, "build", job_input="build", options=options)
+        between = ctx.step(wait_until_suspended, name="pause")
+        second = vm.invoke(ctx, "test", job_input="test", options=options)
+        return [first, between, second]
+
+    @durable_execution
+    def handler(event: Any, context: DurableContext) -> Any:
+        return microvm_session(
+            context, "pipeline", config=holder["config"], handler=pipeline
+        )
+
+    runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05)
+    callbacks = RunnerCallbacks(runner)
+    logger = RecordingLogger()
+    worker = start_microvm_worker(
+        handler=job,
+        create_client=lambda region: callbacks,  # type: ignore[arg-type, return-value]
+        create_microvms_client=lambda region: worker_microvms,  # type: ignore[arg-type, return-value]
+        logger=logger,
+        port=0,
+    )
+    holder["worker"] = worker
+    try:
+        client = StatefulMicrovms(
+            launches=[run_hook_to(worker)],
+            endpoint=url_of(worker),
+        )
+        holder["config"] = MicrovmSessionConfig(
+            image_identifier="img",
+            execution_role_arn="role",
+            timeout=Duration.from_minutes(5),
+            auto_suspend_on_idle=True,
+            auto_suspend_idle_time=Duration.from_seconds(60),
+            client=client,
+        )
+        arn = runner.run_async(input="null", execution_timeout=WAIT)
+        result = runner.wait_for_result(arn, timeout=WAIT)
+        assert worker.idle(WAIT)
+    finally:
+        worker.close()
+        runner.close()
+
+    assert result.status is InvocationStatus.SUCCEEDED, result.error
+    assert json.loads(result.result or "null") == [
+        "build done",
+        "suspended",
+        "test done",
+    ]
+    assert worker_microvms.suspended[0] == "mvm-1"
+    assert client.requests("resume_microvm") == [{"microvmIdentifier": "mvm-1"}]
+    assert client.requests("terminate_microvm") == [{"microvmIdentifier": "mvm-1"}]
+    refusals = logger.count("job refused: the MicroVM is suspending")
+    if resume_hook:
+        assert refusals == 0
+    else:
+        assert refusals >= 1