feat: add tunnel start/stop with PID file tracking
Implements build_autossh_cmd, start_tunnel, and stop_tunnel in process.py. Adds _reap() helper to handle zombie processes in test environments where the parent process is long-lived (SIGKILL leaves zombies unless waitpid is called). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,5 +1,9 @@
|
|||||||
import os
|
import os
|
||||||
|
import shutil
|
||||||
import signal
|
import signal
|
||||||
|
import socket
|
||||||
|
import subprocess
|
||||||
|
import time
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -7,6 +11,8 @@ from typing import Optional
|
|||||||
|
|
||||||
import click
|
import click
|
||||||
|
|
||||||
|
from .config import TunnelConfig
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class TunnelStatus:
|
class TunnelStatus:
|
||||||
@@ -87,3 +93,99 @@ def get_status(config_dir: Path, name: str) -> TunnelStatus:
|
|||||||
if not is_process_alive(pid):
|
if not is_process_alive(pid):
|
||||||
return TunnelStatus(state="stale", pid=pid)
|
return TunnelStatus(state="stale", pid=pid)
|
||||||
return TunnelStatus(state="running", pid=pid, uptime=format_uptime(started_at))
|
return TunnelStatus(state="running", pid=pid, uptime=format_uptime(started_at))
|
||||||
|
|
||||||
|
|
||||||
|
def build_autossh_cmd(tunnel: TunnelConfig) -> list[str]:
|
||||||
|
cmd = [
|
||||||
|
"autossh", "-M", "0", "-N",
|
||||||
|
"-o", "ServerAliveInterval=30",
|
||||||
|
"-o", "ServerAliveCountMax=3",
|
||||||
|
"-i", os.path.expanduser(tunnel.identity_file),
|
||||||
|
"-p", str(tunnel.port),
|
||||||
|
"-R", (
|
||||||
|
f"{tunnel.remote_host}:{tunnel.remote_port}"
|
||||||
|
f":{tunnel.local_host}:{tunnel.local_port}"
|
||||||
|
),
|
||||||
|
f"{tunnel.user}@{tunnel.host}",
|
||||||
|
]
|
||||||
|
if tunnel.ssh_options:
|
||||||
|
for opt in tunnel.ssh_options.split():
|
||||||
|
cmd.append(opt)
|
||||||
|
return cmd
|
||||||
|
|
||||||
|
|
||||||
|
def _is_port_in_use(port: int) -> bool:
|
||||||
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||||
|
try:
|
||||||
|
s.bind(("127.0.0.1", port))
|
||||||
|
return False
|
||||||
|
except OSError:
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def start_tunnel(config_dir: Path, tunnel: TunnelConfig) -> int:
|
||||||
|
if not shutil.which("autossh"):
|
||||||
|
raise click.ClickException("autossh not found. Install it first.")
|
||||||
|
if _is_port_in_use(tunnel.local_port):
|
||||||
|
raise click.ClickException(f"Port {tunnel.local_port} is already in use")
|
||||||
|
|
||||||
|
log_path = config_dir / "logs" / f"{tunnel.name}.log"
|
||||||
|
env = os.environ.copy()
|
||||||
|
env["AUTOSSH_GATETIME"] = "0"
|
||||||
|
|
||||||
|
with open(log_path, "a") as log:
|
||||||
|
proc = subprocess.Popen(
|
||||||
|
build_autossh_cmd(tunnel),
|
||||||
|
stdin=subprocess.DEVNULL,
|
||||||
|
stdout=log,
|
||||||
|
stderr=log,
|
||||||
|
env=env,
|
||||||
|
)
|
||||||
|
|
||||||
|
started_at = datetime.now(timezone.utc)
|
||||||
|
write_pid_file(config_dir, tunnel.name, proc.pid, started_at)
|
||||||
|
|
||||||
|
time.sleep(1)
|
||||||
|
if not is_process_alive(proc.pid):
|
||||||
|
delete_pid_file(config_dir, tunnel.name)
|
||||||
|
raise click.ClickException(
|
||||||
|
f"autossh process for '{tunnel.name}' exited immediately. "
|
||||||
|
f"Check logs: {log_path}"
|
||||||
|
)
|
||||||
|
return proc.pid
|
||||||
|
|
||||||
|
|
||||||
|
def _reap(pid: int) -> None:
|
||||||
|
"""Best-effort waitpid to reap a zombie (only works if we are the parent)."""
|
||||||
|
try:
|
||||||
|
os.waitpid(pid, 0)
|
||||||
|
except (ChildProcessError, PermissionError, OSError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def stop_tunnel(config_dir: Path, name: str) -> None:
|
||||||
|
result = read_pid_file(config_dir, name)
|
||||||
|
if result is None:
|
||||||
|
return
|
||||||
|
pid, _ = result
|
||||||
|
if not is_process_alive(pid):
|
||||||
|
delete_pid_file(config_dir, name)
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
os.kill(pid, signal.SIGTERM)
|
||||||
|
except ProcessLookupError:
|
||||||
|
delete_pid_file(config_dir, name)
|
||||||
|
return
|
||||||
|
for _ in range(50):
|
||||||
|
time.sleep(0.1)
|
||||||
|
if not is_process_alive(pid):
|
||||||
|
_reap(pid)
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
os.kill(pid, signal.SIGKILL)
|
||||||
|
except ProcessLookupError:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
_reap(pid)
|
||||||
|
delete_pid_file(config_dir, name)
|
||||||
|
|||||||
@@ -101,3 +101,88 @@ def test_read_pid_file_corrupt(config_dir):
|
|||||||
pid_path.write_text("only-one-line\n")
|
pid_path.write_text("only-one-line\n")
|
||||||
with pytest.raises(click.ClickException, match="Corrupt PID file"):
|
with pytest.raises(click.ClickException, match="Corrupt PID file"):
|
||||||
read_pid_file(config_dir, "web-service")
|
read_pid_file(config_dir, "web-service")
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Task 4: build_autossh_cmd, start_tunnel, stop_tunnel
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
import shutil
|
||||||
|
import signal
|
||||||
|
import subprocess
|
||||||
|
import time
|
||||||
|
from autossh_mgr.config import TunnelConfig
|
||||||
|
from autossh_mgr.process import start_tunnel, stop_tunnel, build_autossh_cmd
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_autossh_cmd_basic(sample_tunnel):
|
||||||
|
cmd = build_autossh_cmd(sample_tunnel)
|
||||||
|
assert cmd[0] == "autossh"
|
||||||
|
assert "-M" in cmd and "0" in cmd
|
||||||
|
assert "-N" in cmd
|
||||||
|
assert "-i" in cmd
|
||||||
|
assert os.path.expanduser(sample_tunnel.identity_file) in cmd
|
||||||
|
assert "-p" in cmd and str(sample_tunnel.port) in cmd
|
||||||
|
assert "-R" in cmd
|
||||||
|
r_idx = cmd.index("-R")
|
||||||
|
assert cmd[r_idx + 1] == (
|
||||||
|
f"{sample_tunnel.remote_host}:{sample_tunnel.remote_port}"
|
||||||
|
f":{sample_tunnel.local_host}:{sample_tunnel.local_port}"
|
||||||
|
)
|
||||||
|
assert f"{sample_tunnel.user}@{sample_tunnel.host}" in cmd
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_autossh_cmd_with_ssh_options():
|
||||||
|
t = TunnelConfig(
|
||||||
|
name="test", host="h.com", user="u",
|
||||||
|
local_port=80, remote_port=8080,
|
||||||
|
ssh_options="-o StrictHostKeyChecking=no",
|
||||||
|
)
|
||||||
|
cmd = build_autossh_cmd(t)
|
||||||
|
assert "-o" in cmd
|
||||||
|
assert "StrictHostKeyChecking=no" in cmd
|
||||||
|
|
||||||
|
|
||||||
|
def test_start_tunnel_no_autossh(config_dir, sample_tunnel, monkeypatch):
|
||||||
|
monkeypatch.setattr(shutil, "which", lambda _: None)
|
||||||
|
import click
|
||||||
|
with pytest.raises(click.ClickException, match="autossh not found"):
|
||||||
|
start_tunnel(config_dir, sample_tunnel)
|
||||||
|
|
||||||
|
|
||||||
|
def test_start_creates_pid_file(config_dir, sample_tunnel, monkeypatch, tmp_path):
|
||||||
|
# Replace autossh with a real long-lived process for testing
|
||||||
|
fake = tmp_path / "autossh"
|
||||||
|
fake.write_text("#!/bin/sh\nexec sleep 60\n")
|
||||||
|
fake.chmod(0o755)
|
||||||
|
monkeypatch.setenv("PATH", f"{tmp_path}:{os.environ['PATH']}")
|
||||||
|
|
||||||
|
pid = start_tunnel(config_dir, sample_tunnel)
|
||||||
|
try:
|
||||||
|
assert pid > 0
|
||||||
|
result = read_pid_file(config_dir, "web-service")
|
||||||
|
assert result is not None
|
||||||
|
assert result[0] == pid
|
||||||
|
assert is_process_alive(pid)
|
||||||
|
finally:
|
||||||
|
os.kill(pid, signal.SIGKILL)
|
||||||
|
|
||||||
|
|
||||||
|
def test_stop_tunnel(config_dir, sample_tunnel, monkeypatch, tmp_path):
|
||||||
|
fake = tmp_path / "autossh"
|
||||||
|
fake.write_text("#!/bin/sh\nexec sleep 60\n")
|
||||||
|
fake.chmod(0o755)
|
||||||
|
monkeypatch.setenv("PATH", f"{tmp_path}:{os.environ['PATH']}")
|
||||||
|
|
||||||
|
pid = start_tunnel(config_dir, sample_tunnel)
|
||||||
|
assert is_process_alive(pid)
|
||||||
|
stop_tunnel(config_dir, "web-service")
|
||||||
|
time.sleep(0.2)
|
||||||
|
assert not is_process_alive(pid)
|
||||||
|
assert read_pid_file(config_dir, "web-service") is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_stop_tunnel_stale_pid(config_dir):
|
||||||
|
write_pid_file(config_dir, "web-service", 9999999, datetime.now(timezone.utc))
|
||||||
|
stop_tunnel(config_dir, "web-service") # must not raise
|
||||||
|
assert read_pid_file(config_dir, "web-service") is None
|
||||||
|
|||||||
Reference in New Issue
Block a user