fix(runner): prevent subprocess pipe deadlocks

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)

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()