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:
2026-05-21 12:19:06 +08:00
parent 8a2e4b6cc0
commit 5218431349
2 changed files with 187 additions and 0 deletions

View File

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

View File

@@ -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