Skip to content

fix: recover rshell server after command timeout and bound memory #45

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
k-ziomek wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: recover rshell server after command timeout and bound memory #45

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 66 additions & 2 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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/<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.

### 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 '<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)
```

### 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.

Expand All @@ -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 <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
Expand Down
63 changes: 55 additions & 8 deletions mfd_connect/rshell.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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")
Expand All @@ -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")
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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")
Expand Down
77 changes: 60 additions & 17 deletions mfd_connect/rshell_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
Loading
Loading