diff --git a/docs/design/companion-process-manager.md b/docs/design/companion-process-manager.md index 94a0213d..56e3e1fe 100644 --- a/docs/design/companion-process-manager.md +++ b/docs/design/companion-process-manager.md @@ -683,10 +683,10 @@ No per-companion logic in Arbiter. - [x] Preserve and clear `manual_stop` correctly. - [x] Add Unix control socket. - [x] Implement JSON command protocol. -- [ ] Implement `status`. -- [ ] Implement `start`. -- [ ] Implement `stop`. -- [ ] Implement `restart`. +- [x] Implement `status`. +- [x] Implement `start`. +- [x] Implement `stop`. +- [x] Implement `restart`. - [ ] Implement transactional `reread`. - [ ] Add manager spawn/reap logic in Arbiter. - [ ] Add manager shutdown handling in Arbiter. diff --git a/gunicorn/companion/manager.py b/gunicorn/companion/manager.py index 59281138..53e33ce5 100644 --- a/gunicorn/companion/manager.py +++ b/gunicorn/companion/manager.py @@ -11,6 +11,7 @@ import signal import time from typing import TYPE_CHECKING, Callable, Iterable, Union +from gunicorn.companion.control import CommandError from gunicorn.companion.process import CompanionProcess, State if TYPE_CHECKING: @@ -31,6 +32,38 @@ class CompanionManager: self.pid = os.getpid() self.processes = {c.name: CompanionProcess(c) for c in configs} + def handle_command(self, obj: dict) -> dict: + """Route a decoded control command to its action. + + This is the ``dispatch`` the control socket calls. ``status`` returns a + snapshot of every companion; ``start``/``stop``/``restart`` act on the + one named companion and report ``(ok, message)``. Per-companion + commands need a string ``name``, and anything else raises ``CommandError`` so the + socket replies with an error envelope. + """ + cmd = obj["cmd"] + if cmd == "status": + return {"ok": True, "companions": self.status()} + + # Every remaining command acts on one named companion. + name = obj.get("name") + if not isinstance(name, str): + raise CommandError("'%s' requires a 'name'" % cmd) + if cmd == "start": + ok, message = self.start_process(name) + elif cmd == "stop": + ok, message = self.stop_process(name) + elif cmd == "restart": + ok, message = self.restart_process(name) + else: + raise CommandError("unknown command %r" % cmd) + return {"ok": ok, "message": message} + + def status(self, now: float = None) -> list: + """Status entry for every companion, for the ``status`` command.""" + now = now or time.time() + return [proc.status_dict(now) for proc in self.processes.values()] + def spawn_process(self, proc: CompanionProcess) -> int: """Fork one companion. diff --git a/tests/test_companion_manager.py b/tests/test_companion_manager.py index cca1bfef..94e7c130 100644 --- a/tests/test_companion_manager.py +++ b/tests/test_companion_manager.py @@ -8,6 +8,7 @@ from unittest import mock import pytest +from gunicorn.companion.control import CommandError from gunicorn.companion.manager import CompanionManager from gunicorn.companion.process import CompanionConfig, State @@ -126,6 +127,51 @@ def test_reap_no_children(): assert mgr.reap_processes() == [] +def test_status_lists_all_companions(): + mgr = make_manager("rq", "scheduler") + entries = mgr.status(now=100.0) + assert {e["name"] for e in entries} == {"rq", "scheduler"} + assert all("state" in e and "description" in e for e in entries) + + +def test_handle_command_status(): + mgr = make_manager("rq") + resp = mgr.handle_command({"cmd": "status"}) + assert resp["ok"] is True + assert resp["companions"][0]["name"] == "rq" + + +def test_handle_command_start_routes(): + mgr = make_manager("rq") + with mock.patch.object(mgr, "start_process", + return_value=(True, "rq started")) as sp: + resp = mgr.handle_command({"cmd": "start", "name": "rq"}) + sp.assert_called_once_with("rq") + assert resp == {"ok": True, "message": "rq started"} + + +def test_handle_command_stop_and_restart_route(): + mgr = make_manager("rq") + with mock.patch.object(mgr, "stop_process", return_value=(True, "s")) as st, \ + mock.patch.object(mgr, "restart_process", return_value=(True, "r")) as rt: + mgr.handle_command({"cmd": "stop", "name": "rq"}) + mgr.handle_command({"cmd": "restart", "name": "rq"}) + st.assert_called_once_with("rq") + rt.assert_called_once_with("rq") + + +def test_handle_command_missing_name(): + mgr = make_manager("rq") + with pytest.raises(CommandError): + mgr.handle_command({"cmd": "start"}) + + +def test_handle_command_unknown(): + mgr = make_manager("rq") + with pytest.raises(CommandError): + mgr.handle_command({"cmd": "reread"}) + + def test_start_process_stopped_spawns(): mgr = make_manager("rq") proc = mgr.processes["rq"]