121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# 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.
|
|
|
|
"""Cross-platform Process Supervisor and Lifecycle Manager."""
|
|
|
|
import asyncio
|
|
from collections.abc import Sequence
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from artemis.platform import platform
|
|
from artemis.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class ProcessSupervisor:
|
|
"""Manages process spawning, tree cleanup, signals, and UTF-8 stream handling."""
|
|
|
|
@staticmethod
|
|
def prepare_environment(extra_env: dict[str, str] | None = None) -> dict[str, str]:
|
|
"""Prepare child process environment with guaranteed UTF-8 mode."""
|
|
env = os.environ.copy()
|
|
env["PYTHONUNBUFFERED"] = "1"
|
|
env["PYTHONUTF8"] = "1"
|
|
if extra_env:
|
|
env.update(extra_env)
|
|
return env
|
|
|
|
@staticmethod
|
|
async def spawn_async(
|
|
cmd: Sequence[str],
|
|
cwd: Path | str | None = None,
|
|
env: dict[str, str] | None = None,
|
|
capture_output: bool = True,
|
|
) -> asyncio.subprocess.Process:
|
|
"""Spawn an asynchronous subprocess with clean environment and cross-platform flags."""
|
|
prepared_env = ProcessSupervisor.prepare_environment(env)
|
|
cwd_str = str(cwd) if cwd else None
|
|
|
|
stdout_dest = asyncio.subprocess.PIPE if capture_output else None
|
|
stderr_dest = asyncio.subprocess.PIPE if capture_output else None
|
|
|
|
return await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=stdout_dest,
|
|
stderr=stderr_dest,
|
|
cwd=cwd_str,
|
|
env=prepared_env,
|
|
)
|
|
|
|
@staticmethod
|
|
def terminate_tree(pid: int, timeout_seconds: float = 3.0) -> bool:
|
|
"""Recursively terminate a process tree."""
|
|
return platform.process.terminate_process_tree(pid, timeout_seconds=timeout_seconds)
|
|
|
|
@staticmethod
|
|
def terminate_tree_verified(
|
|
pid: int,
|
|
process_created_at: float,
|
|
timeout_seconds: float = 3.0,
|
|
) -> bool:
|
|
"""Terminate a process tree only if PID and creation time still match.
|
|
|
|
Cross-entry task control discovers the worker through a shared lease.
|
|
Verifying creation time prevents a stale lease from terminating an
|
|
unrelated process after the operating system reuses its PID.
|
|
"""
|
|
try:
|
|
import psutil
|
|
except Exception:
|
|
return False
|
|
try:
|
|
process = psutil.Process(pid)
|
|
if process_created_at > 0 and abs(process.create_time() - process_created_at) >= 1.0:
|
|
return False
|
|
except psutil.NoSuchProcess:
|
|
return True
|
|
except Exception:
|
|
return False
|
|
return ProcessSupervisor.terminate_tree(pid, timeout_seconds=timeout_seconds)
|
|
|
|
@staticmethod
|
|
async def stop_process(
|
|
proc: asyncio.subprocess.Process | None,
|
|
timeout_seconds: float = 3.0,
|
|
) -> None:
|
|
"""Safely stop an asyncio Process and kill its children if necessary."""
|
|
if proc is None:
|
|
return
|
|
|
|
if proc.returncode is not None:
|
|
return
|
|
|
|
try:
|
|
if proc.pid:
|
|
ProcessSupervisor.terminate_tree(proc.pid, timeout_seconds=timeout_seconds)
|
|
else:
|
|
proc.terminate()
|
|
await asyncio.wait_for(proc.wait(), timeout=timeout_seconds)
|
|
except Exception as e:
|
|
logger.debug(f"Failed graceful stop for process {proc}: {e}")
|
|
try:
|
|
proc.kill()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
process_supervisor: ProcessSupervisor = ProcessSupervisor()
|