From 3d770207398a0ebc2f8dd5d84ea9f56ad1687773 Mon Sep 17 00:00:00 2001 From: zhangzherui Date: Thu, 23 Jul 2026 03:51:24 +0800 Subject: [PATCH] fix(runner): prevent subprocess pipe deadlocks --- sandbox/runners/base.py | 151 ++++++++++++++++++----- sandbox/tests/runners/test_base.py | 190 +++++++++++++++++++++++++++++ 2 files changed, 312 insertions(+), 29 deletions(-) create mode 100644 sandbox/tests/runners/test_base.py diff --git a/sandbox/runners/base.py b/sandbox/runners/base.py index 0eba1543..e6c5cb4f 100644 --- a/sandbox/runners/base.py +++ b/sandbox/runners/base.py @@ -15,6 +15,7 @@ import asyncio import base64 import os +import signal import subprocess import time import traceback @@ -29,11 +30,95 @@ from sandbox.runners.isolation import tmp_cgroup, tmp_netns, tmp_overlayfs from sandbox.runners.types import CodeRunArgs, CodeRunResult, CommandRunResult, CommandRunStatus from sandbox.utils.common import set_permissions_recursively -from sandbox.utils.execution import cleanup_process, ensure_bash_integrity, get_output_non_blocking, kill_process_tree +from sandbox.utils.execution import cleanup_process, ensure_bash_integrity, kill_process_tree, try_decode logger = structlog.stdlib.get_logger() config = RunConfig.get_instance_sync() +_MAX_CAPTURED_OUTPUT_BYTES = 1024 * 1024 +_OUTPUT_READ_CHUNK_BYTES = 64 * 1024 +_PROCESS_TERMINATION_TIMEOUT = 1.0 + + +async def _drain_stream(stream: Optional[asyncio.StreamReader], captured: bytearray) -> None: + if stream is None: + return + + while chunk := await stream.read(_OUTPUT_READ_CHUNK_BYTES): + remaining = _MAX_CAPTURED_OUTPUT_BYTES - len(captured) + if remaining > 0: + captured.extend(chunk[:remaining]) + + +async def _write_stdin(stream: Optional[asyncio.StreamWriter], stdin: Optional[bytes]) -> None: + if stream is None: + return + + try: + if stdin is not None: + stream.write(stdin) + await stream.drain() + except (BrokenPipeError, ConnectionResetError): + pass + finally: + stream.close() + + +def _terminate_process(process: asyncio.subprocess.Process) -> None: + if os.name == 'posix': + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + except OSError as e: + logger.warning(f'failed to kill process group {process.pid}: {e}') + + if process.returncode is None and psutil.pid_exists(process.pid): + try: + kill_process_tree(process.pid) + except Exception as e: + logger.warning(f'failed to kill process tree {process.pid}: {e}') + + if process.returncode is None: + try: + process.kill() + except OSError as e: + logger.warning(f'failed to kill process {process.pid}: {e}') + + +async def _terminate_and_reap(process: asyncio.subprocess.Process, tasks: List[asyncio.Task], + execution: asyncio.Future) -> None: + current_task = asyncio.current_task() + cancellation_requested = bool(current_task and current_task.cancelling()) + try: + _terminate_process(process) + try: + await asyncio.wait_for(asyncio.shield(execution), timeout=_PROCESS_TERMINATION_TIMEOUT) + except asyncio.CancelledError: + cancellation_requested = True + except Exception: + pass + finally: + for task in tasks: + if not task.done(): + task.cancel() + cleanup = asyncio.gather(*tasks, return_exceptions=True) + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + cancellation_requested = True + cleanup.result() + if not execution.done(): + execution.cancel() + try: + execution.exception() + except asyncio.CancelledError: + pass + + if cancellation_requested: + raise asyncio.CancelledError + async def run_command_bare(command: str | List[str], timeout: float = 10, @@ -44,6 +129,7 @@ async def run_command_bare(command: str | List[str], preexec_fn=None) -> CommandRunResult: try: logger.debug(f'running command {command}') + stdin_bytes = stdin.encode() if stdin is not None else None if use_exec: p = await asyncio.create_subprocess_exec(*command, stdin=subprocess.PIPE, @@ -53,6 +139,7 @@ async def run_command_bare(command: str | List[str], **os.environ, **(extra_env or {}) }, + start_new_session=True, preexec_fn=preexec_fn) else: p = await asyncio.create_subprocess_shell(command, @@ -65,45 +152,51 @@ async def run_command_bare(command: str | List[str], **os.environ, **(extra_env or {}) }, + start_new_session=True, preexec_fn=preexec_fn) - if stdin is not None: - try: - if p.stdin: - p.stdin.write(stdin.encode()) - p.stdin.flush() - else: - logger.warning("Attempted to write to stdin, but stdin is closed.") - except Exception as e: - logger.exception(f"Failed to write to stdin: {e}") - if p.stdin: - try: - p.stdin.close() - except Exception as e: - logger.warning(f"Failed to close stdin: {e}") + stdout = bytearray() + stderr = bytearray() + tasks = [ + asyncio.create_task(_write_stdin(p.stdin, stdin_bytes)), + asyncio.create_task(_drain_stream(p.stdout, stdout)), + asyncio.create_task(_drain_stream(p.stderr, stderr)), + asyncio.create_task(p.wait()), + ] + execution = asyncio.gather(*tasks) start_time = time.time() + completed = False + timed_out = False try: - await asyncio.wait_for(p.wait(), timeout=timeout) + await asyncio.wait_for(asyncio.shield(execution), timeout=timeout) + completed = True execution_time = time.time() - start_time logger.debug(f'stop running command {command}') except asyncio.TimeoutError: - return CommandRunResult(status=CommandRunStatus.TimeLimitExceeded, - execution_time=time.time() - start_time, - stdout=await get_output_non_blocking(p.stdout), - stderr=await get_output_non_blocking(p.stderr)) + timed_out = True + execution_time = time.time() - start_time finally: - if psutil.pid_exists(p.pid): - kill_process_tree(p.pid) - logger.info(f'process killed: {p.pid}') - if config.sandbox.cleanup_process: - cleanup_process() - if config.sandbox.restore_bash: - ensure_bash_integrity() + try: + if not completed: + await _terminate_and_reap(p, tasks, execution) + else: + _terminate_process(p) + finally: + if config.sandbox.cleanup_process: + cleanup_process() + if config.sandbox.restore_bash: + ensure_bash_integrity() + + if timed_out: + return CommandRunResult(status=CommandRunStatus.TimeLimitExceeded, + execution_time=execution_time, + stdout=try_decode(bytes(stdout)), + stderr=try_decode(bytes(stderr))) return CommandRunResult(status=CommandRunStatus.Finished, execution_time=execution_time, return_code=p.returncode, - stdout=await get_output_non_blocking(p.stdout), - stderr=await get_output_non_blocking(p.stderr)) + stdout=try_decode(bytes(stdout)), + stderr=try_decode(bytes(stderr))) except Exception as e: message = f'exception on running command {command}: {e} | {traceback.print_tb(e.__traceback__)}' logger.warning(message) diff --git a/sandbox/tests/runners/test_base.py b/sandbox/tests/runners/test_base.py new file mode 100644 index 00000000..eee585e3 --- /dev/null +++ b/sandbox/tests/runners/test_base.py @@ -0,0 +1,190 @@ +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import os +import shlex +import signal +import sys + +import psutil +import pytest + +import sandbox.runners.base as runner_base +from sandbox.runners.base import run_command_bare +from sandbox.runners.types import CommandRunStatus + +_OUTPUT_CAPTURE_LIMIT = 1024 * 1024 + + +async def _run_with_watchdog(*args, **kwargs): + return await asyncio.wait_for(run_command_bare(*args, **kwargs), timeout=10) + + +async def _assert_process_stopped(pid): + for _ in range(100): + try: + if psutil.Process(pid).status() == psutil.STATUS_ZOMBIE: + return + except psutil.NoSuchProcess: + return + await asyncio.sleep(0.01) + + try: + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + pass + pytest.fail(f'descendant process {pid} is still running') + + +@pytest.mark.parametrize('stream_name', ['stdout', 'stderr']) +@pytest.mark.parametrize(('output_size', 'expected_size'), [(256 * 1024, 256 * 1024), + (2 * 1024 * 1024, _OUTPUT_CAPTURE_LIMIT)]) +async def test_run_command_bare_drains_large_output_with_bounded_capture(stream_name, output_size, expected_size): + command = [sys.executable, '-c', f"import sys; sys.{stream_name}.write('x' * {output_size})"] + + result = await _run_with_watchdog(command, timeout=5, use_exec=True) + + assert result.status == CommandRunStatus.Finished + assert result.return_code == 0 + assert getattr(result, stream_name) == 'x' * expected_size + + +@pytest.mark.skipif(os.name != 'posix', reason='shell quoting is POSIX-specific') +async def test_run_command_bare_drains_large_shell_output(): + script = f"import sys; sys.stdout.write('x' * {2 * _OUTPUT_CAPTURE_LIMIT})" + command = f'{shlex.quote(sys.executable)} -c {shlex.quote(script)}' + + result = await _run_with_watchdog(command, timeout=5) + + assert result.status == CommandRunStatus.Finished + assert result.return_code == 0 + assert result.stdout == 'x' * _OUTPUT_CAPTURE_LIMIT + + +async def test_run_command_bare_preserves_output_on_timeout(): + command = [ + sys.executable, + '-c', + ("import sys, time; " + "sys.stdout.write('stdout before timeout\\n'); sys.stdout.flush(); " + "sys.stderr.write('stderr before timeout\\n'); sys.stderr.flush(); " + "time.sleep(30)"), + ] + + result = await _run_with_watchdog(command, timeout=1, use_exec=True) + + assert result.status == CommandRunStatus.TimeLimitExceeded + assert result.stdout.splitlines() == ['stdout before timeout'] + assert result.stderr.splitlines() == ['stderr before timeout'] + + +@pytest.mark.skipif(os.name != 'posix', reason='process-group semantics are POSIX-specific') +async def test_run_command_bare_kills_descendant_holding_pipes(): + command = [ + sys.executable, + '-c', + ("import subprocess, sys; " + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(30)']); " + "print(child.pid, flush=True)"), + ] + + result = await asyncio.wait_for(run_command_bare(command, timeout=1, use_exec=True), timeout=3) + + assert result.status == CommandRunStatus.TimeLimitExceeded + assert result.execution_time < 3 + await _assert_process_stopped(int(result.stdout.strip())) + + +@pytest.mark.skipif(os.name != 'posix', reason='process-group semantics are POSIX-specific') +async def test_run_command_bare_kills_descendant_after_parent_finishes(): + command = [ + sys.executable, + '-c', + ("import subprocess, sys; " + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(30)'], " + "stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL); " + "print(child.pid, flush=True)"), + ] + + result = await _run_with_watchdog(command, timeout=3, use_exec=True) + + assert result.status == CommandRunStatus.Finished + await _assert_process_stopped(int(result.stdout.strip())) + + +async def test_run_command_bare_passes_stdin(): + command = [sys.executable, '-c', 'import sys; sys.stdout.write(sys.stdin.read())'] + + result = await _run_with_watchdog(command, stdin='hello from stdin', use_exec=True) + + assert result.status == CommandRunStatus.Finished + assert result.stdout == 'hello from stdin' + + +async def test_run_command_bare_does_not_start_process_for_invalid_stdin(monkeypatch): + process_started = False + + async def create_subprocess(*args, **kwargs): + nonlocal process_started + process_started = True + + monkeypatch.setattr(asyncio, 'create_subprocess_exec', create_subprocess) + + result = await run_command_bare([sys.executable], stdin='\ud800', use_exec=True) + + assert result.status == CommandRunStatus.Error + assert 'surrogates not allowed' in result.stderr + assert not process_started + + +async def test_run_command_bare_cleans_up_when_cancelled(monkeypatch): + processes = [] + original_create_subprocess_exec = asyncio.create_subprocess_exec + + async def create_subprocess(*args, **kwargs): + process = await original_create_subprocess_exec(*args, **kwargs) + processes.append(process) + return process + + monkeypatch.setattr(asyncio, 'create_subprocess_exec', create_subprocess) + command = [sys.executable, '-c', 'import time; time.sleep(30)'] + run_task = asyncio.create_task(run_command_bare(command, timeout=30, use_exec=True)) + + for _ in range(100): + if processes: + break + await asyncio.sleep(0.01) + assert processes + + run_task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=3) + + await asyncio.wait_for(processes[0].wait(), timeout=1) + assert processes[0].returncode is not None + + +async def test_terminate_and_reap_preserves_cancellation(monkeypatch): + child_task = asyncio.create_task(asyncio.sleep(30)) + execution = asyncio.gather(child_task) + monkeypatch.setattr(runner_base, '_terminate_process', lambda process: None) + cleanup_task = asyncio.create_task(runner_base._terminate_and_reap(object(), [child_task], execution)) + await asyncio.sleep(0) + + cleanup_task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(cleanup_task, timeout=1) + + assert child_task.cancelled()