From 88e8ef0f3617557d6bc799a3b1146b1492825818 Mon Sep 17 00:00:00 2001 From: Tanmoy Sarkar Date: Tue, 9 Jun 2026 23:22:41 +0530 Subject: [PATCH] feat(companion): Add gunicorn-companion control CLI Add a command-line client for the companion control socket so operators do not have to hand-craft JSON. gunicorn.companion.ctl speaks the manager's newline-delimited JSON protocol: status, start, stop, restart, and reread. The socket path comes from --socket or . Exit status is 0 when the manager reports ok, 1 when it reports a failure, and 2 for a usage error or an unreachable socket. Registered as the gunicorn-companion console script. Co-Authored-By: Claude Opus 4.8 --- gunicorn/companion/ctl.py | 83 +++++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + tests/test_companion_ctl.py | 61 +++++++++++++++++++++++++++ 3 files changed, 145 insertions(+) create mode 100644 gunicorn/companion/ctl.py create mode 100644 tests/test_companion_ctl.py diff --git a/gunicorn/companion/ctl.py b/gunicorn/companion/ctl.py new file mode 100644 index 00000000..6441077b --- /dev/null +++ b/gunicorn/companion/ctl.py @@ -0,0 +1,83 @@ +# +# This file is part of gunicorn released under the MIT license. +# See the NOTICE for more information. + +"""Command-line client for the companion control socket. + +Speaks the newline-delimited JSON protocol the manager's ControlServer serves: +sends one command and prints the manager's reply. Installed as the +``gunicorn-companion`` console script. +""" + +import argparse +import json +import os +import socket +import sys + +# Commands that act on one named companion and so require a name argument. +PER_NAME_COMMANDS = ("start", "stop", "restart") +COMMANDS = ("status", "reread") + PER_NAME_COMMANDS + + +def send_command(socket_path, command): + """Send one command dict to the control socket and return the reply dict.""" + client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + client.connect(socket_path) + client.sendall((json.dumps(command) + "\n").encode("utf-8")) + chunks = [] + while True: + chunk = client.recv(65536) + if not chunk: + break + chunks.append(chunk) + if b"\n" in chunk: + break + finally: + client.close() + return json.loads(b"".join(chunks).decode("utf-8")) + + +def build_parser(): + parser = argparse.ArgumentParser( + prog="gunicorn-companion", + description="Control gunicorn companion processes.") + parser.add_argument( + "-s", "--socket", + default=os.environ.get("GUNICORN_COMPANION_SOCKET"), + help="path to the companion control socket " + "(defaults to $GUNICORN_COMPANION_SOCKET)") + parser.add_argument("command", choices=COMMANDS) + parser.add_argument( + "name", nargs="?", + help="companion name (required for start, stop, restart)") + return parser + + +def run(argv=None): + parser = build_parser() + args = parser.parse_args(argv) + if not args.socket: + parser.error("no control socket; pass --socket or set " + "GUNICORN_COMPANION_SOCKET") + if args.command in PER_NAME_COMMANDS and not args.name: + parser.error("%s requires a companion name" % args.command) + + command = {"cmd": args.command} + if args.name: + command["name"] = args.name + + try: + response = send_command(args.socket, command) + except OSError as error: + print("cannot reach companion socket %s: %s" % (args.socket, error), + file=sys.stderr) + return 2 + + print(json.dumps(response, indent=2)) + return 0 if response.get("ok") else 1 + + +if __name__ == "__main__": + sys.exit(run()) diff --git a/pyproject.toml b/pyproject.toml index eaca1eac..b84366a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,7 @@ testing = [ [project.scripts] # duplicates "python -m gunicorn" handling in __main__.py gunicorn = "gunicorn.app.wsgiapp:run" +gunicorn-companion = "gunicorn.companion.ctl:run" # note the quotes around "paste.server_runner" to escape the dot [project.entry-points."paste.server_runner"] diff --git a/tests/test_companion_ctl.py b/tests/test_companion_ctl.py new file mode 100644 index 00000000..d9969762 --- /dev/null +++ b/tests/test_companion_ctl.py @@ -0,0 +1,61 @@ +# +# This file is part of gunicorn released under the MIT license. +# See the NOTICE for more information. + +from unittest import mock + +import pytest + +from gunicorn.companion import ctl + + +def test_run_status_prints_and_returns_zero(capsys): + with mock.patch.object( + ctl, "send_command", return_value={"ok": True, "companions": []} + ) as send: + code = ctl.run(["--socket", "/tmp/x.sock", "status"]) + assert code == 0 + send.assert_called_once_with("/tmp/x.sock", {"cmd": "status"}) + assert "ok" in capsys.readouterr().out + + +def test_run_per_name_command_sends_name(): + with mock.patch.object( + ctl, "send_command", return_value={"ok": True, "message": "x"} + ) as send: + code = ctl.run(["--socket", "/tmp/x.sock", "stop", "ticker"]) + assert code == 0 + send.assert_called_once_with("/tmp/x.sock", {"cmd": "stop", "name": "ticker"}) + + +def test_run_failure_response_returns_one(): + with mock.patch.object( + ctl, "send_command", return_value={"ok": False, "error": "bad"} + ): + assert ctl.run(["--socket", "/tmp/x.sock", "status"]) == 1 + + +def test_run_per_name_command_requires_name(): + with pytest.raises(SystemExit): + ctl.run(["--socket", "/tmp/x.sock", "stop"]) + + +def test_run_requires_socket(monkeypatch): + monkeypatch.delenv("GUNICORN_COMPANION_SOCKET", raising=False) + with pytest.raises(SystemExit): + ctl.run(["status"]) + + +def test_run_unreachable_socket_returns_two(): + with mock.patch.object(ctl, "send_command", side_effect=OSError("nope")): + assert ctl.run(["--socket", "/tmp/x.sock", "status"]) == 2 + + +def test_send_command_round_trip(): + client = mock.Mock() + client.recv.side_effect = [b'{"ok": true}\n'] + with mock.patch("socket.socket", return_value=client): + result = ctl.send_command("/tmp/x.sock", {"cmd": "status"}) + client.connect.assert_called_once_with("/tmp/x.sock") + assert client.sendall.call_args.args[0] == b'{"cmd": "status"}\n' + assert result == {"ok": True}