feat: Replace Docker monkey-patch with timeout
This commit is contained in:
@@ -247,12 +247,15 @@ class DockerSandbox:
|
||||
# Command execution
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def exec(self, command: str, timeout: int = 120) -> tuple[int, str]:
|
||||
def exec(self, command: str, timeout: int = 120, kill_grace: int = 5) -> tuple[int, str]:
|
||||
"""
|
||||
Run *command* inside the container via the low-level exec API.
|
||||
|
||||
Returns ``(exit_code, combined_stdout_stderr)``.
|
||||
Timeout is enforced via a socket-level timeout on the exec socket.
|
||||
|
||||
Exit code 124 indicates a timeout (137 if escalation to SIGKILL was
|
||||
needed). *command* is passed as a single argv element to ``bash -c``,
|
||||
so shell metacharacters within it need no extra escaping.
|
||||
"""
|
||||
if self._container is None:
|
||||
return 1, "Sandbox container is not running."
|
||||
@@ -261,60 +264,37 @@ class DockerSandbox:
|
||||
if self._working_dir is not None:
|
||||
create_kwargs["workdir"] = self._working_dir
|
||||
|
||||
# TODO(fragile): timeout enforcement relies on private docker-py internals
|
||||
# (frames_iter, demux_adaptor, consume_socket_output from docker.utils.socket)
|
||||
# and monkey-patches select.select / select.poll for the duration of the read
|
||||
# — not thread-safe if multiple exec() calls run concurrently. Replace when
|
||||
# docker-py adds native per-call timeout support.
|
||||
# See https://github.com/docker/docker-py/issues/2651
|
||||
#
|
||||
# On Linux docker-py uses select.poll (not select.select), so both are patched.
|
||||
wrapped = [
|
||||
"timeout",
|
||||
"--kill-after",
|
||||
f"{kill_grace}s",
|
||||
f"{timeout}s",
|
||||
"bash",
|
||||
"-c",
|
||||
command,
|
||||
]
|
||||
|
||||
try:
|
||||
exec_id = self._client.api.exec_create(
|
||||
self._container.id,
|
||||
["bash", "-c", command],
|
||||
wrapped,
|
||||
stdout=True,
|
||||
stderr=True,
|
||||
**create_kwargs,
|
||||
)
|
||||
sock = self._client.api.exec_start(exec_id["Id"], socket=True)
|
||||
sock._sock.settimeout(timeout)
|
||||
|
||||
timeout_ms = timeout * 1000
|
||||
|
||||
class _PollWithTimeout:
|
||||
def __init__(self):
|
||||
self._inner = _original_poll()
|
||||
|
||||
def register(self, fd, eventmask):
|
||||
return self._inner.register(fd, eventmask)
|
||||
|
||||
def poll(self, *args):
|
||||
result = self._inner.poll(timeout_ms)
|
||||
if not result:
|
||||
raise _socket.timeout(f"timed out after {timeout}s")
|
||||
return result
|
||||
|
||||
with ExitStack() as stack:
|
||||
stack.enter_context(patch.object(
|
||||
_select, "select",
|
||||
new=lambda rlist, wlist, xlist: _original_select(
|
||||
rlist, wlist, xlist, timeout
|
||||
),
|
||||
))
|
||||
if _original_poll is not None:
|
||||
stack.enter_context(
|
||||
patch.object(_select, "poll", new=_PollWithTimeout)
|
||||
)
|
||||
gen = (demux_adaptor(*frame) for frame in frames_iter(sock, tty=False))
|
||||
stdout, stderr = consume_socket_output(gen, demux=True)
|
||||
gen = (demux_adaptor(*frame) for frame in frames_iter(sock, tty=False))
|
||||
stdout, stderr = consume_socket_output(gen, demux=True)
|
||||
|
||||
exit_code = self._client.api.exec_inspect(exec_id["Id"])["ExitCode"]
|
||||
if exit_code is None:
|
||||
exit_code = 0
|
||||
output = (stdout or b"") + (stderr or b"")
|
||||
return exit_code, output.decode("utf-8", errors="replace")
|
||||
except _socket.timeout:
|
||||
return 124, f"Command timed out after {timeout}s"
|
||||
|
||||
output = ((stdout or b"") + (stderr or b"")).decode("utf-8", errors="replace")
|
||||
|
||||
if exit_code in (124, 137):
|
||||
return exit_code, output + f"\n[command timed out after {timeout}s]"
|
||||
return exit_code, output
|
||||
except Exception as exc:
|
||||
return 1, f"Error running command in container: {exc}"
|
||||
|
||||
Reference in New Issue
Block a user