diff --git a/README.md b/README.md index 6d129e1..d295c39 100644 --- a/README.md +++ b/README.md @@ -1173,6 +1173,14 @@ If code cannot establish connection, it will start deployment of python using [D * `timeout` - Timeout for command execution. * `command` - Command to be executed. * `ip` - IP address of the EFI Shell target system. + + Responses: + * `200` - command output in the body, `rc` header holds the return code. + * `400` - no command provided. + * `504` - the result did not arrive within `timeout` (+ client poll grace period). The body is empty + and `rc` is `-1`. The command is marked as *abandoned*, so it is never executed late and its result + is discarded on arrival - this keeps the caller and the EFI client in sync after a timeout. + The `X-RShell-Timeout-Reason` header explains *why* the wait failed (see below). * `/post_result` - Endpoint to post results back to the host. Headers fields: * `CommandID` - Unique identifier for the command. @@ -1184,8 +1192,56 @@ If code cannot establish connection, it will start deployment of python using [D * `CommandID` - Unique identifier for the command. Body: * Exception details. - * `/getCommandToExecute` - Endpoint to retrieve commands to be executed on the EFI Shell target system. Returns commandline with generated CommandID. - * `/health/` - Endpoint to check the health status of the connection. + * `/getCommandToExecute` - Endpoint to retrieve commands to be executed on the EFI Shell target system. Returns commandline with generated CommandID. Commands abandoned after a timeout are skipped. Every poll also refreshes the client liveness marker. + * `/health/` - Reports `200` only when the client is known **and** polled within `CLIENT_LIVENESS_TIMEOUT_SECONDS`. A client that was seen once but went silent is reported as `503`, so a dead client no longer looks healthy for the rest of the run. + * `/disconnect_client/` - Removes the client and drops everything still queued for it. + +### Diagnosing a timeout + +The EFI client runs every command through a blocking `os.system()` call and asks for new work only +after the previous command finished. A tool that hangs on the DUT - for example one that waits for a +key press because it was not started in batch/silent mode - therefore stops the polling loop for good. +From the caller side that is indistinguishable from a slow command: **every** later command times out. + +To make that visible the server tracks when each client last polled and which command it picked up, +and returns the conclusion in the `X-RShell-Timeout-Reason` header, which `RShellConnection` logs: + +| Situation | Reported reason | +| --- | --- | +| Client never contacted the server | `has never contacted the server - check that rshell_client.py is running on the DUT` | +| Client fetched this command and never came back | `picked this command up Ns ago and never reported back ... run it manually with output redirected to a file to check` | +| Client stopped polling while running an earlier command | `stopped polling Ns ago ... still blocked executing '' - the command most likely hangs on the DUT` | +| Client keeps polling, command is just slow | `is alive and polling ... consider passing a bigger timeout` | + +The second and third rows point at the DUT, not at the server - no server-side change can make a +hung EFI process return. + +### Timeouts + +`execute_command(timeout=...)` is forwarded to the server. When no timeout is given, the server falls +back to `600` seconds and the connection logs that the fallback was applied. Set `default_timeout` on +the connection for commands that legitimately run longer: + +```python +from mfd_connect import RShellConnection + +conn = RShellConnection(ip="10.10.10.10", server_ip="10.10.10.1", default_timeout=2500) +``` + +### Server resource limits + +The server keeps its in-memory state bounded, so a long test session cannot exhaust RAM: + +| Constant | Default | Meaning | +| --- | --- | --- | +| `STALE_OUTPUT_TTL_SECONDS` | `600` | Age after which an uncollected output or abandoned ID is evicted. | +| `MAX_STORED_OUTPUTS` | `512` | Hard cap on results waiting to be collected. | +| `MAX_ABANDONED_COMMAND_IDS` | `512` | Hard cap on remembered abandoned command IDs. | +| `MAX_PENDING_COMMANDS_PER_CLIENT` | `256` | Hard cap on commands queued for a single client. | +| `CLIENT_LIVENESS_TIMEOUT_SECONDS` | `300` | Silence after which `/health` reports the client as gone. | + +Waiting for a result is event driven - the caller is woken up as soon as the output is posted, +and all shared state is guarded by a lock because Werkzeug serves requests in threads. `rshell.py` is a Connection class that calls RESTful API endpoints provided by `rshell_server.py` to execute commands on the EFI Shell target system. If required, starts `rshell_server.py` on the host machine. @@ -1194,6 +1250,14 @@ RShell server can be started manually using the following command: python -m mfd_connect.rshell_server ``` +`rshell_client.py` runs on the DUT and accepts an optional source IP and source port: +```bash +rshell_client.py [source_ip] [source_port] +``` +The source port defaults to `80`; override it when that fixed local port collides with sockets left +in `TIME_WAIT` by earlier connections. A transient network error no longer terminates the client - +it retries instead, so a single glitch cannot silence the DUT for the rest of the run. + ## OS supported: * LNX * WINDOWS diff --git a/mfd_connect/rshell.py b/mfd_connect/rshell.py index 3b366b3..6793dfe 100644 --- a/mfd_connect/rshell.py +++ b/mfd_connect/rshell.py @@ -39,6 +39,12 @@ # can be adjusted based on requirements and observed behavior of platforms. PLATFORM_POWER_TRANSITION_DELAY_SECONDS = 10 +# Timeout applied by the RShell server when the caller does not provide one. +# Kept in sync with rshell_server.execute_command so the fallback can be logged explicitly. +SERVER_DEFAULT_TIMEOUT_SECONDS = 600 +# Header carrying the server side explanation of a 504. +TIMEOUT_REASON_HEADER = "X-RShell-Timeout-Reason" + class RShellConnection(Connection): """RShell Connection Class.""" @@ -50,6 +56,7 @@ def __init__( model: "BaseModel | None" = None, cache_system_data: bool = True, connection_timeout: int = 60, + default_timeout: int | None = None, ): """ Initialize RShellConnection. @@ -58,8 +65,12 @@ def __init__( :param server_ip: The IP address of the server to connect to (optional). :param model: The Pydantic model to use for the connection (optional). :param cache_system_data: Whether to cache system data (default: True). + :param connection_timeout: Time to wait for the RShell client to show up. + :param default_timeout: Timeout used by execute_command when the caller passes none. + When left as None the server applies its own default + (``SERVER_DEFAULT_TIMEOUT_SECONDS``). """ - super().__init__(model=model, cache_system_data=cache_system_data) + super().__init__(model=model, default_timeout=default_timeout, cache_system_data=cache_system_data) self._ip = ip self.server_ip = server_ip self.server_process: LocalProcess | None = None @@ -76,10 +87,14 @@ def wait_for_connection(self, connection_timeout: int) -> None: logger.log(level=log_levels.MODULE_DEBUG, msg="Checking RShell server health") try: status_code = requests.get( - f"http://{self.server_ip}/health/{self._ip}", proxies={"no_proxy": "*"} + f"http://{self.server_ip}/health/{self._ip}", + proxies={"no_proxy": "*"}, ).status_code except requests.RequestException as e: - logger.log(level=log_levels.MODULE_DEBUG, msg=f"RShell server health check failed with error: {e}") + logger.log( + level=log_levels.MODULE_DEBUG, + msg=f"RShell server health check failed with error: {e}", + ) status_code = None if status_code == 200: logger.log(level=log_levels.MODULE_DEBUG, msg="RShell server is healthy") @@ -96,7 +111,10 @@ def disconnect(self, stop_client: bool = False, stop_server: bool = False) -> No :param stop_client: Whether to stop the RShell client (default: False). """ - requests.post(f"http://{self.server_ip}/disconnect_client/{self._ip}", proxies={"no_proxy": "*"}) + requests.post( + f"http://{self.server_ip}/disconnect_client/{self._ip}", + proxies={"no_proxy": "*"}, + ) if stop_client: logger.log(level=log_levels.MODULE_DEBUG, msg="Stopping RShell client") self.execute_command("end") @@ -200,14 +218,40 @@ def execute_command( level=log_levels.MODULE_DEBUG, msg="Custom exceptions are not supported for RShellConnection and will be ignored.", ) - timeout_string = f" with timeout {timeout} seconds" if timeout is not None else "" - logger.log(level=log_levels.CMD, msg=f"Executing >{self._ip}> '{command}',{timeout_string}") + effective_timeout = timeout if timeout is not None else self.default_timeout + if effective_timeout is None: + logger.log( + level=log_levels.MODULE_DEBUG, + msg=f"No timeout given for '{command}'; the RShell server will apply its default of " + f"{SERVER_DEFAULT_TIMEOUT_SECONDS} seconds. Pass 'timeout' or set 'default_timeout' " + f"on the connection for commands that legitimately run longer.", + ) + timeout_string = f" with timeout {effective_timeout} seconds" if effective_timeout is not None else "" + logger.log( + level=log_levels.CMD, + msg=f"Executing >{self._ip}> '{command}',{timeout_string}", + ) response = requests.post( f"http://{self.server_ip}/execute_command", - data={"command": command, "timeout": timeout, "ip": self._ip}, + data={"command": command, "timeout": effective_timeout, "ip": self._ip}, proxies={"no_proxy": "*"}, ) + if response.status_code == 504: + reason = response.headers.get(TIMEOUT_REASON_HEADER) + waited = effective_timeout if effective_timeout is not None else SERVER_DEFAULT_TIMEOUT_SECONDS + message = ( + f"RShell server timed out after {waited}s waiting for the result of '{command}'. " + f"The command was dropped so it will not be executed later by the client." + ) + if reason: + message = f"{message} Reason: {reason}." + logger.log(level=log_levels.MODULE_DEBUG, msg=message) + elif response.status_code >= 500: + logger.log( + level=log_levels.MODULE_DEBUG, + msg=f"RShell server returned an internal error ({response.status_code}) for '{command}'.", + ) completed_process = ConnectionCompletedProcess( args=command, stdout=response.text, @@ -336,7 +380,10 @@ def stop_server(self) -> None: break time.sleep(1) else: - logger.log(level=log_levels.MODULE_DEBUG, msg="RShell server did not stop within timeout") + logger.log( + level=log_levels.MODULE_DEBUG, + msg="RShell server did not stop within timeout", + ) raise RuntimeError("RShell server did not stop within timeout") logger.log(level=log_levels.MODULE_DEBUG, msg="RShell server stopped") diff --git a/mfd_connect/rshell_client.py b/mfd_connect/rshell_client.py index a305c3f..5408e63 100644 --- a/mfd_connect/rshell_client.py +++ b/mfd_connect/rshell_client.py @@ -5,9 +5,18 @@ Make sure that the Python UEFI interpreter is compiled with Socket module support. + +Usage:: + + rshell_client.py [source_ip] [source_port] + +The client polls the server for work, runs one command at a time and posts the result back. +Commands are executed with a blocking ``os.system()`` call, so a tool that waits for input on +the DUT stops the whole loop - every later command then times out on the server side. Keep +that in mind when adding tools: always run them in a non interactive/batch mode. """ -__version__ = "1.1.0" +__version__ = "1.2.0" try: import httplib as client @@ -23,6 +32,16 @@ source_address = sys.argv[2] else: source_address = None +# Local port the outgoing connection is bound to. Only relevant together with source_address, +# which exists so that the server sees the expected client IP. Override it when the fixed +# port collides with sockets left in TIME_WAIT by previous connections. +if len(sys.argv) > 3: + source_port = int(sys.argv[3]) +else: + source_port = 80 + +# How long to wait before retrying after a failed server interaction. +RETRY_WAIT_SECONDS = 5 os_name = os.name @@ -41,6 +60,16 @@ def _sleep(interval): # noqa: ANN001, ANN202 time.sleep = _sleep +def _close(connection): # noqa: ANN001, ANN202 + """Close a connection without letting a broken socket kill the client.""" + if connection is None: + return + try: + connection.close() + except Exception as exp: # noqa: BLE001 + print("Ignoring error while closing the connection:", exp) + + def _get_command(): # noqa: ANN202 """Get the command from server to execute on client machine.""" # construct the list of tests by interacting with server @@ -61,27 +90,37 @@ def _get_command(): # noqa: ANN202 while True: # Connect to server - source_address_parameter = (source_address, 80) if source_address else None - conn = client.HTTPConnection(http_server, source_address=source_address_parameter) - # get the command from server - _command = _get_command() + source_address_parameter = (source_address, source_port) if source_address else None + conn = None + try: + conn = client.HTTPConnection(http_server, source_address=source_address_parameter) + # get the command from server + _command = _get_command() + except Exception as exp: # noqa: BLE001 + # A transient network error must not end the client. If it did, the DUT would stop + # asking for work and every following command would time out on the server. + print("Failed to get a command from the server:", exp) + _close(conn) + time.sleep(RETRY_WAIT_SECONDS) + continue + if not _command: - conn.close() - time.sleep(5) + _close(conn) + time.sleep(RETRY_WAIT_SECONDS) continue cmd_str, _id = _command cmd_str = cmd_str.decode("utf-8") cmd_name = cmd_str.split(" ")[0] if cmd_name == "end": print("No more commands available to run") - conn.close() + _close(conn) exit(0) print("Executing", cmd_str) if cmd_name.startswith("reset"): print("Reset command received, shutting down the platform") os.system(cmd_str) # execute reset command on machine - conn.close() + _close(conn) exit(0) non_echo = False @@ -120,15 +159,19 @@ def _get_command(): # noqa: ANN202 if non_echo and f: f.close() os.system("del " + out) - except Exception as exp: - conn.request( - "POST", - "exception", - body=cmd + str(exp), - headers={"Content-Type": "text/plain", "Connection": "keep-alive", "CommandID": _id}, - ) + except Exception as exp: # noqa: BLE001 + try: + conn.request( + "POST", + "exception", + body=cmd + str(exp), + headers={"Content-Type": "text/plain", "Connection": "keep-alive", "CommandID": _id}, + ) + except Exception as report_exp: # noqa: BLE001 + # Reporting failed too - stay alive so the next poll can still reach the server. + print("Failed to report the error to the server:", report_exp) print("output posted to server") - conn.close() + _close(conn) print("closed the connection") time.sleep(1) diff --git a/mfd_connect/rshell_server/rshell_server.py b/mfd_connect/rshell_server/rshell_server.py index 36e1646..5b0e23c 100644 --- a/mfd_connect/rshell_server/rshell_server.py +++ b/mfd_connect/rshell_server/rshell_server.py @@ -5,31 +5,235 @@ This script implements a RESTful server using Flask to manage command execution on connected RShell clients. + +Flow: + 1. ``/execute_command`` - caller queues a command and blocks until its result arrives. + 2. ``/getCommandToExecute`` - the EFI client polls for the next command to run. + 3. ``/post_result`` - the EFI client returns the command output. + 4. ``/exception`` - the EFI client reports a failure instead of an output. + +Because the EFI client polls in a slow loop (and some commands take minutes), a waiter may +give up before its result arrives. Such a command is marked as *abandoned* so that: + * it is skipped when the client asks for the next command (it is never executed late), and + * its result is dropped on arrival instead of being stored forever. + +This keeps the caller and the EFI client in sync after a timeout and keeps memory bounded. + +The server also tracks client liveness. The EFI client executes commands with a blocking +``os.system()`` call, so a command that hangs on the DUT (for example a tool waiting for a +key press) stops the polling loop for good. From the caller side this is indistinguishable +from a slow command, therefore every timeout is reported together with a diagnosis telling +whether the client kept polling or went silent while executing an earlier command. """ +import threading import time -from collections import namedtuple -from queue import Queue +from collections import OrderedDict +from queue import Empty, Queue +from typing import NamedTuple from uuid import uuid4 from flask import Flask, Response, request -__version__ = "1.1.0" +__version__ = "1.3.0" + +# How long the EFI client sleeps between polls - added to the caller timeout as a grace period. +CLIENT_LOOP_WAIT_SECONDS = 5 +# How long an output that nobody collected is kept before it is evicted. +STALE_OUTPUT_TTL_SECONDS = 600 +# Hard caps protecting the server against unbounded memory growth. +MAX_STORED_OUTPUTS = 512 +MAX_ABANDONED_COMMAND_IDS = 512 +MAX_PENDING_COMMANDS_PER_CLIENT = 256 +# Longest single wait inside get_output() - keeps the waiting loop responsive. +OUTPUT_WAIT_SLICE_SECONDS = 1.0 +# How long a client may stay silent before /health reports it as gone. The client does not poll +# while it executes a command, so this has to comfortably exceed a normal command duration. +CLIENT_LIVENESS_TIMEOUT_SECONDS = 300 +# Header carrying the human readable explanation of a 504 back to RShellConnection. +TIMEOUT_REASON_HEADER = "X-RShell-Timeout-Reason" + # Global command queue -output_object = namedtuple("OutputObject", ["output", "rc"]) -command_object = namedtuple("CommandObject", ["command_id", "str"]) +class OutputObject(NamedTuple): + """Store command output together with its return code.""" + + output: str + rc: int + + +class CommandObject(NamedTuple): + """Store queued command metadata sent to a specific client.""" + + command_id: str + str: str -output_queue: dict[str, output_object] = dict() + +class RunningCommand(NamedTuple): + """Store the command a client picked up and has not reported back yet.""" + + command_id: str + command: str + started_at: float + + +output_object = OutputObject +command_object = CommandObject + +# Results waiting to be collected by their /execute_command caller. +output_queue: "OrderedDict[str, OutputObject]" = OrderedDict() +output_queue_timestamps: dict[str, float] = dict() +# Commands whose caller already gave up - results must not be stored and they must not run. +abandoned_command_ids: "OrderedDict[str, float]" = OrderedDict() +# Per client (IP) queue of commands waiting to be picked up. command_dict_queue: dict[str, Queue] = dict() clients: list = [] +# Last time (monotonic) each client asked for work - used to tell a slow command from a dead client. +client_last_seen: dict[str, float] = dict() +# Command each client picked up last and has not reported back yet. +client_running_command: dict[str, RunningCommand] = dict() + +# Guards every structure above. Re-entrant so helpers can be called with the lock already held. +_state_lock = threading.RLock() +_output_available = threading.Condition(_state_lock) app = Flask(__name__) -def get_output(command_id: str, timeout: float = 600) -> output_object: +def _cleanup_stale_outputs(now: float | None = None, ttl: int = STALE_OUTPUT_TTL_SECONDS) -> None: + """ + Remove orphaned command outputs that have been kept longer than the configured TTL. + + Also enforces the hard caps on stored outputs and abandoned command IDs. + + :param now: Reference time (``time.monotonic()`` based). Defaults to the current time. + :param ttl: Maximum age, in seconds, of an uncollected output. + """ + with _state_lock: + current_time = time.monotonic() if now is None else now + + stale_outputs = [cid for cid, created in list(output_queue_timestamps.items()) if current_time - created >= ttl] + for command_id in stale_outputs: + output_queue.pop(command_id, None) + output_queue_timestamps.pop(command_id, None) + + stale_abandoned = [cid for cid, created in list(abandoned_command_ids.items()) if current_time - created >= ttl] + for command_id in stale_abandoned: + abandoned_command_ids.pop(command_id, None) + + while len(output_queue) > MAX_STORED_OUTPUTS: + oldest_id, _ = output_queue.popitem(last=False) + output_queue_timestamps.pop(oldest_id, None) + + while len(abandoned_command_ids) > MAX_ABANDONED_COMMAND_IDS: + abandoned_command_ids.popitem(last=False) + + +def _abandon_command(command_id: str) -> None: + """ + Mark a command as no longer awaited, so it is neither executed late nor stored on arrival. + + :param command_id: The ID of the command whose caller gave up. + """ + with _state_lock: + abandoned_command_ids[command_id] = time.monotonic() + output_queue.pop(command_id, None) + output_queue_timestamps.pop(command_id, None) + _cleanup_stale_outputs() + + +def _store_output(command_id: str, output: str, rc: int) -> bool: + """ + Persist command output together with its insertion timestamp and wake up the waiter. + + Results of abandoned commands are dropped instead of being kept forever. + + :param command_id: The ID of the command the output belongs to. + :param output: The command output. + :param rc: The return code of the command. + :return: True when the output was stored, False when it was dropped as abandoned. """ - Retrieve the output for a given command ID. + with _output_available: + if abandoned_command_ids.pop(command_id, None) is not None: + print(f"Dropping output of abandoned command {command_id} - caller already gave up") + return False + output_queue[command_id] = output_object(output=output, rc=rc) + output_queue_timestamps[command_id] = time.monotonic() + _cleanup_stale_outputs() + _output_available.notify_all() + return True + + +def _mark_client_finished(ip_address: str, command_id: str) -> None: + """ + Record that a client reported back, so it is no longer considered busy nor silent. + + :param ip_address: The IP address the result came from. + :param command_id: The ID of the command the client reported. + """ + with _state_lock: + client_last_seen[ip_address] = time.monotonic() + running = client_running_command.get(ip_address) + if running is not None and running.command_id == command_id: + client_running_command.pop(ip_address, None) + + +def diagnose_timeout(ip_address: str, command: str, waiting_since: float, command_id: str | None = None) -> str: + """ + Explain why a command timed out, based on whether the EFI client kept polling. + + The EFI client runs commands through a blocking ``os.system()`` call, so a tool that hangs + on the DUT stops the polling loop permanently. Distinguishing that from a genuinely slow + command is impossible for the caller, hence this explicit diagnosis. + + :param ip_address: The IP address of the client the command was queued for. + :param command: The command that timed out. + :param waiting_since: Monotonic timestamp of the moment the caller started waiting. + :param command_id: The ID of the command that timed out, when known. + :return: Human readable reason, safe to send as an HTTP header value. + """ + with _state_lock: + last_seen = client_last_seen.get(ip_address) + running = client_running_command.get(ip_address) + + if last_seen is None: + return ( + f"the RShell client {ip_address} has never contacted the server - " + f"check that rshell_client.py is running on the DUT" + ) + + now = time.monotonic() + + if running is not None and command_id is not None and running.command_id == command_id: + busy_for = now - running.started_at + return ( + f"the RShell client {ip_address} picked this command up {busy_for:.0f}s ago and never reported " + f"back - '{command}' is either still running or hangs on the DUT (for example waiting for " + f"input); run it manually with output redirected to a file to check" + ) + + if last_seen < waiting_since: + idle_for = now - last_seen + reason = ( + f"the RShell client {ip_address} stopped polling {idle_for:.0f}s ago, so it never asked " + f"for this command" + ) + if running is not None: + busy_for = now - running.started_at + reason += f"; it is still blocked executing '{running.command}' started {busy_for:.0f}s ago" + return reason + " - the command most likely hangs on the DUT (for example waiting for input)" + + return ( + f"the RShell client {ip_address} is alive and polling, but '{command}' did not finish in time - " + f"consider passing a bigger timeout" + ) + + +def get_output(command_id: str, timeout: float = 600) -> OutputObject: + """ + Retrieve the output for a given command ID, waiting until it arrives. + + The wait is event driven - the caller is woken up as soon as the result is posted. :param command_id: The ID of the command to retrieve output for. :param timeout: The maximum time to wait for output (in seconds). @@ -38,14 +242,18 @@ def get_output(command_id: str, timeout: float = 600) -> output_object: """ print("Getting output for command ID:", command_id) print(f"Waiting for output {timeout} seconds") - timeout = timeout + 5 # add time for client loop waiting - while timeout > 0: - result = output_queue.get(command_id, None) - if result is not None: - return result - time.sleep(1) - timeout -= 1 - raise TimeoutError("Command timed out") + deadline = time.monotonic() + timeout + CLIENT_LOOP_WAIT_SECONDS + with _output_available: + while True: + result = output_queue.pop(command_id, None) + if result is not None: + output_queue_timestamps.pop(command_id, None) + return result + remaining = deadline - time.monotonic() + if remaining <= 0: + _abandon_command(command_id) + raise TimeoutError("Command timed out") + _output_available.wait(min(remaining, OUTPUT_WAIT_SLICE_SECONDS)) def add_command_to_queue(command: str, ip_address: str) -> str: @@ -58,19 +266,39 @@ def add_command_to_queue(command: str, ip_address: str) -> str: """ print("Adding command to queue:", command) _id = str(uuid4().int) - if command_dict_queue.get(ip_address) is None: - command_dict_queue[ip_address] = Queue() - command_dict_queue[ip_address].put(command_object(command_id=_id, str=command)) + with _state_lock: + client_queue = command_dict_queue.get(ip_address) + if client_queue is None: + client_queue = Queue() + command_dict_queue[ip_address] = client_queue + while client_queue.qsize() >= MAX_PENDING_COMMANDS_PER_CLIENT: + try: + dropped = client_queue.get_nowait() + except Empty: + break + print(f"Dropping queued command {dropped.command_id} for {ip_address} - queue limit reached") + abandoned_command_ids[dropped.command_id] = time.monotonic() + client_queue.put(command_object(command_id=_id, str=command)) return _id @app.route("/health/", methods=["GET"]) def health_check(ip: str) -> Response: - """Health check endpoint.""" - if ip in clients: - return Response("OK", status=200) - else: + """ + Report whether the RShell client is connected and still polling. + + A client that was seen once but went silent is reported as gone, otherwise a dead client + would look healthy for the rest of the run. + """ + with _state_lock: + last_seen = client_last_seen.get(ip) + connected = ip in clients and last_seen is not None + if not connected or last_seen is None: return Response("Client not connected", status=503) + silent_for = time.monotonic() - last_seen + if silent_for <= CLIENT_LIVENESS_TIMEOUT_SECONDS: + return Response("OK", status=200) + return Response(f"Client stopped polling {silent_for:.0f}s ago", status=503) @app.route("/getCommandToExecute", methods=["GET"]) @@ -78,23 +306,37 @@ def get_command_to_execute() -> Response: """ Get the next command to execute for the connected client. + Commands whose caller already timed out are skipped, so the client never falls behind. + Every poll also refreshes the client liveness marker - the client only asks for work once + the previous command finished. + :return: The next command to execute. """ ip_address = str(request.remote_addr) - if ip_address not in clients: - print(f"Client connected: {ip_address}") - clients.append(ip_address) - client_queue = command_dict_queue.get(ip_address, Queue()) - if not client_queue.empty(): - command_object = client_queue.get() - return Response( - command_object.str, - status=200, - mimetype="text/plain", - headers={"CommandID": command_object.command_id}, - ) - else: - return Response("No more elements left in the queue", status=204) + with _state_lock: + if ip_address not in clients: + print(f"Client connected: {ip_address}") + clients.append(ip_address) + client_last_seen[ip_address] = time.monotonic() + client_running_command.pop(ip_address, None) + client_queue = command_dict_queue.get(ip_address) + while client_queue is not None and not client_queue.empty(): + queued_command = client_queue.get() + if abandoned_command_ids.pop(queued_command.command_id, None) is not None: + print(f"Skipping abandoned command {queued_command.command_id} for {ip_address}") + continue + client_running_command[ip_address] = RunningCommand( + command_id=queued_command.command_id, + command=queued_command.str, + started_at=time.monotonic(), + ) + return Response( + queued_command.str, + status=200, + mimetype="text/plain", + headers={"CommandID": queued_command.command_id}, + ) + return Response("No more elements left in the queue", status=204) @app.route("/exception", methods=["POST"]) @@ -110,7 +352,8 @@ def post_exception() -> Response: command_id = str(request.headers.get("CommandID")) print("CommandID: ", command_id) print(str(read_data, encoding="utf-8")) - output_queue[command_id] = output_object(output=str(read_data, encoding="utf-8"), rc=-1) + _mark_client_finished(str(request.remote_addr), command_id) + _store_output(command_id, str(read_data, encoding="utf-8"), rc=-1) return Response("Exception received", status=200) @@ -127,36 +370,67 @@ def execute_command() -> Response: timeout = int(request.form.get("timeout", 600)) command = request.form.get("command") ip_address = str(request.form.get("ip")) - if command: - _id = add_command_to_queue(command, ip_address) - if command == "end": - return Response("No more commands available to run", status=200) - if command.startswith("reset"): - return Response("Reset command sent", status=200) + if not command: + return Response("No command provided", status=400) + + _id = add_command_to_queue(command, ip_address) + if command == "end": + return Response("No more commands available to run", status=200) + if command.startswith("reset"): + return Response("Reset command sent", status=200) + + waiting_since = time.monotonic() + try: process = get_output(_id, timeout) + except TimeoutError: + # Return a clean gateway timeout instead of a Flask HTML 500 page, which the caller + # would otherwise store verbatim as the command stdout. The body stays empty on + # purpose - the explanation travels in a header so it never pollutes stdout. + reason = diagnose_timeout(ip_address, command, waiting_since, command_id=_id) + print(f"Command {_id} timed out after {timeout}s - marked as abandoned; {reason}") return Response( - process.output.encode("utf-8"), - status=200, + b"", + status=504, headers={ "Content-type": "text/plain", "CommandID": _id, - "rc": process.rc, + "rc": "-1", + TIMEOUT_REASON_HEADER: reason, }, ) - else: - return Response("No command provided", status=400) + + return Response( + process.output.encode("utf-8"), + status=200, + headers={ + "Content-type": "text/plain", + "CommandID": _id, + "rc": str(process.rc), + }, + ) @app.route("/disconnect_client/", methods=["POST"]) def disconnect_client(ip_address: str) -> Response: """ - Disconnect a client from the server. + Disconnect a client from the server and drop everything queued for it. :param ip_address: The IP address of the client to disconnect. """ - if ip_address in clients: - clients.remove(ip_address) - print(f"Client disconnected: {ip_address}") + with _state_lock: + if ip_address in clients: + clients.remove(ip_address) + client_last_seen.pop(ip_address, None) + client_running_command.pop(ip_address, None) + client_queue = command_dict_queue.pop(ip_address, None) + while client_queue is not None and not client_queue.empty(): + try: + pending = client_queue.get_nowait() + except Empty: + break + abandoned_command_ids[pending.command_id] = time.monotonic() + _cleanup_stale_outputs() + print(f"Client disconnected: {ip_address}") return Response("Client disconnected", status=200) @@ -168,7 +442,8 @@ def post_result() -> Response: rc = int(request.headers.get("rc", -1)) print("CommandID: ", command_id) print(str(read_data, encoding="utf-8")) - output_queue[command_id] = output_object(output=str(read_data, encoding="utf-8"), rc=rc) + _mark_client_finished(str(request.remote_addr), command_id) + _store_output(command_id, str(read_data, encoding="utf-8"), rc=rc) return Response("Results received", status=200) diff --git a/tests/unit/test_mfd_connect/test_rshell.py b/tests/unit/test_mfd_connect/test_rshell.py index 9a507d1..84e73ad 100644 --- a/tests/unit/test_mfd_connect/test_rshell.py +++ b/tests/unit/test_mfd_connect/test_rshell.py @@ -16,7 +16,12 @@ from mfd_connect.base import ConnectionCompletedProcess from mfd_connect.exceptions import ConnectionCalledProcessError, OsNotSupported -from mfd_connect.rshell import PLATFORM_POWER_TRANSITION_DELAY_SECONDS, RShellConnection +from mfd_connect.rshell import ( + PLATFORM_POWER_TRANSITION_DELAY_SECONDS, + SERVER_DEFAULT_TIMEOUT_SECONDS, + TIMEOUT_REASON_HEADER, + RShellConnection, +) class TestRShellConnection: @@ -30,6 +35,8 @@ def rshell(self): conn.server_ip = "127.0.0.1" conn.server_process = None conn.cache_system_data = False + # __init__ is mocked out, so the attribute normally set by Connection.__init__ is added here. + conn._default_timeout = None return conn def test_init_local_server_start(self, mocker): @@ -222,7 +229,7 @@ class _FakeBaseModel: def test_execute_command_with_all_unsupported_args_and_skip_logging(self, rshell, mocker): post = mocker.patch("mfd_connect.rshell.requests.post") - post.return_value = Mock(text="out", headers={"rc": "7"}) + post.return_value = Mock(status_code=200, text="out", headers={"rc": "7"}) result = rshell.execute_command( "echo hello", @@ -249,7 +256,10 @@ def test_execute_command_with_all_unsupported_args_and_skip_logging(self, rshell ) def test_execute_command_logs_stdout_and_default_rc(self, rshell, mocker): - mocker.patch("mfd_connect.rshell.requests.post", return_value=Mock(text="stdout", headers={})) + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="stdout", headers={}), + ) result = rshell.execute_command("echo hi") @@ -257,13 +267,95 @@ def test_execute_command_logs_stdout_and_default_rc(self, rshell, mocker): assert result.stdout == "stdout" def test_execute_command_no_stdout(self, rshell, mocker): - mocker.patch("mfd_connect.rshell.requests.post", return_value=Mock(text="", headers={"rc": "0"})) + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="", headers={"rc": "0"}), + ) result = rshell.execute_command("echo hi") assert result.return_code == 0 assert result.stdout == "" + def test_execute_command_logs_server_timeout(self, rshell, mocker, caplog): + """A 504 from the server means the command was dropped - it must be visible in the logs.""" + caplog.set_level(0) + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=504, text="", headers={"rc": "-1"}), + ) + + result = rshell.execute_command("FS0:\\Tools\\nvmupdate64e.efi /i /l") + + assert result.return_code == -1 + assert result.stdout == "" + assert "timed out" in caplog.text + + def test_execute_command_logs_server_internal_error(self, rshell, mocker, caplog): + caplog.set_level(0) + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=500, text="", headers={}), + ) + + result = rshell.execute_command("echo hi") + + assert result.return_code == -1 + assert "internal error" in caplog.text + + def test_execute_command_logs_server_timeout_reason(self, rshell, mocker, caplog): + """The server explains WHY it timed out - that diagnosis must reach the log.""" + caplog.set_level(0) + reason = "the RShell client 10.10.10.10 stopped polling 620s ago, so it never asked for this command" + mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock( + status_code=504, + text="", + headers={"rc": "-1", TIMEOUT_REASON_HEADER: reason}, + ), + ) + + rshell.execute_command("FS0:\\Tools\\nvmupdate64e.efi /i /l") + + assert reason in caplog.text + + def test_execute_command_without_timeout_warns_about_server_default(self, rshell, mocker, caplog): + """A command sent without a timeout is silently capped by the server - make that visible.""" + caplog.set_level(0) + post = mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="out", headers={"rc": "0"}), + ) + + rshell.execute_command("FS0:\\Tools\\nvmupdate64e.efi /i /l") + + assert post.call_args.kwargs["data"]["timeout"] is None + assert str(SERVER_DEFAULT_TIMEOUT_SECONDS) in caplog.text + + def test_execute_command_uses_connection_default_timeout(self, rshell, mocker): + """default_timeout must be applied when the caller does not pass one.""" + rshell._default_timeout = 2500 + post = mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="out", headers={"rc": "0"}), + ) + + rshell.execute_command("FS0:\\Tools\\nvmupdate64e.efi /i /l") + + assert post.call_args.kwargs["data"]["timeout"] == 2500 + + def test_execute_command_explicit_timeout_wins_over_default(self, rshell, mocker): + rshell._default_timeout = 2500 + post = mocker.patch( + "mfd_connect.rshell.requests.post", + return_value=Mock(status_code=200, text="out", headers={"rc": "0"}), + ) + + rshell.execute_command("ver", timeout=20) + + assert post.call_args.kwargs["data"]["timeout"] == 20 + def test_path_python_312_plus(self, rshell, monkeypatch, mocker): monkeypatch.setattr(sys, "version_info", (3, 12, 0)) factory = mocker.patch("mfd_connect.rshell.custom_path_factory", return_value="cp") diff --git a/tests/unit/test_mfd_connect/test_rshell_server.py b/tests/unit/test_mfd_connect/test_rshell_server.py index 13dfaea..e54151e 100644 --- a/tests/unit/test_mfd_connect/test_rshell_server.py +++ b/tests/unit/test_mfd_connect/test_rshell_server.py @@ -5,6 +5,8 @@ import importlib.util import runpy import sys +import threading +import time from pathlib import Path import pytest @@ -30,41 +32,64 @@ class TestRShellServerScript: def server_module(self): module = _load_server_module() module.output_queue.clear() + module.output_queue_timestamps.clear() + module.abandoned_command_ids.clear() module.command_dict_queue.clear() module.clients.clear() + module.client_last_seen.clear() + module.client_running_command.clear() return module def test_get_output_success(self, server_module): command_id = "cmd1" expected = server_module.output_object(output="hello", rc=0) server_module.output_queue[command_id] = expected + server_module.output_queue_timestamps[command_id] = 123.0 result = server_module.get_output(command_id, timeout=0) assert result == expected + assert command_id not in server_module.output_queue + assert command_id not in server_module.output_queue_timestamps def test_get_output_timeout(self, server_module): with pytest.raises(TimeoutError, match="Command timed out"): server_module.get_output("missing", timeout=-5) - def test_get_output_waits_then_returns(self, server_module, monkeypatch): - class _QueueProbe: - def __init__(self): - self.count = 0 + def test_get_output_timeout_marks_command_as_abandoned(self, server_module): + """A caller that gives up must mark its command so it is not executed/stored later.""" + with pytest.raises(TimeoutError): + server_module.get_output("gone", timeout=-5) - def get(self, _command_id, _default=None): - self.count += 1 - if self.count == 1: - return None - return server_module.output_object(output="later", rc=4) + assert "gone" in server_module.abandoned_command_ids - monkeypatch.setattr(server_module, "output_queue", _QueueProbe()) - monkeypatch.setattr(server_module.time, "sleep", lambda _x: None) + def test_get_output_is_woken_up_by_posted_result(self, server_module): + """Waiting is event driven - the waiter returns as soon as the result is stored.""" - result = server_module.get_output("cmd-later", timeout=0) + def _post_later(): + time.sleep(0.1) + server_module._store_output("cmd-later", "later", 4) + + threading.Thread(target=_post_later, daemon=True).start() + + started = time.monotonic() + result = server_module.get_output("cmd-later", timeout=10) + elapsed = time.monotonic() - started assert result.output == "later" assert result.rc == 4 + assert elapsed < 2, "waiter should be notified instead of polling" + + def test_store_output_drops_result_of_abandoned_command(self, server_module): + """Late results must not be kept in memory once nobody waits for them anymore.""" + server_module._abandon_command("dead-cmd") + + stored = server_module._store_output("dead-cmd", "late output", 0) + + assert stored is False + assert "dead-cmd" not in server_module.output_queue + assert "dead-cmd" not in server_module.output_queue_timestamps + assert "dead-cmd" not in server_module.abandoned_command_ids def test_add_command_to_queue_new_and_existing_queue(self, server_module): first_id = server_module.add_command_to_queue("echo 1", "10.0.0.1") @@ -74,6 +99,14 @@ def test_add_command_to_queue_new_and_existing_queue(self, server_module): queue_obj = server_module.command_dict_queue["10.0.0.1"] assert queue_obj.qsize() == 2 + def test_add_command_to_queue_is_bounded(self, server_module): + """Per-client queue must not grow without limits.""" + limit = server_module.MAX_PENDING_COMMANDS_PER_CLIENT + for index in range(limit + 25): + server_module.add_command_to_queue(f"echo {index}", "10.0.0.9") + + assert server_module.command_dict_queue["10.0.0.9"].qsize() <= limit + def test_health_check_endpoint(self, server_module): client = server_module.app.test_client() @@ -81,10 +114,130 @@ def test_health_check_endpoint(self, server_module): assert response_not_connected.status_code == 503 server_module.clients.append("10.0.0.1") + server_module.client_last_seen["10.0.0.1"] = time.monotonic() response_connected = client.get("/health/10.0.0.1") assert response_connected.status_code == 200 assert response_connected.get_data(as_text=True) == "OK" + def test_health_check_reports_client_that_stopped_polling(self, server_module): + """A client that was seen once but went silent must not look healthy forever.""" + client = server_module.app.test_client() + server_module.clients.append("10.0.0.1") + server_module.client_last_seen["10.0.0.1"] = ( + time.monotonic() - server_module.CLIENT_LIVENESS_TIMEOUT_SECONDS - 10 + ) + + response = client.get("/health/10.0.0.1") + + assert response.status_code == 503 + assert "stopped polling" in response.get_data(as_text=True) + + def test_health_check_requires_a_poll_before_reporting_healthy(self, server_module): + """Being present in 'clients' is not enough - the client has to have polled.""" + client = server_module.app.test_client() + server_module.clients.append("10.0.0.1") + + assert client.get("/health/10.0.0.1").status_code == 503 + + def test_poll_refreshes_liveness_and_tracks_running_command(self, server_module): + client = server_module.app.test_client() + command_id = server_module.add_command_to_queue("FS0:\\Tools\\nvmupdate64e.efi /i /l", "1.2.3.4") + + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + + assert "1.2.3.4" in server_module.client_last_seen + running = server_module.client_running_command["1.2.3.4"] + assert running.command_id == command_id + assert running.command == "FS0:\\Tools\\nvmupdate64e.efi /i /l" + + def test_posting_result_clears_running_command(self, server_module): + client = server_module.app.test_client() + command_id = server_module.add_command_to_queue("ver", "1.2.3.4") + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + + client.post( + "/post_result", + data=b"out", + headers={"CommandID": command_id, "rc": "0"}, + environ_base={"REMOTE_ADDR": "1.2.3.4"}, + ) + + assert "1.2.3.4" not in server_module.client_running_command + + def test_diagnose_timeout_for_unknown_client(self, server_module): + reason = server_module.diagnose_timeout("9.9.9.9", "ver", time.monotonic()) + + assert "never contacted the server" in reason + + def test_diagnose_timeout_detects_client_stuck_on_earlier_command(self, server_module): + """The exact situation from the field: the client blocks in os.system() and stops polling.""" + now = time.monotonic() + server_module.client_last_seen["1.2.3.4"] = now - 700 + server_module.client_running_command["1.2.3.4"] = server_module.RunningCommand( + command_id="old", + command="FS0:\\Tools\\nvmupdate64e.efi /i /l", + started_at=now - 700, + ) + + reason = server_module.diagnose_timeout("1.2.3.4", "ls FS0:\\Tools\\nut_upd1.cfg", waiting_since=now - 600) + + assert "stopped polling" in reason + assert "nvmupdate64e.efi /i /l" in reason + assert "hangs on the DUT" in reason + + def test_diagnose_timeout_for_slow_but_alive_client(self, server_module): + now = time.monotonic() + server_module.client_last_seen["1.2.3.4"] = now + + reason = server_module.diagnose_timeout("1.2.3.4", "slow command", waiting_since=now - 600) + + assert "alive and polling" in reason + assert "bigger timeout" in reason + + def test_diagnose_timeout_when_client_took_this_command_and_never_returned(self, server_module): + """First timeout of a hanging command: the client fetched it but never reported back.""" + now = time.monotonic() + server_module.client_last_seen["1.2.3.4"] = now - 600 + server_module.client_running_command["1.2.3.4"] = server_module.RunningCommand( + command_id="cid-1", + command="FS0:\\Tools\\nvmupdate64e.efi /i /l", + started_at=now - 600, + ) + + reason = server_module.diagnose_timeout( + "1.2.3.4", + "FS0:\\Tools\\nvmupdate64e.efi /i /l", + waiting_since=now - 605, + command_id="cid-1", + ) + + assert "picked this command up" in reason + assert "never reported back" in reason + + def test_execute_command_timeout_sends_reason_header_and_empty_body(self, server_module): + """The body must stay empty so the caller never stores an error page as stdout.""" + client = server_module.app.test_client() + + response = client.post("/execute_command", data={"command": "ver", "timeout": "-5", "ip": "1.2.3.4"}) + + assert response.status_code == 504 + assert response.get_data() == b"" + assert response.headers["rc"] == "-1" + assert response.headers[server_module.TIMEOUT_REASON_HEADER] + + def test_disconnect_client_clears_liveness_state(self, server_module): + """A reconnecting client must not inherit the liveness state of the previous session.""" + client = server_module.app.test_client() + server_module.add_command_to_queue("ver", "1.2.3.4") + client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + assert "1.2.3.4" in server_module.client_last_seen + + client.post("/disconnect_client/1.2.3.4") + + assert "1.2.3.4" not in server_module.client_last_seen + assert "1.2.3.4" not in server_module.client_running_command + assert client.get("/health/1.2.3.4").status_code == 503 + def test_get_command_to_execute_endpoint(self, server_module): client = server_module.app.test_client() @@ -99,6 +252,29 @@ def test_get_command_to_execute_endpoint(self, server_module): assert response_with_command.get_data(as_text=True) == "echo hi" assert response_with_command.headers["CommandID"] == command_id + def test_get_command_to_execute_skips_abandoned_commands(self, server_module): + """Abandoned commands must never reach the client, otherwise it stays one command behind.""" + client = server_module.app.test_client() + abandoned_id = server_module.add_command_to_queue("slow command", "1.2.3.4") + server_module._abandon_command(abandoned_id) + live_id = server_module.add_command_to_queue("echo alive", "1.2.3.4") + + response = client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + + assert response.status_code == 200 + assert response.headers["CommandID"] == live_id + assert response.get_data(as_text=True) == "echo alive" + assert abandoned_id not in server_module.abandoned_command_ids + + def test_get_command_to_execute_returns_204_when_only_abandoned(self, server_module): + client = server_module.app.test_client() + abandoned_id = server_module.add_command_to_queue("slow command", "1.2.3.4") + server_module._abandon_command(abandoned_id) + + response = client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": "1.2.3.4"}) + + assert response.status_code == 204 + def test_post_exception_endpoint(self, server_module): client = server_module.app.test_client() response = client.post("/exception", data=b"boom", headers={"CommandID": "cid-1"}) @@ -106,6 +282,7 @@ def test_post_exception_endpoint(self, server_module): assert response.status_code == 200 assert server_module.output_queue["cid-1"].output == "boom" assert server_module.output_queue["cid-1"].rc == -1 + assert "cid-1" in server_module.output_queue_timestamps def test_execute_command_endpoint_paths(self, server_module, monkeypatch): client = server_module.app.test_client() @@ -140,17 +317,70 @@ def test_execute_command_endpoint_paths(self, server_module, monkeypatch): assert response_normal.headers["Content-type"].startswith("text/plain") assert response_normal.headers["CommandID"] + def test_execute_command_returns_gateway_timeout_instead_of_html_error(self, server_module): + """A timeout must not return a Flask HTML 500 page, which the caller stores as stdout.""" + client = server_module.app.test_client() + + response = client.post( + "/execute_command", + data={ + "command": "FS0:\\Tools\\nvmupdate64e.efi /i /l", + "timeout": "-5", + "ip": "1.1.1.1", + }, + ) + + assert response.status_code == 504 + assert response.get_data(as_text=True) == "" + assert response.headers["rc"] == "-1" + assert "" not in response.get_data(as_text=True) + + def test_timed_out_command_does_not_desync_next_command(self, server_module): + """After a timeout the next command must succeed - the server has to re-sync.""" + client = server_module.app.test_client() + ip = "10.102.23.150" + + timed_out = client.post( + "/execute_command", + data={"command": "slow", "timeout": "-5", "ip": ip}, + ) + assert timed_out.status_code == 504 + + # The stale command must not be handed out to the client anymore. + assert client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}).status_code == 204 + + next_id = server_module.add_command_to_queue("ver", ip) + handed_out = client.get("/getCommandToExecute", environ_base={"REMOTE_ADDR": ip}) + assert handed_out.status_code == 200 + assert handed_out.headers["CommandID"] == next_id + + server_module._store_output(next_id, "UEFI Shell", 0) + assert server_module.get_output(next_id, timeout=1).output == "UEFI Shell" + def test_disconnect_client_endpoint(self, server_module): client = server_module.app.test_client() server_module.clients.append("2.2.2.2") + server_module.add_command_to_queue("echo hi", "2.2.2.2") response_existing = client.post("/disconnect_client/2.2.2.2") assert response_existing.status_code == 200 assert "2.2.2.2" not in server_module.clients + assert "2.2.2.2" not in server_module.command_dict_queue response_missing = client.post("/disconnect_client/8.8.8.8") assert response_missing.status_code == 200 + def test_disconnect_client_abandons_pending_commands(self, server_module): + """Pending commands of a disconnected client must not be executed after reconnect.""" + client = server_module.app.test_client() + server_module.clients.append("3.3.3.3") + pending_id = server_module.add_command_to_queue("echo hi", "3.3.3.3") + + client.post("/disconnect_client/3.3.3.3") + + assert pending_id in server_module.abandoned_command_ids + assert server_module._store_output(pending_id, "late", 0) is False + def test_post_result_endpoint(self, server_module): client = server_module.app.test_client() @@ -158,6 +388,7 @@ def test_post_result_endpoint(self, server_module): assert response_default_rc.status_code == 200 assert server_module.output_queue["cmd-a"].output == "output-a" assert server_module.output_queue["cmd-a"].rc == -1 + assert "cmd-a" in server_module.output_queue_timestamps response_given_rc = client.post( "/post_result", @@ -167,6 +398,63 @@ def test_post_result_endpoint(self, server_module): assert response_given_rc.status_code == 200 assert server_module.output_queue["cmd-b"].output == "output-b" assert server_module.output_queue["cmd-b"].rc == 3 + assert "cmd-b" in server_module.output_queue_timestamps + + def test_cleanup_stale_outputs_removes_expired_entries_only(self, server_module): + # Timestamps simulate time.monotonic() values (seconds since arbitrary boot reference). + # now=4000 s, ttl=3600 s -> stale (t=10) is evicted, fresh (t=500) is kept. + server_module.output_queue["stale"] = server_module.output_object(output="old", rc=-1) + server_module.output_queue_timestamps["stale"] = 10.0 + server_module.output_queue["fresh"] = server_module.output_object(output="new", rc=0) + server_module.output_queue_timestamps["fresh"] = 500.0 + + server_module._cleanup_stale_outputs(now=4000.0, ttl=3600) + + assert "stale" not in server_module.output_queue + assert "stale" not in server_module.output_queue_timestamps + assert server_module.output_queue["fresh"].output == "new" + + def test_cleanup_stale_outputs_expires_abandoned_ids(self, server_module): + server_module.abandoned_command_ids["old"] = 10.0 + server_module.abandoned_command_ids["recent"] = 3900.0 + + server_module._cleanup_stale_outputs(now=4000.0, ttl=600) + + assert "old" not in server_module.abandoned_command_ids + assert "recent" in server_module.abandoned_command_ids + + def test_output_queue_is_hard_capped(self, server_module): + """Even without TTL expiry the stored outputs must stay bounded.""" + limit = server_module.MAX_STORED_OUTPUTS + for index in range(limit + 120): + server_module._store_output(f"cid-{index}", "payload", 0) + + assert len(server_module.output_queue) <= limit + assert len(server_module.output_queue_timestamps) == len(server_module.output_queue) + + def test_concurrent_access_is_thread_safe(self, server_module): + """Werkzeug serves requests in threads - shared state must not raise or corrupt.""" + errors = [] + + def _worker(worker_id): + try: + for index in range(30): + command_id = f"w{worker_id}-{index}" + server_module._store_output(command_id, "data", 0) + server_module.get_output(command_id, timeout=1) + server_module.add_command_to_queue(f"cmd{index}", f"10.0.1.{worker_id}") + server_module._cleanup_stale_outputs() + except Exception as exc: # noqa: BLE001 + errors.append(f"{type(exc).__name__}: {exc}") + + threads = [threading.Thread(target=_worker, args=(worker_id,)) for worker_id in range(12)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not errors + assert not server_module.output_queue def test_run_function_starts_flask(self, server_module, monkeypatch): """Test that the run() function starts Flask with correct host and port."""