From d5b21d3c1ba19f668a790503dc70074c79be22ac Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 16 Oct 2026 09:39:45 -0700 Subject: [PATCH 1/4] feat(microvm-worker): act on lifecycle hooks The worker now acts on the suspend, resume, and terminate hooks, and suspends its own MicroVM when a session asks for it. - suspend: answer 200, then refuse jobs with 503 until the resume hook or for 30 seconds. A run hook job that arrives then is deferred. - terminate: set each running job's cancelled event, fail its callback with MicrovmTerminatedError, and answer after the reports, at most 5 seconds. Later jobs get 503. - The three hooks act only on Lambda's own calls: a local Host and no x-amzn-requestid. A forwarded request gets 200 and is ignored. - autoSuspendIdleSeconds: after that many idle seconds the worker calls SuspendMicrovm with the run hook's MicroVM ID, and refuses jobs until the resume hook or for 30 seconds. A rejected first attempt ends the refusal at once. An unknown outcome keeps it. The worker now needs boto3 1.43.35 for the lambda-microvms client. --- .../README.md | 2 +- .../microvm/config.py | 6 +- .../tests/microvm/worker_composed_test.py | 134 ++ .../README.md | 69 +- .../pyproject.toml | 7 +- .../__init__.py | 2 + .../callback_reporter.py | 4 +- .../self_suspend.py | 158 ++ .../worker.py | 836 +++++++- .../tests/lifecycle_hooks_test.py | 1771 +++++++++++++++++ .../tests/self_suspend_test.py | 144 ++ pyproject.toml | 2 +- 12 files changed, 3082 insertions(+), 53 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py create mode 100644 packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py create mode 100644 packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 8886d99ec..9ce81da73 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/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 up to 30 seconds. ### Session permissions diff --git a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py b/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py index 47cbe0db4..27d41573a 100644 --- a/packages/aws-durable-execution-sdk-python-extras/src/aws_durable_execution_sdk_python_extras/microvm/config.py +++ b/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 diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index d86415077..5cc9b1094 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -19,6 +19,8 @@ import http.client import json +import threading +import time from collections.abc import Callable, Iterator from typing import Any @@ -36,6 +38,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 ( @@ -284,3 +287,134 @@ 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("127.0.0.1", 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() + + +def test_resumes_a_worker_that_suspended_itself_between_two_session_jobs( + fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> 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. The worker's timers run 10 times faster, so the + session's 10-second idle time takes 1 second. + """ + original_timer = worker_module._start_timer # noqa: SLF001 + monkeypatch.setattr( + worker_module, + "_start_timer", + lambda delay, action: original_timer(delay / 10, 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" + 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) + 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] + port=0, + ) + holder["worker"] = worker + try: + client = StatefulMicrovms( + launches=[run_hook_to(worker)], + endpoint=f"http://127.0.0.1:{worker.port}", + ) + 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(10), + 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"}] diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index 10c91c80c..6b08b7663 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -35,7 +35,7 @@ if __name__ == "__main__": worker.wait() ``` -The worker needs `handler`, `routes`, or both. A handler is a plain function. An `async def` function, or an object whose `__call__` is `async def`, raises `TypeError`. A handler that returns an awaitable fails its job with a `TypeError`, because the worker cannot run it. The handler's return value becomes the durable function's result, and must be JSON-serializable. An exception fails the durable operation with the exception's class name and message. +The worker needs `handler`, `routes`, or both. A handler is a plain function. An `async def` function, or an object whose `__call__` is `async def`, raises `TypeError`. A handler that returns an awaitable fails its job with a `TypeError`, because the worker cannot run it. The handler's return value becomes the durable function's result, and must be JSON-serializable. An exception fails the durable operation with the exception's class name and message. `context.cancelled` is set when the durable function stops waiting, for example after a heartbeat timeout, and when Lambda calls the `terminate` hook. Python cannot stop a running thread, so the handler should check the event and stop its work. `start_microvm_worker` serves in a daemon thread and returns. So the entry point ends with `worker.wait()`, which blocks until `worker.close()`. @@ -49,20 +49,30 @@ The image needs these hook settings in `CreateMicrovmImage`: { "hooks": { "port": 8080, - "microvmHooks": { "run": "ENABLED", "runTimeoutInSeconds": 10 }, + "microvmHooks": { + "run": "ENABLED", + "runTimeoutInSeconds": 10, + "resume": "ENABLED", + "resumeTimeoutInSeconds": 10, + "suspend": "ENABLED", + "suspendTimeoutInSeconds": 10, + "terminate": "ENABLED", + "terminateTimeoutInSeconds": 10 + }, "microvmImageHooks": { "ready": "ENABLED" } } } ``` - The `run` hook delivers the job, or tells the worker that jobs arrive over HTTP. So it must be enabled. +- The `resume` hook tells a session's worker that its MicroVM runs again after a suspend. It is optional, and recommended for sessions. See [Suspending when idle](#suspending-when-idle). +- The `suspend` hook tells the worker that Lambda is about to freeze the MicroVM. The worker then refuses new jobs until the `resume` hook. It is optional, and recommended for sessions. +- The `terminate` hook tells the worker that Lambda is about to end the MicroVM. The worker then fails the callback of each running job, so the durable function fails the job at once instead of at its heartbeat timeout or its timeout. It is optional, and recommended. Give it at least 10 seconds, because the worker reports before it answers, for at most 5 seconds. A hook timeout left unset ended the hook after about 2 seconds in testing. - The service requires the `ready` hook whenever a lifecycle hook is enabled. The worker answers it with HTTP 200. - `port` must match the worker's port. The worker listens on 8080 by default, which is also the port that the MicroVM endpoint forwards to. Pass `port=` for another port. The worker listens on IPv6 and IPv4 at once when the host supports it. Lambda calls the hooks on `localhost`, and that name can resolve to either family. -The worker answers the `resume`, `suspend`, and `terminate` hooks with HTTP 200, and does nothing else for them yet. A later change acts on them. - ## What the worker does On `POST /aws/lambda-microvms/runtime/v1/run` the worker: @@ -88,11 +98,54 @@ The worker creates a Lambda client for each job, after the job arrives, not at i A document that does not match the contract gets HTTP 400 with the validation message. If the document named a callback, the worker also fails that callback, so the durable function fails at once instead of at its timeout. A request whose body cannot be read gets HTTP 400 with a generic message, and the error is logged. The worker accepts a `run` hook body of at most 64 KB and a job request body of at most 6 MB, with a `Content-Length` or with chunked transfer encoding. In a chunked body, the chunk-size lines, their extensions, and the trailers have a limit of their own, of the same size. A larger body gets HTTP 400. With `Expect: 100-continue`, the 400 replaces `100 Continue`, so the client does not send the body. Without it, the worker reads and drops the body after the answer, for at most 10 seconds and 24 MB. Without that read, the connection would close with unread bytes, and the client would get a TCP reset instead of the 400. A body that is larger, or that is still arriving after 10 seconds, still gets the reset. Other lifecycle hooks get HTTP 200. Other paths get HTTP 404. -`worker.idle(timeout)` waits until every job started so far has ended, and returns `False` when the timeout ends first. `worker.close()` stops the server. Running jobs continue, so call `idle()` first to let them report their outcome. +`worker.idle(timeout)` waits until every job started so far has ended, and a running `SuspendMicrovm` call has returned. It returns `False` when the timeout ends first. `worker.close()` stops the server and the idle time. Running jobs continue, so call `idle()` first to let them report their outcome. + +## Suspending when idle + +A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets it only with `auto_suspend_on_idle=True`, because suspending is off by default. The worker then suspends its own MicroVM when it has run no job for that many seconds: + +1. The idle time starts when the `run` hook arrives without a job, when a job ends, and when the `resume` lifecycle hook arrives, if the image enables it. +2. A job that starts cancels the idle time. So the worker never suspends the MicroVM during a job. +3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The worker waits for it for at most 10 seconds. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. +4. From the moment the idle time ends, the worker answers every job with HTTP 503 and does not start it. A job accepted then would freeze with the MicroVM, and nothing would resume it. After a 503, the durable function checks the state, resumes the MicroVM, and delivers the job again. +5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. So enable the `resume` hook in the image. Without it, the first job after a resume can wait up to 30 seconds. +6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection, because the call sends no request. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. +7. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a throttle, a conflict, expired credentials, a clock offset, a server error, and a dropped connection. A rejection after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. + +`worker.close()` stops the idle time. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. + +The worker counts only jobs. It cannot see background threads or processes, or inbound traffic that is not a job. A session with such work leaves `auto_suspend_on_idle` off, and the payload then has no `autoSuspendIdleSeconds`. + +The worker creates a `lambda-microvms` client for each suspend, and closes it after the call. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. + +## Suspend and terminate hooks + +Lambda waits for the answer to a `suspend` or `terminate` hook, up to the hook's timeout. The network and the execution role's credentials work while the hook runs. These facts come from probes in us-east-1, which the JavaScript SDK ran. + +On the `suspend` hook the worker: + +1. Answers HTTP 200 at once. A suspend hook that ran past its timeout terminated the MicroVM instead of suspending it. +2. Refuses every job with HTTP 503, as after its own suspend. A job accepted now would freeze with the MicroVM. The refusal ends when the `resume` hook arrives, or after 30 seconds. +3. Leaves running jobs alone. They freeze with the MicroVM, and continue after a resume. +4. Defers a job from a `run` hook whose body was still arriving. The durable function cannot deliver a `run` hook job again, so the worker does not refuse it. It starts the job when the refusal ends. + +Lambda called the hook after `SuspendMicrovm` from outside, after the worker's own `SuspendMicrovm`, and when the idle policy suspended the MicroVM. + +On the `terminate` hook the worker: + +1. Sets each running job's `context.cancelled`, and stops its heartbeats. +2. Fails each running job's callback with `MicrovmTerminatedError`. The worker reports `MicrovmTerminatedError` as the callback's error type. The `microvm` operation then fails with a `MicrovmJobFailedError` whose message contains `(MicrovmTerminatedError)`. The worker does not report the handler's outcome when it returns later. +3. Skips a job whose callback a heartbeat already found gone. A report would fail the same way. +4. Answers HTTP 200 when the reports end, or after 5 seconds. A report took 20 to 1,700 milliseconds in testing. The answer also waits for a job that is reporting its own outcome, for a `run` hook request whose body is still arriving, and for every callback failure report in flight, such as the report for an invalid request. Work that begins during the wait is included. A second `terminate` hook that arrives during the wait waits for the same work. Lambda can end the MicroVM as soon as it gets the answer, so a report sent after it can be lost. +5. Refuses every later job with HTTP 503, and never suspends the MicroVM again. A job from a `run` hook whose body was still arriving is not started. Its callback fails with `MicrovmTerminatedError`, as does a job deferred by an earlier `suspend` hook. + +Lambda called the hook after `TerminateMicrovm` on a running MicroVM, and at the end of `maximumDurationInSeconds`. It did not call the hook when it terminated a suspended MicroVM, after `TerminateMicrovm` or at the end of the idle policy's `suspendedDurationSeconds`. The process is frozen then. A job of such a MicroVM fails at its heartbeat timeout or its timeout. + +The MicroVM endpoint forwards a request to a hook path from any caller that holds an auth token for the MicroVM. Lambda's own hook calls carried `Host: localhost:8080` and no `x-amzn-requestid` header. A forwarded request carried the endpoint's host name and an `x-amzn-requestid` header. So the worker acts on `suspend`, `terminate`, and `resume` only when the request has a `localhost`, `127.0.0.1`, or `[::1]` host, with any port, and no `x-amzn-requestid`. It answers any other request to these hooks with 200, and logs a warning. The Lambda documentation does not state these headers. If Lambda changed them, the worker would ignore `suspend` and `terminate`. A refusal after a suspend would then end after 30 seconds instead of at the `resume` hook. Without the check on `resume`, a forwarded `resume` between a `suspend` hook and the freeze could end the refusal early, and a job accepted then would freeze with the MicroVM. The `run` hook does not use the check, because a false rejection there would lose every job. ## Lower-level building blocks -- `MicrovmWorkerListener.create(...)` takes the same arguments as `start_microvm_worker`, except `port`. Its `request_handler_class` serves the hooks and routes from your own `http.server.ThreadingHTTPServer`. Set the server's backlog to 511, as `start_microvm_worker` does. A session can send many jobs at once, and the default backlog of 5 resets the connections beyond it. The server's constructor already listens, so set `request_queue_size = 511` as a class attribute of a `ThreadingHTTPServer` subclass, not on the server object. +- `MicrovmWorkerListener.create(...)` takes the same arguments as `start_microvm_worker`, except `port`. Call its `close()` when you stop serving it, so that it no longer suspends the MicroVM. Its `request_handler_class` serves the hooks and routes from your own `http.server.ThreadingHTTPServer`. Set the server's backlog to 511, as `start_microvm_worker` does. A session can send many jobs at once, and the default backlog of 5 resets the connections beyond it. The server's constructor already listens, so set `request_queue_size = 511` as a class attribute of a `ThreadingHTTPServer` subclass, not on the server object. - `RunHookRequest.from_dict` and `MicrovmJobRequest.from_dict` validate the two documents that deliver a job: the `run` lifecycle hook body, and the body of an HTTP job request. - `CallbackReporter` sends heartbeats, and completes the callback with a result or an error. @@ -147,3 +200,7 @@ The worker logs to the standard library logger `aws_durable_execution_sdk_python ## Permissions The MicroVM's execution role needs `lambda:SendDurableExecutionCallbackSuccess`, `lambda:SendDurableExecutionCallbackFailure`, and `lambda:SendDurableExecutionCallbackHeartbeat` on the durable function's ARN (`arn:aws:lambda:::function::*`). + +A session that suspends when idle also needs `lambda:SuspendMicrovm` on the MicroVM image (`arn:aws:lambda:::microvm-image:`). The action authorizes on the image, not on the MicroVM. So a MicroVM with this permission can suspend any MicroVM from the same image, even during a job. That job then waits until its heartbeat timeout or its timeout. Use one image per trust boundary. + +The worker needs boto3 1.43.35 or later, the first release with the `lambda-microvms` client. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/pyproject.toml b/packages/aws-durable-execution-sdk-python-microvm-worker/pyproject.toml index 76be6bbda..64190f4e7 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/pyproject.toml @@ -22,10 +22,11 @@ classifiers = [ ] # The worker runs inside the MicroVM, not in the durable function. So it does # not depend on the durable execution SDK. It needs boto3 for the durable -# callback APIs. The MicroVM image installs it, because no Lambda runtime -# provides it there. +# callback APIs, and for SuspendMicrovm when a session suspends when idle. The +# MicroVM image installs it, because no Lambda runtime provides it there. dependencies = [ - "boto3>=1.42.1", + # botocore 1.43.35 is the first release with the lambda-microvms client. + "boto3>=1.43.35", ] [project.urls] diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/__init__.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/__init__.py index f8bff7164..85772f14e 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/__init__.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/__init__.py @@ -37,6 +37,7 @@ MICROVM_JOB_PATH, MicrovmJobContext, MicrovmJobHandler, + MicrovmTerminatedError, MicrovmWorker, MicrovmWorkerListener, start_microvm_worker, @@ -57,6 +58,7 @@ "MicrovmJobHandler", "MicrovmJobRequest", "MicrovmRunHookPayload", + "MicrovmTerminatedError", "MicrovmWorker", "MicrovmWorkerListener", "MicrovmWorkerLogger", diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/callback_reporter.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/callback_reporter.py index 2cf8c96bd..8c5291d60 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/callback_reporter.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/callback_reporter.py @@ -88,7 +88,7 @@ # signature rejected because of a clock offset, and throttling. The credential # provider refreshes expired credentials. So a later attempt can succeed, # although the service answered 403. -_TRANSIENT_4XX_CODES = _EXPIRED_CREDENTIAL_CODES | frozenset( +TRANSIENT_4XX_CODES = _EXPIRED_CREDENTIAL_CODES | frozenset( { # Clock skew. "AuthFailure", @@ -248,7 +248,7 @@ def is_permanent_error(error: BaseException) -> bool: """ if is_terminal_callback_error(error): return True - if error_code(error) in _TRANSIENT_4XX_CODES: + if error_code(error) in TRANSIENT_4XX_CODES: return False status = _http_status(error) return status is not None and 400 <= status < 500 and status not in {408, 409, 429} diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py new file mode 100644 index 000000000..a954c8769 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The worker's own SuspendMicrovm call, for a session that suspends when idle.""" + +from __future__ import annotations + +import threading +from concurrent.futures import Future +from typing import TYPE_CHECKING, Any + +import boto3 +from botocore.config import Config +from botocore.exceptions import ClientError, NoCredentialsError + +from aws_durable_execution_sdk_python_microvm_worker.callback_reporter import ( + TRANSIENT_4XX_CODES, + error_code, +) +from aws_durable_execution_sdk_python_microvm_worker.logger import safe_get + + +if TYPE_CHECKING: + from mypy_boto3_lambda_microvms import ( + LambdaMicroVMsClient as LambdaMicrovmsClient, + ) +else: + LambdaMicrovmsClient = Any + + +SUSPEND_CALL_TIMEOUT_SECONDS = 10.0 +"""The longest time the worker waits for its own SuspendMicrovm call. + +The call measured 75 to 94 milliseconds from inside a MicroVM in us-east-1. +""" + +SUSPEND_GRACE_SECONDS = 30.0 +"""How long the worker refuses jobs after a suspend, when no ``resume`` hook +arrives first. + +This applies after the worker's own SuspendMicrovm call returned, after a +call whose outcome is unknown, and after the ``suspend`` hook. The call +returns before the MicroVM freezes, and the worker cannot see the freeze. A +job accepted in between would freeze with the MicroVM, and nothing would +resume it. So the worker answers such a job with HTTP 503, and the durable +function resumes the MicroVM and delivers the job again. The ``resume`` hook +ends the refusal at once. This period covers an image that does not enable +that hook, and a suspend that the service did not carry out. It fits inside +the durable function's default delivery retry window. +""" + +_CONNECT_TIMEOUT_SECONDS = 5 + +_CLIENT_ATTEMPTS = 3 +"""The attempts of the default client, the first attempt included. + +The JavaScript worker's default client makes 3 attempts too. +""" + +# Error codes that show that the service did nothing. +_REJECTION_CODES = frozenset( + {"AccessDeniedException", "ValidationException", "ResourceNotFoundException"} +) + +# 4xx error codes that leave open whether the service suspended the MicroVM. +_OPEN_OUTCOME_CODES = TRANSIENT_4XX_CODES | {"ConflictException"} + + +def create_microvms_client(region: str) -> LambdaMicrovmsClient: + """Create the default ``lambda-microvms`` client for one call. + + Each call gets its own session, because boto3's default session is not + safe to use from several threads at once. + """ + return boto3.session.Session().client( + "lambda-microvms", + region_name=region, + config=Config( + connect_timeout=_CONNECT_TIMEOUT_SECONDS, + read_timeout=SUSPEND_CALL_TIMEOUT_SECONDS, + retries={ + "mode": "standard", + # total_max_attempts counts the first attempt too. + "total_max_attempts": _CLIENT_ATTEMPTS, + }, + ), + ) + + +def suspend_microvm( + client: LambdaMicrovmsClient, microvm_id: str, timeout: float +) -> None: + """Call SuspendMicrovm, and stop waiting after ``timeout`` seconds. + + botocore has no deadline for a whole call. Its read timeout limits one + socket read, and it covers neither credential resolution nor its own + retries. So the call runs in a daemon thread, and this function waits + for the thread. A call that ends after the bound is ignored. The caller + closes a client that it created, which releases the call's connection. + + Raises: + TimeoutError: When the call takes longer than ``timeout``. + Exception: The call's own error. + """ + future: Future[None] = Future() + + def run() -> None: + try: + client.suspend_microvm(microvmIdentifier=microvm_id) + future.set_result(None) + except BaseException as error: # noqa: BLE001 + future.set_exception(error) + + threading.Thread(target=run, name="microvm-suspend-call", daemon=True).start() + try: + future.result(timeout) + except TimeoutError as error: + if future.done(): + raise + msg = f"the SuspendMicrovm call took longer than {timeout} seconds" + raise TimeoutError(msg) from error + + +def is_rejected_call(error: BaseException) -> bool: + """Whether a failed SuspendMicrovm call shows that the service did nothing. + + 1. A botocore retry inside the call means that an earlier attempt may + have suspended the MicroVM before its answer was lost. So any error + after a retry leaves the outcome open. + 2. Missing credentials stop the call before any request. + 3. A known rejection code shows that the service did nothing. + 4. A throttle, a conflict, expired credentials, and a clock offset leave + the outcome open, whatever their status. Some services answer a + throttle with HTTP 400. + 5. Any other 4xx answer except 408, 409, and 429 shows that the service + did nothing. + 6. Any other error, such as a timeout, a dropped connection, or a server + error, leaves the outcome open. + """ + if isinstance(error, NoCredentialsError): + return True + if not isinstance(error, ClientError): + return False + metadata = safe_get(lambda: error.response.get("ResponseMetadata", {})) or {} + retries = safe_get(lambda: metadata.get("RetryAttempts")) + if isinstance(retries, int) and retries > 0: + return False + code = error_code(error) + if code in _REJECTION_CODES: + return True + if code in _OPEN_OUTCOME_CODES: + return False + status = safe_get(lambda: metadata.get("HTTPStatusCode")) + return ( + isinstance(status, int) + and 400 <= status < 500 + and status not in {408, 409, 429} + ) diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index 8ffbe8305..6db36cd77 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -50,12 +50,23 @@ RunHookRequest, loads_strict, ) +from aws_durable_execution_sdk_python_microvm_worker.self_suspend import ( + SUSPEND_CALL_TIMEOUT_SECONDS, + SUSPEND_GRACE_SECONDS, + create_microvms_client, + is_rejected_call, + suspend_microvm, +) if TYPE_CHECKING: from mypy_boto3_lambda import LambdaClient + from mypy_boto3_lambda_microvms import ( + LambdaMicroVMsClient as LambdaMicrovmsClient, + ) else: LambdaClient = Any + LambdaMicrovmsClient = Any HOOK_PATH_PREFIX = "/aws/lambda-microvms/runtime/v1/" @@ -115,6 +126,78 @@ _DISCARD_READ_BYTES = 64 * 1024 +TERMINATE_REPORT_BUDGET_SECONDS = 5.0 +"""How long the worker holds the ``terminate`` hook open while it fails the +callbacks of running jobs. + +Lambda waits for the hook's answer, up to the image's +``terminateTimeoutInSeconds``, and then terminates the MicroVM. The network +and the credentials still work while the hook runs. So the worker reports +first and answers after. A failure report took 20 to 1,700 milliseconds in +testing. This budget leaves room for a slow report, and fits in the +10-second hook timeout that the README recommends. Reports still running +after it continue until the MicroVM ends. +""" + +_LOCAL_HOST = re.compile(r"(localhost|127\.0\.0\.1|\[::1\])(:[0-9]+)?") + + +class MicrovmTerminatedError(Exception): + """The error that the worker reports for a job that the ``terminate`` hook + stopped. + + This covers a job that was running, a ``run`` hook job whose body arrived + after the hook, and a ``run`` hook job that a ``suspend`` hook deferred. + The worker sets a running job's ``context.cancelled`` too. + + The worker reports the class name as the callback's error type. The + ``microvm`` operation in the durable function then fails at once with a + ``MicrovmJobFailedError`` whose message contains + ``(MicrovmTerminatedError)``, instead of at the job's heartbeat timeout or + its timeout. + """ + + def __init__(self) -> None: + super().__init__("the MicroVM was terminated before the job finished") + + +def _start_timer(delay: float, action: Callable[[], None]) -> Callable[[], None]: + """Run ``action`` in a daemon thread after ``delay`` seconds. + + Returns the function that cancels the timer. It does not stop an action + that has already begun. So each action checks that it is still current + before it acts. Tests replace this function. + """ + timer = threading.Timer(delay, action) + timer.daemon = True + timer.name = "microvm-worker-timer" + timer.start() + return timer.cancel + + +@dataclass(frozen=True, kw_only=True) +class _AutoSuspend: + """The run hook's request to suspend the MicroVM when idle.""" + + idle_seconds: float + region: str + microvm_id: str + """The run hook's MicroVM ID. SuspendMicrovm uses it, never an ID from a + job request.""" + + +@dataclass(kw_only=True) +class _ActiveJob: + """A job whose handler has not settled yet, with what the ``terminate`` + hook needs to fail it. The listener's lock guards every field.""" + + region: str + cancelled: threading.Event + heartbeats: Heartbeats | None = None + """Set when the job's thread starts its heartbeats.""" + terminated: bool = False + """Set by the ``terminate`` hook, which then fails the callback.""" + @dataclass(frozen=True, kw_only=True) class MicrovmJobContext: @@ -131,13 +214,17 @@ class MicrovmJobContext: region: str """The Region of the durable function.""" cancelled: threading.Event - """Set when the durable function stops accepting the result. - - This happens, for example, after a callback timeout or a heartbeat - timeout. The worker learns it from a heartbeat answer. So the event is set - only for a job with a heartbeat timeout, and only while the handler runs. - Python cannot stop a running thread. So the handler should check the - event, or wait on it, and stop its work when it is set. + """Set when the durable function stops accepting the result, and when + Lambda calls the ``terminate`` hook. + + The durable function stops accepting the result, for example, after a + callback timeout or a heartbeat timeout. The worker learns it from a + heartbeat answer. So that case sets the event only for a job with a + heartbeat timeout, and only while the handler runs. The ``terminate`` + hook sets it for every running job, and the worker fails the job's + callback with :class:`MicrovmTerminatedError`. Python cannot stop a + running thread. So the handler should check the event, or wait on it, + and stop its work when it is set. """ logger: MicrovmWorkerLogger """The worker's logger.""" @@ -187,6 +274,56 @@ class MicrovmWorkerListener: document still names a callback, the listener fails that callback. So the durable function fails at once instead of at its timeout. + A session's ``run`` hook payload can set ``autoSuspendIdleSeconds``. The + listener then suspends its own MicroVM when no job has run for that many + seconds: + + 1. The idle time starts when the ``run`` hook arrives without a job, when + a job ends, and when the ``resume`` lifecycle hook arrives. + 2. A job that starts cancels the idle time. So the listener never + suspends the MicroVM during a job. + 3. When the idle time ends, the listener calls SuspendMicrovm with the + ``run`` hook's MicroVM ID. The call returns before the MicroVM + freezes. + 4. From the moment the idle time ends, the listener answers every job + with HTTP 503 and does not start it. A job accepted then would freeze + with the MicroVM. After a 503, the durable function resumes the + MicroVM and delivers the job again. + 5. The refusal ends when the ``resume`` hook arrives, or 30 seconds after + the call returned. The idle time then starts again. + 6. A call that the service rejected on its first attempt, for example + for a missing permission, ends the refusal at once. The MicroVM keeps + running, and the listener does not retry the call until a job ends. + 7. A call whose outcome is unknown, such as a timeout, a throttle, a + conflict, a server error, or a rejection after a retried attempt, is + handled like a successful call. The service may have suspended the + MicroVM anyway. + + When the image enables the ``suspend`` and ``terminate`` hooks, the + listener acts on them: + + - ``suspend``: the listener answers HTTP 200 at once, because a suspend + hook that runs past its timeout terminates the MicroVM. It then + refuses jobs with HTTP 503, as after its own suspend, until the + ``resume`` hook or for 30 seconds. Running jobs freeze with the + MicroVM. A ``run`` hook job that arrives during the refusal starts when + the refusal ends, because the durable function cannot deliver it + again. + - ``terminate``: the listener sets each running job's + ``context.cancelled``, and fails its callback with + :class:`MicrovmTerminatedError`. It answers HTTP 200 when the reports + end, or after 5 seconds. The answer also waits for jobs that report + their own outcome, for ``run`` hook requests in flight, and for the + failure reports of invalid requests. From then on it refuses jobs with + HTTP 503, and fails the callback of a ``run`` hook job instead of + starting it. + + The listener acts on these two hooks, and on ``resume``, only when the + request's ``Host`` is ``localhost``, ``127.0.0.1``, or ``[::1]``, with any + port, and it has no ``x-amzn-requestid`` header, as Lambda's own calls + do. A request that a caller of the MicroVM endpoint sent to one of these + hook paths gets HTTP 200 and is ignored. + Every other lifecycle hook gets HTTP 200. This includes the ``ready`` image hook, which the service requires whenever the ``run`` hook is enabled. Other paths get HTTP 404. @@ -200,29 +337,74 @@ def __init__( create_client: Callable[[str], LambdaClient] | None, heartbeat_interval_seconds: float | None, logger: MicrovmWorkerLogger, + create_microvms_client: Callable[[str], LambdaMicrovmsClient] | None = None, ) -> None: self._handler = handler self._routes = dict(routes) self._create_client = create_client + self._create_microvms_client = create_microvms_client self._heartbeat_interval_seconds = heartbeat_interval_seconds self._logger = logger # Guards every field below. Each request and each job runs in its own - # thread. The condition wakes idle() when a job or a report ends. + # thread. The condition wakes idle() and the terminate hook when a + # job, a report, a run hook, or a suspend call ends. self._lock = threading.Lock() self._changed = threading.Condition(self._lock) # The callback IDs of the jobs that run now. self._jobs: set[str] = set() # The callback IDs of jobs that ended, oldest first. self._finished: OrderedDict[str, None] = OrderedDict() - # Callback failure reports in flight, for invalid requests. + # Callback failure reports in flight, from every place that fails a + # callback. The terminate hook waits for them, because Lambda can end + # the MicroVM as soon as it gets the answer. self._reports = 0 # Set by the run hook. A job request that arrives first sets it too. self._microvm_id = _UNKNOWN_MICROVM_ID # Lambda sends one run hook per MicroVM. The MicroVM endpoint also # forwards a request to the hook path from any caller that holds an - # auth token for the MicroVM. So the listener acts only on the first - # valid run hook, and ignores every later one. + # auth token for the MicroVM. SuspendMicrovm authorizes on the image, + # so a later run hook could make the worker suspend another MicroVM + # of the same image. So the listener acts only on the first valid run + # hook, and ignores every later one. self._run_hook_accepted = False + # Run hook requests that are still being handled. The terminate hook + # waits for them, because a run hook whose body is still arriving can + # name a job whose callback it then fails. + self._run_hooks = 0 + # The jobs whose handlers have not settled yet, by callback ID. + self._active: dict[str, _ActiveJob] = {} + # The jobs whose handlers have settled and whose outcome is being + # reported. The terminate hook waits for them. + self._reporting: set[str] = set() + # Set from the run hook payload of a session that suspends when idle. + self._auto_suspend: _AutoSuspend | None = None + # Each timer action carries the generation it was started with, and + # acts only when it is still current. Cancelling a timer increments + # the generation. + self._cancel_idle_timer: Callable[[], None] | None = None + self._idle_generation = 0 + self._cancel_grace_timer: Callable[[], None] | None = None + self._grace_generation = 0 + # Set from the moment the idle time ends, or a suspend hook arrives, + # until the MicroVM runs again. While set, the listener refuses jobs. + self._suspending = False + # Counts the suspend hooks from Lambda. A suspend call of the + # worker's own that the service rejected must not end a refusal that + # a suspend hook started meanwhile. + self._suspend_hooks = 0 + # SuspendMicrovm calls in flight. idle() waits for them. + self._suspend_calls = 0 + # Run hook jobs that arrived during a refusal. The durable function + # cannot deliver a run hook job again. So the listener starts it when + # the refusal ends, instead of refusing it. + self._deferred: list[tuple[MicrovmJobDocument, str]] = [] + # Set by the terminate hook. The MicroVM is about to end, so no job + # can start again. + self._terminating = False + # Set when the first terminate hook has finished its work. Every + # later terminate hook waits for it too. + self._terminated: threading.Event | None = None + self._closed = False @classmethod def create( @@ -231,6 +413,7 @@ def create( handler: MicrovmJobHandler | None = None, routes: Mapping[str, MicrovmJobHandler] | None = None, create_client: Callable[[str], LambdaClient] | None = None, + create_microvms_client: Callable[[str], LambdaMicrovmsClient] | None = None, heartbeat_interval_seconds: float | None = None, logger: MicrovmWorkerLogger | None = None, ) -> MicrovmWorkerListener: @@ -249,6 +432,13 @@ def create( client that this function returns. Defaults to a client of its own per job, from the default credential chain. The listener closes that client when the job ends. + create_microvms_client: Creates the ``lambda-microvms`` client + that suspends this MicroVM when it is idle. The listener + calls it only when a session asks the worker to suspend when + idle, once per suspend. It never closes a client that this + function returns. Defaults to a client of its own per + suspend, from the default credential chain. The listener + closes that client after the call. heartbeat_interval_seconds: The heartbeat interval. Defaults to a third of the job's heartbeat timeout, at most 15 minutes. It is cut to a third of the job's heartbeat timeout. @@ -287,6 +477,7 @@ def create( handler=handler, routes=routes, create_client=create_client, + create_microvms_client=create_microvms_client, heartbeat_interval_seconds=heartbeat_interval_seconds, # A logger that raises must not end a request thread or a job # thread. A job could then never report its outcome. @@ -366,7 +557,8 @@ def handle_expect_100(self, request: BaseHTTPRequestHandler) -> bool: def idle(self, timeout: float | None = None) -> bool: """Wait until every job started so far has ended. - The wait also covers the failure reports of invalid requests. + The wait also covers the failure reports of invalid requests, and a + running SuspendMicrovm call. That call takes at most 10 seconds. Args: timeout: The longest wait in seconds. ``None`` waits until the @@ -377,17 +569,62 @@ def idle(self, timeout: float | None = None) -> bool: """ with self._changed: return self._changed.wait_for( - lambda: not self._jobs and self._reports == 0, timeout + lambda: ( + not self._jobs and self._reports == 0 and self._suspend_calls == 0 + ), + timeout, ) + def close(self) -> None: + """Stop the idle time. The listener then never suspends the MicroVM again. + + Running jobs continue. A refusal that a suspend started still ends on + the ``resume`` hook or after 30 seconds. + """ + with self._lock: + self._closed = True + self._stop_idle_time() + def _dispatch(self, exchange: _Exchange) -> None: path = exchange.path if path.startswith(HOOK_PATH_PREFIX) and exchange.method == "POST": hook = path[len(HOOK_PATH_PREFIX) :] if hook == "run": - self._handle_run_hook(exchange) + with self._lock: + self._run_hooks += 1 + try: + self._handle_run_hook(exchange) + finally: + with self._changed: + self._run_hooks -= 1 + self._changed.notify_all() + return + if hook in {"suspend", "terminate", "resume"} and not _is_lambda_hook_call( + exchange + ): + # A caller of the endpoint, not Lambda, sent this request. + # Acting on it would refuse or fail jobs that Lambda did not + # stop, or accept jobs while the MicroVM is about to freeze. + with self._lock: + microvm_id = self._microvm_id + self._logger.warning( + "lifecycle hook ignored: it did not come from Lambda", + extra={"hook": hook, "microvmId": microvm_id}, + ) + exchange.respond(200, {}) return self._logger.info("lifecycle hook", extra={"hook": hook}) + if hook == "terminate": + self._handle_terminate_hook(exchange) + return + # The state changes before the answer. A request thread could + # otherwise accept a job between the answer and the change. The + # change only takes the lock, so the answer is not delayed. A + # suspend hook that runs past its timeout terminates the MicroVM. + if hook == "suspend": + self._handle_suspend_hook() + elif hook == "resume": + self._handle_resume_hook() exchange.respond(200, {}) return handler = self._job_handler(path) @@ -396,6 +633,110 @@ def _dispatch(self, exchange: _Exchange) -> None: return exchange.respond(404, {"error": "not found"}) + def _handle_suspend_hook(self) -> None: + """Refuse jobs until the ``resume`` hook, or for 30 seconds. + + Lambda freezes the MicroVM after the answer. A job accepted now would + freeze with it. + """ + with self._lock: + self._suspend_hooks += 1 + self._stop_idle_time() + self._suspending = True + self._start_grace() + + def _handle_resume_hook(self) -> None: + """Accept jobs again, and start the idle time. + + A resume without a job would otherwise leave the MicroVM running + until the next job ends. + """ + with self._lock: + if self._suspending: + deferred = self._end_suspending() + else: + deferred = [] + self._start_idle_time() + self._launch_deferred(deferred) + + def _handle_terminate_hook(self, exchange: _Exchange) -> None: + """Fail the callback of every running job, then answer the hook. + + A second terminate hook can arrive while the first one waits. The + first one has already taken the running jobs, so the second would + find nothing to wait for and answer early. So every terminate hook + waits for the first one's work. + """ + with self._lock: + first = self._terminated is None + if self._terminated is None: + self._terminated = threading.Event() + terminated = self._terminated + if first: + try: + self._terminate() + finally: + terminated.set() + else: + terminated.wait() + exchange.respond(200, {}) + + def _terminate(self) -> None: + """Stop the running jobs, fail their callbacks, and wait for the reports. + + Lambda terminates the MicroVM after the answer, or at the hook's + timeout. A running job cannot finish. So the listener sets each job's + ``cancelled`` event, stops its heartbeats, and fails its callback + with :class:`MicrovmTerminatedError`. The durable function then fails + the job at once, instead of at its heartbeat timeout or its timeout. + """ + with self._lock: + self._terminating = True + self._stop_idle_time() + stopped: list[tuple[str, _ActiveJob]] = [] + for callback_id, job in self._active.items(): + if job.cancelled.is_set(): + # A heartbeat found the callback gone. A report would + # fail with the same terminal error. + continue + job.terminated = True + job.cancelled.set() + stopped.append((callback_id, job)) + # The jobs stay registered. Each one's own thread sees the mark, + # and reports nothing. Only the first terminate hook runs this. + deferred = self._deferred + self._deferred = [] + microvm_id = self._microvm_id + for callback_id, job in stopped: + heartbeats = job.heartbeats + self._fail_in_background( + MicrovmTerminatedError(), + callback_id, + job.region, + before=None if heartbeats is None else heartbeats.stop, + ) + for deferred_job, region in deferred: + self._fail_in_background( + MicrovmTerminatedError(), deferred_job.callback_id, region + ) + failed = len(stopped) + len(deferred) + if failed: + self._logger.warning( + "failing running jobs: the MicroVM is terminating", + extra={"microvmId": microvm_id, "jobs": failed}, + ) + # A run hook whose body is still arriving can name a job, and then + # fails its callback. An invalid request fails its callback in the + # background. Either can begin during the wait. The condition is + # checked again on every change, so the wait includes them. + with self._changed: + self._changed.wait_for( + lambda: ( + self._reports == 0 and self._run_hooks == 0 and not self._reporting + ), + TERMINATE_REPORT_BUDGET_SECONDS, + ) + def _job_handler(self, path: str) -> MicrovmJobHandler | None: """The handler for a job route, or ``None`` for another path.""" return self._handler if path == MICROVM_JOB_PATH else self._routes.get(path) @@ -423,17 +764,48 @@ def _handle_run_hook(self, exchange: _Exchange) -> None: job = None if payload is None else payload.job # The checks and the job's registration run under one lock. Requests # run in parallel threads, so two deliveries of one callback ID could - # otherwise both start the job. + # otherwise both start the job. The body is read before these checks. + # So a suspend or terminate hook can arrive while it is still being + # read. with self._lock: later = self._run_hook_accepted if not later: self._run_hook_accepted = True self._microvm_id = parsed.microvm_id + if ( + payload is not None + and payload.auto_suspend_idle_seconds is not None + ): + self._auto_suspend = _AutoSuspend( + idle_seconds=payload.auto_suspend_idle_seconds, + region=payload.region, + microvm_id=parsed.microvm_id, + ) microvm_id = self._microvm_id duplicate = job is not None and self._is_known(job.callback_id) new_job = None if later or duplicate else job - if new_job is not None and self._handler is not None: - self._jobs.add(new_job.callback_id) + # What happens to a new job with a handler: it starts now, it is + # failed because the MicroVM is terminating, or it starts when + # the refusal of a suspend ends. + terminated = False + deferred = False + if ( + payload is not None + and new_job is not None + and self._handler is not None + ): + if self._terminating: + terminated = True + elif self._suspending: + deferred = True + self._deferred.append((new_job, payload.region)) + else: + self._register_job(new_job.callback_id, payload.region) + elif not later and not duplicate and payload is not None: + # No run hook job. A session's HTTP jobs can arrive before + # or after this hook. The idle time does not start while one + # runs, and each one starts it again when it ends. + self._start_idle_time() if payload is not None and new_job is not None and self._handler is None: # The report starts before the answer, so that idle() covers it @@ -446,6 +818,30 @@ def _handle_run_hook(self, exchange: _Exchange) -> None: new_job.callback_id, payload.region, ) + if payload is not None and new_job is not None and terminated: + # The job could not finish. The durable function cannot deliver a + # run hook job again. So the listener fails its callback, as for a + # job that was running when the terminate hook arrived. The report + # starts before the answer, so that the terminate hook waits for + # it. + self._logger.info( + "run hook job not started: the MicroVM is terminating", + extra={"microvmId": microvm_id, "callbackId": new_job.callback_id}, + ) + self._fail_in_background( + MicrovmTerminatedError(), new_job.callback_id, payload.region + ) + exchange.respond(200, {}) + return + if new_job is not None and deferred: + # A job started now would freeze with the MicroVM. The listener + # starts it when the refusal ends. + self._logger.info( + "run hook job deferred: the MicroVM is suspending", + extra={"microvmId": microvm_id, "callbackId": new_job.callback_id}, + ) + exchange.respond(200, {}) + return start_error: Exception | None = None if payload is not None and new_job is not None and self._handler is not None: # The thread starts before the answer, so that idle() covers the @@ -511,9 +907,22 @@ def _handle_job_request( # arrives. self._microvm_id = request.microvm_id microvm_id = self._microvm_id + terminating = self._terminating + suspending = self._suspending duplicate = self._is_known(job.callback_id) - if not duplicate: - self._jobs.add(job.callback_id) + if not (terminating or suspending or duplicate): + self._register_job(job.callback_id, request.region) + if terminating or suspending: + # Terminating: the MicroVM is about to end, and the job could not + # finish. Suspending: the MicroVM is about to freeze. The durable + # function resumes it and delivers the job again. + state = "terminating" if terminating else "suspending" + self._logger.info( + f"job refused: the MicroVM is {state}", + extra={"microvmId": microvm_id, "callbackId": job.callback_id}, + ) + exchange.respond(503, {"error": f"the MicroVM is {state}"}) + return if duplicate: exchange.respond(202, {}) self._logger.warning( @@ -559,10 +968,24 @@ def _is_known(self, callback_id: str) -> bool: """ return callback_id in self._jobs or callback_id in self._finished + def _register_job(self, callback_id: str, region: str) -> None: + """Register a job that is about to start. The caller holds the lock. + + The job is registered for the terminate hook before its thread + starts. So a terminate hook that arrives before the thread runs + still fails the job's callback, and the thread does not run the + handler. + """ + self._jobs.add(callback_id) + self._active[callback_id] = _ActiveJob( + region=region, cancelled=threading.Event() + ) + self._stop_idle_time() + def _start_job( self, handler: MicrovmJobHandler, job: MicrovmJobDocument, region: str ) -> None: - """Run a job in its own thread. The caller registered it in ``_jobs``.""" + """Run a job in its own thread. The caller registered it.""" thread = threading.Thread( target=self._run_job, args=(handler, job, region), @@ -578,6 +1001,8 @@ def _start_job( # sends it only once. with self._changed: self._jobs.discard(job.callback_id) + self._active.pop(job.callback_id, None) + self._start_idle_time() self._changed.notify_all() raise @@ -594,9 +1019,11 @@ def _run_job( finally: with self._changed: self._jobs.discard(job.callback_id) + self._active.pop(job.callback_id, None) self._finished[job.callback_id] = None if len(self._finished) > MAX_FINISHED_JOBS: self._finished.popitem(last=False) + self._start_idle_time() self._changed.notify_all() def _run_job_and_report( @@ -605,6 +1032,10 @@ def _run_job_and_report( callback_id = job.callback_id with self._lock: microvm_id = self._microvm_id + # The job was registered before this thread started, and only + # this thread removes it before the handler settles. The + # terminate hook marks it instead. + active = self._active[callback_id] fields = {"microvmId": microvm_id, "callbackId": callback_id} try: reporter = self._reporter_for(callback_id, region) @@ -617,20 +1048,34 @@ def _run_job_and_report( ) return try: - cancelled = threading.Event() + cancelled = active.cancelled gone: list[BaseException] = [] def on_callback_gone(error: BaseException) -> None: gone.append(error) cancelled.set() - heartbeats = Heartbeats.start( - reporter, - heartbeat_timeout_seconds=job.heartbeat_timeout_seconds, - on_callback_gone=on_callback_gone, - heartbeat_interval_seconds=self._heartbeat_interval_seconds, - logger=self._logger, - ) + # The heartbeats start under the lock. So the terminate hook + # either finds them and stops them, or has already marked the job, + # and the handler does not run. + with self._lock: + started = not active.terminated + if started: + active.heartbeats = Heartbeats.start( + reporter, + heartbeat_timeout_seconds=job.heartbeat_timeout_seconds, + on_callback_gone=on_callback_gone, + heartbeat_interval_seconds=self._heartbeat_interval_seconds, + logger=self._logger, + ) + heartbeats = active.heartbeats + if not started or heartbeats is None: + # The terminate hook arrived before the job started, and + # fails the callback itself. + self._logger.info( + "job not started: the MicroVM is terminating", extra=fields + ) + return try: self._logger.info("job started", extra=fields) context = MicrovmJobContext( @@ -645,8 +1090,25 @@ def on_callback_gone(error: BaseException) -> None: # completion attempt can stall for up to 30 seconds. Without # heartbeats, a short heartbeat timeout would expire before # the retry. - if heartbeats.handler_settled(): - self._report(reporter, outcome, fields) + settled = heartbeats.handler_settled() + with self._lock: + self._active.pop(callback_id, None) + terminated = active.terminated + if settled and not terminated: + self._reporting.add(callback_id) + if terminated: + # The terminate hook has already failed the callback. + self._logger.info( + "job outcome not reported: the MicroVM is terminating", + extra=fields, + ) + elif settled: + try: + self._report(reporter, outcome, fields) + finally: + with self._changed: + self._reporting.discard(callback_id) + self._changed.notify_all() else: # The heartbeat thread decides under a lock, and calls # on_callback_gone after it. So the reason can still be @@ -700,14 +1162,23 @@ def _report( ) def _fail_in_background( - self, error: BaseException, callback_id: str, region: str + self, + error: BaseException, + callback_id: str, + region: str, + *, + before: Callable[[], None] | None = None, ) -> None: - """Fail a callback in its own thread. :meth:`idle` waits for it.""" + """Fail a callback in its own thread. + + :meth:`idle` and the terminate hook wait for it. ``before`` runs + first in that thread, for example to stop the job's heartbeats. + """ with self._lock: self._reports += 1 thread = threading.Thread( target=self._send_failure, - args=(error, callback_id, region), + args=(error, callback_id, region, before), name=f"microvm-failure-{callback_id[:16]}", daemon=True, ) @@ -718,11 +1189,17 @@ def _fail_in_background( raise def _send_failure( - self, error: BaseException, callback_id: str, region: str + self, + error: BaseException, + callback_id: str, + region: str, + before: Callable[[], None] | None = None, ) -> None: fields = {"callbackId": callback_id} reporter: CallbackReporter | None = None try: + if before is not None: + before() reporter = self._reporter_for(callback_id, region) reporter.fail(error) except Exception as report_error: # noqa: BLE001 @@ -764,6 +1241,248 @@ def _reporter_for(self, callback_id: str, region: str) -> CallbackReporter: callback_id, self._create_client(region), logger=self._logger ) + def _start_idle_time(self) -> None: + """Start the idle time, when the session asked for it and no job runs. + + The caller holds the lock. + """ + self._stop_idle_time() + auto_suspend = self._auto_suspend + if ( + self._closed + or self._terminating + or self._suspending + or auto_suspend is None + or self._jobs + ): + return + generation = self._idle_generation + try: + self._cancel_idle_timer = _start_timer( + auto_suspend.idle_seconds, lambda: self._idle_time_ended(generation) + ) + except Exception as error: # noqa: BLE001 + # Without a timer thread, the MicroVM keeps running until a job + # ends and the idle time starts again. + self._logger.error( + "could not start the idle time", extra={"error": describe(error)} + ) + + def _stop_idle_time(self) -> None: + """Cancel the idle time. The caller holds the lock.""" + self._idle_generation += 1 + if self._cancel_idle_timer is not None: + self._cancel_idle_timer() + self._cancel_idle_timer = None + + def _idle_time_ended(self, generation: int) -> None: + with self._lock: + auto_suspend = self._auto_suspend + # A job that started after the timer was set cancels it. This + # check covers a job that started while the timer fired. + if ( + generation != self._idle_generation + or auto_suspend is None + or self._jobs + or self._closed + or self._terminating + or self._suspending + ): + return + self._cancel_idle_timer = None + self._suspending = True + self._suspend_calls += 1 + try: + threading.Thread( + target=self._suspend_self, + args=(auto_suspend,), + name="microvm-suspend", + daemon=True, + ).start() + except Exception as error: # noqa: BLE001 + # No call was made. So jobs can run again at once. + self._logger.error( + "could not start the suspend of the idle MicroVM", + extra={"microvmId": auto_suspend.microvm_id, "error": describe(error)}, + ) + with self._changed: + self._suspend_calls -= 1 + deferred = self._stop_refusing() + self._changed.notify_all() + self._launch_deferred(deferred) + + def _suspend_self(self, auto_suspend: _AutoSuspend) -> None: + try: + self._call_suspend(auto_suspend) + finally: + with self._changed: + self._suspend_calls -= 1 + self._changed.notify_all() + + def _call_suspend(self, auto_suspend: _AutoSuspend) -> None: + target = auto_suspend.microvm_id + fields = {"microvmId": target} + # The listener closes only a client that it created itself. A client + # from create_microvms_client belongs to the caller. + owns_client = self._create_microvms_client is None + try: + client = ( + create_microvms_client(auto_suspend.region) + if self._create_microvms_client is None + else self._create_microvms_client(auto_suspend.region) + ) + except Exception as error: # noqa: BLE001 + # No call was made. So jobs can run again at once. + self._logger.error( + "could not create the Lambda MicroVMs client", + extra={**fields, "error": describe(error)}, + ) + with self._lock: + deferred = self._stop_refusing() + self._launch_deferred(deferred) + return + self._logger.info( + "suspending the idle MicroVM", + extra={**fields, "idleSeconds": auto_suspend.idle_seconds}, + ) + with self._lock: + hooks_before = self._suspend_hooks + try: + suspend_microvm(client, target, SUSPEND_CALL_TIMEOUT_SECONDS) + except Exception as error: # noqa: BLE001 + if is_rejected_call(error): + self._logger.warning( + "could not suspend the idle MicroVM, because the service " + "rejected the call. It keeps running. A missing permission " + "needs lambda:SuspendMicrovm on the MicroVM image in its " + "execution role.", + extra={**fields, "error": describe(error)}, + ) + # The service did nothing. So jobs can run again at once. The + # idle time does not start again, so a missing permission is + # not retried. A suspend hook that arrived during the call + # means that the MicroVM is suspending anyway, so its refusal + # stays. + with self._lock: + deferred = ( + self._stop_refusing() + if self._suspend_hooks == hooks_before + else [] + ) + self._launch_deferred(deferred) + return + # A timeout, a throttle, a conflict, a server error, or a + # rejection after a retried attempt does not show whether the + # service suspended the MicroVM. So the refusal continues, as + # after a successful call. + self._logger.warning( + "the suspend call for the idle MicroVM had an unknown outcome. " + "The worker refuses jobs until the MicroVM resumes, or for 30 " + "seconds.", + extra={**fields, "error": describe(error)}, + ) + finally: + if owns_client: + # Each suspend creates its own client. Closing it releases + # its connections, as the callback reporter does. + try: + client.close() + except Exception as close_error: # noqa: BLE001 + self._logger.error( + "could not close the Lambda MicroVMs client", + extra={**fields, "error": describe(close_error)}, + ) + with self._lock: + # A resume hook that arrived during the call has already ended + # the refusal. The grace timer also runs after close(), so a + # closed listener does not refuse jobs forever. + if self._suspending: + self._start_grace() + + def _start_grace(self) -> None: + """End the refusal after the grace period, unless the ``resume`` hook + ends it first. A timer started earlier is replaced. + + The caller holds the lock. + """ + self._cancel_grace() + generation = self._grace_generation + try: + self._cancel_grace_timer = _start_timer( + SUSPEND_GRACE_SECONDS, lambda: self._grace_ended(generation) + ) + except Exception as error: # noqa: BLE001 + # Without a timer thread, only the resume hook ends the refusal. + self._logger.error( + "could not start the end of the refusal", + extra={"error": describe(error)}, + ) + + def _cancel_grace(self) -> None: + """The caller holds the lock.""" + self._grace_generation += 1 + if self._cancel_grace_timer is not None: + self._cancel_grace_timer() + self._cancel_grace_timer = None + + def _grace_ended(self, generation: int) -> None: + with self._lock: + if generation != self._grace_generation: + return + self._cancel_grace_timer = None + deferred = self._end_suspending() + self._launch_deferred(deferred) + + def _end_suspending(self) -> list[tuple[MicrovmJobDocument, str]]: + """End the refusal of jobs, and start the idle time again. + + The caller holds the lock, and the listener is suspending. The caller + starts the returned jobs after it releases the lock. + """ + deferred = self._stop_refusing() + self._start_idle_time() + return deferred + + def _stop_refusing(self) -> list[tuple[MicrovmJobDocument, str]]: + """End the refusal without starting the idle time. + + The caller holds the lock. It returns the deferred run hook jobs, + registered, for :meth:`_launch_deferred`. + """ + self._cancel_grace() + self._suspending = False + deferred = [ + (job, region) + for job, region in self._deferred + if not self._is_known(job.callback_id) + ] + self._deferred = [] + for job, region in deferred: + self._register_job(job.callback_id, region) + return deferred + + def _launch_deferred(self, deferred: list[tuple[MicrovmJobDocument, str]]) -> None: + """Start the run hook jobs that arrived during a refusal. + + The caller has registered them, and does not hold the lock. + """ + handler = self._handler + if handler is None: # pragma: no cover - only handler jobs are deferred + return + for job, region in deferred: + try: + self._start_job(handler, job, region) + except Exception as error: # noqa: BLE001 + # Lambda sends the run hook once, so no later delivery comes. + # So this thread fails the callback itself. + self._logger.error( + "could not start the run hook job", + extra={"callbackId": job.callback_id, "error": describe(error)}, + ) + with self._lock: + self._reports += 1 + self._send_failure(error, job.callback_id, region) + class MicrovmWorker: """A running worker server. @@ -806,7 +1525,8 @@ def port(self) -> int: return int(self._server.server_address[1]) def idle(self, timeout: float | None = None) -> bool: - """Wait until every job started so far has ended. + """Wait until every job started so far has ended, and a running + SuspendMicrovm call has returned. Returns: ``True`` when no job runs. ``False`` when the timeout ended first. @@ -827,13 +1547,15 @@ def wait(self, timeout: float | None = None) -> bool: return self._closed.wait(timeout) def close(self) -> None: - """Stop accepting requests. Running jobs continue. + """Stop accepting requests, and stop the idle time. Running jobs continue. - Call :meth:`idle` first to let the running jobs report their outcome. + The worker then never suspends the MicroVM again. Call :meth:`idle` + first to let the running jobs report their outcome. """ with self._lock: if self._closed.is_set(): return + self._listener.close() self._server.shutdown() self._server.server_close() self._closed.set() @@ -844,6 +1566,7 @@ def start_microvm_worker( handler: MicrovmJobHandler | None = None, routes: Mapping[str, MicrovmJobHandler] | None = None, create_client: Callable[[str], LambdaClient] | None = None, + create_microvms_client: Callable[[str], LambdaMicrovmsClient] | None = None, heartbeat_interval_seconds: float | None = None, logger: MicrovmWorkerLogger | None = None, port: int = DEFAULT_PORT, @@ -871,6 +1594,7 @@ def start_microvm_worker( handler=handler, routes=routes, create_client=create_client, + create_microvms_client=create_microvms_client, heartbeat_interval_seconds=heartbeat_interval_seconds, logger=logger, ) @@ -998,6 +1722,10 @@ def method(self) -> str: def path(self) -> str: return self._request.path.split("?", 1)[0] + def header(self, name: str) -> str | None: + """A request header, or ``None``. The name is not case-sensitive.""" + return self._request.headers.get(name) + def declared_length(self) -> int | None: """The body length from ``Content-Length``. @@ -1173,6 +1901,40 @@ def _ended_early() -> ConnectionError: return ConnectionError("the connection ended before the request body did") +def _is_lambda_hook_call(exchange: _Exchange) -> bool: + """Whether a lifecycle hook request comes from Lambda, not from a caller + of the MicroVM endpoint. + + The endpoint forwards a request to a hook path from any caller that holds + an auth token for the MicroVM. Probes in us-east-1 found two differences: + + 1. Lambda's own ``run``, ``resume``, ``suspend``, and ``terminate`` calls + carry ``Host: localhost:`` and no ``x-amzn-requestid`` header. + 2. A forwarded request carries the endpoint's host name, and an + ``x-amzn-requestid`` that the endpoint sets. The endpoint rejected a + request that set ``Host: localhost``, and it replaced a request ID + that the caller set. + + The Lambda documentation does not state either behavior. So the worker + uses this check only where a false rejection costs little: + + - ``suspend`` and ``terminate``: the worker would ignore them, which is + what it did before it acted on them. + - ``resume``: the refusal after a suspend would end after its 30-second + limit instead of at once. A forwarded ``resume`` in the time between a + suspend hook and the freeze would otherwise end the refusal early, and + a job accepted then would freeze with the MicroVM. + + The ``run`` hook does not use the check, because a false rejection there + would lose every job. + """ + host = (exchange.header("Host") or "").strip().lower() + return ( + _LOCAL_HOST.fullmatch(host) is not None + and exchange.header("x-amzn-requestid") is None + ) + + class _Server(ThreadingHTTPServer): """The worker's HTTP server.""" diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py new file mode 100644 index 000000000..2fe1b3c58 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -0,0 +1,1771 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the suspend, resume, and terminate hooks, and for auto-suspend. + +The requests go over HTTP to a real worker. The idle time and the refusal +after a suspend run on a fake clock, which each test advances. +""" + +from __future__ import annotations + +import http.client +import json +import socket +import threading +import time +from collections.abc import Callable, Iterator +from concurrent.futures import Future +from dataclasses import dataclass, field +from typing import Any + +import pytest +from botocore.exceptions import ClientError, NoCredentialsError + +from aws_durable_execution_sdk_python_microvm_worker import ( + HOOK_PATH_PREFIX, + MicrovmJobContext, + MicrovmTerminatedError, + MicrovmWorker, + MicrovmWorkerListener, + start_microvm_worker, +) +from aws_durable_execution_sdk_python_microvm_worker import worker as worker_module + + +WAIT = 10.0 +IDLE = 10.0 +GRACE = 30.0 +TERMINATED_MESSAGE = "the MicroVM was terminated before the job finished" + +LAMBDA_HEADERS = {"Host": "localhost:8080"} +FORWARDED_HEADERS = { + "Host": "7ed0892b.lambda-microvm.us-east-1.on.aws", + "x-amzn-requestid": "6f1c0e0a-0000-4000-8000-000000000000", +} + + +def eventually(condition: Callable[[], bool]) -> bool: + deadline = time.monotonic() + WAIT + while not condition(): + if time.monotonic() > deadline: + return False + time.sleep(0.005) + return True + + +def suspend_error(code: str, status: int, retries: int = 0) -> ClientError: + response: Any = { + "Error": {"Code": code, "Message": f"{code} message"}, + "ResponseMetadata": {"HTTPStatusCode": status, "RetryAttempts": retries}, + } + return ClientError(response, "SuspendMicrovm") + + +class FakeTimers: + """Replaces the listener's timers with a clock that the test advances. + + ``advance`` runs every due action in the test's thread, in time order. + The clock counts whole milliseconds, so that sums of seconds compare + exactly. + """ + + def __init__(self) -> None: + self.now = 0 + self.failure: BaseException | None = None + # Fire cancelled timers too. A threading.Timer whose action has + # begun runs it although cancel() was called. + self.fire_cancelled = False + self._lock = threading.Lock() + self._timers: list[_FakeTimer] = [] + + def start(self, delay: float, action: Callable[[], None]) -> Callable[[], None]: + if self.failure is not None: + raise self.failure + with self._lock: + timer = _FakeTimer(due=self.now + round(delay * 1000), action=action) + self._timers.append(timer) + return timer.cancel + + def advance(self, seconds: float) -> None: + target = self.now + round(seconds * 1000) + while True: + with self._lock: + due = [ + t + for t in self._timers + if (self.fire_cancelled or not t.cancelled) and t.due <= target + ] + if not due: + self.now = target + return + timer = min(due, key=lambda t: t.due) + self._timers.remove(timer) + self.now = timer.due + timer.action() + + +@dataclass +class _FakeTimer: + due: int + action: Callable[[], None] + cancelled: bool = False + + def cancel(self) -> None: + self.cancelled = True + + +class FakeMicrovmsClient: + """Records SuspendMicrovm calls. ``hold`` keeps each call open until + ``finish()``. ``failure`` answers every call with an error.""" + + def __init__(self) -> None: + self.suspended: list[str] = [] + self.failure: BaseException | None = None + self.hold = False + self.closed = False + self._release = threading.Event() + self._lock = threading.Lock() + + def suspend_microvm(self, *, microvmIdentifier: str) -> dict[str, Any]: # noqa: N803 + with self._lock: + self.suspended.append(microvmIdentifier) + if self.hold: + self._release.wait(WAIT) + if self.failure is not None: + raise self.failure + return {} + + def finish(self) -> None: + self._release.set() + + def close(self) -> None: + self.closed = True + + +class Jobs: + """A handler whose jobs each run until the test releases them.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._releases: dict[str, threading.Event] = {} + self.contexts: dict[str, MicrovmJobContext] = {} + self.order: list[str] = [] + + def _release(self, callback_id: str) -> threading.Event: + with self._lock: + return self._releases.setdefault(callback_id, threading.Event()) + + def __call__(self, job_input: Any, context: MicrovmJobContext) -> Any: + with self._lock: + self.order.append(context.callback_id) + self.contexts[context.callback_id] = context + self._release(context.callback_id).wait(WAIT) + return "done" + + def started(self) -> list[str]: + with self._lock: + return list(self.order) + + def release(self, callback_id: str) -> None: + self._release(callback_id).set() + + def release_all(self) -> None: + with self._lock: + releases = list(self._releases.values()) + for release in releases: + release.set() + + +class HeldLambdaClient: + """Records callback completions. While ``hold`` is set, each completion + waits until ``finish()`` or ``finish_first()``.""" + + def __init__(self, heartbeat_error: BaseException | None = None) -> None: + self.heartbeat_error = heartbeat_error + self.completion_error: BaseException | None = None + self.hold = False + self._lock = threading.Lock() + self._changed = threading.Condition(self._lock) + self._waiting: list[threading.Event] = [] + self.completions: list[dict[str, Any]] = [] + + def send_durable_execution_callback_heartbeat(self, **kwargs: Any) -> Any: + if self.heartbeat_error is not None: + raise self.heartbeat_error + return {} + + def _complete(self, entry: dict[str, Any]) -> Any: + with self._changed: + gate = None + if self.hold: + gate = threading.Event() + self._waiting.append(gate) + self._changed.notify_all() + if gate is not None: + gate.wait(WAIT) + if self.completion_error is not None: + raise self.completion_error + with self._lock: + self.completions.append(entry) + return {} + + def send_durable_execution_callback_success(self, **kwargs: Any) -> Any: + return self._complete({"kind": "success", "callbackId": kwargs["CallbackId"]}) + + def send_durable_execution_callback_failure(self, **kwargs: Any) -> Any: + error = kwargs["Error"] + return self._complete( + { + "kind": "failure", + "callbackId": kwargs["CallbackId"], + "errorType": error["ErrorType"], + "message": error["ErrorMessage"], + } + ) + + def waiting(self) -> int: + with self._lock: + return len(self._waiting) + + def finish_first(self) -> None: + with self._lock: + gate = self._waiting.pop(0) if self._waiting else None + if gate is not None: + gate.set() + + def finish(self) -> None: + with self._lock: + self.hold = False + gates, self._waiting = self._waiting, [] + for gate in gates: + gate.set() + + def recorded(self) -> list[dict[str, Any]]: + with self._lock: + return list(self.completions) + + def close(self) -> None: + return + + +def terminated_failure(callback_id: str) -> dict[str, Any]: + return { + "kind": "failure", + "callbackId": callback_id, + "errorType": "MicrovmTerminatedError", + "message": TERMINATED_MESSAGE, + } + + +@dataclass +class Harness: + worker: MicrovmWorker + lambda_client: HeldLambdaClient + microvms: FakeMicrovmsClient + timers: FakeTimers + jobs: Jobs + logger: Any + _threads: list[threading.Thread] = field(default_factory=list) + + @property + def listener(self) -> MicrovmWorkerListener: + return self.worker._listener # noqa: SLF001 + + def post(self, path: str, body: Any, headers: dict[str, str] | None = None) -> int: + connection = http.client.HTTPConnection( + "127.0.0.1", self.worker.port, timeout=WAIT + ) + try: + connection.request( + "POST", + path, + body=json.dumps(body).encode(), + headers={"Content-Type": "application/json", **(headers or {})}, + ) + response = connection.getresponse() + response.read() + finally: + connection.close() + return response.status + + def run_hook( + self, + idle_seconds: float | None = None, + microvm_id: str = "mvm-1", + job: dict[str, Any] | None = None, + ) -> int: + payload: dict[str, Any] = {"version": 1, "region": "us-east-1"} + if idle_seconds is not None: + payload["autoSuspendIdleSeconds"] = idle_seconds + if job is not None: + payload["job"] = job + return self.post( + f"{HOOK_PATH_PREFIX}run", + {"microvmId": microvm_id, "runHookPayload": json.dumps(payload)}, + ) + + def hook(self, name: str, headers: dict[str, str] | None = None) -> int: + return self.post(f"{HOOK_PATH_PREFIX}{name}", {}, headers) + + def hook_in_background( + self, name: str, headers: dict[str, str] | None = None + ) -> Future[int]: + future: Future[int] = Future() + + def run() -> None: + try: + future.set_result(self.hook(name, headers)) + except BaseException as error: # noqa: BLE001 + future.set_exception(error) + + thread = threading.Thread(target=run, daemon=True) + self._threads.append(thread) + thread.start() + return future + + def job(self, callback_id: str, **fields: Any) -> int: + return self.post( + "/job", + { + "version": 1, + "region": "us-east-1", + "callbackId": callback_id, + "input": {}, + **fields, + }, + ) + + def finish_job(self, callback_id: str) -> None: + """Release a job, and wait until it has ended.""" + self.jobs.release(callback_id) + assert eventually(lambda: callback_id not in self.listener._jobs) # noqa: SLF001 + + def settle(self) -> None: + """Wait until no SuspendMicrovm call runs.""" + assert eventually(lambda: self.listener._suspend_calls == 0) # noqa: SLF001 + + def advance(self, seconds: float) -> None: + """Advance the clock, and let a suspend call that it started end.""" + self.timers.advance(seconds) + if not self.microvms.hold: + self.settle() + + def started(self, callback_id: str) -> bool: + return callback_id in self.jobs.started() + + +@pytest.fixture +def timers(monkeypatch: pytest.MonkeyPatch) -> FakeTimers: + fake = FakeTimers() + monkeypatch.setattr(worker_module, "_start_timer", fake.start) + return fake + + +@pytest.fixture +def start(logger: Any, timers: FakeTimers) -> Iterator[Callable[..., Harness]]: + harnesses: list[Harness] = [] + + def _start( + *, + lambda_client: HeldLambdaClient | None = None, + create_microvms_client: Any = "fake", + **options: Any, + ) -> Harness: + client = lambda_client or HeldLambdaClient() + microvms = FakeMicrovmsClient() + jobs = Jobs() + if create_microvms_client == "fake": + create_microvms_client = lambda region: microvms # noqa: E731 + + def create_client(region: str) -> Any: + return client + + options.setdefault("logger", logger) + worker = start_microvm_worker( + handler=jobs, + routes={"/job": jobs}, + create_client=create_client, + create_microvms_client=create_microvms_client, + port=0, + **options, + ) + harness = Harness( + worker=worker, + lambda_client=client, + microvms=microvms, + timers=timers, + jobs=jobs, + logger=options["logger"], + ) + harnesses.append(harness) + return harness + + yield _start + for harness in harnesses: + harness.listener.close() + harness.jobs.release_all() + harness.lambda_client.finish() + harness.microvms.finish() + timers.advance(GRACE * 10) + harness.worker.idle(WAIT) + harness.worker.close() + for thread in harness._threads: # noqa: SLF001 + thread.join(WAIT) + + +class SlowCall: + """A request whose body arrives in two parts. + + The worker has the headers and half of the body when the constructor + returns. ``finish()`` sends the rest, and returns the status. + """ + + def __init__(self, harness: Harness, path: str, body: Any) -> None: + data = json.dumps(body).encode() + half = len(data) // 2 + self._socket = socket.create_connection( + ("127.0.0.1", harness.worker.port), timeout=WAIT + ) + head = ( + f"POST {path} HTTP/1.1\r\nHost: localhost:8080\r\n" + f"Content-Type: application/json\r\nContent-Length: {len(data)}\r\n\r\n" + ).encode() + self._socket.sendall(head + data[:half]) + self._rest = data[half:] + + def finish(self) -> int: + try: + self._socket.sendall(self._rest) + chunks = [] + while chunk := self._socket.recv(65536): + chunks.append(chunk) + finally: + self._socket.close() + status_line = b"".join(chunks).split(b"\r\n", 1)[0] + return int(status_line.split()[1]) + + def abandon(self) -> None: + self._socket.close() + + +RUN_HOOK_WITH_JOB = { + "microvmId": "mvm-1", + "runHookPayload": json.dumps( + { + "version": 1, + "region": "us-east-1", + "job": {"callbackId": "cb-run", "input": {}}, + } + ), +} + + +def slow_run_hook(harness: Harness) -> SlowCall: + call = SlowCall(harness, f"{HOOK_PATH_PREFIX}run", RUN_HOOK_WITH_JOB) + assert eventually(lambda: harness.listener._run_hooks == 1) # noqa: SLF001 + return call + + +def answered_within(future: Future[int], seconds: float) -> bool: + try: + future.result(seconds) + except TimeoutError: + return False + return True + + +# --------------------------------------------------------------------------- +# Auto-suspend when idle +# --------------------------------------------------------------------------- + + +def test_suspends_its_own_microvm_after_the_idle_time_when_no_job_arrives( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.advance(IDLE - 0.001) + assert h.microvms.suspended == [] + + h.advance(0.001) + assert h.microvms.suspended == ["mvm-1"] + + +def test_closes_the_default_client_after_each_suspend_and_never_a_provided_one( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + created: list[FakeMicrovmsClient] = [] + + def create_default(region: str) -> FakeMicrovmsClient: + client = FakeMicrovmsClient() + created.append(client) + return client + + monkeypatch.setattr(worker_module, "create_microvms_client", create_default) + h = start(create_microvms_client=None) + h.run_hook(IDLE) + h.advance(IDLE) + + assert [client.suspended for client in created] == [["mvm-1"]] + assert created[0].closed + + provided = start() + provided.run_hook(IDLE) + provided.advance(IDLE) + + assert provided.microvms.suspended == ["mvm-1"] + assert not provided.microvms.closed + + +def test_never_suspends_when_the_run_hook_payload_does_not_ask_for_it( + start: Any, +) -> None: + h = start() + + h.run_hook() + h.advance(IDLE * 10) + + assert h.microvms.suspended == [] + + +def test_does_not_suspend_while_a_job_runs_and_suspends_one_idle_time_after_it( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + assert h.job("cb-1") == 202 + h.advance(IDLE * 4) + assert h.microvms.suspended == [] + + h.finish_job("cb-1") + h.advance(IDLE - 0.001) + assert h.microvms.suspended == [] + h.advance(0.001) + assert h.microvms.suspended == ["mvm-1"] + + +def test_suspends_with_the_run_hook_id_after_a_job_that_arrived_first_with_another( + start: Any, +) -> None: + h = start() + + assert h.job("cb-1", microvmId="mvm-from-job") == 202 + h.run_hook(IDLE) + # The run hook arrived while the job runs, so the idle time waits. + h.advance(IDLE * 2) + assert h.microvms.suspended == [] + + h.finish_job("cb-1") + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1"] + + +def test_keeps_the_first_run_hook_target_when_a_later_run_hook_has_another_id( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.run_hook(microvm_id="mvm-2") + h.advance(IDLE) + + assert h.microvms.suspended == ["mvm-1"] + + +def test_ignores_a_later_run_hook_that_asks_to_suspend_another_microvm( + start: Any, +) -> None: + # The endpoint forwards the hook path from any caller with an auth token. + # SuspendMicrovm authorizes on the image, so the later hook must not + # choose the suspend target. + h = start() + + h.run_hook() + assert h.run_hook(IDLE, microvm_id="mvm-2") == 200 + h.advance(IDLE * 2) + + assert h.microvms.suspended == [] + + +def test_accepts_a_valid_run_hook_after_an_invalid_one(start: Any) -> None: + h = start() + + assert h.post(f"{HOOK_PATH_PREFIX}run", {}) == 400 + h.run_hook(IDLE) + h.advance(IDLE) + + assert h.microvms.suspended == ["mvm-1"] + + +def test_starts_the_idle_time_again_when_a_job_arrives_before_it_ends( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.advance(IDLE / 2) + h.job("cb-1") + h.finish_job("cb-1") + # The first idle time would have ended here. + h.advance(IDLE / 2) + assert h.microvms.suspended == [] + + h.advance(IDLE / 2) + assert h.microvms.suspended == ["mvm-1"] + + +def test_does_not_suspend_while_one_of_two_parallel_jobs_still_runs( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.job("cb-1") + h.job("cb-2") + h.finish_job("cb-1") + h.advance(IDLE * 4) + assert h.microvms.suspended == [] + + h.finish_job("cb-2") + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1"] + + +def test_starts_the_idle_time_after_a_run_hook_job(start: Any) -> None: + h = start() + + h.run_hook(IDLE, job={"callbackId": "cb-run", "input": {}}) + assert eventually(lambda: h.started("cb-run")) + h.advance(IDLE * 2) + assert h.microvms.suspended == [] + + h.finish_job("cb-run") + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1"] + + +def test_refuses_a_job_with_503_while_its_suspend_call_runs(start: Any) -> None: + h = start() + h.microvms.hold = True + + h.run_hook(IDLE) + h.advance(IDLE) + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + + assert h.job("cb-1") == 503 + assert h.jobs.started() == [] + h.microvms.finish() + + +def test_refuses_jobs_after_the_call_returned_until_the_resume_hook_arrives( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.advance(IDLE) + assert h.job("cb-1") == 503 + + assert h.hook("resume") == 200 + assert h.job("cb-1") == 202 + assert eventually(lambda: h.jobs.started() == ["cb-1"]) + + # The idle time starts again after the job. + h.finish_job("cb-1") + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +def test_accepts_jobs_again_30_seconds_after_the_call_without_a_resume_hook( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.advance(IDLE) + h.advance(GRACE - 0.001) + assert h.job("cb-1") == 503 + + h.advance(0.001) + assert h.job("cb-1") == 202 + assert eventually(lambda: h.jobs.started() == ["cb-1"]) + + +def test_starts_the_idle_time_again_at_the_end_of_the_30_second_refusal( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.advance(IDLE) + h.advance(GRACE) + assert h.microvms.suspended == ["mvm-1"] + + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +def test_suspends_again_one_idle_time_after_a_resume_that_brings_no_job( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.advance(IDLE) + h.hook("resume") + h.advance(IDLE - 0.001) + assert h.microvms.suspended == ["mvm-1"] + + h.advance(0.001) + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +def test_accepts_jobs_at_once_when_the_resume_hook_arrives_during_the_call( + start: Any, +) -> None: + h = start() + h.microvms.hold = True + + h.run_hook(IDLE) + h.advance(IDLE) + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + h.hook("resume") + assert h.job("cb-1") == 202 + + # The call returns after the resume. It must not start a new refusal. + h.microvms.finish() + h.settle() + assert h.job("cb-2") == 202 + assert eventually(lambda: sorted(h.jobs.started()) == ["cb-1", "cb-2"]) + + +def test_logs_a_rejected_suspend_accepts_jobs_at_once_and_does_not_retry( + start: Any, logger: Any +) -> None: + h = start() + h.microvms.failure = suspend_error("AccessDeniedException", 403) + + h.run_hook(IDLE) + h.advance(IDLE * 6) + + assert h.microvms.suspended == ["mvm-1"] + warnings = [line for line in logger.lines if line[0] == "warning"] + assert len(warnings) == 1 + assert "lambda:SuspendMicrovm" in warnings[0][1] + assert warnings[0][2]["microvmId"] == "mvm-1" + assert warnings[0][2]["error"]["name"] == "ClientError" + assert h.job("cb-1") == 202 + + +def test_retries_the_suspend_one_idle_time_after_a_job_that_follows_a_rejection( + start: Any, +) -> None: + h = start() + h.microvms.failure = suspend_error("AccessDeniedException", 403) + + h.run_hook(IDLE) + h.advance(IDLE) + h.job("cb-1") + h.finish_job("cb-1") + h.advance(IDLE) + + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +@pytest.mark.parametrize( + "failure", + [ + TimeoutError("timed out"), + suspend_error("ThrottlingException", 429), + suspend_error("ThrottlingException", 400), + suspend_error("ConflictException", 409), + suspend_error("InternalServerException", 500), + suspend_error("RequestTimeoutException", 408), + # The client retried a 5xx. The first attempt may have suspended the + # MicroVM, and the retry then fails validation. + suspend_error("ValidationException", 400, retries=1), + ConnectionResetError("reset"), + ], + ids=[ + "timeout", + "throttle", + "throttle-400", + "conflict", + "server-error", + "request-timeout", + "retried-rejection", + "connection-error", + ], +) +def test_keeps_refusing_jobs_after_an_unknown_outcome( + start: Any, logger: Any, failure: BaseException +) -> None: + h = start() + h.microvms.failure = failure + + h.run_hook(IDLE) + h.advance(IDLE) + assert "unknown outcome" in logger.messages("warning")[0] + assert h.job("cb-1") == 503 + + h.advance(GRACE) + assert h.job("cb-1") == 202 + + +@pytest.mark.parametrize( + "failure", + [NoCredentialsError(), suspend_error("ValidationException", 400)], + ids=["missing-credentials", "first-attempt-rejection"], +) +def test_accepts_jobs_at_once_after_a_rejection( + start: Any, failure: BaseException +) -> None: + h = start() + h.microvms.failure = failure + + h.run_hook(IDLE) + h.advance(IDLE) + + assert h.job("cb-1") == 202 + + +def test_treats_a_call_that_outlasts_its_timeout_as_an_unknown_outcome( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "SUSPEND_CALL_TIMEOUT_SECONDS", 0.05) + h = start() + h.microvms.hold = True + + h.run_hook(IDLE) + h.advance(IDLE) + h.settle() + + assert "unknown outcome" in logger.messages("warning")[0] + assert h.job("cb-1") == 503 + h.microvms.finish() + + +def test_accepts_jobs_at_once_when_the_microvms_client_cannot_be_created( + start: Any, logger: Any +) -> None: + def broken(region: str) -> Any: + raise RuntimeError("no client") + + h = start(create_microvms_client=broken) + + h.run_hook(IDLE) + h.advance(IDLE) + + assert "could not create the Lambda MicroVMs client" in logger.messages("error") + assert h.job("cb-1") == 202 + + +def test_logs_a_default_client_that_cannot_be_closed( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + client = FakeMicrovmsClient() + + def close() -> None: + raise RuntimeError("close failed") + + client.close = close # type: ignore[method-assign] + monkeypatch.setattr(worker_module, "create_microvms_client", lambda r: client) + h = start(create_microvms_client=None) + + h.run_hook(IDLE) + h.advance(IDLE) + + assert "could not close the Lambda MicroVMs client" in logger.messages("error") + assert h.job("cb-1") == 503 + + +def test_tries_the_call_again_one_idle_time_after_an_unknown_outcome_refusal( + start: Any, +) -> None: + h = start() + h.microvms.failure = suspend_error("InternalServerException", 500, retries=2) + + h.run_hook(IDLE) + h.advance(IDLE) + h.advance(GRACE + IDLE - 0.001) + assert h.microvms.suspended == ["mvm-1"] + + h.advance(0.001) + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +def test_starts_the_idle_time_on_a_resume_hook_that_no_suspend_preceded( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.job("cb-1") + h.finish_job("cb-1") + h.advance(IDLE / 2) + # An idle policy can resume the MicroVM. The hook starts the idle time + # again from zero. + h.hook("resume") + h.advance(IDLE / 2) + assert h.microvms.suspended == [] + h.advance(IDLE / 2) + assert h.microvms.suspended == ["mvm-1"] + + +def test_never_suspends_after_close(start: Any) -> None: + h = start() + + h.run_hook(IDLE) + h.listener.close() + h.advance(IDLE * 4) + + assert h.microvms.suspended == [] + + +def test_does_not_start_the_idle_time_when_a_job_ends_after_close( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.job("cb-1") + h.listener.close() + h.finish_job("cb-1") + h.advance(IDLE * 4) + + assert h.microvms.suspended == [] + + +def test_ends_a_refusal_30_seconds_after_a_call_that_returned_after_close( + start: Any, +) -> None: + h = start() + h.microvms.hold = True + + h.run_hook(IDLE) + h.advance(IDLE) + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + h.listener.close() + h.microvms.finish() + h.settle() + assert h.job("cb-1") == 503 + + h.advance(GRACE) + assert h.job("cb-1") == 202 + # The closed listener does not suspend again. + h.finish_job("cb-1") + h.advance(IDLE * 4) + assert h.microvms.suspended == ["mvm-1"] + + +def test_ends_a_refusal_30_seconds_after_the_call_when_close_follows_the_call( + start: Any, +) -> None: + h = start() + + h.run_hook(IDLE) + h.advance(IDLE) + # The call has returned, and its grace timer runs. + h.listener.close() + assert h.job("cb-1") == 503 + + h.advance(GRACE) + assert h.job("cb-1") == 202 + + +def test_idle_waits_for_a_running_suspend_call(start: Any) -> None: + h = start() + h.microvms.hold = True + + h.run_hook(IDLE) + h.advance(IDLE) + + assert not h.worker.idle(0.1) + h.microvms.finish() + assert h.worker.idle(WAIT) + + +# --------------------------------------------------------------------------- +# The terminate hook +# --------------------------------------------------------------------------- + + +def test_terminate_fails_each_running_job_and_sets_its_cancelled_event( + start: Any, +) -> None: + h = start() + h.run_hook() + assert h.job("cb-1") == 202 + assert h.job("cb-2") == 202 + assert eventually(lambda: len(h.jobs.started()) == 2) + + assert h.hook("terminate") == 200 + + assert sorted(h.lambda_client.recorded(), key=lambda c: c["callbackId"]) == [ + terminated_failure("cb-1"), + terminated_failure("cb-2"), + ] + assert h.jobs.contexts["cb-1"].cancelled.is_set() + assert h.jobs.contexts["cb-2"].cancelled.is_set() + + +def test_terminate_reports_nothing_more_when_the_handler_returns_after_it( + start: Any, logger: Any +) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + h.hook("terminate") + + h.finish_job("cb-1") + + assert h.lambda_client.recorded() == [terminated_failure("cb-1")] + assert eventually( + lambda: "job outcome not reported: the MicroVM is terminating" + in logger.messages("info") + ) + + +def test_terminate_answers_after_the_reports_end(start: Any) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + h.lambda_client.hold = True + + answer = h.hook_in_background("terminate") + assert eventually(lambda: h.lambda_client.waiting() == 1) + assert not answered_within(answer, 0.2) + + h.lambda_client.finish() + assert answer.result(WAIT) == 200 + assert h.lambda_client.recorded() == [terminated_failure("cb-1")] + + +def test_terminate_answers_after_the_budget_when_a_report_stalls( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "TERMINATE_REPORT_BUDGET_SECONDS", 0.3) + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + h.lambda_client.hold = True + + started = time.monotonic() + assert h.hook("terminate") == 200 + + assert 0.25 <= time.monotonic() - started < WAIT + assert h.lambda_client.recorded() == [] + + +def test_terminate_waits_for_a_job_that_is_reporting_its_own_outcome( + start: Any, +) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + h.lambda_client.hold = True + h.jobs.release("cb-1") + assert eventually(lambda: h.lambda_client.waiting() == 1) + + answer = h.hook_in_background("terminate") + assert not answered_within(answer, 0.2) + + h.lambda_client.finish() + assert answer.result(WAIT) == 200 + assert h.lambda_client.recorded() == [{"kind": "success", "callbackId": "cb-1"}] + + +def test_holds_a_second_concurrent_terminate_until_the_first_one_ends( + start: Any, +) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + h.lambda_client.hold = True + + first = h.hook_in_background("terminate") + assert eventually(lambda: h.lambda_client.waiting() == 1) + second = h.hook_in_background("terminate") + assert not answered_within(second, 0.2) + assert not first.done() + + h.lambda_client.finish() + assert first.result(WAIT) == 200 + assert second.result(WAIT) == 200 + assert h.lambda_client.recorded() == [terminated_failure("cb-1")] + + +def test_terminate_reports_nothing_for_a_job_whose_callback_a_heartbeat_found_gone( + start: Any, make_error: Any +) -> None: + h = start( + lambda_client=HeldLambdaClient( + heartbeat_error=make_error("CallbackTimeoutException") + ) + ) + h.run_hook() + h.job("cb-gone", heartbeatTimeoutSeconds=30) + assert eventually(lambda: h.started("cb-gone")) + assert eventually(lambda: h.jobs.contexts["cb-gone"].cancelled.is_set()) + + assert h.hook("terminate") == 200 + h.finish_job("cb-gone") + + assert h.lambda_client.recorded() == [] + + +def test_terminate_stops_the_heartbeats_of_a_job_before_it_fails_it( + start: Any, +) -> None: + h = start() + h.run_hook() + h.job("cb-1", heartbeatTimeoutSeconds=30) + assert eventually(lambda: h.started("cb-1")) + heartbeats = h.listener._active["cb-1"].heartbeats # noqa: SLF001 + assert heartbeats is not None + + h.hook("terminate") + + assert heartbeats._stopped.is_set() # noqa: SLF001 + assert h.lambda_client.recorded() == [terminated_failure("cb-1")] + + +def test_terminate_logs_a_failure_report_that_the_service_refuses_at_info( + start: Any, logger: Any, make_error: Any +) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + h.lambda_client.completion_error = make_error("CallbackTimeoutException") + + h.hook("terminate") + + assert "the callback no longer accepts a failure" in logger.messages("info") + assert "could not fail the callback" not in logger.messages("error") + + +def test_terminate_answers_at_once_when_no_job_runs(start: Any) -> None: + h = start() + h.run_hook() + + assert h.hook("terminate") == 200 + assert h.lambda_client.recorded() == [] + + +def test_terminate_fails_a_job_once_when_the_hook_arrives_twice(start: Any) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + + h.hook("terminate") + h.hook("terminate") + + assert h.lambda_client.recorded() == [terminated_failure("cb-1")] + + +def test_refuses_jobs_with_503_after_the_terminate_hook(start: Any) -> None: + h = start() + h.run_hook() + h.hook("terminate") + + assert h.job("cb-1") == 503 + assert not h.started("cb-1") + + +def test_never_suspends_after_the_terminate_hook(start: Any) -> None: + h = start() + h.run_hook(IDLE) + h.hook("terminate") + + h.advance(IDLE * 6) + + assert h.microvms.suspended == [] + + +def test_fails_a_job_whose_thread_had_not_run_when_the_terminate_hook_came( + start: Any, logger: Any +) -> None: + # The job is registered before its thread runs. A terminate hook in + # between fails the callback, and the thread then does not run the + # handler. + entered = threading.Event() + release = threading.Event() + lambda_client = HeldLambdaClient() + calls = 0 + lock = threading.Lock() + + def create_client(region: str) -> HeldLambdaClient: + nonlocal calls + with lock: + calls += 1 + first = calls == 1 + if first: + entered.set() + release.wait(WAIT) + return lambda_client + + h = start(lambda_client=lambda_client) + h.listener._create_client = create_client # noqa: SLF001 + h.run_hook() + h.job("cb-1") + assert entered.wait(WAIT) + + assert h.hook("terminate") == 200 + release.set() + assert eventually(lambda: "cb-1" not in h.listener._jobs) # noqa: SLF001 + + assert not h.started("cb-1") + assert h.lambda_client.recorded() == [terminated_failure("cb-1")] + assert "job not started: the MicroVM is terminating" in logger.messages("info") + + +def test_does_not_start_a_job_whose_thread_ran_after_the_terminate_hook( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # The terminate hook can mark the job between its thread's two checks. + h = start() + original = h.listener._reporter_for # noqa: SLF001 + marked = threading.Event() + + def reporter_for(callback_id: str, region: str) -> Any: + if callback_id == "cb-1" and not marked.is_set(): + marked.set() + with h.listener._lock: # noqa: SLF001 + h.listener._active[callback_id].terminated = True # noqa: SLF001 + return original(callback_id, region) + + monkeypatch.setattr(h.listener, "_reporter_for", reporter_for) + h.run_hook() + h.job("cb-1") + + assert eventually(lambda: "cb-1" not in h.listener._jobs) # noqa: SLF001 + assert not h.started("cb-1") + assert "job not started: the MicroVM is terminating" in logger.messages("info") + + +@pytest.mark.parametrize( + "headers", + [ + FORWARDED_HEADERS, + {**LAMBDA_HEADERS, "x-amzn-requestid": "r-1"}, + {"Host": ""}, + ], + ids=["forwarded", "request-id", "empty-host"], +) +def test_ignores_a_terminate_hook_that_did_not_come_from_lambda( + start: Any, logger: Any, headers: dict[str, str] +) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + + assert h.hook("terminate", headers) == 200 + + assert h.lambda_client.recorded() == [] + assert not h.jobs.contexts["cb-1"].cancelled.is_set() + assert h.job("cb-2") == 202 + ignored = [ + line + for line in logger.lines + if line[1] == "lifecycle hook ignored: it did not come from Lambda" + ] + assert ignored[0][2]["hook"] == "terminate" + h.finish_job("cb-1") + assert {"kind": "success", "callbackId": "cb-1"} in h.lambda_client.recorded() + + +def test_ignores_a_terminate_hook_without_a_host_header(start: Any) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + + connection = http.client.HTTPConnection("127.0.0.1", h.worker.port, timeout=WAIT) + try: + connection.putrequest("POST", f"{HOOK_PATH_PREFIX}terminate", skip_host=True) + connection.putheader("Content-Length", "2") + connection.endheaders(b"{}") + response = connection.getresponse() + response.read() + finally: + connection.close() + + assert response.status == 200 + assert h.lambda_client.recorded() == [] + + +@pytest.mark.parametrize( + "host", ["127.0.0.1:8080", "localhost", "[::1]:8080", "LOCALHOST:8080"] +) +def test_acts_on_a_terminate_hook_with_a_local_host(start: Any, host: str) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + + h.hook("terminate", {"Host": host}) + + assert h.lambda_client.recorded() == [terminated_failure("cb-1")] + + +# --------------------------------------------------------------------------- +# A request whose body is still arriving +# --------------------------------------------------------------------------- + + +def test_fails_a_slow_run_hook_job_when_the_terminate_hook_came_first( + start: Any, +) -> None: + h = start() + run = slow_run_hook(h) + + # Lambda can end the MicroVM as soon as the terminate hook answers. So + # the failure must be reported before that answer. + answer = h.hook_in_background("terminate") + assert not answered_within(answer, 0.2) + assert run.finish() == 200 + + assert answer.result(WAIT) == 200 + assert h.lambda_client.recorded() == [terminated_failure("cb-run")] + assert not h.started("cb-run") + + +def test_waits_for_a_run_hook_that_begins_while_the_terminate_hook_waits( + start: Any, +) -> None: + h = start() + # An HTTP job can arrive before the run hook. Its report holds the + # terminate answer open while the run hook begins. + assert h.job("cb-http") == 202 + assert eventually(lambda: h.started("cb-http")) + h.lambda_client.hold = True + + answer = h.hook_in_background("terminate") + assert eventually(lambda: h.lambda_client.waiting() == 1) + run = slow_run_hook(h) + h.lambda_client.finish() + assert not answered_within(answer, 0.2) + assert run.finish() == 200 + + assert answer.result(WAIT) == 200 + assert sorted(c["callbackId"] for c in h.lambda_client.recorded()) == [ + "cb-http", + "cb-run", + ] + + +INVALID_RUN_HOOK = { + "microvmId": "mvm-1", + "runHookPayload": json.dumps( + { + "version": 99, + "region": "us-east-1", + "job": {"callbackId": "cb-invalid", "input": {}}, + } + ), +} +INVALID_JOB = { + "version": 99, + "region": "us-east-1", + "callbackId": "cb-invalid", + "input": {}, +} + + +@pytest.mark.parametrize( + ("path", "body"), + [(f"{HOOK_PATH_PREFIX}run", INVALID_RUN_HOOK), ("/job", INVALID_JOB)], + ids=["run-hook", "job-request"], +) +def test_an_invalid_request_failure_report_holds_the_terminate_answer( + start: Any, path: str, body: Any +) -> None: + h = start() + if path == "/job": + h.run_hook() + h.lambda_client.hold = True + assert h.post(path, body) == 400 + assert eventually(lambda: h.lambda_client.waiting() == 1) + + answer = h.hook_in_background("terminate") + assert not answered_within(answer, 0.2) + + h.lambda_client.finish() + assert answer.result(WAIT) == 200 + assert [c["callbackId"] for c in h.lambda_client.recorded()] == ["cb-invalid"] + + +def test_an_invalid_request_failure_report_holds_the_answer_for_the_budget_only( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "TERMINATE_REPORT_BUDGET_SECONDS", 0.3) + h = start() + h.lambda_client.hold = True + assert h.post(f"{HOOK_PATH_PREFIX}run", INVALID_RUN_HOOK) == 400 + + started = time.monotonic() + assert h.hook("terminate") == 200 + assert time.monotonic() - started >= 0.25 + + +def test_an_invalid_request_that_begins_while_the_terminate_hook_waits_holds_it( + start: Any, +) -> None: + h = start() + h.run_hook() + assert h.job("cb-http") == 202 + assert eventually(lambda: h.started("cb-http")) + h.lambda_client.hold = True + answer = h.hook_in_background("terminate") + assert eventually(lambda: h.lambda_client.waiting() == 1) + + assert h.post("/job", INVALID_JOB) == 400 + assert eventually(lambda: h.lambda_client.waiting() == 2) + # The running job's report ends. The invalid request's report still runs, + # so the answer must wait. + h.lambda_client.finish_first() + assert eventually(lambda: len(h.lambda_client.recorded()) == 1) + assert not answered_within(answer, 0.2) + + h.lambda_client.finish() + assert answer.result(WAIT) == 200 + assert sorted(c["callbackId"] for c in h.lambda_client.recorded()) == [ + "cb-http", + "cb-invalid", + ] + + +def test_answers_the_terminate_hook_after_the_budget_when_a_run_hook_never_ends( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "TERMINATE_REPORT_BUDGET_SECONDS", 0.3) + h = start() + run = slow_run_hook(h) + + started = time.monotonic() + assert h.hook("terminate") == 200 + assert time.monotonic() - started >= 0.25 + run.abandon() + + +def test_starts_a_slow_run_hook_job_after_the_resume_hook_when_a_suspend_came_first( + start: Any, +) -> None: + h = start() + run = slow_run_hook(h) + + assert h.hook("suspend") == 200 + assert run.finish() == 200 + assert not h.started("cb-run") + + h.hook("resume") + assert eventually(lambda: h.started("cb-run")) + h.finish_job("cb-run") + assert h.lambda_client.recorded() == [{"kind": "success", "callbackId": "cb-run"}] + + +def test_starts_a_deferred_job_30_seconds_after_the_suspend_hook_without_a_resume( + start: Any, +) -> None: + h = start() + run = slow_run_hook(h) + h.hook("suspend") + assert run.finish() == 200 + + h.advance(GRACE - 0.001) + assert not h.started("cb-run") + h.advance(0.001) + assert eventually(lambda: h.started("cb-run")) + + +def test_fails_a_deferred_job_when_the_terminate_hook_follows_the_suspend_hook( + start: Any, +) -> None: + h = start() + run = slow_run_hook(h) + h.hook("suspend") + assert run.finish() == 200 + + h.hook("terminate") + h.advance(GRACE) + + assert not h.started("cb-run") + assert h.lambda_client.recorded() == [terminated_failure("cb-run")] + + +def test_refuses_a_slow_http_job_with_503_when_the_terminate_hook_came_first( + start: Any, +) -> None: + h = start() + h.run_hook() + request = SlowCall( + h, + "/job", + {"version": 1, "region": "us-east-1", "callbackId": "cb-http", "input": {}}, + ) + + h.hook("terminate") + assert request.finish() == 503 + assert not h.started("cb-http") + + +# --------------------------------------------------------------------------- +# The suspend hook +# --------------------------------------------------------------------------- + + +def test_the_suspend_hook_refuses_jobs_with_503_until_the_resume_hook( + start: Any, +) -> None: + h = start() + h.run_hook() + + assert h.hook("suspend") == 200 + assert h.job("cb-1") == 503 + + h.hook("resume") + assert h.job("cb-1") == 202 + + +def test_accepts_jobs_again_30_seconds_after_the_suspend_hook_without_a_resume( + start: Any, +) -> None: + h = start() + h.run_hook() + h.hook("suspend") + + h.advance(GRACE - 0.001) + assert h.job("cb-1") == 503 + h.advance(0.001) + assert h.job("cb-1") == 202 + + +def test_the_suspend_hook_leaves_running_jobs_running(start: Any) -> None: + h = start() + h.run_hook() + h.job("cb-1") + assert eventually(lambda: h.started("cb-1")) + + h.hook("suspend") + h.finish_job("cb-1") + + assert not h.jobs.contexts["cb-1"].cancelled.is_set() + assert h.lambda_client.recorded() == [{"kind": "success", "callbackId": "cb-1"}] + + +def test_keeps_refusing_jobs_after_a_forwarded_resume(start: Any) -> None: + h = start() + h.run_hook() + h.hook("suspend") + + assert h.hook("resume", FORWARDED_HEADERS) == 200 + assert h.job("cb-1") == 503 + + h.hook("resume") + assert h.job("cb-1") == 202 + + +def test_ignores_a_forwarded_suspend_hook(start: Any) -> None: + h = start() + h.run_hook() + + assert h.hook("suspend", FORWARDED_HEADERS) == 200 + assert h.job("cb-1") == 202 + + +def test_the_suspend_hook_stops_the_idle_time_and_the_resume_hook_restarts_it( + start: Any, +) -> None: + h = start() + h.run_hook(IDLE) + h.advance(IDLE / 2) + h.hook("suspend") + + h.advance(IDLE * 2) + assert h.microvms.suspended == [] + + h.hook("resume") + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1"] + + +def test_keeps_the_refusal_when_the_hook_arrives_during_a_rejected_own_call( + start: Any, +) -> None: + # The worker's own SuspendMicrovm call triggers the suspend hook. A + # rejection of a retried or concurrent call must not end the refusal + # that the hook started, because the MicroVM is suspending. + h = start() + h.microvms.hold = True + h.microvms.failure = suspend_error("AccessDeniedException", 403) + h.run_hook(IDLE) + h.advance(IDLE) + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + + h.hook("suspend") + h.microvms.finish() + h.settle() + + assert h.job("cb-1") == 503 + h.hook("resume") + assert h.job("cb-1") == 202 + + +def test_ends_the_refusal_at_once_when_a_rejected_own_call_had_no_hook( + start: Any, +) -> None: + h = start() + h.microvms.failure = suspend_error("AccessDeniedException", 403) + h.run_hook(IDLE) + h.advance(IDLE) + + assert h.job("cb-1") == 202 + + +def test_keeps_one_30_second_refusal_when_the_hook_follows_the_own_call( + start: Any, +) -> None: + h = start() + h.run_hook(IDLE) + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1"] + + h.advance(5) + h.hook("suspend") + + # The hook started a new 30-second refusal 5 seconds after the call's. + h.advance(25) + assert h.job("cb-1") == 503 + h.advance(5) + assert h.job("cb-1") == 202 + + +def test_a_terminated_error_names_its_cause() -> None: + error = MicrovmTerminatedError() + + assert str(error) == TERMINATED_MESSAGE + assert type(error).__name__ == "MicrovmTerminatedError" + + +# --------------------------------------------------------------------------- +# Threads and timers that cannot start +# --------------------------------------------------------------------------- + + +def failing_thread_start(monkeypatch: pytest.MonkeyPatch, prefix: str) -> None: + """Make every thread whose name starts with ``prefix`` fail to start.""" + original = threading.Thread.start + + def start(self: threading.Thread) -> None: + if self.name.startswith(prefix): + msg = "can't start new thread" + raise RuntimeError(msg) + original(self) + + monkeypatch.setattr(threading.Thread, "start", start) + + +def test_logs_an_idle_time_that_cannot_start( + start: Any, logger: Any, timers: FakeTimers +) -> None: + h = start() + timers.failure = RuntimeError("can't start new thread") + + assert h.run_hook(IDLE) == 200 + + assert "could not start the idle time" in logger.messages("error") + timers.failure = None + h.advance(IDLE * 2) + assert h.microvms.suspended == [] + + +def test_accepts_jobs_at_once_when_the_suspend_thread_cannot_start( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + h = start() + h.run_hook(IDLE) + failing_thread_start(monkeypatch, "microvm-suspend") + + h.advance(IDLE) + + assert "could not start the suspend of the idle MicroVM" in logger.messages("error") + assert h.microvms.suspended == [] + monkeypatch.undo() + assert h.job("cb-1") == 202 + + +def test_refuses_jobs_until_the_resume_hook_when_the_refusal_timer_cannot_start( + start: Any, logger: Any, timers: FakeTimers +) -> None: + h = start() + h.run_hook() + timers.failure = RuntimeError("can't start new thread") + + h.hook("suspend") + + assert "could not start the end of the refusal" in logger.messages("error") + timers.failure = None + h.advance(GRACE * 2) + assert h.job("cb-1") == 503 + h.hook("resume") + assert h.job("cb-1") == 202 + + +def test_fails_a_deferred_run_hook_job_whose_thread_cannot_start( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + h = start() + run = slow_run_hook(h) + h.hook("suspend") + assert run.finish() == 200 + failing_thread_start(monkeypatch, "microvm-job-") + + h.hook("resume") + + assert eventually(lambda: len(h.lambda_client.recorded()) == 1) + assert h.lambda_client.recorded()[0]["callbackId"] == "cb-run" + assert h.lambda_client.recorded()[0]["errorType"] == "RuntimeError" + assert "could not start the run hook job" in logger.messages("error") + assert not h.started("cb-run") + + +def test_a_stale_idle_timer_does_not_suspend(start: Any, timers: FakeTimers) -> None: + h = start() + timers.fire_cancelled = True + + h.run_hook(IDLE) + h.job("cb-1") + h.advance(IDLE) + + assert h.microvms.suspended == [] + + +def test_a_stale_refusal_timer_does_not_end_a_later_refusal( + start: Any, timers: FakeTimers +) -> None: + h = start() + timers.fire_cancelled = True + h.run_hook() + + h.hook("suspend") + h.advance(1) + h.hook("resume") + h.advance(9) + h.hook("suspend") + # The first refusal's timer fires 30 seconds after the first hook. + h.advance(GRACE - 9) + + assert h.job("cb-1") == 503 + h.advance(9) + assert h.job("cb-1") == 202 diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py new file mode 100644 index 000000000..d03a9134b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the worker's own SuspendMicrovm call.""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +import pytest +from botocore.exceptions import ( + ClientError, + EndpointConnectionError, + NoCredentialsError, +) + +from aws_durable_execution_sdk_python_microvm_worker.self_suspend import ( + SUSPEND_CALL_TIMEOUT_SECONDS, + create_microvms_client, + is_rejected_call, + suspend_microvm, +) + + +def suspend_error(code: str, status: int, retries: int = 0) -> ClientError: + response: Any = { + "Error": {"Code": code, "Message": f"{code} message"}, + "ResponseMetadata": {"HTTPStatusCode": status, "RetryAttempts": retries}, + } + return ClientError(response, "SuspendMicrovm") + + +@pytest.mark.parametrize( + "error", + [ + suspend_error("AccessDeniedException", 403), + suspend_error("ValidationException", 400), + suspend_error("ResourceNotFoundException", 404), + suspend_error("SomeServiceError", 403), + # The rejection codes count whatever their status. + suspend_error("AccessDeniedException", 500), + NoCredentialsError(), + ], +) +def test_counts_a_rejection_on_the_first_attempt_as_rejected( + error: BaseException, +) -> None: + assert is_rejected_call(error) + + +@pytest.mark.parametrize( + "error", + [ + # An earlier attempt may have suspended the MicroVM. + suspend_error("ValidationException", 400, retries=1), + suspend_error("AccessDeniedException", 403, retries=2), + suspend_error("ThrottlingException", 429), + # Some services answer a throttle with HTTP 400. + suspend_error("ThrottlingException", 400), + suspend_error("ExpiredTokenException", 403), + suspend_error("ConflictException", 409), + suspend_error("ConflictException", 400), + suspend_error("RequestTimeoutException", 408), + suspend_error("InternalServerException", 500), + TimeoutError("timed out"), + EndpointConnectionError(endpoint_url="https://example.com"), + RuntimeError("boom"), + ], +) +def test_counts_any_other_error_as_an_unknown_outcome(error: BaseException) -> None: + assert not is_rejected_call(error) + + +def test_reads_a_client_error_without_metadata_as_an_unknown_outcome() -> None: + error = ClientError({"Error": {"Code": "SomeError"}}, "SuspendMicrovm") # type: ignore[typeddict-item] + + assert not is_rejected_call(error) + + +class _Client: + def __init__(self, answer: Any = None, delay: float = 0.0) -> None: + self.calls: list[dict[str, Any]] = [] + self._answer = answer + self._delay = delay + self.release = threading.Event() + + def suspend_microvm(self, **kwargs: Any) -> dict[str, Any]: + self.calls.append(kwargs) + if self._delay: + self.release.wait(self._delay) + if isinstance(self._answer, BaseException): + raise self._answer + return {} + + +def test_calls_suspend_microvm_with_the_microvm_id() -> None: + client: Any = _Client() + + suspend_microvm(client, "mvm-1", 5.0) + + assert client.calls == [{"microvmIdentifier": "mvm-1"}] + + +def test_raises_the_call_s_own_error() -> None: + error = suspend_error("AccessDeniedException", 403) + + client: Any = _Client(error) + + with pytest.raises(ClientError) as raised: + suspend_microvm(client, "mvm-1", 5.0) + + assert raised.value is error + + +def test_raises_the_call_s_own_timeout_error() -> None: + client: Any = _Client(TimeoutError("read timed out")) + + with pytest.raises(TimeoutError, match="^read timed out$"): + suspend_microvm(client, "mvm-1", 5.0) + + +def test_stops_waiting_for_a_stalled_call_after_the_timeout() -> None: + client: Any = _Client(delay=10.0) + started = time.monotonic() + + with pytest.raises(TimeoutError, match="longer than 0.05 seconds"): + suspend_microvm(client, "mvm-1", 0.05) + + assert time.monotonic() - started < 2.0 + client.release.set() + + +def test_the_default_client_ends_a_read_at_the_call_timeout() -> None: + client = create_microvms_client("us-west-2") + try: + config: Any = client.meta.config + assert client.meta.region_name == "us-west-2" + assert config.read_timeout == SUSPEND_CALL_TIMEOUT_SECONDS + assert config.connect_timeout == 5 + assert config.retries == {"mode": "standard", "total_max_attempts": 3} + finally: + client.close() diff --git a/pyproject.toml b/pyproject.toml index 365a7d25b..5c6655bd1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -126,7 +126,7 @@ dependencies = [ "pytest-cov", "coverage[toml]", "mypy>=1.0.0", - "boto3-stubs[lambda]", + "boto3-stubs[lambda,lambda-microvms]", ] [tool.hatch.envs.dev-microvm-worker.scripts] From dbe368647c4a0bb3a1793e00182755bce97cfebf Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 16 Oct 2026 11:01:41 -0700 Subject: [PATCH 2/4] fix(microvm-worker): bound suspend call attempts Address the review of the lifecycle hooks. - SuspendMicrovm: a before-send handler stops botocore's retries at the 10-second deadline. Before, retries went on in the call's thread, and an attempt started 22.5 seconds after the first. - A throttle, expired credentials, or a clock offset on the first attempt ends the refusal at once, and the idle time starts again. The service answers these before it acts. - Docs: idle() covers the wait, not the call, and the wait without the resume hook is a little over 30 seconds. - A composed test covers an image without the resume hook. --- .../README.md | 2 +- .../tests/microvm/worker_composed_test.py | 50 ++++++- .../README.md | 11 +- .../self_suspend.py | 115 +++++++++++----- .../worker.py | 80 +++++++----- .../tests/lifecycle_hooks_test.py | 68 +++++++++- .../tests/self_suspend_test.py | 123 ++++++++++++++++-- 7 files changed, 363 insertions(+), 86 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/README.md b/packages/aws-durable-execution-sdk-python-extras/README.md index 9ce81da73..bfc9e463e 100644 --- a/packages/aws-durable-execution-sdk-python-extras/README.md +++ b/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) 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 up to 30 seconds. +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 diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index 5cc9b1094..235f89f29 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -304,22 +304,50 @@ def post_hook(worker: MicrovmWorker, hook: str) -> None: 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 + 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. The worker's timers run 10 times faster, so the - session's 10-second idle time takes 1 second. + 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 / 10, action), + lambda delay, action: original_timer(delay / 20, action), ) state = {"value": "RUNNING"} lock = threading.Lock() @@ -336,7 +364,8 @@ def resume_microvm(self, **request: Any) -> dict[str, Any]: super().resume_microvm(**request) with lock: state["value"] = "RUNNING" - post_hook(holder["worker"], "resume") + if resume_hook: + post_hook(holder["worker"], "resume") return {} class WorkerMicrovms: @@ -382,10 +411,12 @@ def handler(event: Any, context: DurableContext) -> Any: 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 @@ -399,7 +430,7 @@ def handler(event: Any, context: DurableContext) -> Any: execution_role_arn="role", timeout=Duration.from_minutes(5), auto_suspend_on_idle=True, - auto_suspend_idle_time=Duration.from_seconds(10), + auto_suspend_idle_time=Duration.from_seconds(60), client=client, ) arn = runner.run_async(input="null", execution_timeout=WAIT) @@ -418,3 +449,8 @@ def handler(event: Any, context: DurableContext) -> Any: 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 diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index 6b08b7663..074978d84 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -98,7 +98,7 @@ The worker creates a Lambda client for each job, after the job arrives, not at i A document that does not match the contract gets HTTP 400 with the validation message. If the document named a callback, the worker also fails that callback, so the durable function fails at once instead of at its timeout. A request whose body cannot be read gets HTTP 400 with a generic message, and the error is logged. The worker accepts a `run` hook body of at most 64 KB and a job request body of at most 6 MB, with a `Content-Length` or with chunked transfer encoding. In a chunked body, the chunk-size lines, their extensions, and the trailers have a limit of their own, of the same size. A larger body gets HTTP 400. With `Expect: 100-continue`, the 400 replaces `100 Continue`, so the client does not send the body. Without it, the worker reads and drops the body after the answer, for at most 10 seconds and 24 MB. Without that read, the connection would close with unread bytes, and the client would get a TCP reset instead of the 400. A body that is larger, or that is still arriving after 10 seconds, still gets the reset. Other lifecycle hooks get HTTP 200. Other paths get HTTP 404. -`worker.idle(timeout)` waits until every job started so far has ended, and a running `SuspendMicrovm` call has returned. It returns `False` when the timeout ends first. `worker.close()` stops the server and the idle time. Running jobs continue, so call `idle()` first to let them report their outcome. +`worker.idle(timeout)` waits until every job started so far has ended, and the wait for a running `SuspendMicrovm` call has ended. It returns `False` when the timeout ends first. `worker.close()` stops the server and the idle time. Running jobs continue, so call `idle()` first to let them report their outcome. ## Suspending when idle @@ -106,17 +106,18 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets 1. The idle time starts when the `run` hook arrives without a job, when a job ends, and when the `resume` lifecycle hook arrives, if the image enables it. 2. A job that starts cancels the idle time. So the worker never suspends the MicroVM during a job. -3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The worker waits for it for at most 10 seconds. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. +3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The worker waits for it for at most 10 seconds. No attempt of the call starts after that, because a `before-send` handler on the client stops botocore's retries at the deadline. An attempt already in flight then runs until its read timeout. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. 4. From the moment the idle time ends, the worker answers every job with HTTP 503 and does not start it. A job accepted then would freeze with the MicroVM, and nothing would resume it. After a 503, the durable function checks the state, resumes the MicroVM, and delivers the job again. -5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. So enable the `resume` hook in the image. Without it, the first job after a resume can wait up to 30 seconds. +5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. So enable the `resume` hook in the image. Without it, the first job after a resume can wait a little over 30 seconds: the 30 seconds, plus up to 4 seconds until the durable function's next delivery attempt. A test measured 32 seconds. 6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection, because the call sends no request. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. -7. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a throttle, a conflict, expired credentials, a clock offset, a server error, and a dropped connection. A rejection after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. +7. A call that the service refused for now on its first attempt is logged as a warning too, and the worker accepts jobs again at once. This covers a throttle, expired credentials, and a clock offset, whatever the HTTP status. The service answers these before it acts. The idle time then starts again, so the worker tries the call again. The JavaScript worker decides these cases by HTTP status. +8. A call whose outcome is unknown is handled like a successful call, because the service may have suspended the MicroVM anyway. This covers a timeout, a conflict, a server error, and a dropped connection. An error after botocore retried an earlier attempt also counts as unknown, because the earlier attempt may have suspended the MicroVM. botocore retries a throttle, so a throttle usually ends here. The refusal continues until the `resume` hook or the 30-second limit. The idle time then starts again, so the worker tries the call again. `worker.close()` stops the idle time. A closed worker never suspends its MicroVM again. A refusal that is already running still ends on the `resume` hook or after 30 seconds. The worker counts only jobs. It cannot see background threads or processes, or inbound traffic that is not a job. A session with such work leaves `auto_suspend_on_idle` off, and the payload then has no `autoSuspendIdleSeconds`. -The worker creates a `lambda-microvms` client for each suspend, and closes it after the call. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. +The worker creates a `lambda-microvms` client for each suspend, and closes it when its wait for the call ends. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. ## Suspend and terminate hooks diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py index a954c8769..b315ddfc1 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py @@ -6,7 +6,9 @@ from __future__ import annotations import threading +import time from concurrent.futures import Future +from enum import Enum from typing import TYPE_CHECKING, Any import boto3 @@ -31,7 +33,8 @@ SUSPEND_CALL_TIMEOUT_SECONDS = 10.0 """The longest time the worker waits for its own SuspendMicrovm call. -The call measured 75 to 94 milliseconds from inside a MicroVM in us-east-1. +No attempt of the call starts after this time. The call measured 75 to 94 +milliseconds from inside a MicroVM in us-east-1. """ SUSPEND_GRACE_SECONDS = 30.0 @@ -57,13 +60,34 @@ The JavaScript worker's default client makes 3 attempts too. """ +_BEFORE_SEND = "before-send.lambda-microvms.SuspendMicrovm" + # Error codes that show that the service did nothing. _REJECTION_CODES = frozenset( {"AccessDeniedException", "ValidationException", "ResourceNotFoundException"} ) -# 4xx error codes that leave open whether the service suspended the MicroVM. -_OPEN_OUTCOME_CODES = TRANSIENT_4XX_CODES | {"ConflictException"} + +class SuspendOutcome(Enum): + """What a failed SuspendMicrovm call shows about the MicroVM.""" + + REJECTED = "rejected" + """The service did nothing, and a later call would most likely fail the + same way, for example for a missing permission.""" + REFUSED_FOR_NOW = "refused_for_now" + """The service did nothing, and a later call can succeed: a throttle, + expired credentials, or a clock offset.""" + UNKNOWN = "unknown" + """The service may have suspended the MicroVM.""" + + +class SuspendDeadlineError(TimeoutError): + """The SuspendMicrovm deadline passed before an attempt could start. + + botocore retries only its own connection and HTTP client errors. So it + does not retry this one. The call can end with it just before the wait + ends, so it is a ``TimeoutError`` too. + """ def create_microvms_client(region: str) -> LambdaMicrovmsClient: @@ -94,15 +118,38 @@ def suspend_microvm( botocore has no deadline for a whole call. Its read timeout limits one socket read, and it covers neither credential resolution nor its own - retries. So the call runs in a daemon thread, and this function waits - for the thread. A call that ends after the bound is ignored. The caller - closes a client that it created, which releases the call's connection. + retries. So the call needs two limits: + + 1. The call runs in a daemon thread, and this function waits for the + thread for at most ``timeout``. + 2. A ``before-send`` handler on the client raises + :class:`SuspendDeadlineError` when an attempt would start after + ``timeout``. botocore does not retry that error. So no attempt starts + after this function has stopped waiting. Without the handler, botocore + went on retrying in the thread, and an attempt started 22.5 seconds + after the first one. + + An attempt that is already in flight at the deadline still runs until + its read timeout. The service receives its request before the deadline. + The handler is removed when the call ends. So a client from the caller + keeps no handler. Raises: TimeoutError: When the call takes longer than ``timeout``. Exception: The call's own error. """ future: Future[None] = Future() + deadline = time.monotonic() + timeout + unique_id = f"aws-durable-execution-microvm-worker-deadline-{id(future)}" + events = safe_get(lambda: client.meta.events) + + def stop_after_deadline(**_: object) -> None: + if time.monotonic() >= deadline: + msg = f"the SuspendMicrovm deadline of {timeout} seconds has passed" + raise SuspendDeadlineError(msg) + + if events is not None: + events.register_first(_BEFORE_SEND, stop_after_deadline, unique_id=unique_id) def run() -> None: try: @@ -110,6 +157,9 @@ def run() -> None: future.set_result(None) except BaseException as error: # noqa: BLE001 future.set_exception(error) + finally: + if events is not None: + events.unregister(_BEFORE_SEND, unique_id=unique_id) threading.Thread(target=run, name="microvm-suspend-call", daemon=True).start() try: @@ -121,38 +171,45 @@ def run() -> None: raise TimeoutError(msg) from error -def is_rejected_call(error: BaseException) -> bool: - """Whether a failed SuspendMicrovm call shows that the service did nothing. +def classify_suspend_error(error: BaseException) -> SuspendOutcome: + """What a failed SuspendMicrovm call shows about the MicroVM. 1. A botocore retry inside the call means that an earlier attempt may have suspended the MicroVM before its answer was lost. So any error - after a retry leaves the outcome open. - 2. Missing credentials stop the call before any request. - 3. A known rejection code shows that the service did nothing. - 4. A throttle, a conflict, expired credentials, and a clock offset leave - the outcome open, whatever their status. Some services answer a - throttle with HTTP 400. - 5. Any other 4xx answer except 408, 409, and 429 shows that the service - did nothing. - 6. Any other error, such as a timeout, a dropped connection, or a server - error, leaves the outcome open. + after a retry is :attr:`SuspendOutcome.UNKNOWN`. + 2. Missing credentials stop the call before any request. They are + :attr:`SuspendOutcome.REJECTED`, as in the JavaScript worker. + 3. A known rejection code is :attr:`SuspendOutcome.REJECTED`. + 4. A conflict is :attr:`SuspendOutcome.UNKNOWN`, whatever its status. It + can mean that the MicroVM is already suspending. + 5. A throttle, expired credentials, or a clock offset on the first + attempt is :attr:`SuspendOutcome.REFUSED_FOR_NOW`, whatever its + status. The service answers them before it acts, and a later call can + succeed. Some services answer a throttle with HTTP 400. + 6. Any other 4xx answer except 408, 409, and 429 is + :attr:`SuspendOutcome.REJECTED`. A 429 is + :attr:`SuspendOutcome.REFUSED_FOR_NOW`. + 7. Any other error, such as a timeout, a dropped connection, a 408, or a + server error, is :attr:`SuspendOutcome.UNKNOWN`. """ if isinstance(error, NoCredentialsError): - return True + return SuspendOutcome.REJECTED if not isinstance(error, ClientError): - return False + return SuspendOutcome.UNKNOWN metadata = safe_get(lambda: error.response.get("ResponseMetadata", {})) or {} retries = safe_get(lambda: metadata.get("RetryAttempts")) if isinstance(retries, int) and retries > 0: - return False + return SuspendOutcome.UNKNOWN code = error_code(error) if code in _REJECTION_CODES: - return True - if code in _OPEN_OUTCOME_CODES: - return False + return SuspendOutcome.REJECTED + if code == "ConflictException": + return SuspendOutcome.UNKNOWN + if code in TRANSIENT_4XX_CODES: + return SuspendOutcome.REFUSED_FOR_NOW status = safe_get(lambda: metadata.get("HTTPStatusCode")) - return ( - isinstance(status, int) - and 400 <= status < 500 - and status not in {408, 409, 429} - ) + if status == 429: + return SuspendOutcome.REFUSED_FOR_NOW + if isinstance(status, int) and 400 <= status < 500 and status not in {408, 409}: + return SuspendOutcome.REJECTED + return SuspendOutcome.UNKNOWN diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index 6db36cd77..742f28cbd 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -53,8 +53,9 @@ from aws_durable_execution_sdk_python_microvm_worker.self_suspend import ( SUSPEND_CALL_TIMEOUT_SECONDS, SUSPEND_GRACE_SECONDS, + SuspendOutcome, + classify_suspend_error, create_microvms_client, - is_rejected_call, suspend_microvm, ) @@ -294,10 +295,15 @@ class MicrovmWorkerListener: 6. A call that the service rejected on its first attempt, for example for a missing permission, ends the refusal at once. The MicroVM keeps running, and the listener does not retry the call until a job ends. - 7. A call whose outcome is unknown, such as a timeout, a throttle, a - conflict, a server error, or a rejection after a retried attempt, is - handled like a successful call. The service may have suspended the - MicroVM anyway. + 7. A call that the service refused for now on its first attempt, such + as a throttle, expired credentials, or a clock offset, ends the + refusal at once too. The idle time then starts again, so the listener + tries the call again. + 8. A call whose outcome is unknown, such as a timeout, a conflict, a + server error, or an error after a retried attempt, is handled like a + successful call. The service may have suspended the MicroVM anyway. + The listener waits for the call for at most 10 seconds, and no + attempt starts after that. When the image enables the ``suspend`` and ``terminate`` hooks, the listener acts on them: @@ -557,8 +563,10 @@ def handle_expect_100(self, request: BaseHTTPRequestHandler) -> bool: def idle(self, timeout: float | None = None) -> bool: """Wait until every job started so far has ended. - The wait also covers the failure reports of invalid requests, and a - running SuspendMicrovm call. That call takes at most 10 seconds. + The wait also covers the failure reports of invalid requests, and the + worker's wait for its own SuspendMicrovm call. That wait takes at most + 10 seconds, and no attempt of the call starts after it. An attempt + already in flight then can run until its read timeout. Args: timeout: The longest wait in seconds. ``None`` waits until the @@ -1350,31 +1358,41 @@ def _call_suspend(self, auto_suspend: _AutoSuspend) -> None: try: suspend_microvm(client, target, SUSPEND_CALL_TIMEOUT_SECONDS) except Exception as error: # noqa: BLE001 - if is_rejected_call(error): - self._logger.warning( - "could not suspend the idle MicroVM, because the service " - "rejected the call. It keeps running. A missing permission " - "needs lambda:SuspendMicrovm on the MicroVM image in its " - "execution role.", - extra={**fields, "error": describe(error)}, - ) - # The service did nothing. So jobs can run again at once. The - # idle time does not start again, so a missing permission is - # not retried. A suspend hook that arrived during the call - # means that the MicroVM is suspending anyway, so its refusal - # stays. - with self._lock: - deferred = ( - self._stop_refusing() - if self._suspend_hooks == hooks_before - else [] + outcome = classify_suspend_error(error) + if outcome is not SuspendOutcome.UNKNOWN: + if outcome is SuspendOutcome.REJECTED: + self._logger.warning( + "could not suspend the idle MicroVM, because the " + "service rejected the call. It keeps running. A " + "missing permission needs lambda:SuspendMicrovm on the " + "MicroVM image in its execution role.", + extra={**fields, "error": describe(error)}, ) + else: + self._logger.warning( + "could not suspend the idle MicroVM, because the " + "service refused the call for now. It keeps running, " + "and the worker tries again after the next idle time.", + extra={**fields, "error": describe(error)}, + ) + # The service did nothing. So jobs can run again at once. A + # suspend hook that arrived during the call means that the + # MicroVM is suspending anyway, so its refusal stays. + with self._lock: + deferred = [] + if self._suspend_hooks == hooks_before: + deferred = self._stop_refusing() + # A rejection, such as a missing permission, is not + # retried until a job ends. A refusal for now is + # retried after the next idle time. + if outcome is SuspendOutcome.REFUSED_FOR_NOW: + self._start_idle_time() self._launch_deferred(deferred) return - # A timeout, a throttle, a conflict, a server error, or a - # rejection after a retried attempt does not show whether the - # service suspended the MicroVM. So the refusal continues, as - # after a successful call. + # A timeout, a conflict, a server error, or an error after a + # retried attempt does not show whether the service suspended + # the MicroVM. So the refusal continues, as after a successful + # call. self._logger.warning( "the suspend call for the idle MicroVM had an unknown outcome. " "The worker refuses jobs until the MicroVM resumes, or for 30 " @@ -1525,8 +1543,8 @@ def port(self) -> int: return int(self._server.server_address[1]) def idle(self, timeout: float | None = None) -> bool: - """Wait until every job started so far has ended, and a running - SuspendMicrovm call has returned. + """Wait until every job started so far has ended, and the wait for a + running SuspendMicrovm call has ended. Returns: ``True`` when no job runs. ``False`` when the timeout ended first. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index 2fe1b3c58..1bcb6b5a9 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -779,24 +779,24 @@ def test_retries_the_suspend_one_idle_time_after_a_job_that_follows_a_rejection( "failure", [ TimeoutError("timed out"), - suspend_error("ThrottlingException", 429), - suspend_error("ThrottlingException", 400), suspend_error("ConflictException", 409), suspend_error("InternalServerException", 500), suspend_error("RequestTimeoutException", 408), # The client retried a 5xx. The first attempt may have suspended the # MicroVM, and the retry then fails validation. suspend_error("ValidationException", 400, retries=1), + # botocore retries a throttle. An earlier attempt may have reached + # the service. + suspend_error("ThrottlingException", 429, retries=2), ConnectionResetError("reset"), ], ids=[ "timeout", - "throttle", - "throttle-400", "conflict", "server-error", "request-timeout", "retried-rejection", + "retried-throttle", "connection-error", ], ) @@ -832,6 +832,66 @@ def test_accepts_jobs_at_once_after_a_rejection( assert h.job("cb-1") == 202 +@pytest.mark.parametrize( + "failure", + [ + suspend_error("ThrottlingException", 429), + suspend_error("ThrottlingException", 400), + suspend_error("ExpiredTokenException", 403), + suspend_error("SignatureDoesNotMatch", 403), + ], + ids=["throttle", "throttle-400", "expired-token", "clock-offset"], +) +def test_accepts_jobs_at_once_and_retries_after_a_refusal_for_now( + start: Any, logger: Any, failure: BaseException +) -> None: + h = start() + h.microvms.failure = failure + + h.run_hook(IDLE) + h.advance(IDLE) + assert "refused the call for now" in logger.messages("warning")[0] + assert h.job("cb-1") == 202 + h.finish_job("cb-1") + + # The idle time started again when the call failed, and once more when + # the job ended. So the next call comes one idle time after the job. + h.advance(IDLE) + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +def test_retries_a_refusal_for_now_one_idle_time_later_without_a_job( + start: Any, +) -> None: + h = start() + h.microvms.failure = suspend_error("ThrottlingException", 429) + + h.run_hook(IDLE) + h.advance(IDLE) + h.advance(IDLE - 0.001) + assert h.microvms.suspended == ["mvm-1"] + + h.advance(0.001) + assert h.microvms.suspended == ["mvm-1", "mvm-1"] + + +def test_keeps_a_hook_refusal_after_a_refusal_for_now_of_the_own_call( + start: Any, +) -> None: + h = start() + h.microvms.hold = True + h.microvms.failure = suspend_error("ThrottlingException", 429) + h.run_hook(IDLE) + h.advance(IDLE) + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + + h.hook("suspend") + h.microvms.finish() + h.settle() + + assert h.job("cb-1") == 503 + + def test_treats_a_call_that_outlasts_its_timeout_as_an_unknown_outcome( start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py index d03a9134b..88d244ab6 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py @@ -5,11 +5,14 @@ from __future__ import annotations +import socket import threading import time from typing import Any +import boto3 import pytest +from botocore.config import Config from botocore.exceptions import ( ClientError, EndpointConnectionError, @@ -18,8 +21,10 @@ from aws_durable_execution_sdk_python_microvm_worker.self_suspend import ( SUSPEND_CALL_TIMEOUT_SECONDS, + SuspendDeadlineError, + SuspendOutcome, + classify_suspend_error, create_microvms_client, - is_rejected_call, suspend_microvm, ) @@ -44,39 +49,58 @@ def suspend_error(code: str, status: int, retries: int = 0) -> ClientError: NoCredentialsError(), ], ) -def test_counts_a_rejection_on_the_first_attempt_as_rejected( +def test_classifies_a_rejection_on_the_first_attempt_as_rejected( error: BaseException, ) -> None: - assert is_rejected_call(error) + assert classify_suspend_error(error) is SuspendOutcome.REJECTED @pytest.mark.parametrize( "error", [ - # An earlier attempt may have suspended the MicroVM. - suspend_error("ValidationException", 400, retries=1), - suspend_error("AccessDeniedException", 403, retries=2), suspend_error("ThrottlingException", 429), # Some services answer a throttle with HTTP 400. suspend_error("ThrottlingException", 400), + suspend_error("TooManyRequestsException", 429), + suspend_error("SomeThrottle", 429), suspend_error("ExpiredTokenException", 403), + suspend_error("SignatureDoesNotMatch", 403), + suspend_error("RequestTimeTooSkewed", 403), + ], +) +def test_classifies_a_first_attempt_that_can_succeed_later_as_refused_for_now( + error: BaseException, +) -> None: + assert classify_suspend_error(error) is SuspendOutcome.REFUSED_FOR_NOW + + +@pytest.mark.parametrize( + "error", + [ + # An earlier attempt may have suspended the MicroVM. + suspend_error("ValidationException", 400, retries=1), + suspend_error("AccessDeniedException", 403, retries=2), + suspend_error("ThrottlingException", 429, retries=2), suspend_error("ConflictException", 409), suspend_error("ConflictException", 400), suspend_error("RequestTimeoutException", 408), suspend_error("InternalServerException", 500), TimeoutError("timed out"), + SuspendDeadlineError("passed"), EndpointConnectionError(endpoint_url="https://example.com"), RuntimeError("boom"), ], ) -def test_counts_any_other_error_as_an_unknown_outcome(error: BaseException) -> None: - assert not is_rejected_call(error) +def test_classifies_any_other_error_as_an_unknown_outcome( + error: BaseException, +) -> None: + assert classify_suspend_error(error) is SuspendOutcome.UNKNOWN def test_reads_a_client_error_without_metadata_as_an_unknown_outcome() -> None: error = ClientError({"Error": {"Code": "SomeError"}}, "SuspendMicrovm") # type: ignore[typeddict-item] - assert not is_rejected_call(error) + assert classify_suspend_error(error) is SuspendOutcome.UNKNOWN class _Client: @@ -142,3 +166,84 @@ def test_the_default_client_ends_a_read_at_the_call_timeout() -> None: assert config.retries == {"mode": "standard", "total_max_attempts": 3} finally: client.close() + + +class _SilentService: + """A local HTTP service that accepts connections and never answers, and + records when each attempt arrives.""" + + def __init__(self) -> None: + self.attempts: list[float] = [] + self._server = socket.create_server(("127.0.0.1", 0)) + self._server.settimeout(0.1) + self._stop = threading.Event() + self._connections: list[socket.socket] = [] + self._thread = threading.Thread(target=self._serve, daemon=True) + self._thread.start() + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self._server.getsockname()[1]}" + + def _serve(self) -> None: + while not self._stop.is_set(): + try: + connection, _ = self._server.accept() + except OSError: + continue + self.attempts.append(time.monotonic()) + self._connections.append(connection) + + def close(self) -> None: + self._stop.set() + self._thread.join(5) + for connection in self._connections: + connection.close() + self._server.close() + + +def test_starts_no_attempt_after_the_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # botocore retries a read timeout. Without the deadline handler, it went + # on retrying in the call's thread after the wait had ended. + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test") + service = _SilentService() + try: + client = boto3.session.Session().client( + "lambda-microvms", + region_name="us-east-1", + endpoint_url=service.url, + config=Config( + connect_timeout=1, + read_timeout=0.3, + retries={"mode": "standard", "total_max_attempts": 5}, + ), + ) + started = time.monotonic() + + # The call ends with the wait's TimeoutError, or with + # SuspendDeadlineError when botocore meets the handler first. + with pytest.raises(TimeoutError): + suspend_microvm(client, "mvm-1", 0.5) + + # The attempt in flight at the deadline ends at its read timeout, + # and botocore then meets the deadline handler. + time.sleep(2.0) + assert service.attempts + assert all(at - started < 0.5 for at in service.attempts) + # The handler is gone once the call has ended. + handlers = client.meta.events._emitter._unique_id_handlers # type: ignore[attr-defined] # noqa: SLF001 + assert not [key for key in handlers if "microvm-worker-deadline" in key] + client.close() + finally: + service.close() + + +def test_skips_the_deadline_handler_for_a_client_without_events() -> None: + client: Any = _Client() + + suspend_microvm(client, "mvm-1", 5.0) + + assert client.calls == [{"microvmIdentifier": "mvm-1"}] From 9cc250395c9a7a3c503748cb0b0fdc39f2094c3a Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Fri, 16 Oct 2026 11:59:23 -0700 Subject: [PATCH 3/4] test(microvm-worker): fix parallel-run flakes - Reach a dual-stack worker on ::1. On macOS, the worker's [::] bind with SO_REUSEADDR succeeds on a port that another process holds on 127.0.0.1. A test that connected to 127.0.0.1 then reached that process. Linux refuses such a bind, so a MicroVM is not affected. - finish_job waits for the lock, so the idle time has started when it returns. - The deadline test waits for the call to end, and records attempts on the client side, so a slow accept on the fake service cannot fail it. 80 runs of the worker tests with 8 workers passed. Before, 15 of 262 runs had a failure. --- .../tests/microvm/worker_composed_test.py | 57 +++++++++++++-- .../tests/lifecycle_hooks_test.py | 47 ++++++++++++- .../tests/self_suspend_test.py | 38 ++++++++-- .../tests/worker_test.py | 70 ++++++++++++------- 4 files changed, 172 insertions(+), 40 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index 235f89f29..af03b4ca4 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -17,8 +17,10 @@ from __future__ import annotations +import functools import http.client import json +import socket import threading import time from collections.abc import Callable, Iterator @@ -57,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.""" @@ -97,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)) @@ -145,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", @@ -267,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", @@ -291,7 +338,7 @@ def handler(event: Any, context: DurableContext) -> Any: def post_hook(worker: MicrovmWorker, hook: str) -> None: """Send a lifecycle hook as Lambda does, with a local Host.""" - connection = http.client.HTTPConnection("127.0.0.1", worker.port, timeout=WAIT) + connection = http.client.HTTPConnection(host_of(worker), worker.port, timeout=WAIT) try: connection.request( "POST", @@ -423,7 +470,7 @@ def handler(event: Any, context: DurableContext) -> Any: try: client = StatefulMicrovms( launches=[run_hook_to(worker)], - endpoint=f"http://127.0.0.1:{worker.port}", + endpoint=url_of(worker), ) holder["config"] = MicrovmSessionConfig( image_identifier="img", diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index 1bcb6b5a9..1847c7cab 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -9,6 +9,7 @@ from __future__ import annotations +import functools import http.client import json import socket @@ -45,6 +46,39 @@ } +@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 eventually(condition: Callable[[], bool]) -> bool: deadline = time.monotonic() + WAIT while not condition(): @@ -274,7 +308,7 @@ def listener(self) -> MicrovmWorkerListener: def post(self, path: str, body: Any, headers: dict[str, str] | None = None) -> int: connection = http.client.HTTPConnection( - "127.0.0.1", self.worker.port, timeout=WAIT + host_of(self.worker), self.worker.port, timeout=WAIT ) try: connection.request( @@ -340,6 +374,11 @@ def finish_job(self, callback_id: str) -> None: """Release a job, and wait until it has ended.""" self.jobs.release(callback_id) assert eventually(lambda: callback_id not in self.listener._jobs) # noqa: SLF001 + # The job's thread removes the job and starts the idle time in one + # locked block. The check above reads without the lock, so it can see + # the removal first. Taking the lock waits for the rest of the block. + with self.listener._lock: # noqa: SLF001 + pass def settle(self) -> None: """Wait until no SuspendMicrovm call runs.""" @@ -425,7 +464,7 @@ def __init__(self, harness: Harness, path: str, body: Any) -> None: data = json.dumps(body).encode() half = len(data) // 2 self._socket = socket.create_connection( - ("127.0.0.1", harness.worker.port), timeout=WAIT + (host_of(harness.worker), harness.worker.port), timeout=WAIT ) head = ( f"POST {path} HTTP/1.1\r\nHost: localhost:8080\r\n" @@ -1351,7 +1390,9 @@ def test_ignores_a_terminate_hook_without_a_host_header(start: Any) -> None: h.job("cb-1") assert eventually(lambda: h.started("cb-1")) - connection = http.client.HTTPConnection("127.0.0.1", h.worker.port, timeout=WAIT) + connection = http.client.HTTPConnection( + host_of(h.worker), h.worker.port, timeout=WAIT + ) try: connection.putrequest("POST", f"{HOOK_PATH_PREFIX}terminate", skip_host=True) connection.putheader("Content-Length", "2") diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py index 88d244ab6..a26173c11 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py @@ -221,6 +221,23 @@ def test_starts_no_attempt_after_the_deadline( retries={"mode": "standard", "total_max_attempts": 5}, ), ) + # Runs after the deadline handler, so it records only the attempts + # that the handler let through. + sent: list[float] = [] + client.meta.events.register( + "before-send.lambda-microvms.SuspendMicrovm", + lambda **_: sent.append(time.monotonic()), + ) + call = client.suspend_microvm + ended = threading.Event() + + def tracked(**kwargs: Any) -> Any: + try: + return call(**kwargs) + finally: + ended.set() + + client.suspend_microvm = tracked # type: ignore[method-assign] started = time.monotonic() # The call ends with the wait's TimeoutError, or with @@ -228,14 +245,21 @@ def test_starts_no_attempt_after_the_deadline( with pytest.raises(TimeoutError): suspend_microvm(client, "mvm-1", 0.5) - # The attempt in flight at the deadline ends at its read timeout, - # and botocore then meets the deadline handler. - time.sleep(2.0) - assert service.attempts - assert all(at - started < 0.5 for at in service.attempts) - # The handler is gone once the call has ended. + # The attempt in flight at the deadline ends at its read timeout. + # botocore then waits its backoff, up to a few seconds, and meets + # the deadline handler, which ends the call. + assert ended.wait(15) handlers = client.meta.events._emitter._unique_id_handlers # type: ignore[attr-defined] # noqa: SLF001 - assert not [key for key in handlers if "microvm-worker-deadline" in key] + # The call's thread removes the handler right after the call ends. + deadline = time.monotonic() + 5 + while [key for key in handlers if "microvm-worker-deadline" in key]: + assert time.monotonic() < deadline + time.sleep(0.01) + assert service.attempts + assert sent + # The handler and this recorder run back to back. The margin covers + # a thread switch between them on a loaded machine. + assert all(at - started < 0.55 for at in sent) client.close() finally: service.close() diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py index 18c87e039..66fe27ec0 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py @@ -6,6 +6,7 @@ from __future__ import annotations import contextlib +import functools import http.client import inspect import io @@ -35,6 +36,39 @@ WAIT = 10.0 +@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 run_payload(*, with_job: bool = True, **job: Any) -> dict[str, Any]: """A run hook payload with one job, whose fields ``job`` replaces.""" payload: dict[str, Any] = {"version": 1, "region": "us-east-1"} @@ -69,9 +103,12 @@ def request( raw: Any = None, headers: dict[str, str] | None = None, chunked: bool = False, + host: str | None = None, ) -> tuple[int, Any]: """Send one request, and return the status and the parsed answer.""" - connection = http.client.HTTPConnection("127.0.0.1", worker.port, timeout=WAIT) + connection = http.client.HTTPConnection( + host or host_of(worker), worker.port, timeout=WAIT + ) try: data = raw if raw is not None else json.dumps(body).encode() connection.request( @@ -90,7 +127,7 @@ def request( def raw_request(worker: MicrovmWorker, data: bytes) -> bytes: """Send bytes, close the sending side, and return the raw answer.""" - with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + with socket.create_connection((host_of(worker), worker.port), timeout=WAIT) as sock: sock.sendall(data) sock.shutdown(socket.SHUT_WR) chunks = [] @@ -588,7 +625,7 @@ def test_answers_expect_100_continue_before_it_reads_the_body(start: Any) -> Non f"Content-Length: {len(data)}\r\n\r\n" ) - with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + with socket.create_connection((host_of(worker), worker.port), timeout=WAIT) as sock: sock.sendall(head.encode()) interim = sock.recv(65536) sock.sendall(data) @@ -644,7 +681,7 @@ def test_answers_400_instead_of_100_continue_to_a_body_over_the_limit( f"Content-Length: {limit + 1}\r\n\r\n" ) - with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + with socket.create_connection((host_of(worker), worker.port), timeout=WAIT) as sock: sock.sendall(head.encode()) answer = b"" while chunk := sock.recv(65536): @@ -674,7 +711,7 @@ def test_answers_100_continue_when_the_declared_size_is_not_over_a_limit( f"{framing}\r\n\r\n" ) - with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + with socket.create_connection((host_of(worker), worker.port), timeout=WAIT) as sock: sock.sendall(head.encode()) interim = sock.recv(65536) @@ -695,7 +732,7 @@ def _send_after_the_answer( The answer, and the error of the first failed write, or ``None``. """ head = f"POST {RUN} HTTP/1.1\r\nHost: localhost\r\nContent-Length: 999999\r\n\r\n" - with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + with socket.create_connection((host_of(worker), worker.port), timeout=WAIT) as sock: sock.sendall(head.encode() + b"x" * 100) answer = b"" while chunk := sock.recv(65536): @@ -1086,7 +1123,7 @@ def test_prints_no_traceback_for_a_client_that_resets_during_the_headers( ) -> None: worker, _ = start(Held()) - with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + with socket.create_connection((host_of(worker), worker.port), timeout=WAIT) as sock: sock.sendall(b"POST /aws/lambda-microvms/runtime/v1/run HTTP/1.1\r\nHo") # The reset must arrive while the server waits for the rest of the # headers. A probe printed the traceback only then. @@ -1368,23 +1405,6 @@ def test_raises_when_the_port_is_taken(start: Any) -> None: start_microvm_worker(handler=Held(), port=worker.port) -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 test_listens_on_ipv6_and_ipv4(start: Any) -> None: if not _has_ipv6_loopback(): pytest.skip("this host has no IPv6 loopback address") @@ -1396,7 +1416,7 @@ def test_listens_on_ipv6_and_ipv4(start: Any) -> None: assert connection.getresponse().status == 200 finally: connection.close() - assert request(worker, f"{HOOK_PATH_PREFIX}ready", {})[0] == 200 + assert request(worker, f"{HOOK_PATH_PREFIX}ready", {}, host="127.0.0.1")[0] == 200 def test_listens_on_ipv4_without_dual_stack( From c1ca4bbdb944ba86cac7b2ca04a6c96dd3777745 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 19 Oct 2026 17:53:20 -0700 Subject: [PATCH 4/4] fix(microvm-worker): wait for late suspend attempt Address the review of #782. - botocore emits before-send before it connects, and cannot cancel an attempt in flight. So an attempt that started before the deadline could reach the service after the refusal ended. The worker now waits for the call to end, and the 30 seconds start after it. - The deadline handler is removed in a finally block in the calling thread. A failed call no longer leaves it on a caller's client. - The resume hook answers before it starts deferred jobs. A job whose thread cannot start fails its callback after the answer. - The IPv4-only test connects to 127.0.0.1. - Tests start the worker again when the picked port reaches another process on the loopback address. --- .../tests/microvm/worker_composed_test.py | 75 +++++++--- .../README.md | 6 +- .../self_suspend.py | 78 ++++------ .../worker.py | 36 +++-- .../tests/lifecycle_hooks_test.py | 92 ++++++++++-- .../tests/self_suspend_test.py | 137 ++++++++++++------ .../tests/worker_test.py | 41 +++++- 7 files changed, 325 insertions(+), 140 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py index af03b4ca4..5fa7dfe6a 100644 --- a/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py +++ b/packages/aws-durable-execution-sdk-python-extras/tests/microvm/worker_composed_test.py @@ -92,6 +92,39 @@ def host_of(worker: MicrovmWorker) -> str: return "127.0.0.1" +def reaches_worker(worker: MicrovmWorker) -> bool: + """Whether a connection to ``host_of(worker)`` reaches this worker. + + The port that the kernel picks can be held by another process on one + loopback address, ``127.0.0.1`` or ``::1``. The worker's ``[::]`` bind + still succeeds on macOS, and a connection then reaches the other + process. So a test checks the ``Server`` header, and starts the worker + again on another port when it is not the worker's. + """ + connection = http.client.HTTPConnection(host_of(worker), worker.port, timeout=5) + try: + connection.request("GET", "/__reachability_check__") + response = connection.getresponse() + response.read() + server = response.getheader("Server") or "" + except OSError: + return False + finally: + connection.close() + return server.startswith("aws-durable-execution-sdk-python-microvm-worker") + + +def start_reachable_worker(start: Callable[[], MicrovmWorker]) -> MicrovmWorker: + """Start a worker with ``start`` until a test connection reaches it.""" + for _ in range(10): + worker = start() + if reaches_worker(worker): + return worker + worker.close() + msg = "every port that the kernel picked was held by another process" + raise RuntimeError(msg) + + def url_of(worker: MicrovmWorker) -> str: """The worker's endpoint URL for the durable function.""" host = host_of(worker) @@ -181,13 +214,17 @@ def handler(event: Any, context: DurableContext) -> Any: runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05) closers.append(runner.close) callbacks = RunnerCallbacks(runner) - worker = start_microvm_worker( - handler=job, - routes=fields.pop("routes", None), - # The double has the three callback methods that the worker calls. - create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] - heartbeat_interval_seconds=fields.pop("heartbeat_interval", None), - port=0, + routes = fields.pop("routes", None) + heartbeat_interval = fields.pop("heartbeat_interval", None) + worker = start_reachable_worker( + lambda: start_microvm_worker( + handler=job, + routes=routes, + # The double has the three callback methods that the worker calls. + create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] + heartbeat_interval_seconds=heartbeat_interval, + port=0, + ) ) closers.append(worker.close) client = fake_client( @@ -306,10 +343,12 @@ def handler(event: Any, context: DurableContext) -> Any: runner = DurableFunctionTestRunner(handler=handler, poll_interval=0.05) callbacks = RunnerCallbacks(runner) - worker = start_microvm_worker( - handler=job, - create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] - port=0, + worker = start_reachable_worker( + lambda: start_microvm_worker( + handler=job, + create_client=lambda region: callbacks, # type: ignore[arg-type, return-value] + port=0, + ) ) try: client = fake_client( @@ -459,12 +498,14 @@ def handler(event: Any, context: DurableContext) -> Any: 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, + worker = start_reachable_worker( + lambda: 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: diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md index 074978d84..885bbedda 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -98,7 +98,7 @@ The worker creates a Lambda client for each job, after the job arrives, not at i A document that does not match the contract gets HTTP 400 with the validation message. If the document named a callback, the worker also fails that callback, so the durable function fails at once instead of at its timeout. A request whose body cannot be read gets HTTP 400 with a generic message, and the error is logged. The worker accepts a `run` hook body of at most 64 KB and a job request body of at most 6 MB, with a `Content-Length` or with chunked transfer encoding. In a chunked body, the chunk-size lines, their extensions, and the trailers have a limit of their own, of the same size. A larger body gets HTTP 400. With `Expect: 100-continue`, the 400 replaces `100 Continue`, so the client does not send the body. Without it, the worker reads and drops the body after the answer, for at most 10 seconds and 24 MB. Without that read, the connection would close with unread bytes, and the client would get a TCP reset instead of the 400. A body that is larger, or that is still arriving after 10 seconds, still gets the reset. Other lifecycle hooks get HTTP 200. Other paths get HTTP 404. -`worker.idle(timeout)` waits until every job started so far has ended, and the wait for a running `SuspendMicrovm` call has ended. It returns `False` when the timeout ends first. `worker.close()` stops the server and the idle time. Running jobs continue, so call `idle()` first to let them report their outcome. +`worker.idle(timeout)` waits until every job started so far has ended, and a running `SuspendMicrovm` call has returned. It returns `False` when the timeout ends first. `worker.close()` stops the server and the idle time. Running jobs continue, so call `idle()` first to let them report their outcome. ## Suspending when idle @@ -106,7 +106,7 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets 1. The idle time starts when the `run` hook arrives without a job, when a job ends, and when the `resume` lifecycle hook arrives, if the image enables it. 2. A job that starts cancels the idle time. So the worker never suspends the MicroVM during a job. -3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. The worker waits for it for at most 10 seconds. No attempt of the call starts after that, because a `before-send` handler on the client stops botocore's retries at the deadline. An attempt already in flight then runs until its read timeout. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. +3. When the idle time ends, the worker calls `SuspendMicrovm` with the MicroVM ID from the `run` hook, never one from a job request. The call returned in under 100 milliseconds in testing, before the MicroVM froze. No attempt of the call starts after 10 seconds, because a `before-send` handler on the client stops botocore's retries at that deadline. botocore cannot cancel an attempt in flight, and it opens the connection after `before-send`. So an attempt that started before the deadline can still reach the service after it. The worker waits for that attempt to end, and refuses jobs until then. The client's connect and read timeouts bound the attempt: 5 and 10 seconds for the default client. A client from `create_microvms_client` should use similar timeouts. botocore's own defaults are 60 seconds each, so with them the refusal can last minutes. The durable function resumes the MicroVM before the next job, and the process continues with its memory and files. 4. From the moment the idle time ends, the worker answers every job with HTTP 503 and does not start it. A job accepted then would freeze with the MicroVM, and nothing would resume it. After a 503, the durable function checks the state, resumes the MicroVM, and delivers the job again. 5. The refusal ends when the `resume` lifecycle hook arrives, or 30 seconds after the call returned, if no hook arrives. The idle time then starts again. So enable the `resume` hook in the image. Without it, the first job after a resume can wait a little over 30 seconds: the 30 seconds, plus up to 4 seconds until the durable function's next delivery attempt. A test measured 32 seconds. 6. A call that the service rejected on its first attempt, such as `AccessDeniedException` for a missing permission, is logged as a warning, and the MicroVM keeps running. Missing credentials count as a rejection, because the call sends no request. The worker accepts jobs again at once and does not retry the call. The next job that ends starts the idle time again. @@ -117,7 +117,7 @@ A session's run hook payload can set `autoSuspendIdleSeconds`. The session sets The worker counts only jobs. It cannot see background threads or processes, or inbound traffic that is not a job. A session with such work leaves `auto_suspend_on_idle` off, and the payload then has no `autoSuspendIdleSeconds`. -The worker creates a `lambda-microvms` client for each suspend, and closes it when its wait for the call ends. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. +The worker creates a `lambda-microvms` client for each suspend, and closes it after the call. Pass `create_microvms_client` to provide the client. The worker never closes a client that it returns. ## Suspend and terminate hooks diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py index b315ddfc1..1ed3157b4 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/self_suspend.py @@ -5,9 +5,7 @@ from __future__ import annotations -import threading import time -from concurrent.futures import Future from enum import Enum from typing import TYPE_CHECKING, Any @@ -31,9 +29,9 @@ SUSPEND_CALL_TIMEOUT_SECONDS = 10.0 -"""The longest time the worker waits for its own SuspendMicrovm call. +"""No attempt of the worker's own SuspendMicrovm call starts after this time. -No attempt of the call starts after this time. The call measured 75 to 94 +An attempt in flight then runs until it ends. The call measured 75 to 94 milliseconds from inside a MicroVM in us-east-1. """ @@ -85,8 +83,7 @@ class SuspendDeadlineError(TimeoutError): """The SuspendMicrovm deadline passed before an attempt could start. botocore retries only its own connection and HTTP client errors. So it - does not retry this one. The call can end with it just before the wait - ends, so it is a ``TimeoutError`` too. + does not retry this one. """ @@ -112,63 +109,52 @@ def create_microvms_client(region: str) -> LambdaMicrovmsClient: def suspend_microvm( - client: LambdaMicrovmsClient, microvm_id: str, timeout: float + client: LambdaMicrovmsClient, microvm_id: str, deadline_seconds: float ) -> None: - """Call SuspendMicrovm, and stop waiting after ``timeout`` seconds. + """Call SuspendMicrovm in this thread, and start no attempt after + ``deadline_seconds``. botocore has no deadline for a whole call. Its read timeout limits one socket read, and it covers neither credential resolution nor its own - retries. So the call needs two limits: - - 1. The call runs in a daemon thread, and this function waits for the - thread for at most ``timeout``. - 2. A ``before-send`` handler on the client raises - :class:`SuspendDeadlineError` when an attempt would start after - ``timeout``. botocore does not retry that error. So no attempt starts - after this function has stopped waiting. Without the handler, botocore - went on retrying in the thread, and an attempt started 22.5 seconds - after the first one. - - An attempt that is already in flight at the deadline still runs until - its read timeout. The service receives its request before the deadline. - The handler is removed when the call ends. So a client from the caller - keeps no handler. + retries. So a ``before-send`` handler on the client raises + :class:`SuspendDeadlineError` when an attempt would start after the + deadline. botocore does not retry that error. Without the handler, + botocore went on retrying, and an attempt started 22.5 seconds after the + first one. + + botocore emits ``before-send`` before it opens the connection, and it + cannot cancel an attempt in flight. So an attempt that started before + the deadline can still reach the service after it. This function + returns only when that attempt has ended, so that the caller keeps + refusing jobs until then. The client's connect and read timeouts bound + the attempt: 5 and 10 seconds for the default client. + + The handler is removed when the call ends, also when it raises. So a + client from the caller keeps no handler. Raises: - TimeoutError: When the call takes longer than ``timeout``. + SuspendDeadlineError: When an attempt would start after the + deadline. Exception: The call's own error. """ - future: Future[None] = Future() - deadline = time.monotonic() + timeout - unique_id = f"aws-durable-execution-microvm-worker-deadline-{id(future)}" + deadline = time.monotonic() + deadline_seconds + unique_id = f"aws-durable-execution-microvm-worker-deadline-{id(object())}" events = safe_get(lambda: client.meta.events) def stop_after_deadline(**_: object) -> None: if time.monotonic() >= deadline: - msg = f"the SuspendMicrovm deadline of {timeout} seconds has passed" + msg = ( + f"the SuspendMicrovm deadline of {deadline_seconds} seconds has passed" + ) raise SuspendDeadlineError(msg) if events is not None: events.register_first(_BEFORE_SEND, stop_after_deadline, unique_id=unique_id) - - def run() -> None: - try: - client.suspend_microvm(microvmIdentifier=microvm_id) - future.set_result(None) - except BaseException as error: # noqa: BLE001 - future.set_exception(error) - finally: - if events is not None: - events.unregister(_BEFORE_SEND, unique_id=unique_id) - - threading.Thread(target=run, name="microvm-suspend-call", daemon=True).start() try: - future.result(timeout) - except TimeoutError as error: - if future.done(): - raise - msg = f"the SuspendMicrovm call took longer than {timeout} seconds" - raise TimeoutError(msg) from error + client.suspend_microvm(microvmIdentifier=microvm_id) + finally: + if events is not None: + events.unregister(_BEFORE_SEND, unique_id=unique_id) def classify_suspend_error(error: BaseException) -> SuspendOutcome: diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py index 208425e94..872bbe249 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -304,8 +304,8 @@ class MicrovmWorkerListener: 8. A call whose outcome is unknown, such as a timeout, a conflict, a server error, or an error after a retried attempt, is handled like a successful call. The service may have suspended the MicroVM anyway. - The listener waits for the call for at most 10 seconds, and no - attempt starts after that. + No attempt of the call starts after 10 seconds. An attempt already in + flight runs until it ends, and the refusal lasts at least until then. When the image enables the ``suspend`` and ``terminate`` hooks, the listener acts on them: @@ -565,10 +565,10 @@ def handle_expect_100(self, request: BaseHTTPRequestHandler) -> bool: def idle(self, timeout: float | None = None) -> bool: """Wait until every job started so far has ended. - The wait also covers the failure reports of invalid requests, and the - worker's wait for its own SuspendMicrovm call. That wait takes at most - 10 seconds, and no attempt of the call starts after it. An attempt - already in flight then can run until its read timeout. + The wait also covers the failure reports of invalid requests, and a + running SuspendMicrovm call. No attempt of that call starts after 10 + seconds. With the default client, an attempt in flight then ends + within 15 seconds. Args: timeout: The longest wait in seconds. ``None`` waits until the @@ -631,11 +631,16 @@ def _dispatch(self, exchange: _Exchange) -> None: # otherwise accept a job between the answer and the change. The # change only takes the lock, so the answer is not delayed. A # suspend hook that runs past its timeout terminates the MicroVM. + deferred: list[tuple[MicrovmJobDocument, str]] = [] if hook == "suspend": self._handle_suspend_hook() elif hook == "resume": - self._handle_resume_hook() + deferred = self._handle_resume_hook() exchange.respond(200, {}) + # The deferred jobs are registered, and start after the answer. A + # job whose thread cannot start fails its callback in this thread, + # which can take minutes. The resume hook must not wait for that. + self._launch_deferred(deferred) return handler = self._job_handler(path) if handler is not None and exchange.method == "POST": @@ -655,19 +660,18 @@ def _handle_suspend_hook(self) -> None: self._suspending = True self._start_grace() - def _handle_resume_hook(self) -> None: + def _handle_resume_hook(self) -> list[tuple[MicrovmJobDocument, str]]: """Accept jobs again, and start the idle time. A resume without a job would otherwise leave the MicroVM running - until the next job ends. + until the next job ends. Returns the deferred run hook jobs, + registered. The caller starts them after the answer. """ with self._lock: if self._suspending: - deferred = self._end_suspending() - else: - deferred = [] - self._start_idle_time() - self._launch_deferred(deferred) + return self._end_suspending() + self._start_idle_time() + return [] def _handle_terminate_hook(self, exchange: _Exchange) -> None: """Fail the callback of every running job, then answer the hook. @@ -1569,8 +1573,8 @@ def port(self) -> int: return int(self._server.server_address[1]) def idle(self, timeout: float | None = None) -> bool: - """Wait until every job started so far has ended, and the wait for a - running SuspendMicrovm call has ended. + """Wait until every job started so far has ended, and a running + SuspendMicrovm call has returned. Returns: ``True`` when no job runs. ``False`` when the timeout ended first. diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py index a66a99239..175306726 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/lifecycle_hooks_test.py @@ -79,6 +79,39 @@ def host_of(worker: MicrovmWorker) -> str: return "127.0.0.1" +def reaches_worker(worker: MicrovmWorker) -> bool: + """Whether a connection to ``host_of(worker)`` reaches this worker. + + The port that the kernel picks can be held by another process on one + loopback address, ``127.0.0.1`` or ``::1``. The worker's ``[::]`` bind + still succeeds on macOS, and a connection then reaches the other + process. So a test checks the ``Server`` header, and starts the worker + again on another port when it is not the worker's. + """ + connection = http.client.HTTPConnection(host_of(worker), worker.port, timeout=5) + try: + connection.request("GET", "/__reachability_check__") + response = connection.getresponse() + response.read() + server = response.getheader("Server") or "" + except OSError: + return False + finally: + connection.close() + return server.startswith("aws-durable-execution-sdk-python-microvm-worker") + + +def start_reachable_worker(start: Callable[[], MicrovmWorker]) -> MicrovmWorker: + """Start a worker with ``start`` until a test connection reaches it.""" + for _ in range(10): + worker = start() + if reaches_worker(worker): + return worker + worker.close() + msg = "every port that the kernel picked was held by another process" + raise RuntimeError(msg) + + def eventually(condition: Callable[[], bool]) -> bool: deadline = time.monotonic() + WAIT while not condition(): @@ -421,13 +454,15 @@ def create_client(region: str) -> Any: return client options.setdefault("logger", logger) - worker = start_microvm_worker( - handler=jobs, - routes={"/job": jobs}, - create_client=create_client, - create_microvms_client=create_microvms_client, - port=0, - **options, + worker = start_reachable_worker( + lambda: start_microvm_worker( + handler=jobs, + routes={"/job": jobs}, + create_client=create_client, + create_microvms_client=create_microvms_client, + port=0, + **options, + ) ) harness = Harness( worker=worker, @@ -931,20 +966,31 @@ def test_keeps_a_hook_refusal_after_a_refusal_for_now_of_the_own_call( assert h.job("cb-1") == 503 -def test_treats_a_call_that_outlasts_its_timeout_as_an_unknown_outcome( - start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +def test_keeps_refusing_jobs_until_a_call_past_its_deadline_has_ended( + start: Any, monkeypatch: pytest.MonkeyPatch ) -> None: + # botocore cannot cancel an attempt in flight, so the attempt can reach + # the service after the deadline. The refusal and its 30 seconds start + # only when the call has ended. monkeypatch.setattr(worker_module, "SUSPEND_CALL_TIMEOUT_SECONDS", 0.05) h = start() h.microvms.hold = True h.run_hook(IDLE) h.advance(IDLE) - h.settle() - - assert "unknown outcome" in logger.messages("warning")[0] + assert eventually(lambda: h.microvms.suspended == ["mvm-1"]) + # The deadline passes while the call is held. + time.sleep(0.3) + h.advance(GRACE * 3) assert h.job("cb-1") == 503 + assert not h.worker.idle(0.05) + h.microvms.finish() + h.settle() + h.advance(GRACE - 0.001) + assert h.job("cb-1") == 503 + h.advance(0.001) + assert h.job("cb-1") == 202 def test_accepts_jobs_at_once_when_the_microvms_client_cannot_be_created( @@ -1895,3 +1941,25 @@ def test_terminate_waits_for_and_does_not_repeat_a_heartbeat_start_failure( ("cb-1", "RuntimeError") ] assert not h.started("cb-1") + + +def test_answers_the_resume_hook_before_it_reports_a_deferred_job_that_cannot_start( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # A deferred job whose thread cannot start fails its callback, which can + # take minutes. Lambda allows the resume hook about 10 seconds. + h = start() + run = slow_run_hook(h) + h.hook("suspend") + assert run.finish() == 200 + h.lambda_client.hold = True + failing_thread_start(monkeypatch, "microvm-job-") + + answer = h.hook_in_background("resume") + + assert answer.result(WAIT) == 200 + assert eventually(lambda: h.lambda_client.waiting() == 1) + monkeypatch.undo() + h.lambda_client.finish() + assert eventually(lambda: len(h.lambda_client.recorded()) == 1) + assert h.lambda_client.recorded()[0]["errorType"] == "RuntimeError" diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py index a26173c11..c6ba6ddaa 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/self_suspend_test.py @@ -10,6 +10,8 @@ import time from typing import Any +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + import boto3 import pytest from botocore.config import Config @@ -145,15 +147,15 @@ def test_raises_the_call_s_own_timeout_error() -> None: suspend_microvm(client, "mvm-1", 5.0) -def test_stops_waiting_for_a_stalled_call_after_the_timeout() -> None: - client: Any = _Client(delay=10.0) +def test_returns_only_when_a_call_in_flight_at_the_deadline_has_ended() -> None: + # botocore cannot cancel an attempt in flight. So the call returns only + # when the attempt ends, and the worker refuses jobs until then. + client: Any = _Client(delay=0.3) started = time.monotonic() - with pytest.raises(TimeoutError, match="longer than 0.05 seconds"): - suspend_microvm(client, "mvm-1", 0.05) + suspend_microvm(client, "mvm-1", 0.05) - assert time.monotonic() - started < 2.0 - client.release.set() + assert time.monotonic() - started >= 0.29 def test_the_default_client_ends_a_read_at_the_call_timeout() -> None: @@ -202,25 +204,36 @@ def close(self) -> None: self._server.close() -def test_starts_no_attempt_after_the_deadline( - monkeypatch: pytest.MonkeyPatch, -) -> None: - # botocore retries a read timeout. Without the deadline handler, it went - # on retrying in the call's thread after the wait had ended. +def _microvms_client(url: str, read_timeout: float = 5.0) -> Any: + return boto3.session.Session().client( + "lambda-microvms", + region_name="us-east-1", + endpoint_url=url, + config=Config( + connect_timeout=1, + read_timeout=read_timeout, + retries={"mode": "standard", "total_max_attempts": 5}, + ), + ) + + +def _deadline_handlers(client: Any) -> list[str]: + handlers = client.meta.events._emitter._unique_id_handlers # noqa: SLF001 + return [key for key in handlers if "microvm-worker-deadline" in key] + + +@pytest.fixture +def credentials(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test") + + +def test_starts_no_attempt_after_the_deadline(credentials: None) -> None: + # botocore retries a read timeout. Without the deadline handler, it went + # on retrying after the deadline. service = _SilentService() try: - client = boto3.session.Session().client( - "lambda-microvms", - region_name="us-east-1", - endpoint_url=service.url, - config=Config( - connect_timeout=1, - read_timeout=0.3, - retries={"mode": "standard", "total_max_attempts": 5}, - ), - ) + client = _microvms_client(service.url, read_timeout=0.3) # Runs after the deadline handler, so it records only the attempts # that the handler let through. sent: list[float] = [] @@ -228,38 +241,74 @@ def test_starts_no_attempt_after_the_deadline( "before-send.lambda-microvms.SuspendMicrovm", lambda **_: sent.append(time.monotonic()), ) - call = client.suspend_microvm - ended = threading.Event() - - def tracked(**kwargs: Any) -> Any: - try: - return call(**kwargs) - finally: - ended.set() - - client.suspend_microvm = tracked # type: ignore[method-assign] started = time.monotonic() - # The call ends with the wait's TimeoutError, or with - # SuspendDeadlineError when botocore meets the handler first. - with pytest.raises(TimeoutError): + # The attempt in flight at the deadline ends at its read timeout. + # botocore then waits its backoff, and meets the deadline handler. + with pytest.raises(SuspendDeadlineError): suspend_microvm(client, "mvm-1", 0.5) - # The attempt in flight at the deadline ends at its read timeout. - # botocore then waits its backoff, up to a few seconds, and meets - # the deadline handler, which ends the call. - assert ended.wait(15) - handlers = client.meta.events._emitter._unique_id_handlers # type: ignore[attr-defined] # noqa: SLF001 - # The call's thread removes the handler right after the call ends. - deadline = time.monotonic() + 5 - while [key for key in handlers if "microvm-worker-deadline" in key]: - assert time.monotonic() < deadline - time.sleep(0.01) assert service.attempts assert sent # The handler and this recorder run back to back. The margin covers # a thread switch between them on a loaded machine. assert all(at - started < 0.55 for at in sent) + assert _deadline_handlers(client) == [] + client.close() + finally: + service.close() + + +class _AnsweringService: + """A local HTTP service that answers every request with ``200 {}``.""" + + def __init__(self) -> None: + self.requests = 0 + service = self + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: # noqa: N802 + service.requests += 1 + self.rfile.read(int(self.headers.get("Content-Length", "0"))) + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", "2") + self.end_headers() + self.wfile.write(b"{}") + + def log_message(self, *args: Any) -> None: + return + + self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=self._server.serve_forever, daemon=True).start() + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self._server.server_address[1]}" + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + + +def test_a_caller_client_keeps_no_expired_handler_after_a_failed_call( + credentials: None, +) -> None: + # A client from create_microvms_client can serve every suspend. An + # expired handler left on it would fail every later call before it + # sends. + service = _AnsweringService() + try: + client = _microvms_client(service.url) + + with pytest.raises(SuspendDeadlineError): + suspend_microvm(client, "mvm-1", 0) + assert _deadline_handlers(client) == [] + assert service.requests == 0 + + suspend_microvm(client, "mvm-1", 5.0) + assert service.requests == 1 + assert _deadline_handlers(client) == [] client.close() finally: service.close() diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py index 6a88e454e..899b602ee 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py @@ -69,6 +69,39 @@ def host_of(worker: MicrovmWorker) -> str: return "127.0.0.1" +def reaches_worker(worker: MicrovmWorker) -> bool: + """Whether a connection to ``host_of(worker)`` reaches this worker. + + The port that the kernel picks can be held by another process on one + loopback address, ``127.0.0.1`` or ``::1``. The worker's ``[::]`` bind + still succeeds on macOS, and a connection then reaches the other + process. So a test checks the ``Server`` header, and starts the worker + again on another port when it is not the worker's. + """ + connection = http.client.HTTPConnection(host_of(worker), worker.port, timeout=5) + try: + connection.request("GET", "/__reachability_check__") + response = connection.getresponse() + response.read() + server = response.getheader("Server") or "" + except OSError: + return False + finally: + connection.close() + return server.startswith("aws-durable-execution-sdk-python-microvm-worker") + + +def start_reachable_worker(start: Callable[[], MicrovmWorker]) -> MicrovmWorker: + """Start a worker with ``start`` until a test connection reaches it.""" + for _ in range(10): + worker = start() + if reaches_worker(worker): + return worker + worker.close() + msg = "every port that the kernel picked was held by another process" + raise RuntimeError(msg) + + def run_payload(*, with_job: bool = True, **job: Any) -> dict[str, Any]: """A run hook payload with one job, whose fields ``job`` replaces.""" payload: dict[str, Any] = {"version": 1, "region": "us-east-1"} @@ -185,7 +218,9 @@ def _start( client = client if client is not None else fake_client() options.setdefault("create_client", lambda region: client) options.setdefault("logger", logger) - worker = start_microvm_worker(handler=handler, port=0, **options) + worker = start_reachable_worker( + lambda: start_microvm_worker(handler=handler, port=0, **options) + ) workers.append(worker) return worker, client @@ -1473,7 +1508,9 @@ def test_listens_on_ipv4_without_dual_stack( monkeypatch.setattr(worker_module.socket, "has_dualstack_ipv6", lambda: False) worker = start_microvm_worker(handler=Held(), port=0) try: - assert request(worker, f"{HOOK_PATH_PREFIX}ready", {})[0] == 200 + assert ( + request(worker, f"{HOOK_PATH_PREFIX}ready", {}, host="127.0.0.1")[0] == 200 + ) finally: worker.close()