diff --git a/core/libs/commonwealth/src/commonwealth/utils/DHCPDiscovery.py b/core/libs/commonwealth/src/commonwealth/utils/DHCPDiscovery.py index 60aeb5235c..b3e2b8919d 100644 --- a/core/libs/commonwealth/src/commonwealth/utils/DHCPDiscovery.py +++ b/core/libs/commonwealth/src/commonwealth/utils/DHCPDiscovery.py @@ -4,6 +4,7 @@ import sys from typing import List +from commonwealth.utils.general import run_subprocess from loguru import logger @@ -40,15 +41,11 @@ async def discover_dhcp_servers(iface: str, timeout: int = 15) -> List[str]: # Run nmap asynchronously logger.info(f"Running nmap command: {' '.join(cmd)}") - process = await asyncio.create_subprocess_exec( - *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE - ) + returncode, stdout_bytes, stderr_bytes = await run_subprocess(cmd) + output = stdout_bytes.decode() + error = stderr_bytes.decode() - stdout, stderr = await process.communicate() - output = stdout.decode() - error = stderr.decode() - - if process.returncode != 0: + if returncode != 0: logger.info(f"nmap output: {output}") logger.info(f"nmap error: {error}") raise DHCPDiscoveryError(f"nmap failed: {error}") diff --git a/core/libs/commonwealth/src/commonwealth/utils/general.py b/core/libs/commonwealth/src/commonwealth/utils/general.py index f67d53c0ce..22d5462dcf 100644 --- a/core/libs/commonwealth/src/commonwealth/utils/general.py +++ b/core/libs/commonwealth/src/commonwealth/utils/general.py @@ -6,7 +6,7 @@ from enum import Enum from functools import cache from pathlib import Path -from typing import Any, AsyncGenerator +from typing import Any, AsyncGenerator, List, Optional, Tuple import psutil from commonwealth.utils.commands import load_file @@ -222,6 +222,27 @@ def file_is_open(path: Path) -> bool: return _file_is_open_logic_lsof(result.returncode, result.stdout.strip(), result.stderr.strip()) +async def run_subprocess(command: List[str], merge_stderr: bool = False) -> Tuple[Optional[int], bytes, bytes]: + """Run a command asynchronously and wait for it to finish. + + Args: + command: Command and arguments to execute. + merge_stderr: If True, redirect stderr into stdout (mirrors `subprocess.STDOUT`). + The returned stderr is empty bytes in this case. + + Returns: + Tuple of (returncode, stdout, stderr) as raw bytes. Callers that expect text output + are responsible for decoding it themselves. + """ + process = await asyncio.create_subprocess_exec( + *command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT if merge_stderr else asyncio.subprocess.PIPE, + ) + stdout, stderr = await process.communicate() + return process.returncode, stdout, stderr or b"" + + async def file_is_open_async(path: Path) -> bool: cmd = _file_is_open_command(path) diff --git a/core/libs/commonwealth/src/commonwealth/utils/tests/test_general.py b/core/libs/commonwealth/src/commonwealth/utils/tests/test_general.py new file mode 100644 index 0000000000..7d6ceca6e8 --- /dev/null +++ b/core/libs/commonwealth/src/commonwealth/utils/tests/test_general.py @@ -0,0 +1,41 @@ +import sys + +import pytest + +from ..general import run_subprocess + + +@pytest.mark.asyncio +async def test_run_subprocess_captures_stdout_and_stderr() -> None: + command = [ + sys.executable, + "-c", + "import sys; sys.stdout.write('out'); sys.stderr.write('err')", + ] + returncode, stdout, stderr = await run_subprocess(command) + + assert returncode == 0 + assert stdout == b"out" + assert stderr == b"err" + + +@pytest.mark.asyncio +async def test_run_subprocess_returns_nonzero_returncode() -> None: + command = [sys.executable, "-c", "import sys; sys.exit(3)"] + returncode, _stdout, _stderr = await run_subprocess(command) + + assert returncode == 3 + + +@pytest.mark.asyncio +async def test_run_subprocess_merge_stderr_folds_into_stdout() -> None: + command = [ + sys.executable, + "-c", + "import sys; sys.stdout.write('out'); sys.stderr.write('err')", + ] + returncode, stdout, stderr = await run_subprocess(command, merge_stderr=True) + + assert returncode == 0 + assert stdout == b"outerr" + assert stderr == b"" diff --git a/core/services/disk_usage/main.py b/core/services/disk_usage/main.py index 024c3747f6..8a32d70e0d 100755 --- a/core/services/disk_usage/main.py +++ b/core/services/disk_usage/main.py @@ -13,6 +13,7 @@ from typing import Any, AsyncGenerator, Dict, List, Optional from commonwealth.utils.apis import GenericErrorHandlingRoute, PrettyJSONResponse +from commonwealth.utils.general import run_subprocess from commonwealth.utils.logs import InterceptHandler, init_logger from commonwealth.utils.sentry_config import init_sentry_async from commonwealth.utils.streaming import streamer @@ -214,16 +215,10 @@ async def collect_disk_usage( if depth >= 0: args.extend(["-d", str(depth)]) - process = await asyncio.create_subprocess_exec( - *args, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - text=False, - ) - stdout_bytes, stderr_bytes = await process.communicate() - if process.returncode not in (0, 1): + returncode, stdout_bytes, stderr_bytes = await run_subprocess(args) + if returncode not in (0, 1): stderr = stderr_bytes.decode("utf-8", "ignore") - logger.warning(f"du command returned {process.returncode}: {stderr}") + logger.warning(f"du command returned {returncode}: {stderr}") entries = parse_du_output(stdout_bytes) tree = build_tree(entries, path, min_size_bytes) @@ -355,25 +350,20 @@ async def run_single_speed_test(size_bytes: int) -> DiskSpeedResult: logger.info(f"Running disk speed test: {' '.join(args)}") - process = await asyncio.create_subprocess_exec( - *args, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.STDOUT, - ) - stdout_bytes, _ = await process.communicate() + returncode, stdout_bytes, _stderr_bytes = await run_subprocess(args, merge_stderr=True) output = stdout_bytes.decode("utf-8", "ignore") logger.debug(f"disktest output: {output}") - if process.returncode != 0: - logger.warning(f"disktest returned {process.returncode}: {output}") + if returncode != 0: + logger.warning(f"disktest returned {returncode}: {output}") return DiskSpeedResult( write_speed_mbps=None, read_speed_mbps=None, bytes_tested=size_bytes, seed="", success=False, - error=f"disktest failed with return code {process.returncode}: {output}", + error=f"disktest failed with return code {returncode}: {output}", ) write_speed, read_speed, seed = parse_disktest_speed(output) diff --git a/core/services/helper/main.py b/core/services/helper/main.py index aba7a0eea5..664ecc6f0e 100755 --- a/core/services/helper/main.py +++ b/core/services/helper/main.py @@ -29,6 +29,7 @@ get_cpu_type, local_hardware_identifier, local_unique_identifier, + run_subprocess, ) from commonwealth.utils.logs import InterceptHandler, init_logger from commonwealth.utils.sentry_config import init_sentry_async @@ -582,11 +583,8 @@ def software_id() -> Any: @version(1, 0) async def ping(host: str, interface_addr: Optional[str] = None) -> bool: iface = ["-I", interface_addr] if interface_addr else [] - process = await asyncio.create_subprocess_exec( - "ping", "-c", "1", *iface, host, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE - ) - await process.communicate() - return process.returncode == 0 + returncode, _stdout, _stderr = await run_subprocess(["ping", "-c", "1", *iface, host]) + return returncode == 0 async def periodic() -> None: diff --git a/core/services/recorder_extractor/main.py b/core/services/recorder_extractor/main.py index 041f514eee..22e8861c16 100755 --- a/core/services/recorder_extractor/main.py +++ b/core/services/recorder_extractor/main.py @@ -13,7 +13,7 @@ from aiocache import cached from commonwealth.utils.apis import GenericErrorHandlingRoute, PrettyJSONResponse -from commonwealth.utils.general import file_is_open_async +from commonwealth.utils.general import file_is_open_async, run_subprocess from commonwealth.utils.logs import InterceptHandler, init_logger from commonwealth.utils.sentry_config import init_sentry_async from fastapi import APIRouter, FastAPI, HTTPException, status @@ -122,21 +122,15 @@ async def check_and_recover_mcap(mcap_path: Path) -> None: logger.info(f"Running mcap doctor on {mcap_path}") # Run mcap doctor doctor_cmd = [mcap_binary, "doctor", str(mcap_path)] - doctor_proc = await asyncio.create_subprocess_exec( - *doctor_cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - text=False, - ) - stdout_bytes, stderr_bytes = await doctor_proc.communicate() + doctor_returncode, stdout_bytes, stderr_bytes = await run_subprocess(doctor_cmd) stdout = stdout_bytes.decode("utf-8", "ignore") stderr = stderr_bytes.decode("utf-8", "ignore") - if doctor_proc.returncode == 0: + if doctor_returncode == 0: logger.info(f"mcap doctor passed for {mcap_path}: {stdout.strip()}") return - logger.warning(f"mcap doctor failed for {mcap_path} (code={doctor_proc.returncode}): {stderr.strip()}") + logger.warning(f"mcap doctor failed for {mcap_path} (code={doctor_returncode}): {stderr.strip()}") logger.info(f"Attempting to recover {mcap_path}") # Create a temporary file path in the same directory as the mcap file @@ -147,19 +141,13 @@ async def check_and_recover_mcap(mcap_path: Path) -> None: tmp_path = Path(tmpfile.name) recover_cmd = [mcap_binary, "recover", str(mcap_path), "-o", str(tmp_path)] - recover_proc = await asyncio.create_subprocess_exec( - *recover_cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - text=False, - ) - _, recover_stderr_bytes = await recover_proc.communicate() + recover_returncode, _stdout_bytes, recover_stderr_bytes = await run_subprocess(recover_cmd) recover_stderr = recover_stderr_bytes.decode("utf-8", "ignore") # Check if recovery succeeded - if recover_proc.returncode != 0: + if recover_returncode != 0: logger.error( - f"mcap recover command failed for {mcap_path} (code={recover_proc.returncode}): {recover_stderr.strip()}", + f"mcap recover command failed for {mcap_path} (code={recover_returncode}): {recover_stderr.strip()}", ) return @@ -199,16 +187,10 @@ async def build_thumbnail_bytes(path: Path) -> bytes: """ # 1) Discover duration (nanoseconds) using gst-discoverer discover_cmd = ["gst-discoverer-1.0", f"file://{path}"] - discover_proc = await asyncio.create_subprocess_exec( - *discover_cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - text=False, - ) - stdout_bytes, stderr_bytes = await discover_proc.communicate() + discover_returncode, stdout_bytes, stderr_bytes = await run_subprocess(discover_cmd) stdout = stdout_bytes.decode("utf-8", "ignore") stderr = stderr_bytes.decode("utf-8", "ignore") - if discover_proc.returncode != 0: + if discover_returncode != 0: logger.error(f"gst-discoverer-1.0 failed for {path}: {stderr.strip()}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -240,16 +222,10 @@ async def build_thumbnail_bytes(path: Path) -> bytes: f"Thumbnail target: duration_ns={duration_ns} target_ns={target_ns} target_sec={target_sec:.3f} file={path}" ) logger.info(f"Thumbnail command: {' '.join(play_cmd)}") - play_proc = await asyncio.create_subprocess_exec( - *play_cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - text=False, - ) - stdout_bytes, stderr_bytes = await play_proc.communicate() + play_returncode, stdout_bytes, stderr_bytes = await run_subprocess(play_cmd) stderr = stderr_bytes.decode("utf-8", "ignore") - if play_proc.returncode != 0 or not stdout: - logger.error(f"gst-play-1.0 failed for {path} (code={play_proc.returncode}): {stderr}") + if play_returncode != 0 or not stdout_bytes: + logger.error(f"gst-play-1.0 failed for {path} (code={play_returncode}): {stderr}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to generate thumbnail.", @@ -290,20 +266,14 @@ async def extract_mcap_recordings() -> None: processing_mcap_files.add(mcap_relative) try: async with thumbnail_lock: - process = await asyncio.create_subprocess_exec( - *command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - text=False, - ) - stdout_bytes, stderr_bytes = await process.communicate() + extract_returncode, stdout_bytes, stderr_bytes = await run_subprocess(command) stdout = stdout_bytes.decode("utf-8", "ignore") stderr = stderr_bytes.decode("utf-8", "ignore") finally: processing_mcap_files.discard(mcap_relative) - if process.returncode != 0: + if extract_returncode != 0: logger.error( - f"MCAP extract failed for {mcap_path} (code={process.returncode}): {stderr}", + f"MCAP extract failed for {mcap_path} (code={extract_returncode}): {stderr}", ) else: logger.info(f"MCAP extract completed for {mcap_path}: {stdout.strip()}")