mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 16:21:50 +03:00
fix(gateway): prevent lifecycle races
This commit is contained in:
@@ -14,6 +14,7 @@ from rich.console import Console
|
||||
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.gateway import (
|
||||
GatewayAlreadyRunningError,
|
||||
GatewayInstance,
|
||||
GatewayRuntime,
|
||||
GatewayStartOptions,
|
||||
@@ -240,14 +241,7 @@ def create_gateway_app(
|
||||
unconfigured_provider_error = None
|
||||
if validate_startup_config is not None:
|
||||
unconfigured_provider_error = validate_startup_config(cfg)
|
||||
if unconfigured_provider_error is None:
|
||||
run_gateway(
|
||||
cfg,
|
||||
port=port,
|
||||
webui_bundle_mode=interactive_build_mode(),
|
||||
gateway_instance=instance,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
run_gateway(
|
||||
cfg,
|
||||
port=port,
|
||||
@@ -255,6 +249,10 @@ def create_gateway_app(
|
||||
unconfigured_provider_error=unconfigured_provider_error,
|
||||
gateway_instance=instance,
|
||||
)
|
||||
except GatewayAlreadyRunningError as exc:
|
||||
console.print("[yellow]Gateway is already running.[/yellow]")
|
||||
print_status(exc.status)
|
||||
raise typer.Exit(1) from None
|
||||
|
||||
@gateway_app.command("status")
|
||||
def gateway_status( # pyright: ignore[reportUnusedFunction]
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Lightweight background runtime for the nanobot gateway."""
|
||||
|
||||
from nanobot.gateway.runtime import (
|
||||
GatewayAlreadyRunningError,
|
||||
GatewayClientLease,
|
||||
GatewayInstance,
|
||||
GatewayRuntime,
|
||||
@@ -12,6 +13,7 @@ from nanobot.gateway.runtime import (
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"GatewayAlreadyRunningError",
|
||||
"GatewayClientLease",
|
||||
"GatewayInstance",
|
||||
"GatewayRuntime",
|
||||
|
||||
+117
-72
@@ -66,6 +66,14 @@ class RuntimeResult(ProcessResult):
|
||||
promoted: bool = False
|
||||
|
||||
|
||||
class GatewayAlreadyRunningError(RuntimeError):
|
||||
"""Raised when a foreground gateway tries to replace a live instance."""
|
||||
|
||||
def __init__(self, status: GatewayStatus) -> None:
|
||||
super().__init__("gateway_already_running")
|
||||
self.status = status
|
||||
|
||||
|
||||
def build_gateway_command(python_executable: str, options: GatewayStartOptions) -> list[str]:
|
||||
"""Build a foreground gateway command for process supervisors."""
|
||||
command = [
|
||||
@@ -188,12 +196,16 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
|
||||
def _build_child_command(self, options: ProcessStartOptions) -> list[str]:
|
||||
return build_gateway_command(self.python_executable, options)
|
||||
|
||||
def _transition_lock(self) -> FileLock:
|
||||
"""Serialize long lifecycle transitions without blocking child cleanup."""
|
||||
return FileLock(f"{self.paths.state_path}.transition.lock")
|
||||
|
||||
def start_background(self, options: ProcessStartOptions) -> RuntimeResult:
|
||||
"""Start the gateway detached from the current terminal."""
|
||||
lease = GatewayClientLease(self, kind="gateway-background")
|
||||
while True:
|
||||
lease.wait_for_shutdown()
|
||||
with self._lifecycle_lock():
|
||||
with self._transition_lock(), self._lifecycle_lock():
|
||||
promoted = lease._try_mark_persistent_locked()
|
||||
if promoted is None:
|
||||
continue
|
||||
@@ -202,12 +214,17 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
|
||||
|
||||
def start_on_demand(self, options: ProcessStartOptions) -> RuntimeResult:
|
||||
"""Atomically reuse a gateway or start one owned by local client leases."""
|
||||
with self._lifecycle_lock():
|
||||
status = self.status()
|
||||
if status.running:
|
||||
return RuntimeResult(False, "gateway_already_running", status)
|
||||
GatewayClientLease(self, kind="gateway-start")._mark_ephemeral_locked()
|
||||
return self._start_background(options)
|
||||
lease = GatewayClientLease(self, kind="gateway-start")
|
||||
while True:
|
||||
lease.wait_for_shutdown()
|
||||
with self._transition_lock(), self._lifecycle_lock():
|
||||
if lease._shutdown_pending_locked():
|
||||
continue
|
||||
status = self.status()
|
||||
if status.running:
|
||||
return RuntimeResult(False, "gateway_already_running", status)
|
||||
lease._mark_ephemeral_locked()
|
||||
return self._start_background(options)
|
||||
|
||||
def _start_background(self, options: ProcessStartOptions) -> RuntimeResult:
|
||||
result = super()._start_background(options)
|
||||
@@ -221,10 +238,14 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
|
||||
|
||||
def stop(self, *, timeout_s: int = 20) -> RuntimeResult:
|
||||
"""Stop the gateway recorded by this runtime."""
|
||||
with self._lifecycle_lock():
|
||||
with self._transition_lock():
|
||||
result = self._stop(timeout_s=timeout_s)
|
||||
if result.ok or result.message in {"gateway_not_running", "gateway_state_stale"}:
|
||||
GatewayClientLease(self, kind="gateway-stop")._clear_locked()
|
||||
with self._lifecycle_lock():
|
||||
if result.ok or result.message in {
|
||||
"gateway_not_running",
|
||||
"gateway_state_stale",
|
||||
}:
|
||||
GatewayClientLease(self, kind="gateway-stop")._clear_locked()
|
||||
return self._result(result)
|
||||
|
||||
def status(self, *, reason: str | None = None) -> GatewayStatus:
|
||||
@@ -260,55 +281,69 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
|
||||
self._release_current_process()
|
||||
|
||||
def _claim_current_process(self, options: ProcessStartOptions) -> GatewayLaunchMode:
|
||||
with self._lifecycle_lock():
|
||||
state = self._read_state() or {}
|
||||
pid = os.getpid()
|
||||
launch_mode = (
|
||||
"background"
|
||||
if state.get("pid") == pid and state.get("launch_mode") == "background"
|
||||
else "foreground"
|
||||
)
|
||||
state.update(
|
||||
{
|
||||
"pid": pid,
|
||||
"identity": self._process_identity(pid),
|
||||
"started_at": datetime.now(UTC).isoformat(),
|
||||
"platform": self.platform_name,
|
||||
"port": options.port,
|
||||
"workspace": options.workspace,
|
||||
"config_path": options.config_path,
|
||||
"command": self._build_child_command(options),
|
||||
"log_path": str(self.paths.log_path),
|
||||
"launch_mode": launch_mode,
|
||||
}
|
||||
)
|
||||
self._write_state(state)
|
||||
if launch_mode == "foreground":
|
||||
GatewayClientLease(
|
||||
self,
|
||||
kind="gateway-foreground",
|
||||
)._try_mark_persistent_locked()
|
||||
return launch_mode
|
||||
lease = GatewayClientLease(self, kind="gateway-foreground")
|
||||
pid = os.getpid()
|
||||
while True:
|
||||
lease.wait_for_shutdown()
|
||||
with self._transition_lock(), self._lifecycle_lock():
|
||||
current = self.status()
|
||||
if current.running and current.pid != pid:
|
||||
raise GatewayAlreadyRunningError(current)
|
||||
if lease._shutdown_pending_locked():
|
||||
continue
|
||||
state = self._read_state() or {}
|
||||
launch_mode: GatewayLaunchMode = (
|
||||
"background"
|
||||
if state.get("pid") == pid and state.get("launch_mode") == "background"
|
||||
else "foreground"
|
||||
)
|
||||
state.update(
|
||||
{
|
||||
"pid": pid,
|
||||
"identity": self._process_identity(pid),
|
||||
"started_at": datetime.now(UTC).isoformat(),
|
||||
"platform": self.platform_name,
|
||||
"port": options.port,
|
||||
"workspace": options.workspace,
|
||||
"config_path": options.config_path,
|
||||
"command": self._build_child_command(options),
|
||||
"log_path": str(self.paths.log_path),
|
||||
"launch_mode": launch_mode,
|
||||
}
|
||||
)
|
||||
self._write_state(state)
|
||||
if launch_mode == "foreground":
|
||||
lease._try_mark_persistent_locked()
|
||||
return launch_mode
|
||||
|
||||
def _release_current_process(self) -> None:
|
||||
with self._lifecycle_lock():
|
||||
state = self._read_state()
|
||||
if state and self._record_matches_process(state, os.getpid()):
|
||||
self._clear_state()
|
||||
GatewayClientLease(self, kind="gateway-exit")._finish_shutdown_locked()
|
||||
GatewayClientLease(
|
||||
self,
|
||||
kind="gateway-exit",
|
||||
)._finish_shutdown_locked()
|
||||
|
||||
def restart(self, options: ProcessStartOptions, *, timeout_s: int = 20) -> RuntimeResult:
|
||||
"""Restart an existing gateway without creating a new persistent instance."""
|
||||
with self._lifecycle_lock():
|
||||
status = self.status()
|
||||
if not status.running:
|
||||
return RuntimeResult(False, "gateway_not_running", status)
|
||||
if status.launch_mode == "foreground":
|
||||
return RuntimeResult(False, "gateway_foreground_restart_required", status)
|
||||
with self._transition_lock():
|
||||
with self._lifecycle_lock():
|
||||
status = self.status()
|
||||
if not status.running:
|
||||
return RuntimeResult(False, "gateway_not_running", status)
|
||||
if status.launch_mode == "foreground":
|
||||
return RuntimeResult(
|
||||
False,
|
||||
"gateway_foreground_restart_required",
|
||||
status,
|
||||
)
|
||||
stop_result = self._stop(timeout_s=timeout_s)
|
||||
if not stop_result.ok:
|
||||
return self._result(stop_result)
|
||||
return self._start_background(options)
|
||||
with self._lifecycle_lock():
|
||||
return self._start_background(options)
|
||||
|
||||
def _result(self, result: ProcessResult) -> RuntimeResult:
|
||||
status = result.status
|
||||
@@ -335,6 +370,7 @@ class GatewayClientLease:
|
||||
self.state_path = state_path.with_name(
|
||||
f"{state_path.stem}.clients{state_path.suffix}"
|
||||
)
|
||||
self.transition_lock = FileLock(f"{state_path}.transition.lock")
|
||||
self.lifecycle_lock = FileLock(f"{state_path}.lock")
|
||||
self.lock = FileLock(f"{self.state_path}.lock")
|
||||
self._acquired = False
|
||||
@@ -343,7 +379,7 @@ class GatewayClientLease:
|
||||
"""Register this client before it starts or attaches to the gateway."""
|
||||
while True:
|
||||
self.wait_for_shutdown()
|
||||
with self.lifecycle_lock, self.lock:
|
||||
with self.transition_lock, self.lifecycle_lock, self.lock:
|
||||
state = self._live_state()
|
||||
if state.get("stopping"):
|
||||
continue
|
||||
@@ -358,7 +394,7 @@ class GatewayClientLease:
|
||||
|
||||
def mark_ephemeral(self) -> None:
|
||||
"""Mark a gateway started by a client for last-client shutdown."""
|
||||
with self.lifecycle_lock:
|
||||
with self.transition_lock, self.lifecycle_lock:
|
||||
self._mark_ephemeral_locked()
|
||||
|
||||
def _mark_ephemeral_locked(self) -> None:
|
||||
@@ -371,7 +407,7 @@ class GatewayClientLease:
|
||||
"""Keep an explicitly backgrounded gateway alive; return whether it was promoted."""
|
||||
while True:
|
||||
self.wait_for_shutdown()
|
||||
with self.lifecycle_lock:
|
||||
with self.transition_lock, self.lifecycle_lock:
|
||||
promoted = self._try_mark_persistent_locked()
|
||||
if promoted is not None:
|
||||
return promoted
|
||||
@@ -388,7 +424,7 @@ class GatewayClientLease:
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Forget leases after an explicit gateway stop."""
|
||||
with self.lifecycle_lock:
|
||||
with self.transition_lock, self.lifecycle_lock:
|
||||
self._clear_locked()
|
||||
|
||||
def _clear_locked(self) -> None:
|
||||
@@ -407,7 +443,7 @@ class GatewayClientLease:
|
||||
|
||||
def begin_orphan_shutdown(self) -> bool:
|
||||
"""Commit shutdown only while an on-demand gateway still has no clients."""
|
||||
with self.lifecycle_lock, self.lock:
|
||||
with self.transition_lock, self.lifecycle_lock, self.lock:
|
||||
state = self._live_state()
|
||||
if not bool(state.get("auto_stop")) or self._clients(state):
|
||||
self._write_or_clear(state)
|
||||
@@ -420,26 +456,31 @@ class GatewayClientLease:
|
||||
"""Release this client and stop an ephemeral gateway when it was the last."""
|
||||
if not self._acquired:
|
||||
return False
|
||||
with self.lifecycle_lock:
|
||||
with self.lock:
|
||||
state = self._live_state()
|
||||
clients = self._clients(state)
|
||||
clients.pop(self.token, None)
|
||||
self._acquired = False
|
||||
should_stop = not clients and bool(state.get("auto_stop"))
|
||||
self._write_or_clear(state)
|
||||
if not should_stop:
|
||||
return False
|
||||
result = self.runtime._stop(timeout_s=timeout_s)
|
||||
stopped = result.ok or result.message in {
|
||||
"gateway_not_running",
|
||||
"gateway_state_stale",
|
||||
}
|
||||
if stopped:
|
||||
self._clear_locked()
|
||||
else:
|
||||
self._mark_ephemeral_locked()
|
||||
return stopped
|
||||
while True:
|
||||
self.wait_for_shutdown()
|
||||
with self.transition_lock:
|
||||
with self.lifecycle_lock, self.lock:
|
||||
state = self._live_state()
|
||||
if state.get("stopping"):
|
||||
continue
|
||||
clients = self._clients(state)
|
||||
clients.pop(self.token, None)
|
||||
self._acquired = False
|
||||
should_stop = not clients and bool(state.get("auto_stop"))
|
||||
self._write_or_clear(state)
|
||||
if not should_stop:
|
||||
return False
|
||||
result = self.runtime._stop(timeout_s=timeout_s)
|
||||
stopped = result.ok or result.message in {
|
||||
"gateway_not_running",
|
||||
"gateway_state_stale",
|
||||
}
|
||||
with self.lifecycle_lock:
|
||||
if stopped:
|
||||
self._clear_locked()
|
||||
else:
|
||||
self._mark_ephemeral_locked()
|
||||
return stopped
|
||||
|
||||
def wait_for_shutdown(self, *, timeout_s: float = 20) -> None:
|
||||
"""Wait until a committed orphan shutdown can no longer accept clients."""
|
||||
@@ -457,6 +498,10 @@ class GatewayClientLease:
|
||||
raise RuntimeError("gateway is still shutting down; try again shortly")
|
||||
time.sleep(0.05)
|
||||
|
||||
def _shutdown_pending_locked(self) -> bool:
|
||||
with self.lock:
|
||||
return bool(self._live_state().get("stopping"))
|
||||
|
||||
def _finish_shutdown_locked(self) -> None:
|
||||
with self.lock:
|
||||
state = self._live_state()
|
||||
|
||||
@@ -9,6 +9,7 @@ from typer.testing import CliRunner
|
||||
from nanobot.cli.gateway import _resolved_config_selector, create_gateway_app
|
||||
from nanobot.config.schema import Config
|
||||
from nanobot.gateway import (
|
||||
GatewayAlreadyRunningError,
|
||||
GatewayInstance,
|
||||
GatewayRuntimePaths,
|
||||
GatewayStartOptions,
|
||||
@@ -113,6 +114,7 @@ def _test_app(
|
||||
tmp_path: Path,
|
||||
config: Config | None = None,
|
||||
startup_error: str | None = None,
|
||||
run_error: Exception | None = None,
|
||||
):
|
||||
app = typer.Typer()
|
||||
fake_runtime = FakeRuntime(tmp_path)
|
||||
@@ -133,6 +135,8 @@ def _test_app(
|
||||
unconfigured_provider_error: str | None = None,
|
||||
gateway_instance: GatewayInstance | None = None,
|
||||
) -> None:
|
||||
if run_error is not None:
|
||||
raise run_error
|
||||
run_calls.append(
|
||||
(config, port, webui_bundle_mode, unconfigured_provider_error, gateway_instance)
|
||||
)
|
||||
@@ -176,6 +180,27 @@ def test_gateway_default_still_runs_foreground(tmp_path):
|
||||
)
|
||||
|
||||
|
||||
def test_gateway_foreground_reports_a_competing_live_instance(tmp_path):
|
||||
status = GatewayStatus(
|
||||
running=True,
|
||||
pid=12345,
|
||||
state_path=tmp_path / "gateway.json",
|
||||
log_path=tmp_path / "gateway.log",
|
||||
reason="running",
|
||||
launch_mode="foreground",
|
||||
)
|
||||
app, _runtime, _service, _calls, _prepare_calls = _test_app(
|
||||
tmp_path,
|
||||
run_error=GatewayAlreadyRunningError(status),
|
||||
)
|
||||
|
||||
result = runner.invoke(app, ["gateway"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "Gateway is already running" in result.output
|
||||
assert "PID: 12345" in result.output
|
||||
|
||||
|
||||
def test_config_workspace_does_not_split_the_foreground_instance(tmp_path: Path) -> None:
|
||||
config = Config()
|
||||
config.agents.defaults.workspace = str(tmp_path / "configured-workspace")
|
||||
|
||||
@@ -5,6 +5,7 @@ import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -40,6 +41,66 @@ def _paths(tmp_path: Path) -> GatewayRuntimePaths:
|
||||
return GatewayRuntimePaths.for_instance(data_dir=tmp_path)
|
||||
|
||||
|
||||
_FOREGROUND_CHILD = r"""
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from nanobot.gateway import (
|
||||
GatewayAlreadyRunningError,
|
||||
GatewayRuntime,
|
||||
GatewayRuntimePaths,
|
||||
GatewayStartOptions,
|
||||
)
|
||||
|
||||
root = Path(sys.argv[1])
|
||||
duration = float(sys.argv[2])
|
||||
runtime = GatewayRuntime(paths=GatewayRuntimePaths.for_instance(data_dir=root))
|
||||
|
||||
def stop(*_args):
|
||||
raise SystemExit(0)
|
||||
|
||||
signal.signal(signal.SIGTERM, stop)
|
||||
try:
|
||||
with runtime.foreground_instance(GatewayStartOptions(port=18790)):
|
||||
print("claimed", flush=True)
|
||||
time.sleep(duration)
|
||||
except GatewayAlreadyRunningError:
|
||||
print("occupied", flush=True)
|
||||
raise SystemExit(17)
|
||||
"""
|
||||
|
||||
|
||||
def _foreground_child(runtime_dir: Path, duration_s: float) -> subprocess.Popen[str]:
|
||||
env = os.environ.copy()
|
||||
root = str(Path(__file__).resolve().parents[2])
|
||||
env["PYTHONPATH"] = os.pathsep.join(filter(None, (root, env.get("PYTHONPATH"))))
|
||||
return subprocess.Popen(
|
||||
[sys.executable, "-c", _FOREGROUND_CHILD, str(runtime_dir), str(duration_s)],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
env=env,
|
||||
start_new_session=os.name != "nt",
|
||||
)
|
||||
|
||||
|
||||
def _wait_for_claim(runtime: GatewayRuntime, process: subprocess.Popen[str]) -> None:
|
||||
deadline = time.monotonic() + 3
|
||||
while time.monotonic() < deadline and process.poll() is None:
|
||||
try:
|
||||
state = json.loads(runtime.paths.state_path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
time.sleep(0.01)
|
||||
continue
|
||||
if state.get("pid") == process.pid:
|
||||
return
|
||||
time.sleep(0.01)
|
||||
stdout, stderr = process.communicate(timeout=3)
|
||||
pytest.fail(f"gateway claim failed: stdout={stdout!r}, stderr={stderr!r}")
|
||||
|
||||
|
||||
def test_paths_use_stable_instance_suffix_for_custom_selectors(tmp_path):
|
||||
default_paths = GatewayRuntimePaths.for_instance(data_dir=tmp_path)
|
||||
first_paths = GatewayRuntimePaths.for_instance(
|
||||
@@ -200,6 +261,51 @@ def test_foreground_gateway_clears_its_state_after_an_error(tmp_path, monkeypatc
|
||||
assert not runtime.paths.state_path.exists()
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="POSIX signal escalation regression")
|
||||
def test_stop_allows_a_foreground_gateway_to_release_before_timeout(tmp_path):
|
||||
runtime = GatewayRuntime(paths=_paths(tmp_path))
|
||||
child = _foreground_child(tmp_path, 30)
|
||||
try:
|
||||
_wait_for_claim(runtime, child)
|
||||
|
||||
result = runtime.stop(timeout_s=1)
|
||||
child.wait(timeout=3)
|
||||
|
||||
assert result.ok is True
|
||||
assert child.returncode == 0
|
||||
assert not runtime.paths.state_path.exists()
|
||||
finally:
|
||||
if child.poll() is None:
|
||||
child.kill()
|
||||
child.wait(timeout=3)
|
||||
|
||||
|
||||
def test_competing_foreground_claim_preserves_the_live_gateway(tmp_path):
|
||||
runtime = GatewayRuntime(paths=_paths(tmp_path))
|
||||
first = _foreground_child(tmp_path, 30)
|
||||
try:
|
||||
_wait_for_claim(runtime, first)
|
||||
second = _foreground_child(tmp_path, 0)
|
||||
try:
|
||||
second.wait(timeout=3)
|
||||
assert second.stdout is not None
|
||||
assert second.stdout.read().strip() == "occupied"
|
||||
assert second.returncode == 17
|
||||
finally:
|
||||
if second.poll() is None:
|
||||
second.kill()
|
||||
second.wait(timeout=3)
|
||||
|
||||
state = json.loads(runtime.paths.state_path.read_text(encoding="utf-8"))
|
||||
assert state["pid"] == first.pid
|
||||
assert runtime.status().pid == first.pid
|
||||
finally:
|
||||
if first.poll() is None:
|
||||
first.terminate()
|
||||
first.wait(timeout=3)
|
||||
runtime.status()
|
||||
|
||||
|
||||
def test_stop_reaps_an_owned_child_without_consuming_the_shutdown_timeout(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
|
||||
Reference in New Issue
Block a user