feat: muxterm process supervision (start, restart on crash, stop)
This commit is contained in:
Ken
@@ -0,0 +1,165 @@
|
||||
"""
|
||||
muxterm process supervision — start, monitor, restart on crash, stop.
|
||||
|
||||
Module state:
|
||||
_muxterm_process — the currently running muxterm subprocess (or None)
|
||||
_restart_task — the background monitor/restart task (or None)
|
||||
|
||||
Public API:
|
||||
start_muxterm(secret, port, binary_path, auto_restart)
|
||||
stop_muxterm()
|
||||
|
||||
Internal:
|
||||
_find_muxterm_binary(binary_path)
|
||||
_monitor_and_restart(secret, port, binary_path)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import shutil
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_muxterm_process: asyncio.subprocess.Process | None = None
|
||||
_restart_task: asyncio.Task | None = None # type: ignore[type-arg]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Binary discovery
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _find_muxterm_binary(binary_path: str | None = None) -> str:
|
||||
"""Locate the muxterm binary.
|
||||
|
||||
Checks explicit *binary_path* first, then falls back to
|
||||
``shutil.which('muxterm')``. Raises :exc:`FileNotFoundError` if
|
||||
the binary cannot be found.
|
||||
"""
|
||||
if binary_path is not None:
|
||||
return binary_path
|
||||
found = shutil.which("muxterm")
|
||||
if found is None:
|
||||
raise FileNotFoundError(
|
||||
"muxterm binary not found on PATH; pass binary_path explicitly"
|
||||
)
|
||||
return found
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Start
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def start_muxterm(
|
||||
secret: str,
|
||||
port: int = 7682,
|
||||
binary_path: str | None = None,
|
||||
auto_restart: bool = True,
|
||||
) -> None:
|
||||
"""Start the muxterm process.
|
||||
|
||||
Finds the binary via :func:`_find_muxterm_binary`, spawns it with
|
||||
``--addr 127.0.0.1:<port> --secret <secret>``, and optionally starts
|
||||
the background monitor/restart loop.
|
||||
"""
|
||||
global _muxterm_process, _restart_task
|
||||
|
||||
binary = _find_muxterm_binary(binary_path)
|
||||
|
||||
_muxterm_process = await asyncio.create_subprocess_exec(
|
||||
binary,
|
||||
"--addr",
|
||||
f"127.0.0.1:{port}",
|
||||
"--secret",
|
||||
secret,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
logger.info("muxterm started (pid=%s, port=%s)", _muxterm_process.pid, port)
|
||||
|
||||
if auto_restart:
|
||||
_restart_task = asyncio.create_task(
|
||||
_monitor_and_restart(secret, port, binary_path)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Monitor / restart
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _monitor_and_restart(
|
||||
secret: str,
|
||||
port: int,
|
||||
binary_path: str | None,
|
||||
) -> None:
|
||||
"""Wait for process exit; restart on non-zero exit codes.
|
||||
|
||||
Normal exit (rc=0) ends the loop. On crash, waits 1 s before
|
||||
restarting, with 5 s backoff on repeated failures.
|
||||
"""
|
||||
global _muxterm_process
|
||||
|
||||
backoff = 1.0
|
||||
while True:
|
||||
if _muxterm_process is None:
|
||||
return
|
||||
|
||||
rc = await _muxterm_process.wait()
|
||||
|
||||
if rc == 0:
|
||||
logger.info("muxterm exited cleanly (rc=0)")
|
||||
return
|
||||
|
||||
logger.warning("muxterm exited with rc=%s, restarting in %.0fs", rc, backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
|
||||
try:
|
||||
binary = _find_muxterm_binary(binary_path)
|
||||
_muxterm_process = await asyncio.create_subprocess_exec(
|
||||
binary,
|
||||
"--addr",
|
||||
f"127.0.0.1:{port}",
|
||||
"--secret",
|
||||
secret,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
logger.info(
|
||||
"muxterm restarted (pid=%s, port=%s)", _muxterm_process.pid, port
|
||||
)
|
||||
# Reset backoff on successful spawn, increase on repeated failures
|
||||
backoff = min(backoff + 1.0, 5.0)
|
||||
except Exception:
|
||||
logger.exception("failed to restart muxterm, retrying in 5s")
|
||||
backoff = 5.0
|
||||
await asyncio.sleep(backoff)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stop
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def stop_muxterm() -> None:
|
||||
"""Stop the muxterm process and cancel the restart task.
|
||||
|
||||
Terminates with a 5 s grace period, then kills if still alive.
|
||||
"""
|
||||
global _muxterm_process, _restart_task
|
||||
|
||||
if _restart_task is not None:
|
||||
_restart_task.cancel()
|
||||
_restart_task = None
|
||||
|
||||
if _muxterm_process is not None:
|
||||
_muxterm_process.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(_muxterm_process.wait(), timeout=5.0)
|
||||
except (asyncio.TimeoutError, ProcessLookupError):
|
||||
_muxterm_process.kill()
|
||||
_muxterm_process = None
|
||||
@@ -0,0 +1,150 @@
|
||||
"""
|
||||
Tests for muxplex/muxterm.py — muxterm process supervision (start, restart, stop).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import muxplex.muxterm as muxterm_mod
|
||||
from muxplex.muxterm import (
|
||||
_find_muxterm_binary,
|
||||
start_muxterm,
|
||||
stop_muxterm,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# autouse fixture — reset module-level state between tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_muxterm_state():
|
||||
"""Reset module-level muxterm state before every test."""
|
||||
muxterm_mod._muxterm_process = None
|
||||
muxterm_mod._restart_task = None
|
||||
yield
|
||||
# Teardown: clean up again
|
||||
muxterm_mod._muxterm_process = None
|
||||
muxterm_mod._restart_task = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _find_muxterm_binary tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_find_muxterm_binary_explicit_path(tmp_path):
|
||||
"""Explicit binary path is returned as-is."""
|
||||
binary = tmp_path / "muxterm"
|
||||
binary.touch()
|
||||
binary.chmod(0o755)
|
||||
result = _find_muxterm_binary(str(binary))
|
||||
assert result == str(binary)
|
||||
|
||||
|
||||
def test_find_muxterm_binary_uses_which():
|
||||
"""Falls back to shutil.which('muxterm') when no explicit path."""
|
||||
with patch("shutil.which", return_value="/usr/bin/muxterm") as mock_which:
|
||||
result = _find_muxterm_binary()
|
||||
mock_which.assert_called_once_with("muxterm")
|
||||
assert result == "/usr/bin/muxterm"
|
||||
|
||||
|
||||
def test_find_muxterm_binary_raises_when_not_found():
|
||||
"""Raises FileNotFoundError if muxterm is not found."""
|
||||
with patch("shutil.which", return_value=None):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
_find_muxterm_binary()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# start_muxterm tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_start_muxterm_spawns_process():
|
||||
"""start_muxterm calls create_subprocess_exec with correct args and stores process."""
|
||||
mock_proc = MagicMock()
|
||||
mock_proc.returncode = None
|
||||
|
||||
with (
|
||||
patch(
|
||||
"muxplex.muxterm._find_muxterm_binary",
|
||||
return_value="/usr/local/bin/muxterm",
|
||||
),
|
||||
patch(
|
||||
"asyncio.create_subprocess_exec",
|
||||
new=AsyncMock(return_value=mock_proc),
|
||||
) as mock_create,
|
||||
):
|
||||
await start_muxterm("test-secret", port=7682, auto_restart=False)
|
||||
|
||||
# Verify the binary path was passed as first arg
|
||||
mock_create.assert_called_once()
|
||||
call_args = mock_create.call_args
|
||||
assert call_args[0][0] == "/usr/local/bin/muxterm"
|
||||
assert "--addr" in call_args[0]
|
||||
assert "127.0.0.1:7682" in call_args[0]
|
||||
assert "--secret" in call_args[0]
|
||||
assert "test-secret" in call_args[0]
|
||||
|
||||
# Verify process stored in module state
|
||||
assert muxterm_mod._muxterm_process is mock_proc
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stop_muxterm tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_stop_muxterm_terminates_process():
|
||||
"""stop_muxterm terminates the process and sets state to None."""
|
||||
mock_proc = MagicMock()
|
||||
mock_proc.terminate = MagicMock()
|
||||
mock_proc.kill = MagicMock()
|
||||
mock_proc.wait = AsyncMock()
|
||||
|
||||
muxterm_mod._muxterm_process = mock_proc
|
||||
|
||||
await stop_muxterm()
|
||||
|
||||
mock_proc.terminate.assert_called_once()
|
||||
assert muxterm_mod._muxterm_process is None
|
||||
|
||||
|
||||
async def test_stop_muxterm_cancels_restart_task():
|
||||
"""stop_muxterm cancels the restart task if active."""
|
||||
mock_proc = MagicMock()
|
||||
mock_proc.terminate = MagicMock()
|
||||
mock_proc.kill = MagicMock()
|
||||
mock_proc.wait = AsyncMock()
|
||||
|
||||
mock_task = MagicMock()
|
||||
mock_task.cancel = MagicMock()
|
||||
|
||||
muxterm_mod._muxterm_process = mock_proc
|
||||
muxterm_mod._restart_task = mock_task
|
||||
|
||||
await stop_muxterm()
|
||||
|
||||
mock_task.cancel.assert_called_once()
|
||||
assert muxterm_mod._restart_task is None
|
||||
|
||||
|
||||
async def test_stop_muxterm_kills_on_timeout():
|
||||
"""stop_muxterm kills the process if terminate doesn't finish in time."""
|
||||
mock_proc = MagicMock()
|
||||
mock_proc.terminate = MagicMock()
|
||||
mock_proc.kill = MagicMock()
|
||||
mock_proc.wait = AsyncMock(side_effect=asyncio.TimeoutError)
|
||||
|
||||
muxterm_mod._muxterm_process = mock_proc
|
||||
|
||||
await stop_muxterm()
|
||||
|
||||
mock_proc.terminate.assert_called_once()
|
||||
mock_proc.kill.assert_called_once()
|
||||
assert muxterm_mod._muxterm_process is None
|
||||
Reference in New Issue
Block a user