From e95851755dd56a96cf39a758714c3f7c0800eb58 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 12 Oct 2026 09:15:11 -0700 Subject: [PATCH 1/4] feat(microvm-worker): add HTTP listener and server The worker now answers the MicroVM lifecycle hooks and the job routes. start_microvm_worker serves them on port 8080. - The run hook gets HTTP 200 at once, and its job runs after the answer. Only the first valid run hook is used. - A job request gets HTTP 202, and runs the route's handler. MICROVM_JOB_PATH runs the default handler. - Each job runs in its own thread, with heartbeats and a callback reporter. A second delivery of a callback ID runs nothing. - An invalid document gets HTTP 400, and fails the callback that it names. A body over the size limit gets its 400 in place of 100 Continue, or the worker drops the rest of the body after the answer, so the client reads the 400 instead of a TCP reset. - The server listens on IPv6 and IPv4, reads chunked bodies with strict chunk sizes, and answers Expect: 100-continue. - An async handler is refused at create time. A handler that returns an awaitable fails its job with TypeError. The closed-callback warning now uses the wording of the JS worker (aws/aws-durable-execution-sdk-js#984), and says that the durable function no longer waits for the callback in either case. The suspend, terminate, and resume hooks get HTTP 200 and no action. A later change acts on them. --- .../README.md | 90 +- .../__init__.py | 22 +- .../callback_reporter.py | 9 +- .../worker.py | 1138 +++++++++++++++ .../tests/callback_reporter_test.py | 3 +- .../tests/worker_test.py | 1295 +++++++++++++++++ 6 files changed, 2546 insertions(+), 11 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py create mode 100644 packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py 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 c83ad8615..750f2a6ee 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -6,13 +6,95 @@ This package runs inside an AWS Lambda MicroVM. It receives jobs from a durable The package does not depend on the durable execution SDK, because it runs in the MicroVM and not in the durable function. The job document is the same in every SDK language. So a durable function in any language can send jobs to this worker. -The package currently contains these parts: +## Usage +```python +import logging + +from aws_durable_execution_sdk_python_microvm_worker import start_microvm_worker + + +def run_build(job_input, context): + # context.cancelled is set when the durable function stops waiting. + return {"passed": build(job_input["repo"], stop=context.cancelled)} + + +def clone_build(job_input, context): + return clone_and_build(job_input) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + worker = start_microvm_worker( + # Jobs that name no route. A small job arrives in the run hook, and a + # large one over HTTP. The handler receives both the same way. + handler=run_build, + # Jobs sent to a named route. + routes={"/clone-build": clone_build}, + ) + 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. + +`start_microvm_worker` serves in a daemon thread and returns. So the entry point ends with `worker.wait()`, which blocks until `worker.close()`. + +## Image configuration + +Start the worker when the image's process starts. Lambda snapshots the process after the `ready` hook, so the worker must already be listening. + +The image needs these hook settings in `CreateMicrovmImage`: + +```json +{ + "hooks": { + "port": 8080, + "microvmHooks": { "run": "ENABLED", "runTimeoutInSeconds": 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 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: + +1. Validates the run hook payload. +2. Answers HTTP 200 at once, because a slow or failed `run` hook can terminate the MicroVM. +3. Runs `handler` for the payload's job, when it has one. A payload without a job belongs to a MicroVM whose jobs arrive over HTTP. + +The worker acts only on the first valid `run` hook. The MicroVM endpoint forwards a request to a hook path from any caller that holds an auth token for the MicroVM. So a later `run` hook gets HTTP 200, and is ignored. + +On `POST ` the worker validates the job request, answers HTTP 202 at once, and runs the route's handler. `POST /durable-execution/v1/job` (`MICROVM_JOB_PATH`) runs `handler`. That path is reserved, so a route cannot use it. A job request can arrive before the `run` hook. Until the `run` hook arrives, the worker takes the MicroVM ID for `context.microvm_id` from the first job request that has a `microvmId`. The `run` hook's ID then replaces it for later jobs. Without either, `context.microvm_id` is `"unknown"`. + +Each job runs in its own thread. For every job, the worker: + +1. Sends heartbeats while the job runs and while its outcome is reported, when the job has a heartbeat timeout. See [Heartbeats](#heartbeats). +2. Sets `context.cancelled` when a heartbeat finds the callback gone while the handler runs. Python cannot stop a running thread, so the handler should check the event and stop its work. The worker then does not report the outcome. +3. Completes the callback with the handler's result, or fails it with the handler's exception. A `SystemExit` from the handler is reported as a failure too. +4. Reports a result that is over 256 KB, or that is not JSON-serializable, as a failure. So the durable function learns why at once, instead of at its timeout. + +A second delivery of a job is answered like the first and ignored, while the job runs and after it ended. The durable function can deliver the same job again: it retries a delivery whose answer it did not receive, and its request step runs again on replay when its result was not checkpointed. The worker remembers the callback IDs of the last 1,000 finished jobs. So a job runs once, unless more than 1,000 other jobs ended in between, or the worker process restarted. + +The worker creates a Lambda client for each job, after the job arrives, not at image build. A client created before the snapshot would carry build-time state into every MicroVM. A client kept across jobs would keep connections that a suspend of the MicroVM can leave dead. The worker closes each job's default client when the job ends. It never closes a client from `create_client`. + +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. 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. 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. + +## 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`. - `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. -- The `heartbeats` module sends a job's heartbeats on a schedule. The listener in a later change starts them for each job. - -The HTTP listener for the lifecycle hooks and the job routes comes in a later change. ## The job documents 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 c23eadf1f..f8bff7164 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 @@ -6,9 +6,9 @@ This package is experimental. Its API can change in any release. The exports match the parts of the JavaScript worker's ``index.ts`` that -this package has so far. The listener API comes in a later change. -``CancelScope`` and ``CallCancelledError`` are extra, because the public -``heartbeat()`` signature needs them, and Python has no ``AbortSignal``. +this package has so far. ``CancelScope`` and ``CallCancelledError`` are +extra, because the public ``heartbeat()`` signature needs them, and Python +has no ``AbortSignal``. """ from aws_durable_execution_sdk_python_microvm_worker.__about__ import __version__ @@ -32,22 +32,38 @@ MicrovmRunHookPayload, RunHookRequest, ) +from aws_durable_execution_sdk_python_microvm_worker.worker import ( + HOOK_PATH_PREFIX, + MICROVM_JOB_PATH, + MicrovmJobContext, + MicrovmJobHandler, + MicrovmWorker, + MicrovmWorkerListener, + start_microvm_worker, +) __all__ = [ + "HOOK_PATH_PREFIX", "MAX_CALLBACK_RESULT_BYTES", + "MICROVM_JOB_PATH", "SUPPORTED_PAYLOAD_VERSION", "CallCancelledError", "CallbackReporter", "CancelScope", "InvalidRunHookPayloadError", + "MicrovmJobContext", "MicrovmJobDocument", + "MicrovmJobHandler", "MicrovmJobRequest", "MicrovmRunHookPayload", + "MicrovmWorker", + "MicrovmWorkerListener", "MicrovmWorkerLogger", "ResultSerializationError", "ResultTooLargeError", "RunHookRequest", "__version__", "is_terminal_callback_error", + "start_microvm_worker", ] 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 3116053e7..2cf8c96bd 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 @@ -460,9 +460,12 @@ def _with_retry(self, call: Callable[[], object]) -> None: error ) == CLOSED_CALLBACK_CODE: self._logger.warning( - "the callback is already complete or timed out. An " - "earlier attempt whose answer never arrived probably " - "reported the outcome.", + "the callback is closed. The service gives the same " + "answer for a callback that is already complete and " + "for one that timed out. So either an earlier attempt " + "whose answer never arrived reported the outcome, or " + "the callback timed out first. The durable function " + "no longer waits for this callback in either case.", extra={"callbackId": self.callback_id, "attempt": attempt}, ) return 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 new file mode 100644 index 000000000..b832a16ba --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/src/aws_durable_execution_sdk_python_microvm_worker/worker.py @@ -0,0 +1,1138 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""The HTTP listener that receives jobs inside the MicroVM. + +Lambda calls the MicroVM's lifecycle hooks over HTTP. The durable function +sends a job in the ``run`` hook, or in an HTTP request to a job route. This +module answers both, runs each job in its own thread, and reports the job's +outcome to its callback. +""" + +from __future__ import annotations + +import contextlib +import inspect +import json +import os +import re +import socket +import socketserver +import threading +import time +from collections import OrderedDict +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import cached_property +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import TYPE_CHECKING, Any + +from aws_durable_execution_sdk_python_microvm_worker.callback_reporter import ( + CallbackReporter, + ResultSerializationError, + ResultTooLargeError, + is_terminal_callback_error, +) +from aws_durable_execution_sdk_python_microvm_worker.heartbeats import ( + Heartbeats, + validate_heartbeat_interval, +) +from aws_durable_execution_sdk_python_microvm_worker.logger import ( + MicrovmWorkerLogger, + describe, + safe_logger, +) +from aws_durable_execution_sdk_python_microvm_worker.payload import ( + InvalidRunHookPayloadError, + MicrovmJobDocument, + MicrovmJobRequest, + RunHookRequest, + loads_strict, +) + + +if TYPE_CHECKING: + from mypy_boto3_lambda import LambdaClient +else: + LambdaClient = Any + + +HOOK_PATH_PREFIX = "/aws/lambda-microvms/runtime/v1/" +"""Lambda sends the lifecycle hooks to paths under this prefix.""" + +MICROVM_JOB_PATH = "/durable-execution/v1/job" +"""The route that receives a job over HTTP when the job names no route. + +The durable function sends a job here when the job is too large for the +``run`` hook, and for an invoke without a path. The worker runs ``handler`` +for it. The value must equal ``DEFAULT_MICROVM_JOB_PATH`` in the durable +function's MicroVM operations. +""" + +DEFAULT_PORT = 8080 +"""The default port. It is also the port that the MicroVM endpoint forwards to.""" + +MAX_HOOK_BODY_BYTES = 64 * 1024 +"""The largest ``run`` hook body the worker reads. The payload is at most 16 KB.""" + +MAX_JOB_BODY_BYTES = 6 * 1024 * 1024 +"""The largest job request body the worker reads.""" + +MAX_FINISHED_JOBS = 1_000 +"""How many finished callback IDs the worker remembers. + +The durable function can deliver a job again after it finished. Its request +step runs again on replay when its result was not checkpointed, and it +resends a job whose 202 answer was lost. Each callback ID belongs to one job. +So the worker ignores a callback ID that it already ran. A callback ID is a +short string, so 1,000 of them use well under 1 MB. +""" + +_REQUEST_TIMEOUT_SECONDS = 60.0 +"""How long the worker waits for each read of a request. + +A client that stops sending would otherwise hold its request thread forever. +Lambda allows the ``run`` hook at most 60 seconds. +""" + +_SERVER_POLL_SECONDS = 0.1 +"""How often the server checks for a shutdown request. It bounds ``close()``.""" + +_UNKNOWN_MICROVM_ID = "unknown" + +_CONTENT_LENGTH = re.compile(r"[0-9]+") + +_CHUNK_SIZE = re.compile(rb"[0-9A-Fa-f]+") + +_MAX_CHUNK_LINE_BYTES = 1024 + +_DISCARD_SECONDS = 10.0 +"""How long the worker reads and drops a body that it answered with 400.""" + +_MAX_DISCARD_BYTES = 4 * MAX_JOB_BODY_BYTES +"""How many bytes of such a body the worker reads and drops at most.""" + +_DISCARD_READ_BYTES = 64 * 1024 + + +@dataclass(frozen=True, kw_only=True) +class MicrovmJobContext: + """What the job handler receives besides its input.""" + + callback_id: str + """The callback ID that the worker completes with the handler's result.""" + microvm_id: str + """The MicroVM identifier. + + The worker learns it from the ``run`` hook, or from a job request that + arrives first. It is ``"unknown"`` until then. + """ + 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. + """ + logger: MicrovmWorkerLogger + """The worker's logger.""" + + +MicrovmJobHandler = Callable[[Any, MicrovmJobContext], Any] +"""The job to run. + +The handler receives the job's input and a :class:`MicrovmJobContext`. Its +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. +""" + + +class MicrovmWorkerListener: + """Answers the MicroVM lifecycle hooks and the job routes. + + Create it with :meth:`create`. :func:`start_microvm_worker` creates one + and serves it. Use this class directly to serve the hooks from your own + ``http.server`` server, through :attr:`request_handler_class`. + + On ``POST .../run`` the listener: + + 1. Validates the body and its run hook payload. + 2. Answers HTTP 200 at once. Lambda allows the ``run`` hook at most 60 + seconds, and a slow or failed hook can terminate the MicroVM. So the + job starts after the answer. + 3. Runs ``handler`` for the payload's job, when the payload has one. A + payload without a job belongs to a MicroVM whose jobs arrive over + HTTP. + + On ``POST `` the listener validates the job request, answers HTTP + 202, and runs the route's handler. ``POST`` to :data:`MICROVM_JOB_PATH` + runs ``handler``. + + For every job, the listener sends heartbeats until the job's outcome is + reported, when the job sets a heartbeat timeout. When the handler + returns, the listener completes the callback with the result. When the + handler raises, the listener fails the callback with the exception. + + A second delivery of a callback ID is answered like the first and + ignored, while that job runs and after it ended. The listener remembers + the last 1,000 finished callback IDs. + + A document that does not match the contract gets HTTP 400. If the + document still names a callback, the listener fails that callback. So + the durable function fails at once instead of at its timeout. + + 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. + """ + + def __init__( + self, + *, + handler: MicrovmJobHandler | None, + routes: Mapping[str, MicrovmJobHandler], + create_client: Callable[[str], LambdaClient] | None, + heartbeat_interval_seconds: float | None, + logger: MicrovmWorkerLogger, + ) -> None: + self._handler = handler + self._routes = dict(routes) + self._create_client = create_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. + 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. + 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. + self._run_hook_accepted = False + + @classmethod + def create( + cls, + *, + handler: MicrovmJobHandler | None = None, + routes: Mapping[str, MicrovmJobHandler] | None = None, + create_client: Callable[[str], LambdaClient] | None = None, + heartbeat_interval_seconds: float | None = None, + logger: MicrovmWorkerLogger | None = None, + ) -> MicrovmWorkerListener: + """Create a listener. + + Args: + handler: The job to run when a job names no route. The durable + function sends such a job in the ``run`` hook, or over HTTP + to :data:`MICROVM_JOB_PATH` when the job is too large for the + ``run`` hook. + routes: The jobs to run when a job arrives over HTTP on a named + route, keyed by path. + create_client: Creates the Lambda client for a Region. The + listener calls it for each job, after the job arrives, and to + fail the callback of an invalid request. It never closes a + 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. + 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. + logger: The worker's logger. Defaults to the package's standard + library logger. + + Raises: + TypeError: When there is neither a handler nor a route, or a + handler is not a plain function. An ``async def`` function + is not supported. + ValueError: When a route path is not valid, or + ``heartbeat_interval_seconds`` is not above 0 and at most + 900. + """ + routes = dict(routes or {}) + if handler is None and not routes: + msg = "the MicroVM worker needs a handler or routes" + raise TypeError(msg) + for job_handler in [handler, *routes.values()]: + _check_handler(job_handler) + for path in routes: + if not path.startswith("/") or path.startswith(HOOK_PATH_PREFIX): + msg = ( + f'route "{path}" must start with "/" and must not use the ' + "lifecycle hook prefix" + ) + raise ValueError(msg) + if path == MICROVM_JOB_PATH: + msg = ( + f'route "{path}" is reserved for jobs that name no route; ' + "set handler instead" + ) + raise ValueError(msg) + validate_heartbeat_interval(heartbeat_interval_seconds) + return cls( + handler=handler, + routes=routes, + create_client=create_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. + logger=safe_logger(logger), + ) + + @cached_property + def request_handler_class(self) -> type[BaseHTTPRequestHandler]: + """A request handler class for ``http.server`` that serves this listener. + + Pass it to a ``ThreadingHTTPServer``. The listener needs a thread per + request, because it answers each request before its job ends. + """ + listener = self + + class _Handler(_RequestHandler): + def handle_exchange(self) -> None: + listener.handle_request(self) + + def handle_expect_100(self) -> bool: + return listener.handle_expect_100(self) + + return _Handler + + def handle_request(self, request: BaseHTTPRequestHandler) -> None: + """Handle one request that ``http.server`` received. + + :attr:`request_handler_class` calls it for every request. An + unexpected error is logged, and gets HTTP 500 when no answer was sent. + """ + exchange = _Exchange(request, self._logger) + try: + self._dispatch(exchange) + except Exception as error: + self._logger.error( + "request handling failed", extra={"error": describe(error)} + ) + if not exchange.responded: + exchange.respond(500, {"error": "internal error"}) + finally: + exchange.discard_unread_body() + + def handle_expect_100(self, request: BaseHTTPRequestHandler) -> bool: + """Answer a request that sent ``Expect: 100-continue``. + + :attr:`request_handler_class` calls it before it reads the body. + + 1. A request whose ``Content-Length`` is over the path's limit gets + its HTTP 400 now, instead of ``100 Continue``. RFC 9110 section + 10.1.1 allows a final answer in place of ``100 Continue``. The + client then does not send the body. + 2. Every other request gets ``100 Continue``. + + Returns: + Whether the base class goes on to read the request. + """ + exchange = _Exchange(request, self._logger) + limit = self._body_limit(exchange) + length = exchange.declared_length() + if limit is not None and length is not None and length > limit: + self._reject(exchange.too_large(limit), exchange) + exchange.discard_unread_body() + return False + return BaseHTTPRequestHandler.handle_expect_100(request) + + 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. + + Args: + timeout: The longest wait in seconds. ``None`` waits until the + jobs end. + + Returns: + ``True`` when no job runs. ``False`` when the timeout ended first. + """ + with self._changed: + return self._changed.wait_for( + lambda: not self._jobs and self._reports == 0, timeout + ) + + 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) + return + self._logger.info("lifecycle hook", extra={"hook": hook}) + exchange.respond(200, {}) + return + handler = self._job_handler(path) + if handler is not None and exchange.method == "POST": + self._handle_job_request(handler, exchange) + return + exchange.respond(404, {"error": "not found"}) + + 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) + + def _body_limit(self, exchange: _Exchange) -> int | None: + """The largest body that the request's path reads. + + ``None`` for a request whose body the worker does not read. + """ + if exchange.method != "POST": + return None + if exchange.path == f"{HOOK_PATH_PREFIX}run": + return MAX_HOOK_BODY_BYTES + if self._job_handler(exchange.path) is not None: + return MAX_JOB_BODY_BYTES + return None + + def _handle_run_hook(self, exchange: _Exchange) -> None: + try: + parsed = RunHookRequest.from_dict(exchange.read_json(MAX_HOOK_BODY_BYTES)) + except Exception as error: # noqa: BLE001 + self._reject(error, exchange) + return + payload = parsed.payload + 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. + with self._lock: + later = self._run_hook_accepted + if not later: + self._run_hook_accepted = True + self._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) + + 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 + # once the caller has the answer. + self._fail_in_background( + InvalidRunHookPayloadError( + "the run hook delivered a job, and this worker has no " + "handler for run hook jobs" + ), + new_job.callback_id, + payload.region, + ) + exchange.respond(200, {}) + if payload is not None and new_job is not None and self._handler is not None: + self._start_job(self._handler, new_job, payload.region) + return + + if later: + self._logger.warning( + "later run hook ignored", + extra={"microvmId": microvm_id, "hookMicrovmId": parsed.microvm_id}, + ) + elif payload is None: + self._logger.warning( + "run hook without runHookPayload", extra={"microvmId": microvm_id} + ) + elif job is not None and duplicate: + self._logger.warning( + "duplicate job ignored", + extra={"microvmId": microvm_id, "callbackId": job.callback_id}, + ) + + def _handle_job_request( + self, handler: MicrovmJobHandler, exchange: _Exchange + ) -> None: + try: + request = MicrovmJobRequest.from_dict( + exchange.read_json(MAX_JOB_BODY_BYTES) + ) + except Exception as error: # noqa: BLE001 + self._reject(error, exchange) + return + job = request.job + with self._lock: + if ( + self._microvm_id == _UNKNOWN_MICROVM_ID + and request.microvm_id is not None + ): + # The endpoint can accept a request before Lambda has sent the + # run hook. The run hook sets the same identifier when it + # arrives. + self._microvm_id = request.microvm_id + microvm_id = self._microvm_id + duplicate = self._is_known(job.callback_id) + if not duplicate: + self._jobs.add(job.callback_id) + exchange.respond(202, {}) + if duplicate: + self._logger.warning( + "duplicate job ignored", + extra={"microvmId": microvm_id, "callbackId": job.callback_id}, + ) + return + self._start_job(handler, job, request.region) + + def _reject(self, error: BaseException, exchange: _Exchange) -> None: + """Answer 400 to a request that the worker could not accept. + + Only the worker's own validation messages are sent back. They describe + the request, not the worker. Any other error, such as a connection + that failed while the body was read, is logged, and the caller gets a + generic message. So no internal error text reaches the HTTP response. + """ + if not isinstance(error, InvalidRunHookPayloadError): + self._logger.error( + "could not read the request", extra={"error": describe(error)} + ) + exchange.respond(400, {"error": "invalid request"}) + return + message = str(error) + self._logger.error("invalid request", extra={"error": message}) + if error.callback_id is not None: + region = error.region or os.environ.get("AWS_REGION") + if region: + # The report starts before the answer, so that idle() covers + # it once the caller has the answer. + self._fail_in_background(error, error.callback_id, region) + exchange.respond(400, {"error": message}) + + def _is_known(self, callback_id: str) -> bool: + """Whether a job with this callback ID runs now, or ran recently. + + The caller holds the lock. + """ + return callback_id in self._jobs or callback_id in self._finished + + def _start_job( + self, handler: MicrovmJobHandler, job: MicrovmJobDocument, region: str + ) -> None: + """Run a job in its own thread. The caller registered it in ``_jobs``.""" + thread = threading.Thread( + target=self._run_job, + args=(handler, job, region), + name=f"microvm-job-{job.callback_id[:16]}", + daemon=True, + ) + try: + thread.start() + except BaseException: + # The job never ran. So a later delivery may run it. + with self._changed: + self._jobs.discard(job.callback_id) + self._changed.notify_all() + raise + + def _run_job( + self, handler: MicrovmJobHandler, job: MicrovmJobDocument, region: str + ) -> None: + try: + self._run_job_and_report(handler, job, region) + except Exception as error: # noqa: BLE001 + self._logger.error( + "job failed inside the worker", + extra={"callbackId": job.callback_id, "error": describe(error)}, + ) + finally: + with self._changed: + self._jobs.discard(job.callback_id) + self._finished[job.callback_id] = None + if len(self._finished) > MAX_FINISHED_JOBS: + self._finished.popitem(last=False) + self._changed.notify_all() + + def _run_job_and_report( + self, handler: MicrovmJobHandler, job: MicrovmJobDocument, region: str + ) -> None: + callback_id = job.callback_id + with self._lock: + microvm_id = self._microvm_id + fields = {"microvmId": microvm_id, "callbackId": callback_id} + try: + reporter = self._reporter_for(callback_id, region) + except Exception as error: # noqa: BLE001 + # Without a client, the worker cannot report anything. The durable + # function fails the job at its heartbeat timeout or its timeout. + self._logger.error( + "could not create the Lambda client for the job", + extra={**fields, "error": describe(error)}, + ) + return + try: + cancelled = threading.Event() + 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, + ) + try: + self._logger.info("job started", extra=fields) + context = MicrovmJobContext( + callback_id=callback_id, + microvm_id=microvm_id, + region=region, + cancelled=cancelled, + logger=self._logger, + ) + outcome = _Outcome.of(lambda: handler(job.input, context)) + # Heartbeats keep running while the outcome is reported. A + # 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) + else: + # The heartbeat thread decides under a lock, and calls + # on_callback_gone after it. So the reason can still be + # missing here. + self._logger.warning( + "job outcome not reported: the callback is gone", + extra={ + **fields, + "reason": describe(gone[0]) if gone else None, + }, + ) + finally: + heartbeats.stop() + finally: + reporter.close() + + def _report( + self, + reporter: CallbackReporter, + outcome: _Outcome, + fields: Mapping[str, object], + ) -> None: + try: + if outcome.error is None: + try: + reporter.succeed(outcome.value) + except (ResultTooLargeError, ResultSerializationError) as error: + # The result was never sent. Report it as a failure, so the + # durable function learns why instead of waiting for its + # timeout. + reporter.fail(error) + self._logger.info( + "job failed: its result cannot be sent", + extra={**fields, "error": describe(error)}, + ) + return + self._logger.info("job succeeded", extra=fields) + return + reporter.fail(outcome.error) + self._logger.info( + "job failed", extra={**fields, "error": describe(outcome.error)} + ) + except Exception as error: # noqa: BLE001 + self._logger.error( + "job outcome could not be reported", + extra={ + **fields, + "terminal": is_terminal_callback_error(error), + "error": describe(error), + }, + ) + + def _fail_in_background( + self, error: BaseException, callback_id: str, region: str + ) -> None: + """Fail a callback in its own thread. :meth:`idle` waits for it.""" + with self._lock: + self._reports += 1 + thread = threading.Thread( + target=self._send_failure, + args=(error, callback_id, region), + name=f"microvm-failure-{callback_id[:16]}", + daemon=True, + ) + try: + thread.start() + except BaseException: + self._report_ended() + raise + + def _send_failure( + self, error: BaseException, callback_id: str, region: str + ) -> None: + fields = {"callbackId": callback_id} + reporter: CallbackReporter | None = None + try: + reporter = self._reporter_for(callback_id, region) + reporter.fail(error) + except Exception as report_error: # noqa: BLE001 + if is_terminal_callback_error(report_error): + # The callback has already ended, for example at its timeout. + # The durable function no longer waits for this report. + self._logger.info( + "the callback no longer accepts a failure", + extra={**fields, "error": describe(report_error)}, + ) + else: + self._logger.error( + "could not fail the callback", + extra={**fields, "error": describe(report_error)}, + ) + finally: + try: + if reporter is not None: + reporter.close() + except Exception as close_error: # noqa: BLE001 + self._logger.error( + "could not close the Lambda client", + extra={**fields, "error": describe(close_error)}, + ) + self._report_ended() + + def _report_ended(self) -> None: + with self._changed: + self._reports -= 1 + self._changed.notify_all() + + def _reporter_for(self, callback_id: str, region: str) -> CallbackReporter: + # Each job gets its own client. A client kept across jobs would keep + # its connections across a suspend of the MicroVM, and those can be + # dead after the resume. + if self._create_client is None: + return CallbackReporter.create(callback_id, region, logger=self._logger) + return CallbackReporter( + callback_id, self._create_client(region), logger=self._logger + ) + + +class MicrovmWorker: + """A running worker server. + + Create it with :func:`start_microvm_worker`. + """ + + def __init__( + self, + *, + listener: MicrovmWorkerListener, + server: ThreadingHTTPServer, + ) -> None: + self._listener = listener + self._server = server + self._lock = threading.Lock() + self._closed = threading.Event() + + @classmethod + def start( + cls, listener: MicrovmWorkerListener, *, port: int = DEFAULT_PORT + ) -> MicrovmWorker: + """Serve ``listener`` on ``port``, in a background thread. + + Raises: + OSError: When the port cannot be bound. + """ + server = _create_server(port, listener.request_handler_class) + threading.Thread( + target=server.serve_forever, + kwargs={"poll_interval": _SERVER_POLL_SECONDS}, + name="microvm-worker-server", + daemon=True, + ).start() + return cls(listener=listener, server=server) + + @property + def port(self) -> int: + """The port the server listens on.""" + return int(self._server.server_address[1]) + + def idle(self, timeout: float | None = None) -> bool: + """Wait until every job started so far has ended. + + Returns: + ``True`` when no job runs. ``False`` when the timeout ended first. + """ + return self._listener.idle(timeout) + + def wait(self, timeout: float | None = None) -> bool: + """Block until :meth:`close` is called. + + The server runs in a daemon thread, so the process would otherwise + end when the main thread ends. Call this method at the end of the + image's entry point. + + Returns: + ``True`` when the worker is closed. ``False`` when the timeout + ended first. + """ + return self._closed.wait(timeout) + + def close(self) -> None: + """Stop accepting requests. Running jobs continue. + + Call :meth:`idle` first to let the running jobs report their outcome. + """ + with self._lock: + if self._closed.is_set(): + return + self._server.shutdown() + self._server.server_close() + self._closed.set() + + +def start_microvm_worker( + *, + handler: MicrovmJobHandler | None = None, + routes: Mapping[str, MicrovmJobHandler] | None = None, + create_client: Callable[[str], LambdaClient] | None = None, + heartbeat_interval_seconds: float | None = None, + logger: MicrovmWorkerLogger | None = None, + port: int = DEFAULT_PORT, +) -> MicrovmWorker: + """Start an HTTP server for the MicroVM lifecycle hooks and the job routes. + + Start it when the image's process starts. Lambda snapshots the process + after the ``ready`` hook, so the worker must already be listening. The + image's hook port must match ``port``. The server runs in a daemon thread. + So end the entry point with :meth:`MicrovmWorker.wait`. + + See :meth:`MicrovmWorkerListener.create` for the other arguments, and + :class:`MicrovmWorkerListener` for the request handling. + + Args: + port: The port to listen on. Defaults to 8080, the MicroVM default. + 0 picks a free port. + + Raises: + TypeError: As :meth:`MicrovmWorkerListener.create`. + ValueError: As :meth:`MicrovmWorkerListener.create`. + OSError: When the port cannot be bound. + """ + listener = MicrovmWorkerListener.create( + handler=handler, + routes=routes, + create_client=create_client, + heartbeat_interval_seconds=heartbeat_interval_seconds, + logger=logger, + ) + return MicrovmWorker.start(listener, port=port) + + +def _check_handler(handler: object) -> None: + if handler is None: + return + if not callable(handler): + msg = f"a job handler must be callable. Got {type(handler).__name__}." + raise TypeError(msg) + # An async function returns a coroutine, which the worker would report as + # a result that is not JSON-serializable. So it is refused here. An object + # whose __call__ is async is an async function too, and + # iscoroutinefunction does not detect it on the object itself. + if inspect.iscoroutinefunction(handler) or inspect.iscoroutinefunction( + getattr(type(handler), "__call__", None) + ): + msg = "a job handler must be a plain function, not an async function" + raise TypeError(msg) + + +@dataclass(frozen=True, kw_only=True) +class _Outcome: + """What the handler returned, or what it raised.""" + + value: Any = None + error: BaseException | None = None + + @classmethod + def of(cls, run: Callable[[], Any]) -> _Outcome: + """Run the handler, and keep its result or its exception. + + ``BaseException`` is kept too. A handler that calls ``sys.exit()`` + raises ``SystemExit``, which would otherwise end the job thread + without a report. + + A handler can still return an awaitable, for example a plain function + that calls an async function and returns its coroutine. The worker + cannot run it. So the outcome is a ``TypeError`` that names the + cause, and a coroutine is closed, so that Python does not warn that + it was never awaited. + """ + try: + value = run() + except BaseException as error: # noqa: BLE001 + return cls(error=error) + if inspect.isawaitable(value): + if inspect.iscoroutine(value): + value.close() + return cls( + error=TypeError( + "the job handler returned an awaitable. A job handler must " + "be a plain function that returns its result." + ) + ) + return cls(value=value) + + +class _RequestHandler(BaseHTTPRequestHandler): + """Passes every request to :meth:`handle_exchange`. + + The listener's subclass implements it. + """ + + timeout = _REQUEST_TIMEOUT_SECONDS + # HTTP/1.1 makes the base class answer "Expect: 100-continue". Under + # HTTP/1.0 it does not, and a client such as curl then waited 1 second + # before it sent the body. Every answer closes its connection, so a body + # that the worker did not read cannot be parsed as the next request. + protocol_version = "HTTP/1.1" + + def handle_exchange(self) -> None: # pragma: no cover + raise NotImplementedError + + def do_POST(self) -> None: # noqa: N802 + self.handle_exchange() + + def do_GET(self) -> None: # noqa: N802 + self.handle_exchange() + + def do_PUT(self) -> None: # noqa: N802 + self.handle_exchange() + + def do_PATCH(self) -> None: # noqa: N802 + self.handle_exchange() + + def do_DELETE(self) -> None: # noqa: N802 + self.handle_exchange() + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + # The base class writes every request to stderr. The worker logs what + # matters through its own logger. + return + + +class _Exchange: + """One request and its answer.""" + + def __init__( + self, request: BaseHTTPRequestHandler, logger: MicrovmWorkerLogger + ) -> None: + self._request = request + self._logger = logger + self.responded = False + # Set when the worker answers before it has read the whole body. + self._body_unread = False + + @property + def method(self) -> str: + return self._request.command + + @property + def path(self) -> str: + return self._request.path.split("?", 1)[0] + + def declared_length(self) -> int | None: + """The body length from ``Content-Length``. + + ``None`` when the body is chunked, or the header is absent or not + valid. Reading the body reports an invalid header. + """ + headers = self._request.headers + if "chunked" in headers.get("Transfer-Encoding", "").lower(): + return None + raw = headers.get("Content-Length") + if raw is None or not _CONTENT_LENGTH.fullmatch(raw.strip()): + return None + return int(raw) + + def read_json(self, max_bytes: int) -> Any: + """Read the body, and parse it as JSON. + + Raises: + InvalidRunHookPayloadError: When the body is larger than + ``max_bytes``, its framing is not valid, or it is not JSON. + OSError: When the connection fails or ends before the body does. + """ + body = self._read_body(max_bytes) + try: + return loads_strict(body) + except ValueError as error: + msg = "request body is not valid JSON" + raise InvalidRunHookPayloadError(msg) from error + + def too_large(self, max_bytes: int) -> InvalidRunHookPayloadError: + """Return the error for a body over ``max_bytes``, and remember that + the worker answers before it reads the rest of the body.""" + self._body_unread = True + return InvalidRunHookPayloadError(f"request body exceeds {max_bytes} bytes") + + def respond(self, status: int, body: Mapping[str, object]) -> None: + """Send the answer. A client that has gone away is logged.""" + self.responded = True + data = json.dumps(body, separators=(",", ":")).encode() + request = self._request + try: + request.send_response(status) + request.send_header("Content-Type", "application/json") + request.send_header("Content-Length", str(len(data))) + # The header also makes the base class close the connection. + request.send_header("Connection", "close") + request.end_headers() + request.wfile.write(data) + except OSError as error: + self._logger.warning( + "could not send the response", + extra={"status": status, "error": describe(error)}, + ) + + def discard_unread_body(self) -> None: + """Read and drop the rest of a body that the worker did not read. + + 1. The worker answers a body over the limit without reading it. + 2. The connection then closes with unread bytes in its receive + buffer. So the kernel sends a TCP reset instead of a normal close. + 3. A client that is still sending gets the reset before it reads the + answer. A probe showed this in 20 of 20 tries, so no client saw + the 400. + 4. The durable function fails at once on a 400, and retries a reset. + So the reset turns a clear rejection into retries. + + So after the answer, the worker ends its side of the connection, and + reads and drops input until the client closes. The read stops after + :data:`_DISCARD_SECONDS` or :data:`_MAX_DISCARD_BYTES`. A client that + sends more still gets the reset. + """ + if not self._body_unread: + return + connection = self._request.connection + try: + connection.shutdown(socket.SHUT_WR) + deadline = time.monotonic() + _DISCARD_SECONDS + discarded = 0 + while discarded < _MAX_DISCARD_BYTES: + remaining = deadline - time.monotonic() + if remaining <= 0: + return + connection.settimeout(remaining) + data = connection.recv(_DISCARD_READ_BYTES) + if not data: + return + discarded += len(data) + except OSError: + # The client has gone, or the time ran out in recv. + return + + def _read_body(self, max_bytes: int) -> bytes: + headers = self._request.headers + if "chunked" in headers.get("Transfer-Encoding", "").lower(): + return self._read_chunked(max_bytes) + if headers.get("Content-Length") is None: + return b"" + length = self.declared_length() + if length is None: + msg = "request Content-Length is not valid" + raise InvalidRunHookPayloadError(msg) + if length > max_bytes: + raise self.too_large(max_bytes) + return self._read_exactly(length) + + def _read_chunked(self, max_bytes: int) -> bytes: + """Read a body with chunked transfer encoding. + + ``http.server`` does not decode it, and Node's server does. A durable + function in another language, or Lambda's own hook client, can send + one. So the worker decodes it. + + RFC 9112 defines a chunk size as one or more hex digits. ``int(x, 16)`` + also accepts a ``0x`` prefix, a sign, and underscores. So the size is + matched against the RFC form first, as ``Content-Length`` is. + """ + chunks: list[bytes] = [] + size = 0 + while True: + line = self._request.rfile.readline(_MAX_CHUNK_LINE_BYTES) + if not line.endswith(b"\n"): + raise _malformed_chunks() if line else _ended_early() + token = line.split(b";", 1)[0].strip() + if not _CHUNK_SIZE.fullmatch(token): + raise _malformed_chunks() + chunk_size = int(token, 16) + if chunk_size == 0: + self._skip_trailers() + return b"".join(chunks) + size += chunk_size + if size > max_bytes: + raise self.too_large(max_bytes) + chunks.append(self._read_exactly(chunk_size)) + if self._read_exactly(2) != b"\r\n": + raise _malformed_chunks() + + def _skip_trailers(self) -> None: + while True: + line = self._request.rfile.readline(_MAX_CHUNK_LINE_BYTES) + if line in {b"\r\n", b"\n", b""}: + return + + def _read_exactly(self, length: int) -> bytes: + data = self._request.rfile.read(length) + if len(data) < length: + raise _ended_early() + return data + + +def _malformed_chunks() -> InvalidRunHookPayloadError: + return InvalidRunHookPayloadError("request body has malformed chunked encoding") + + +def _ended_early() -> ConnectionError: + return ConnectionError("the connection ended before the request body did") + + +class _Server(ThreadingHTTPServer): + """The worker's HTTP server.""" + + def server_bind(self) -> None: + # HTTPServer.server_bind looks up the host's fully qualified name, + # which can wait for DNS. The worker never uses the name. So the bind + # skips the lookup. + socketserver.TCPServer.server_bind(self) + self.server_name = "localhost" + self.server_port = int(self.server_address[1]) + + +class _DualStackServer(_Server): + """Listens on IPv6 and IPv4 at once. + + Lambda calls the hooks on ``localhost``. That name can resolve to ``::1`` + or to ``127.0.0.1``. A server on one address family would refuse the + other, and a refused ``run`` hook loses the job. Node's server listens on + both by default. + """ + + address_family = socket.AF_INET6 + + def server_bind(self) -> None: + with contextlib.suppress(OSError, AttributeError): + self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0) + super().server_bind() + + +def _create_server( + port: int, handler_class: type[BaseHTTPRequestHandler] +) -> ThreadingHTTPServer: + if socket.has_dualstack_ipv6(): + return _DualStackServer(("::", port), handler_class) + return _Server(("0.0.0.0", port), handler_class) # noqa: S104 diff --git a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/callback_reporter_test.py b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/callback_reporter_test.py index 92f72f22f..ae3a32c94 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/tests/callback_reporter_test.py +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/callback_reporter_test.py @@ -239,7 +239,8 @@ def test_closed_callback_after_uncertain_attempt_is_delivered( assert len(logger.lines) == 1 level, message, extra = logger.lines[0] assert level == "warning" - assert "already complete or timed out" in message + assert "already complete" in message + assert "timed out first" in message assert extra["callbackId"] == "cb-1" 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 new file mode 100644 index 000000000..b1df372e0 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/tests/worker_test.py @@ -0,0 +1,1295 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the HTTP listener and the worker server.""" + +from __future__ import annotations + +import http.client +import inspect +import io +import json +import socket +import threading +import time +from collections.abc import Callable, Iterator +from http.server import ThreadingHTTPServer +from typing import Any + +import pytest + +from aws_durable_execution_sdk_python_microvm_worker import ( + HOOK_PATH_PREFIX, + MICROVM_JOB_PATH, + MicrovmJobContext, + MicrovmWorker, + MicrovmWorkerListener, + start_microvm_worker, +) +from aws_durable_execution_sdk_python_microvm_worker import worker as worker_module + + +RUN = f"{HOOK_PATH_PREFIX}run" +WAIT = 10.0 + + +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"} + if with_job: + payload["job"] = {"callbackId": "cb-1", "input": {"repo": "org/app"}, **job} + return payload + + +def run_body(payload: dict[str, Any] | None, microvm_id: str = "mvm-1") -> Any: + body: dict[str, Any] = {"microvmId": microvm_id} + if payload is not None: + body["runHookPayload"] = json.dumps(payload) + return body + + +def job_body(**fields: Any) -> dict[str, Any]: + return { + "version": 1, + "region": "us-east-1", + "callbackId": "cb-1", + "input": {"repo": "org/app"}, + **fields, + } + + +def request( + worker: MicrovmWorker, + path: str, + body: Any = None, + *, + method: str = "POST", + raw: Any = None, + headers: dict[str, str] | None = None, + chunked: bool = False, +) -> 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) + try: + data = raw if raw is not None else json.dumps(body).encode() + connection.request( + method, + path, + body=data, + headers={"Content-Type": "application/json", **(headers or {})}, + encode_chunked=chunked, + ) + response = connection.getresponse() + text = response.read() + finally: + connection.close() + return response.status, json.loads(text) if text else None + + +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: + sock.sendall(data) + sock.shutdown(socket.SHUT_WR) + chunks = [] + while chunk := sock.recv(65536): + chunks.append(chunk) + return b"".join(chunks) + + +def eventually(condition: Callable[[], bool]) -> bool: + """Wait until ``condition`` holds. + + The worker writes some log lines after it has sent the answer. So a test + that has the answer can still be ahead of the line. + """ + deadline = time.monotonic() + WAIT + while not condition(): + if time.monotonic() > deadline: + return False + time.sleep(0.005) + return True + + +def completions(client: Any) -> list[tuple[str, dict[str, Any]]]: + return [(name, kwargs) for name, kwargs in client.calls if name != "heartbeat"] + + +class Held: + """A handler that runs until the test releases it.""" + + def __init__(self, result: Any = "ok") -> None: + self.started = threading.Event() + self.release = threading.Event() + self.contexts: list[MicrovmJobContext] = [] + self.inputs: list[Any] = [] + self.result = result + + def __call__(self, job_input: Any, context: MicrovmJobContext) -> Any: + self.inputs.append(job_input) + self.contexts.append(context) + self.started.set() + assert self.release.wait(WAIT) + return self.result + + +@pytest.fixture +def start( + fake_client: Any, logger: Any +) -> Iterator[Callable[..., tuple[MicrovmWorker, Any]]]: + """Start a worker on a free port with a fake client, and close it after.""" + workers: list[MicrovmWorker] = [] + + def _start( + handler: Any = None, *, client: Any = None, **options: Any + ) -> tuple[MicrovmWorker, Any]: + 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) + workers.append(worker) + return worker, client + + yield _start + for worker in workers: + worker.idle(WAIT) + worker.close() + + +def test_answers_the_run_hook_before_the_job_ends_then_reports_the_result( + start: Any, +) -> None: + handler = Held({"passed": True}) + worker, client = start(handler) + + assert request(worker, RUN, run_body(run_payload())) == (200, {}) + assert handler.started.wait(WAIT) + assert completions(client) == [] + handler.release.set() + assert worker.idle(WAIT) + + assert handler.inputs == [{"repo": "org/app"}] + context = handler.contexts[0] + assert (context.callback_id, context.microvm_id, context.region) == ( + "cb-1", + "mvm-1", + "us-east-1", + ) + assert completions(client) == [ + ("success", {"CallbackId": "cb-1", "Result": b'{"passed":true}'}) + ] + assert not client.closed + + +def test_sends_heartbeats_while_the_job_runs_and_none_after(start: Any) -> None: + handler = Held() + worker, client = start(handler, heartbeat_interval_seconds=0.02) + + request(worker, RUN, run_body(run_payload(heartbeatTimeoutSeconds=30))) + assert handler.started.wait(WAIT) + deadline = time.monotonic() + WAIT + while client.names().count("heartbeat") < 3 and time.monotonic() < deadline: + time.sleep(0.01) + handler.release.set() + assert worker.idle(WAIT) + beats = client.names().count("heartbeat") + time.sleep(0.1) + + assert beats >= 3 + assert client.names().count("heartbeat") == beats + # Heartbeats continue while the outcome is reported. So a heartbeat can + # be recorded after the success call, and only the completions are + # compared. + assert [name for name, _ in completions(client)] == ["success"] + + +def test_sends_no_heartbeats_without_a_heartbeat_timeout(start: Any) -> None: + worker, client = start(lambda job_input, context: "ok") + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + assert client.names() == ["success"] + + +def test_sets_cancelled_and_reports_nothing_when_the_callback_is_gone( + start: Any, fake_client: Any, make_error: Any, logger: Any +) -> None: + client = fake_client(make_error("CallbackTimeoutException")) + seen: list[bool] = [] + + def handler(job_input: Any, context: MicrovmJobContext) -> str: + seen.append(context.cancelled.wait(WAIT)) + return "late" + + worker, _ = start(handler, client=client, heartbeat_interval_seconds=0.01) + + request(worker, RUN, run_body(run_payload(heartbeatTimeoutSeconds=30))) + assert worker.idle(WAIT) + + assert seen == [True] + assert completions(client) == [] + assert "job outcome not reported: the callback is gone" in logger.messages( + "warning" + ) + + +def test_reports_the_handler_error_as_a_failure(start: Any) -> None: + def handler(job_input: Any, context: MicrovmJobContext) -> None: + msg = "3 tests failed" + raise RuntimeError(msg) + + worker, client = start(handler) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + assert completions(client) == [ + ( + "failure", + { + "CallbackId": "cb-1", + "Error": { + "ErrorType": "RuntimeError", + "ErrorMessage": "3 tests failed", + }, + }, + ) + ] + + +def test_reports_system_exit_from_the_handler_as_a_failure(start: Any) -> None: + def handler(job_input: Any, context: MicrovmJobContext) -> None: + raise SystemExit(3) + + worker, client = start(handler) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + assert completions(client)[0][1]["Error"]["ErrorType"] == "SystemExit" + + +@pytest.mark.parametrize( + ("result", "error_type"), + [ + ({1, 2}, "ResultSerializationError"), + ("x" * (256 * 1024), "ResultTooLargeError"), + ], +) +def test_reports_a_result_that_cannot_be_sent_as_a_failure( + start: Any, logger: Any, result: Any, error_type: str +) -> None: + worker, client = start(lambda job_input, context: result) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + [(name, kwargs)] = completions(client) + assert name == "failure" + assert kwargs["Error"]["ErrorType"] == error_type + assert "job failed: its result cannot be sent" in logger.messages("info") + + +def test_logs_an_outcome_that_cannot_be_reported( + start: Any, fake_client: Any, make_error: Any, logger: Any +) -> None: + client = fake_client(make_error("AccessDeniedException", 403)) + worker, _ = start(lambda job_input, context: "ok", client=client) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + [line] = [line for line in logger.lines if line[0] == "error"] + assert line[1] == "job outcome could not be reported" + assert line[2]["terminal"] is False + assert line[2]["callbackId"] == "cb-1" + + +def test_logs_a_job_whose_client_cannot_be_created_and_runs_later_jobs( + start: Any, fake_client: Any, logger: Any +) -> None: + client = fake_client() + attempts: list[str] = [] + + def create_client(region: str) -> Any: + attempts.append(region) + if len(attempts) == 1: + msg = "no client" + raise RuntimeError(msg) + return client + + handler = Held() + handler.release.set() + worker, _ = start(handler, routes={"/build": handler}, create_client=create_client) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + assert request(worker, "/build", job_body(callbackId="cb-2"))[0] == 202 + assert worker.idle(WAIT) + + assert "could not create the Lambda client for the job" in logger.messages("error") + assert [kwargs["CallbackId"] for _, kwargs in completions(client)] == ["cb-2"] + + +def test_reports_the_outcome_when_the_logger_raises(start: Any) -> None: + class RaisingLogger: + def info(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError + + warning = error = info + + worker, client = start(lambda job_input, context: "ok", logger=RaisingLogger()) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + assert client.names() == ["success"] + + +def test_logs_an_unexpected_job_error( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + worker, _ = start(lambda job_input, context: "ok") + + def broken(*args: Any) -> None: + msg = "broken" + raise RuntimeError(msg) + + monkeypatch.setattr(MicrovmWorkerListener, "_run_job_and_report", broken) + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + assert "job failed inside the worker" in logger.messages("error") + + +def test_ignores_every_run_hook_after_the_first(start: Any, logger: Any) -> None: + handler = Held() + handler.release.set() + worker, client = start(handler) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + later = run_body(run_payload(callbackId="cb-2"), microvm_id="mvm-2") + assert request(worker, RUN, later) == (200, {}) + assert worker.idle(WAIT) + + assert [kwargs["CallbackId"] for _, kwargs in completions(client)] == ["cb-1"] + assert eventually(lambda: "later run hook ignored" in logger.messages("warning")) + + +def test_answers_400_to_an_invalid_payload_and_fails_the_named_callback( + start: Any, logger: Any +) -> None: + handler = Held() + worker, client = start(handler) + + status, body = request(worker, RUN, run_body({**run_payload(), "version": 2})) + assert worker.idle(WAIT) + + assert status == 400 + assert "version 2 is not supported" in body["error"] + [(name, kwargs)] = completions(client) + assert name == "failure" + assert kwargs["CallbackId"] == "cb-1" + assert kwargs["Error"]["ErrorType"] == "InvalidRunHookPayloadError" + assert not handler.started.is_set() + + +def test_answers_400_to_an_empty_microvm_id_and_fails_the_job(start: Any) -> None: + worker, client = start(Held()) + + status, _ = request(worker, RUN, run_body(run_payload(), microvm_id="")) + assert worker.idle(WAIT) + + assert status == 400 + assert [name for name, _ in completions(client)] == ["failure"] + + +def test_fails_a_callback_with_the_environment_region_when_the_request_has_none( + start: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + client = fake_client() + regions: list[str] = [] + + def create_client(region: str) -> Any: + regions.append(region) + return client + + monkeypatch.setenv("AWS_REGION", "eu-west-1") + worker, _ = start(routes={"/build": Held()}, create_client=create_client) + + body = job_body() + del body["region"] + assert request(worker, "/build", body)[0] == 400 + assert worker.idle(WAIT) + + assert regions == ["eu-west-1"] + assert [name for name, _ in completions(client)] == ["failure"] + + +def test_fails_no_callback_without_any_region( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("AWS_REGION", raising=False) + worker, client = start(routes={"/build": Held()}) + + body = job_body() + del body["region"] + assert request(worker, "/build", body)[0] == 400 + assert worker.idle(WAIT) + + assert client.calls == [] + + +def test_logs_a_failure_report_that_the_service_rejects( + start: Any, fake_client: Any, make_error: Any, logger: Any +) -> None: + client = fake_client(make_error("AccessDeniedException", 403)) + worker, _ = start(routes={"/build": Held()}, client=client) + + request(worker, "/build", job_body(version=2)) + assert worker.idle(WAIT) + + assert "could not fail the callback" in logger.messages("error") + + +def test_logs_a_failure_report_for_a_closed_callback_as_info( + start: Any, fake_client: Any, make_error: Any, logger: Any +) -> None: + client = fake_client(make_error("CallbackTimeoutException")) + worker, _ = start(routes={"/build": Held()}, client=client) + + request(worker, "/build", job_body(version=2)) + assert worker.idle(WAIT) + + assert "the callback no longer accepts a failure" in logger.messages("info") + + +def test_logs_a_failure_report_whose_client_cannot_be_closed( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + worker, _ = start(routes={"/build": Held()}) + + def broken_close(self: Any) -> None: + raise RuntimeError + + monkeypatch.setattr(worker_module.CallbackReporter, "close", broken_close) + request(worker, "/build", job_body(version=2)) + assert worker.idle(WAIT) + + assert "could not close the Lambda client" in logger.messages("error") + + +def test_creates_a_default_client_per_job_and_closes_it( + start: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + created: list[Any] = [] + + def create_client(region: str) -> Any: + client = fake_client() + created.append(client) + return client + + monkeypatch.setattr( + "aws_durable_execution_sdk_python_microvm_worker.callback_reporter._create_client", + create_client, + ) + worker = start_microvm_worker(handler=lambda job_input, context: "ok", port=0) + try: + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + finally: + worker.close() + + assert len(created) == 1 + assert created[0].closed + assert created[0].names() == ["success"] + + +@pytest.mark.parametrize( + ("raw", "message"), + [ + (b"{not json", "request body is not valid JSON"), + (b"NaN", "request body is not valid JSON"), + (b"", "request body is not valid JSON"), + ], +) +def test_answers_400_to_a_body_that_is_not_json( + start: Any, raw: bytes, message: str +) -> None: + worker, client = start(Held()) + + assert request(worker, RUN, raw=raw) == (400, {"error": message}) + assert worker.idle(WAIT) + assert client.calls == [] + + +def test_answers_400_to_a_body_without_a_content_length(start: Any) -> None: + worker, _ = start(Held()) + + answer = raw_request( + worker, f"POST {RUN} HTTP/1.1\r\nHost: localhost\r\n\r\n".encode() + ) + + assert answer.startswith(b"HTTP/1.1 400") + assert b"request body is not valid JSON" in answer + + +def test_answers_400_to_an_invalid_content_length(start: Any) -> None: + worker, _ = start(Held()) + + answer = raw_request( + worker, + f"POST {RUN} HTTP/1.1\r\nHost: localhost\r\nContent-Length: 1_0\r\n\r\n".encode(), + ) + + assert answer.startswith(b"HTTP/1.1 400") + assert b"Content-Length is not valid" in answer + + +def test_answers_400_with_a_generic_message_when_the_body_ends_early( + start: Any, logger: Any +) -> None: + worker, _ = start(Held()) + + answer = raw_request( + worker, + f"POST {RUN} HTTP/1.1\r\nHost: localhost\r\nContent-Length: 50\r\n\r\n{{".encode(), + ) + + assert answer.startswith(b"HTTP/1.1 400") + assert answer.endswith(b'{"error":"invalid request"}') + assert "could not read the request" in logger.messages("error") + + +def test_answers_400_to_a_body_over_the_limit_without_reading_it(start: Any) -> None: + worker, _ = start(Held()) + size = worker_module.MAX_HOOK_BODY_BYTES + 1 + + answer = raw_request( + worker, + f"POST {RUN} HTTP/1.1\r\nHost: localhost\r\nContent-Length: {size}\r\n\r\n".encode(), + ) + + assert answer.startswith(b"HTTP/1.1 400") + assert b"request body exceeds 65536 bytes" in answer + + +def test_answers_expect_100_continue_before_it_reads_the_body(start: Any) -> None: + worker, client = start(routes={"/build": lambda job_input, context: "ok"}) + data = json.dumps(job_body()).encode() + head = ( + "POST /build HTTP/1.1\r\nHost: localhost\r\nExpect: 100-continue\r\n" + f"Content-Length: {len(data)}\r\n\r\n" + ) + + with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + sock.sendall(head.encode()) + interim = sock.recv(65536) + sock.sendall(data) + answer = b"" + while chunk := sock.recv(65536): + answer += chunk + assert worker.idle(WAIT) + + assert interim.startswith(b"HTTP/1.1 100 Continue") + assert answer.startswith(b"HTTP/1.1 202") + assert b"Connection: close" in answer + assert client.names() == ["success"] + + +def test_answers_400_to_a_body_over_the_limit_that_the_client_sends( + start: Any, +) -> None: + worker, client = start(routes={"/build": Held()}) + data = b"x" * (worker_module.MAX_JOB_BODY_BYTES + 1) + + status, body = request(worker, "/build", raw=data) + + assert (status, body) == (400, {"error": "request body exceeds 6291456 bytes"}) + assert client.calls == [] + + +def test_answers_400_to_a_chunked_body_over_the_limit_that_the_client_sends( + start: Any, +) -> None: + worker, _ = start(Held()) + data = b"x" * (worker_module.MAX_HOOK_BODY_BYTES + 1) + + status, body = request( + worker, RUN, raw=iter([data[:1000], data[1000:]]), chunked=True + ) + + assert (status, body) == (400, {"error": "request body exceeds 65536 bytes"}) + + +@pytest.mark.parametrize( + ("path", "limit"), + [ + ("/build", worker_module.MAX_JOB_BODY_BYTES), + (RUN, worker_module.MAX_HOOK_BODY_BYTES), + ], +) +def test_answers_400_instead_of_100_continue_to_a_body_over_the_limit( + start: Any, logger: Any, path: str, limit: int +) -> None: + worker, _ = start(Held(), routes={"/build": Held()}) + head = ( + f"POST {path} HTTP/1.1\r\nHost: localhost\r\nExpect: 100-continue\r\n" + f"Content-Length: {limit + 1}\r\n\r\n" + ) + + with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + sock.sendall(head.encode()) + answer = b"" + while chunk := sock.recv(65536): + answer += chunk + + assert answer.startswith(b"HTTP/1.1 400") + assert b"100 Continue" not in answer + assert answer.endswith(f'{{"error":"request body exceeds {limit} bytes"}}'.encode()) + assert "invalid request" in logger.messages("error") + + +@pytest.mark.parametrize( + ("method", "path", "framing"), + [ + ("POST", "/other", "Content-Length: 99999999"), + ("POST", f"{HOOK_PATH_PREFIX}ready", "Content-Length: 99999999"), + ("GET", "/build", "Content-Length: 99999999"), + ("POST", "/build", "Transfer-Encoding: chunked"), + ], +) +def test_answers_100_continue_when_the_declared_size_is_not_over_a_limit( + start: Any, method: str, path: str, framing: str +) -> None: + worker, _ = start(Held(), routes={"/build": Held()}) + head = ( + f"{method} {path} HTTP/1.1\r\nHost: localhost\r\nExpect: 100-continue\r\n" + f"{framing}\r\n\r\n" + ) + + with socket.create_connection(("127.0.0.1", worker.port), timeout=WAIT) as sock: + sock.sendall(head.encode()) + interim = sock.recv(65536) + + assert interim.startswith(b"HTTP/1.1 100 Continue") + + +def test_stops_discarding_a_rejected_body_after_the_time_limit( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "_DISCARD_SECONDS", 0.2) + worker, _ = start(Held()) + 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: + sock.sendall(head.encode() + b"x" * 100) + began = time.monotonic() + answer = b"" + while chunk := sock.recv(65536): + answer += chunk + elapsed = time.monotonic() - began + + assert answer.startswith(b"HTTP/1.1 400") + assert elapsed < WAIT + + +class _EndlessConnection: + """A socket whose peer never stops sending.""" + + def __init__(self) -> None: + self.received = 0 + self.shut = False + + def shutdown(self, how: int) -> None: + self.shut = True + + def settimeout(self, timeout: float) -> None: ... + + def recv(self, size: int) -> bytes: + self.received += size + return b"x" * size + + +def _rejected_exchange(connection: _EndlessConnection, logger: Any) -> Any: + request = type("Request", (), {"connection": connection})() + exchange = worker_module._Exchange(request, logger) # noqa: SLF001 + exchange.too_large(1) + return exchange + + +def test_stops_discarding_a_rejected_body_after_the_byte_limit( + logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "_MAX_DISCARD_BYTES", 200_000) + connection = _EndlessConnection() + + _rejected_exchange(connection, logger).discard_unread_body() + + assert connection.shut + assert 200_000 <= connection.received < 200_000 + 65536 + + +def test_stops_discarding_a_rejected_body_when_the_client_is_gone( + logger: Any, +) -> None: + class GoneConnection(_EndlessConnection): + def shutdown(self, how: int) -> None: + raise ConnectionResetError + + connection = GoneConnection() + + _rejected_exchange(connection, logger).discard_unread_body() + + assert connection.received == 0 + + +def test_stops_discarding_a_rejected_body_when_the_time_is_up( + logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "_DISCARD_SECONDS", 0) + connection = _EndlessConnection() + + _rejected_exchange(connection, logger).discard_unread_body() + + assert connection.received == 0 + + +def test_accepts_a_chunked_body(start: Any) -> None: + worker, client = start(lambda job_input, context: job_input) + data = json.dumps(run_body(run_payload())).encode() + + status, _ = request(worker, RUN, raw=iter([data[:10], data[10:]]), chunked=True) + assert worker.idle(WAIT) + + assert status == 200 + assert completions(client)[0][1]["Result"] == b'{"repo":"org/app"}' + + +def test_accepts_a_chunked_body_with_extensions_and_trailers(start: Any) -> None: + worker, client = start(routes={"/build": lambda job_input, context: job_input}) + data = json.dumps(job_body()).encode() + head = ( + "POST /build HTTP/1.1\r\nHost: localhost\r\nTransfer-Encoding: chunked\r\n\r\n" + ) + chunked = ( + f"{len(data):x};name=value\r\n".encode() + + data + + b"\r\n0\r\nX-Trailer: 1\r\n\r\n" + ) + + answer = raw_request(worker, head.encode() + chunked) + assert worker.idle(WAIT) + + assert answer.startswith(b"HTTP/1.1 202") + assert completions(client)[0][1]["Result"] == b'{"repo":"org/app"}' + + +@pytest.mark.parametrize( + ("chunks", "expected"), + [ + (b"zz\r\n{}\r\n0\r\n\r\n", b"malformed chunked encoding"), + (b"-2\r\n{}\r\n0\r\n\r\n", b"malformed chunked encoding"), + (b"0x2\r\n{}\r\n0\r\n\r\n", b"malformed chunked encoding"), + (b"+2\r\n{}\r\n0\r\n\r\n", b"malformed chunked encoding"), + (b"0_2\r\n{}\r\n0\r\n\r\n", b"malformed chunked encoding"), + (b"2\r\n{}XX0\r\n\r\n", b"malformed chunked encoding"), + (b"2" * 2000, b"malformed chunked encoding"), + (b"2\r\n{}\r\n", b'{"error":"invalid request"}'), + (b"20000\r\n", b"request body exceeds 65536 bytes"), + ], +) +def test_answers_400_to_a_malformed_chunked_body( + start: Any, chunks: bytes, expected: bytes +) -> None: + worker, _ = start(Held()) + head = ( + f"POST {RUN} HTTP/1.1\r\nHost: localhost\r\nTransfer-Encoding: chunked\r\n\r\n" + ) + + answer = raw_request(worker, head.encode() + chunks) + + assert answer.startswith(b"HTTP/1.1 400") + assert expected in answer + + +def test_answers_404_outside_the_hook_paths_and_routes(start: Any) -> None: + worker, _ = start(Held(), routes={"/build": Held()}) + + assert request(worker, "/other", {}) == (404, {"error": "not found"}) + assert request(worker, "/build", method="GET", raw=b"") == ( + 404, + {"error": "not found"}, + ) + assert request(worker, RUN, method="PUT", raw=b"") == (404, {"error": "not found"}) + for method in ["PATCH", "DELETE"]: + assert request(worker, "/build", method=method, raw=b"")[0] == 404 + + +def test_logs_a_failure_report_whose_client_cannot_be_created( + start: Any, logger: Any +) -> None: + def create_client(region: str) -> Any: + msg = "no client" + raise RuntimeError(msg) + + worker, _ = start(routes={"/build": Held()}, create_client=create_client) + + assert request(worker, "/build", job_body(version=2))[0] == 400 + assert worker.idle(WAIT) + + assert "could not fail the callback" in logger.messages("error") + + +@pytest.mark.parametrize("hook", ["ready", "resume", "suspend", "terminate", "other"]) +def test_answers_200_to_every_other_lifecycle_hook( + start: Any, logger: Any, hook: str +) -> None: + worker, client = start(Held()) + + assert request(worker, f"{HOOK_PATH_PREFIX}{hook}", {}) == (200, {}) + assert client.calls == [] + assert ("info", "lifecycle hook", {"hook": hook}) in logger.lines + + +def test_answers_a_run_hook_without_a_payload_with_200(start: Any, logger: Any) -> None: + worker, client = start(Held()) + + assert request(worker, RUN, run_body(None)) == (200, {}) + assert worker.idle(WAIT) + + assert client.calls == [] + assert eventually( + lambda: "run hook without runHookPayload" in logger.messages("warning") + ) + + +def test_answers_a_run_hook_without_a_job_with_200(start: Any, logger: Any) -> None: + worker, client = start(Held()) + + assert request(worker, RUN, run_body(run_payload(with_job=False))) == (200, {}) + assert worker.idle(WAIT) + + assert client.calls == [] + assert logger.messages("warning") == [] + + +def test_answers_a_job_request_with_202_and_reports_the_route_result( + start: Any, +) -> None: + worker, client = start(routes={"/build": lambda job_input, context: "built"}) + + assert request(worker, "/build?x=1", job_body()) == (202, {}) + assert worker.idle(WAIT) + + assert completions(client) == [ + ("success", {"CallbackId": "cb-1", "Result": b'"built"'}) + ] + + +def test_runs_handler_for_a_job_on_the_default_job_path(start: Any) -> None: + worker, client = start(lambda job_input, context: "default") + + assert request(worker, MICROVM_JOB_PATH, job_body()) == (202, {}) + assert worker.idle(WAIT) + + assert completions(client)[0][1]["Result"] == b'"default"' + + +def test_answers_404_on_the_default_job_path_without_a_handler(start: Any) -> None: + worker, _ = start(routes={"/build": Held()}) + + assert request(worker, MICROVM_JOB_PATH, job_body())[0] == 404 + + +def test_fails_a_run_hook_job_when_the_worker_has_only_routes(start: Any) -> None: + worker, client = start(routes={"/build": Held()}) + + assert request(worker, RUN, run_body(run_payload())) == (200, {}) + assert worker.idle(WAIT) + + [(name, kwargs)] = completions(client) + assert name == "failure" + assert kwargs["Error"]["ErrorType"] == "InvalidRunHookPayloadError" + assert "no handler for run hook jobs" in kwargs["Error"]["ErrorMessage"] + + +def test_answers_400_to_an_invalid_job_request_and_fails_the_named_callback( + start: Any, +) -> None: + worker, client = start(routes={"/build": Held()}) + + status, body = request(worker, "/build", job_body(heartbeatTimeoutSeconds=0)) + assert worker.idle(WAIT) + + assert status == 400 + assert "heartbeatTimeoutSeconds" in body["error"] + assert [name for name, _ in completions(client)] == ["failure"] + + +def test_takes_the_microvm_id_from_a_job_request_before_the_run_hook( + start: Any, +) -> None: + handler = Held() + handler.release.set() + worker, _ = start(handler, routes={"/build": handler}) + + request(worker, "/build", job_body(microvmId="mvm-early")) + assert worker.idle(WAIT) + request(worker, "/build", job_body(callbackId="cb-2", microvmId="mvm-other")) + assert worker.idle(WAIT) + request(worker, RUN, run_body(run_payload(callbackId="cb-3"), microvm_id="mvm-1")) + assert worker.idle(WAIT) + request(worker, "/build", job_body(callbackId="cb-4", microvmId="mvm-later")) + assert worker.idle(WAIT) + + assert [context.microvm_id for context in handler.contexts] == [ + "mvm-early", + "mvm-early", + "mvm-1", + "mvm-1", + ] + + +def test_uses_unknown_until_the_microvm_id_arrives(start: Any) -> None: + handler = Held() + handler.release.set() + worker, _ = start(routes={"/build": handler}) + + request(worker, "/build", job_body()) + assert worker.idle(WAIT) + + assert handler.contexts[0].microvm_id == "unknown" + + +def test_runs_a_job_once_when_its_request_arrives_twice( + start: Any, logger: Any +) -> None: + handler = Held() + worker, client = start(routes={"/build": handler}) + + assert request(worker, "/build", job_body())[0] == 202 + assert handler.started.wait(WAIT) + assert request(worker, "/build", job_body())[0] == 202 + handler.release.set() + assert worker.idle(WAIT) + assert request(worker, "/build", job_body())[0] == 202 + assert worker.idle(WAIT) + + assert len(handler.inputs) == 1 + assert len(completions(client)) == 1 + assert eventually( + lambda: logger.messages("warning").count("duplicate job ignored") == 2 + ) + + +def test_runs_a_job_once_when_parallel_requests_deliver_it(start: Any) -> None: + handler = Held() + handler.release.set() + worker, client = start(routes={"/build": handler}) + barrier = threading.Barrier(8) + + def deliver() -> None: + barrier.wait(WAIT) + request(worker, "/build", job_body()) + + threads = [threading.Thread(target=deliver) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(WAIT) + assert worker.idle(WAIT) + + assert len(handler.inputs) == 1 + assert len(completions(client)) == 1 + + +def test_ignores_a_run_hook_job_that_already_arrived_over_http( + start: Any, logger: Any +) -> None: + handler = Held() + handler.release.set() + worker, client = start(handler) + + request(worker, MICROVM_JOB_PATH, job_body()) + assert worker.idle(WAIT) + assert request(worker, RUN, run_body(run_payload())) == (200, {}) + assert worker.idle(WAIT) + + assert len(handler.inputs) == 1 + assert eventually(lambda: "duplicate job ignored" in logger.messages("warning")) + + +def test_forgets_the_oldest_finished_job( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "MAX_FINISHED_JOBS", 2) + handler = Held() + handler.release.set() + worker, _ = start(routes={"/build": handler}) + + for callback_id in ["cb-1", "cb-2", "cb-3", "cb-2", "cb-1"]: + request(worker, "/build", job_body(callbackId=callback_id)) + assert worker.idle(WAIT) + + # cb-1 was forgotten when cb-3 ended. cb-2 was still remembered. + assert len(handler.inputs) == 4 + + +def test_answers_500_to_an_unexpected_error( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + worker, _ = start(Held()) + + def broken(self: Any, exchange: Any) -> None: + raise RuntimeError + + monkeypatch.setattr(MicrovmWorkerListener, "_dispatch", broken) + + assert request(worker, RUN, {}) == (500, {"error": "internal error"}) + assert "request handling failed" in logger.messages("error") + + +def test_sends_no_second_answer_after_an_unexpected_error( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + worker, _ = start(Held()) + + def broken(self: Any, exchange: Any) -> None: + exchange.respond(202, {}) + raise RuntimeError + + monkeypatch.setattr(MicrovmWorkerListener, "_dispatch", broken) + + assert request(worker, RUN, {}) == (202, {}) + + +def test_starts_the_job_when_the_answer_cannot_be_sent( + logger: Any, fake_client: Any +) -> None: + client = fake_client() + handler = Held() + handler.release.set() + listener = MicrovmWorkerListener.create( + routes={"/build": handler}, create_client=lambda region: client, logger=logger + ) + + class BrokenPipe: + def write(self, data: bytes) -> None: + raise BrokenPipeError + + data = json.dumps(job_body()).encode() + + class Request: + command = "POST" + path = "/build" + headers = {"Content-Length": str(len(data))} + rfile = io.BytesIO(data) + wfile = BrokenPipe() + + def send_response(self, status: int) -> None: ... + def send_header(self, name: str, value: str) -> None: ... + def end_headers(self) -> None: ... + + listener.handle_request(Request()) # type: ignore[arg-type] + assert listener.idle(WAIT) + + assert len(handler.inputs) == 1 + assert "could not send the response" in logger.messages("warning") + + +def test_forgets_a_job_whose_thread_cannot_start( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + handler = Held() + handler.release.set() + worker, _ = start(routes={"/build": handler}) + original = threading.Thread.start + failures = [True] + + def flaky_start(self: threading.Thread) -> None: + if self.name.startswith("microvm-job-") and failures: + failures.pop() + msg = "can't start new thread" + raise RuntimeError(msg) + original(self) + + monkeypatch.setattr(threading.Thread, "start", flaky_start) + assert request(worker, "/build", job_body())[0] == 202 + assert worker.idle(WAIT) + assert request(worker, "/build", job_body())[0] == 202 + assert worker.idle(WAIT) + + assert len(handler.inputs) == 1 + + +def test_counts_no_report_whose_thread_cannot_start( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + worker, _ = start(routes={"/build": Held()}) + original = threading.Thread.start + + def failing_start(self: threading.Thread) -> None: + if self.name.startswith("microvm-failure-"): + msg = "can't start new thread" + raise RuntimeError(msg) + original(self) + + monkeypatch.setattr(threading.Thread, "start", failing_start) + + assert request(worker, "/build", job_body(version=2))[0] == 500 + assert worker.idle(WAIT) + + +def test_idle_returns_false_while_a_job_runs(start: Any) -> None: + handler = Held() + worker, _ = start(handler) + + request(worker, RUN, run_body(run_payload())) + assert handler.started.wait(WAIT) + + assert worker.idle(0.05) is False + handler.release.set() + assert worker.idle(WAIT) is True + + +def test_close_ends_wait_and_can_run_twice(start: Any) -> None: + worker, _ = start(Held()) + + assert worker.wait(0.01) is False + worker.close() + worker.close() + + assert worker.wait(0) is True + with pytest.raises(OSError): + request(worker, RUN, {}) + + +def test_raises_when_the_port_is_taken(start: Any) -> None: + worker, _ = start(Held()) + + with pytest.raises(OSError): + start_microvm_worker(handler=Held(), port=worker.port) + + +def test_listens_on_ipv6_and_ipv4(start: Any) -> None: + if not socket.has_dualstack_ipv6(): + pytest.skip("this host has no dual-stack IPv6") + worker, _ = start(Held()) + + connection = http.client.HTTPConnection("::1", worker.port, timeout=WAIT) + try: + connection.request("POST", f"{HOOK_PATH_PREFIX}ready", body=b"{}") + assert connection.getresponse().status == 200 + finally: + connection.close() + assert request(worker, f"{HOOK_PATH_PREFIX}ready", {})[0] == 200 + + +def test_listens_on_ipv4_without_dual_stack( + fake_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + 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 + finally: + worker.close() + + +def test_listener_serves_a_server_of_the_caller(logger: Any, fake_client: Any) -> None: + client = fake_client() + listener = MicrovmWorkerListener.create( + handler=lambda job_input, context: "ok", + create_client=lambda region: client, + logger=logger, + ) + server = ThreadingHTTPServer(("127.0.0.1", 0), listener.request_handler_class) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + connection = http.client.HTTPConnection( + "127.0.0.1", server.server_address[1], timeout=WAIT + ) + connection.request("POST", RUN, body=json.dumps(run_body(run_payload()))) + assert connection.getresponse().status == 200 + connection.close() + assert listener.idle(WAIT) + finally: + server.shutdown() + server.server_close() + + assert listener.request_handler_class is listener.request_handler_class + assert client.names() == ["success"] + + +async def _async_handler(job_input: Any, context: MicrovmJobContext) -> None: + return None # pragma: no cover + + +class _AsyncCallable: + async def __call__(self, job_input: Any, context: MicrovmJobContext) -> None: + return None # pragma: no cover + + +class _Awaitable: + def __await__(self) -> Iterator[None]: + yield None # pragma: no cover + + +def test_fails_a_job_whose_handler_returns_a_coroutine_and_closes_it( + start: Any, +) -> None: + coroutine = _async_handler(None, None) # type: ignore[arg-type] + worker, client = start(lambda job_input, context: coroutine) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + # A closed coroutine does not warn that it was never awaited. + assert inspect.getcoroutinestate(coroutine) == inspect.CORO_CLOSED + [(name, kwargs)] = completions(client) + assert name == "failure" + assert kwargs["Error"]["ErrorType"] == "TypeError" + assert "returned an awaitable" in kwargs["Error"]["ErrorMessage"] + + +def test_fails_a_job_whose_handler_returns_another_awaitable(start: Any) -> None: + worker, client = start(lambda job_input, context: _Awaitable()) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + assert completions(client)[0][1]["Error"]["ErrorType"] == "TypeError" + + +@pytest.mark.parametrize( + ("options", "error", "message"), + [ + ({}, TypeError, "needs a handler or routes"), + ({"routes": {}}, TypeError, "needs a handler or routes"), + ({"handler": "run"}, TypeError, "must be callable. Got str"), + ({"handler": _async_handler}, TypeError, "not an async function"), + ({"handler": _AsyncCallable()}, TypeError, "not an async function"), + ({"routes": {"/build": _async_handler}}, TypeError, "not an async function"), + ({"routes": {"build": Held()}}, ValueError, 'must start with "/"'), + ( + {"routes": {f"{HOOK_PATH_PREFIX}run": Held()}}, + ValueError, + "lifecycle hook prefix", + ), + ({"routes": {MICROVM_JOB_PATH: Held()}}, ValueError, "is reserved"), + ( + {"handler": Held(), "heartbeat_interval_seconds": 0}, + ValueError, + "heartbeat_interval_seconds", + ), + ], +) +def test_create_rejects_invalid_options( + options: dict[str, Any], error: type[Exception], message: str +) -> None: + with pytest.raises(error, match=message): + MicrovmWorkerListener.create(**options) From 7462618e264dcf8bb624e3647e8a2d856642e51e Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Mon, 12 Oct 2026 10:40:57 -0700 Subject: [PATCH 2/4] test(microvm-worker): skip IPv6 test without ::1 The CI containers support IPv6 sockets, but have no ::1 address. So socket.has_dualstack_ipv6() is true there, and a connection to ::1 fails with EADDRNOTAVAIL. The test now binds ::1 first, and skips when the bind fails. The worker's dual-stack server still started there, and every other test reached it over 127.0.0.1. Two parametrized tests now set short IDs. The default ID of the 256 KB result case was a 262 KB line, and the GitHub log view stopped at it. --- .../tests/worker_test.py | 35 +++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) 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 b1df372e0..b6cc38c55 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 @@ -276,6 +276,8 @@ def handler(job_input: Any, context: MicrovmJobContext) -> None: ({1, 2}, "ResultSerializationError"), ("x" * (256 * 1024), "ResultTooLargeError"), ], + # The default test ID would hold the 256 KB result. + ids=["not-json", "over-256-kb"], ) def test_reports_a_result_that_cannot_be_sent_as_a_failure( start: Any, logger: Any, result: Any, error_type: str @@ -800,6 +802,18 @@ def test_accepts_a_chunked_body_with_extensions_and_trailers(start: Any) -> None (b"2\r\n{}\r\n", b'{"error":"invalid request"}'), (b"20000\r\n", b"request body exceeds 65536 bytes"), ], + # The default test IDs would hold the raw bytes, one of them 2,000 long. + ids=[ + "not-hex", + "negative", + "hex-prefix", + "plus-sign", + "underscore", + "no-crlf-after-chunk", + "size-line-too-long", + "ends-early", + "over-the-limit", + ], ) def test_answers_400_to_a_malformed_chunked_body( start: Any, chunks: bytes, expected: bytes @@ -1174,9 +1188,26 @@ def test_raises_when_the_port_is_taken(start: Any) -> None: start_microvm_worker(handler=Held(), port=worker.port) -def test_listens_on_ipv6_and_ipv4(start: Any) -> None: +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(): - pytest.skip("this host has no dual-stack 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") worker, _ = start(Held()) connection = http.client.HTTPConnection("::1", worker.port, timeout=WAIT) From c014618c895d2dd8321fe7124a0f97fb52636f4b Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Tue, 13 Oct 2026 09:21:30 -0700 Subject: [PATCH 3/4] fix(microvm-worker): address listener review feedback - The server listens with a backlog of 511, as Node does. With the default of 5, a burst of 50 parallel job requests got 10 to 13 resets on Linux. - A job request starts its thread before the 202. A thread that cannot start now gets HTTP 500, which the durable function retries. After a 202, it does not deliver the job again. - Chunk-size lines, their extensions, and trailers count against a limit of their own. Before, they made the worker read without bound. - The Server header no longer names the Python version. - A client that resets the connection during the headers no longer prints a traceback to stderr. - A coroutine whose close() raises still fails its job with TypeError. - setsockopt(IPV6_V6ONLY, 0) raises instead of failing silently. A new test starts IPv6 sockets IPv6-only, as Windows and the BSDs do. - The time-limit test for dropping a rejected body now detects the limit, and a second test checks that it does not end early. --- .../README.md | 2 +- .../worker.py | 80 +++++-- .../tests/worker_test.py | 203 ++++++++++++++++-- 3 files changed, 257 insertions(+), 28 deletions(-) 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 750f2a6ee..bc786b8ee 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -92,7 +92,7 @@ A document that does not match the contract gets HTTP 400 with the validation me ## 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`. +- `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 `request_queue_size` 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. - `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. 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 b832a16ba..6d8937db8 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 @@ -18,6 +18,7 @@ import re import socket import socketserver +import sys import threading import time from collections import OrderedDict @@ -296,8 +297,13 @@ def create( def request_handler_class(self) -> type[BaseHTTPRequestHandler]: """A request handler class for ``http.server`` that serves this listener. - Pass it to a ``ThreadingHTTPServer``. The listener needs a thread per - request, because it answers each request before its job ends. + Pass it to a ``ThreadingHTTPServer``. A single-threaded server serves + one request at a time. So one slow client, which the listener waits + for up to 60 seconds per read, would hold back every other request, + including the ``run`` hook. Also set the server's + ``request_queue_size`` to 511, as :func:`start_microvm_worker` does. + A session can send many jobs at once, and the default backlog of 5 + resets the connections beyond it. """ listener = self @@ -478,14 +484,19 @@ def _handle_job_request( duplicate = self._is_known(job.callback_id) if not duplicate: self._jobs.add(job.callback_id) - exchange.respond(202, {}) if duplicate: + exchange.respond(202, {}) self._logger.warning( "duplicate job ignored", extra={"microvmId": microvm_id, "callbackId": job.callback_id}, ) return + # The thread starts before the answer. A thread that cannot start + # raises, and the caller gets HTTP 500, which the durable function + # retries. After a 202, the durable function does not deliver the job + # again. self._start_job(handler, job, request.region) + exchange.respond(202, {}) def _reject(self, error: BaseException, exchange: _Exchange) -> None: """Answer 400 to a request that the worker could not accept. @@ -877,7 +888,11 @@ def of(cls, run: Callable[[], Any]) -> _Outcome: return cls(error=error) if inspect.isawaitable(value): if inspect.iscoroutine(value): - value.close() + # close() raises when the coroutine has started and ignores + # GeneratorExit, for example with an await in a finally + # block. The TypeError below must still be reported. + with contextlib.suppress(Exception): + value.close() return cls( error=TypeError( "the job handler returned an awaitable. A job handler must " @@ -923,6 +938,12 @@ def log_message(self, format: str, *args: Any) -> None: # noqa: A002 # matters through its own logger. return + def version_string(self) -> str: + # The base class names the exact Python version in the Server header. + # Every caller of the MicroVM endpoint would learn it, and no caller + # needs it. + return "aws-durable-execution-sdk-python-microvm-worker" + class _Exchange: """One request and its answer.""" @@ -1059,11 +1080,27 @@ def _read_chunked(self, max_bytes: int) -> bytes: RFC 9112 defines a chunk size as one or more hex digits. ``int(x, 16)`` also accepts a ``0x`` prefix, a sign, and underscores. So the size is matched against the RFC form first, as ``Content-Length`` is. + + The size lines, their extensions, and the trailers do not count as + body data. Without a limit of their own, they could make the worker + read without bound. Node fixed the same issue in CVE-2024-22019. So + they get their own limit of ``max_bytes``, and the data limit keeps + its meaning. """ chunks: list[bytes] = [] size = 0 - while True: + framing = 0 + + def read_framing_line() -> bytes: + nonlocal framing line = self._request.rfile.readline(_MAX_CHUNK_LINE_BYTES) + framing += len(line) + if framing > max_bytes: + raise self.too_large(max_bytes) + return line + + while True: + line = read_framing_line() if not line.endswith(b"\n"): raise _malformed_chunks() if line else _ended_early() token = line.split(b";", 1)[0].strip() @@ -1071,7 +1108,9 @@ def _read_chunked(self, max_bytes: int) -> bytes: raise _malformed_chunks() chunk_size = int(token, 16) if chunk_size == 0: - self._skip_trailers() + # The trailer lines end at an empty line. + while read_framing_line() not in {b"\r\n", b"\n", b""}: + pass return b"".join(chunks) size += chunk_size if size > max_bytes: @@ -1080,12 +1119,6 @@ def _read_chunked(self, max_bytes: int) -> bytes: if self._read_exactly(2) != b"\r\n": raise _malformed_chunks() - def _skip_trailers(self) -> None: - while True: - line = self._request.rfile.readline(_MAX_CHUNK_LINE_BYTES) - if line in {b"\r\n", b"\n", b""}: - return - def _read_exactly(self, length: int) -> bytes: data = self._request.rfile.read(length) if len(data) < length: @@ -1104,6 +1137,12 @@ def _ended_early() -> ConnectionError: class _Server(ThreadingHTTPServer): """The worker's HTTP server.""" + # socketserver listens with a backlog of 5, and Node's server with 511. + # A session can send many jobs at once. A connection beyond the backlog + # waits for a SYN retransmit, or gets a TCP reset. So the worker uses + # Node's backlog. The kernel caps it at net.core.somaxconn. + request_queue_size = 511 + def server_bind(self) -> None: # HTTPServer.server_bind looks up the host's fully qualified name, # which can wait for DNS. The worker never uses the name. So the bind @@ -1112,6 +1151,15 @@ def server_bind(self) -> None: self.server_name = "localhost" self.server_port = int(self.server_address[1]) + def handle_error(self, request: Any, client_address: Any) -> None: + # A client that resets the connection while it sends the request line + # or the headers ends here, before the listener gets the request. The + # base class prints a traceback to stderr. A client that went away + # needs none. + if isinstance(sys.exception(), ConnectionError): + return + super().handle_error(request, client_address) + class _DualStackServer(_Server): """Listens on IPv6 and IPv4 at once. @@ -1125,8 +1173,12 @@ class _DualStackServer(_Server): address_family = socket.AF_INET6 def server_bind(self) -> None: - with contextlib.suppress(OSError, AttributeError): - self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0) + # _create_server uses this class only when has_dualstack_ipv6() is + # true. That check makes this same call on a new socket. So the call + # fails only on a fault, and the worker then fails to start. A + # suppressed failure would leave a server that refuses IPv4 on a host + # whose IPv6 sockets start IPv6-only. + self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0) super().server_bind() 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 b6cc38c55..0e39dfab8 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 @@ -5,14 +5,16 @@ from __future__ import annotations +import contextlib import http.client import inspect import io import json import socket +import struct import threading import time -from collections.abc import Callable, Iterator +from collections.abc import Callable, Generator, Iterator from http.server import ThreadingHTTPServer from typing import Any @@ -679,23 +681,56 @@ def test_answers_100_continue_when_the_declared_size_is_not_over_a_limit( assert interim.startswith(b"HTTP/1.1 100 Continue") -def test_stops_discarding_a_rejected_body_after_the_time_limit( - start: Any, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(worker_module, "_DISCARD_SECONDS", 0.2) - worker, _ = start(Held()) - head = f"POST {RUN} HTTP/1.1\r\nHost: localhost\r\nContent-Length: 999999\r\n\r\n" +def _send_after_the_answer( + worker: MicrovmWorker, seconds: float +) -> tuple[bytes, BaseException | None]: + """Get a 400 for an oversized body, then keep sending for ``seconds``. + The worker shuts down its side right after the answer. So the client + reads the answer and end of file at once. The worker keeps reading and + dropping input until its time limit, and then closes the socket. So a + write after the limit fails. + + Returns: + 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: sock.sendall(head.encode() + b"x" * 100) - began = time.monotonic() answer = b"" while chunk := sock.recv(65536): answer += chunk - elapsed = time.monotonic() - began + deadline = time.monotonic() + seconds + try: + while time.monotonic() < deadline: + sock.sendall(b"x" * 4096) + time.sleep(0.01) + except ConnectionError as error: + return answer, error + return answer, None + + +def test_stops_discarding_a_rejected_body_after_the_time_limit( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(worker_module, "_DISCARD_SECONDS", 0.2) + worker, _ = start(Held()) + + answer, error = _send_after_the_answer(worker, 1.5) + + assert answer.startswith(b"HTTP/1.1 400") + assert isinstance(error, ConnectionError) + + +def test_keeps_discarding_a_rejected_body_until_the_time_limit(start: Any) -> None: + # The default limit is 10 seconds. So writes during the first second must + # succeed. A limit that ended too early would fail them. + worker, _ = start(Held()) + + answer, error = _send_after_the_answer(worker, 1.0) assert answer.startswith(b"HTTP/1.1 400") - assert elapsed < WAIT + assert error is None class _EndlessConnection: @@ -801,6 +836,14 @@ def test_accepts_a_chunked_body_with_extensions_and_trailers(start: Any) -> None (b"2" * 2000, b"malformed chunked encoding"), (b"2\r\n{}\r\n", b'{"error":"invalid request"}'), (b"20000\r\n", b"request body exceeds 65536 bytes"), + ( + b"0\r\n" + b"X: y\r\n" * 20_000 + b"\r\n", + b"request body exceeds 65536 bytes", + ), + ( + b"".join(b"1;" + b"e" * 1000 + b"\r\nx\r\n" for _ in range(70)), + b"request body exceeds 65536 bytes", + ), ], # The default test IDs would hold the raw bytes, one of them 2,000 long. ids=[ @@ -813,6 +856,8 @@ def test_accepts_a_chunked_body_with_extensions_and_trailers(start: Any) -> None "size-line-too-long", "ends-early", "over-the-limit", + "trailers-over-the-limit", + "extensions-over-the-limit", ], ) def test_answers_400_to_a_malformed_chunked_body( @@ -999,6 +1044,109 @@ def test_runs_a_job_once_when_its_request_arrives_twice( ) +def test_accepts_a_burst_of_parallel_job_requests(start: Any) -> None: + # socketserver's default backlog of 5 reset 10 to 13 of these 50 + # connections. The worker listens with Node's backlog of 511. + worker, _ = start(routes={"/build": lambda job_input, context: job_input}) + barrier = threading.Barrier(50) + statuses: list[int] = [] + + def deliver(k: int) -> None: + barrier.wait(WAIT) + statuses.append(request(worker, "/build", job_body(callbackId=f"cb-{k}"))[0]) + + threads = [threading.Thread(target=deliver, args=(k,)) for k in range(50)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(WAIT) + assert worker.idle(WAIT) + + assert statuses == [202] * 50 + + +def test_names_no_python_version_in_the_server_header(start: Any) -> None: + worker, _ = start(Held()) + + answer = raw_request( + worker, + f"POST {HOOK_PATH_PREFIX}ready HTTP/1.1\r\nHost: localhost\r\n" + "Content-Length: 2\r\n\r\n{}".encode(), + ) + + assert b"Server: aws-durable-execution-sdk-python-microvm-worker\r\n" in answer + assert b"Python/" not in answer + + +def test_prints_no_traceback_for_a_client_that_resets_during_the_headers( + start: Any, capsys: pytest.CaptureFixture[str] +) -> None: + worker, _ = start(Held()) + + with socket.create_connection(("127.0.0.1", 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. + time.sleep(0.1) + # A zero linger time makes close() send a TCP reset. + sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) + # The request thread handles the reset after the close. + time.sleep(0.3) + assert request(worker, f"{HOOK_PATH_PREFIX}ready", {})[0] == 200 + + assert "Traceback" not in capsys.readouterr().err + + +def test_reports_other_server_errors_as_the_base_class_does( + start: Any, capsys: pytest.CaptureFixture[str] +) -> None: + worker, _ = start(Held()) + server = worker._server # noqa: SLF001 + + try: + msg = "unexpected" + raise ValueError(msg) + except ValueError: + server.handle_error(None, ("127.0.0.1", 0)) + + assert "ValueError: unexpected" in capsys.readouterr().err + + +class _IPv6OnlyByDefault(socket.socket): + """Starts each new IPv6 socket IPv6-only. + + Windows and the BSDs do this, and so does Linux with + ``net.ipv6.bindv6only=1``. + """ + + def __init__( + self, + family: int = -1, + type: int = -1, # noqa: A002 + proto: int = -1, + fileno: int | None = None, + ) -> None: + super().__init__(family, type, proto, fileno) + # accept() passes fileno for each connection. Only a new socket + # starts with the system default. + if fileno is None and self.family == socket.AF_INET6: + self.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1) + + +def test_listens_on_ipv4_when_ipv6_sockets_start_ipv6_only( + start: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + # Linux starts an IPv6 socket dual-stack by default. So there the server + # accepts IPv4 without its IPV6_V6ONLY call, and the test above passes + # either way. Here every new IPv6 socket starts IPv6-only. + if not socket.has_dualstack_ipv6(): + pytest.skip("this host has no dual-stack IPv6") + monkeypatch.setattr(socket, "socket", _IPv6OnlyByDefault) + worker, _ = start(Held()) + + assert request(worker, f"{HOOK_PATH_PREFIX}ready", {})[0] == 200 + + def test_runs_a_job_once_when_parallel_requests_deliver_it(start: Any) -> None: handler = Held() handler.release.set() @@ -1131,7 +1279,9 @@ def flaky_start(self: threading.Thread) -> None: original(self) monkeypatch.setattr(threading.Thread, "start", flaky_start) - assert request(worker, "/build", job_body())[0] == 202 + # The thread starts before the answer. So a thread that cannot start gets + # HTTP 500, which the durable function retries, and the retry runs the job. + assert request(worker, "/build", job_body())[0] == 500 assert worker.idle(WAIT) assert request(worker, "/build", job_body())[0] == 202 assert worker.idle(WAIT) @@ -1266,8 +1416,8 @@ async def __call__(self, job_input: Any, context: MicrovmJobContext) -> None: class _Awaitable: - def __await__(self) -> Iterator[None]: - yield None # pragma: no cover + def __await__(self) -> Generator[None, None, None]: + yield None def test_fails_a_job_whose_handler_returns_a_coroutine_and_closes_it( @@ -1296,6 +1446,33 @@ def test_fails_a_job_whose_handler_returns_another_awaitable(start: Any) -> None assert completions(client)[0][1]["Error"]["ErrorType"] == "TypeError" +async def _ignores_generator_exit() -> None: + try: + await _Awaitable() + finally: + await _Awaitable() + + +def test_fails_a_job_whose_started_coroutine_cannot_be_closed(start: Any) -> None: + # The coroutine has started and awaits in a finally block. So close() + # raises RuntimeError: coroutine ignored GeneratorExit. The job must still + # report its TypeError. + coroutine = _ignores_generator_exit() + coroutine.send(None) + worker, client = start(lambda job_input, context: coroutine) + + request(worker, RUN, run_body(run_payload())) + assert worker.idle(WAIT) + + [(name, kwargs)] = completions(client) + assert name == "failure" + assert kwargs["Error"]["ErrorType"] == "TypeError" + # Run the finally block to its end, so the coroutine closes before the + # garbage collector meets it. + with contextlib.suppress(GeneratorExit, StopIteration): + coroutine.send(None) + + @pytest.mark.parametrize( ("options", "error", "message"), [ From ce8e39442ef908546d4bfb2885c307cf3dd03f46 Mon Sep 17 00:00:00 2001 From: Pooya Paridel Date: Tue, 13 Oct 2026 10:01:21 -0700 Subject: [PATCH 4/4] fix(microvm-worker): fail a run hook job that cannot start - The run hook starts the job's thread before its 200. When the thread cannot start, the request thread fails the job's callback. Lambda sends the run hook once, so the job would otherwise wait for its heartbeat timeout or its timeout. - The framing limit of a chunked body has its own message, so a rejection names its cause. - The docstring and the README say to set request_queue_size as a class attribute. The server's constructor already listens, so an attribute set on the server object changes nothing. - Two comments no longer quote probe counts. --- .../README.md | 4 +- .../worker.py | 67 +++++++++++++++---- .../tests/worker_test.py | 36 +++++++++- 3 files changed, 88 insertions(+), 19 deletions(-) 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 bc786b8ee..10c91c80c 100644 --- a/packages/aws-durable-execution-sdk-python-microvm-worker/README.md +++ b/packages/aws-durable-execution-sdk-python-microvm-worker/README.md @@ -86,13 +86,13 @@ A second delivery of a job is answered like the first and ignored, while the job The worker creates a Lambda client for each job, after the job arrives, not at image build. A client created before the snapshot would carry build-time state into every MicroVM. A client kept across jobs would keep connections that a suspend of the MicroVM can leave dead. The worker closes each job's default client when the job ends. It never closes a client from `create_client`. -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. 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. Other lifecycle hooks get HTTP 200. Other paths get HTTP 404. +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. ## 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 `request_queue_size` 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. +- `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. - `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. 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 6d8937db8..8ffbe8305 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 @@ -300,10 +300,16 @@ def request_handler_class(self) -> type[BaseHTTPRequestHandler]: Pass it to a ``ThreadingHTTPServer``. A single-threaded server serves one request at a time. So one slow client, which the listener waits for up to 60 seconds per read, would hold back every other request, - including the ``run`` hook. Also set the server's - ``request_queue_size`` to 511, as :func:`start_microvm_worker` does. - A session can send many jobs at once, and the default backlog of 5 - resets the connections beyond it. + including the ``run`` hook. + + Also give the server a backlog of 511, as :func:`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 setting ``request_queue_size`` on a server object + changes nothing. Set it as a class attribute of a subclass instead:: + + class Server(ThreadingHTTPServer): + request_queue_size = 511 """ listener = self @@ -440,9 +446,33 @@ def _handle_run_hook(self, exchange: _Exchange) -> None: new_job.callback_id, payload.region, ) - exchange.respond(200, {}) + start_error: Exception | None = None if payload is not None and new_job is not None and self._handler is not None: - self._start_job(self._handler, new_job, payload.region) + # The thread starts before the answer, so that idle() covers the + # job once the caller has the answer. + try: + self._start_job(self._handler, new_job, payload.region) + except Exception as error: # noqa: BLE001 + start_error = error + with self._lock: + self._reports += 1 + exchange.respond(200, {}) + if payload is not None and new_job is not None and start_error is not None: + # Lambda sends the run hook once. So no later delivery comes, and + # the job would wait for its heartbeat timeout or its timeout. A + # new thread cannot start either. So this request thread, whose + # answer has gone out, fails the callback itself. + self._logger.error( + "could not start the run hook job", + extra={ + "microvmId": microvm_id, + "callbackId": new_job.callback_id, + "error": describe(start_error), + }, + ) + self._send_failure(start_error, new_job.callback_id, payload.region) + return + if new_job is not None: return if later: @@ -542,7 +572,10 @@ def _start_job( try: thread.start() except BaseException: - # The job never ran. So a later delivery may run it. + # The job never ran. A job request then answers 500, so the + # durable function delivers the job again, and the retry may run + # it. The run hook fails the callback instead, because Lambda + # sends it only once. with self._changed: self._jobs.discard(job.callback_id) self._changed.notify_all() @@ -994,11 +1027,16 @@ def read_json(self, max_bytes: int) -> Any: msg = "request body is not valid JSON" raise InvalidRunHookPayloadError(msg) from error - def too_large(self, max_bytes: int) -> InvalidRunHookPayloadError: - """Return the error for a body over ``max_bytes``, and remember that - the worker answers before it reads the rest of the body.""" + def too_large( + self, max_bytes: int, part: str = "request body" + ) -> InvalidRunHookPayloadError: + """Return the error for a request part over ``max_bytes``. + + The worker answers such a request before it reads the rest of the + body. So this method also marks the body unread. + """ self._body_unread = True - return InvalidRunHookPayloadError(f"request body exceeds {max_bytes} bytes") + return InvalidRunHookPayloadError(f"{part} exceeds {max_bytes} bytes") def respond(self, status: int, body: Mapping[str, object]) -> None: """Send the answer. A client that has gone away is logged.""" @@ -1026,8 +1064,7 @@ def discard_unread_body(self) -> None: 2. The connection then closes with unread bytes in its receive buffer. So the kernel sends a TCP reset instead of a normal close. 3. A client that is still sending gets the reset before it reads the - answer. A probe showed this in 20 of 20 tries, so no client saw - the 400. + answer. So the client never sees the 400. 4. The durable function fails at once on a 400, and retries a reset. So the reset turns a clear rejection into retries. @@ -1096,7 +1133,9 @@ def read_framing_line() -> bytes: line = self._request.rfile.readline(_MAX_CHUNK_LINE_BYTES) framing += len(line) if framing > max_bytes: - raise self.too_large(max_bytes) + raise self.too_large( + max_bytes, "the chunked framing of the request body" + ) return line while True: 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 0e39dfab8..18c87e039 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 @@ -824,6 +824,9 @@ def test_accepts_a_chunked_body_with_extensions_and_trailers(start: Any) -> None assert completions(client)[0][1]["Result"] == b'{"repo":"org/app"}' +_FRAMING_TOO_LARGE = b"the chunked framing of the request body exceeds 65536 bytes" + + @pytest.mark.parametrize( ("chunks", "expected"), [ @@ -838,11 +841,11 @@ def test_accepts_a_chunked_body_with_extensions_and_trailers(start: Any) -> None (b"20000\r\n", b"request body exceeds 65536 bytes"), ( b"0\r\n" + b"X: y\r\n" * 20_000 + b"\r\n", - b"request body exceeds 65536 bytes", + _FRAMING_TOO_LARGE, ), ( b"".join(b"1;" + b"e" * 1000 + b"\r\nx\r\n" for _ in range(70)), - b"request body exceeds 65536 bytes", + _FRAMING_TOO_LARGE, ), ], # The default test IDs would hold the raw bytes, one of them 2,000 long. @@ -1045,7 +1048,7 @@ def test_runs_a_job_once_when_its_request_arrives_twice( def test_accepts_a_burst_of_parallel_job_requests(start: Any) -> None: - # socketserver's default backlog of 5 reset 10 to 13 of these 50 + # On Linux, socketserver's default backlog of 5 resets some of these # connections. The worker listens with Node's backlog of 511. worker, _ = start(routes={"/build": lambda job_input, context: job_input}) barrier = threading.Barrier(50) @@ -1289,6 +1292,33 @@ def flaky_start(self: threading.Thread) -> None: assert len(handler.inputs) == 1 +def test_fails_a_run_hook_job_whose_thread_cannot_start( + start: Any, logger: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + worker, client = start(Held()) + original = threading.Thread.start + + def failing_start(self: threading.Thread) -> None: + if self.name.startswith("microvm-job-"): + msg = "can't start new thread" + raise RuntimeError(msg) + original(self) + + monkeypatch.setattr(threading.Thread, "start", failing_start) + # Lambda sends the run hook once. So the worker answers 200, and fails + # the callback, instead of leaving the job to its heartbeat timeout. + assert request(worker, RUN, run_body(run_payload())) == (200, {}) + assert worker.idle(WAIT) + + [(name, kwargs)] = completions(client) + assert name == "failure" + assert kwargs["Error"] == { + "ErrorType": "RuntimeError", + "ErrorMessage": "can't start new thread", + } + assert "could not start the run hook job" in logger.messages("error") + + def test_counts_no_report_whose_thread_cannot_start( start: Any, monkeypatch: pytest.MonkeyPatch ) -> None: