fix: recover rshell server after command timeout and bound memory

README.md

@@ -1173,6 +1173,17 @@ 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.
+    * `413` - command output exceeds 16 MiB.
+    * `429` - a client or command admission limit was reached; existing commands are not evicted.
+    * `503` - pending results exceed the storage limit.
+    * `504` - the result did not arrive within `timeout` (+ client poll grace period). The body is empty
+      and `rc` is `-1`. Queued commands are removed and late results are discarded.
+      A command already picked up by the client may still execute; a timeout does not stop it.
+      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 +1195,78 @@ 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/<ip>` - 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/<ip>` - 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/<ip>` - Removes the client and drops everything still queued for it.
+
+Results are removed after collection. Unknown, duplicate and abandoned results are not retained.
+The server accepts up to 256 clients, 512 active commands and 256 queued commands per client.
+Stored result payloads are limited to 64 MiB in total and 16 MiB per result, measured as UTF-8 bytes.
+Overflow is reported without truncating results or evicting another caller's result.
+Disconnecting clears the client's pending results and wakes its waiting callers.
+Idle client records are reclaimed during subsequent requests; actively awaited commands keep their client records.
+Unfinished-command diagnostics are retained separately, up to 256 records. Records without an active waiter
+expire one hour after pickup; the oldest such records may be evicted earlier to enforce the cap.
+Receiving the matching result, a new poll, or an explicit disconnect clears the diagnostic record.
+
+HTTP errors other than `504` raise `requests.HTTPError` in `RShellConnection.execute_command()`.
+Server-generated HTTP errors have an empty body, `rc: -1` and an `X-RShell-Error` header;
+they are not returned as successful command output.
+
+### 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:
+
+The header value is escaped to single-line ASCII and limited to 1024 characters.
+
+| 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 '<command>' - 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)
+```
+
+The fire-and-forget commands `end` and `reset*` return immediately, but must be picked up within
+the requested timeout (default 600 seconds) plus the 5-second client poll grace period.
+Undelivered commands are removed during subsequent requests, releasing their command and client slots.
+
+### 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. |
+| `RUNNING_COMMAND_TTL_SECONDS` | `3600` | Age since pickup after which an unfinished-command diagnostic without an active waiter is evicted. |
+| `MAX_TIMEOUT_REASON_LENGTH` | `1024` | Maximum length of the ASCII-escaped timeout reason header. |
+
+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 +1275,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 <server_ip> [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

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,42 @@ 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"Any queued command was removed; a command already picked up by the client may still be running."
+            )
+            if reason:
+                message = f"{message} Reason: {reason}."
+            logger.log(level=log_levels.MODULE_DEBUG, msg=message)
+        elif response.status_code >= 400:
+            logger.log(
+                level=log_levels.MODULE_DEBUG,
+                msg=f"RShell request failed ({response.status_code}): "
+                f"{response.headers.get('X-RShell-Error', response.reason)}",
+            )
+            response.raise_for_status()
         completed_process = ConnectionCompletedProcess(
             args=command,
             stdout=response.text,
@@ -336,7 +382,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")

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 <server_ip> [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,20 @@ 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.
+
+    :param connection: HTTP connection to close, or None when no connection was created.
+    """
+    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
@@ -50,6 +83,9 @@ def _get_command():  # noqa: ANN202
     _id = rsp.getheader("CommandID")
     if status == 204:
         return None
+    if status != 200:
+        rsp.read()
+        raise RuntimeError("RShell command request failed with HTTP {}".format(status))
 
     print("Waiting for command from server: ")
     data_received = rsp.read()
@@ -61,27 +97,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 +166,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)