zulip/scripts/lib/supervisor.py

80 lines
2.4 KiB
Python

import socket
import time
from http.client import HTTPConnection
from typing import Optional, Union
from xmlrpc import client
from typing_extensions import override
class UnixStreamHTTPConnection(HTTPConnection):
@override
def connect(self) -> None:
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
connected = False
for i in range(2):
try:
self.sock.connect(self.host)
connected = True
break
except FileNotFoundError:
# Backoff and retry
time.sleep(2**i)
if not connected:
raise Exception(
"Failed to connect to supervisor -- check that it is running, by running 'service supervisor status'"
)
class UnixStreamTransport(client.Transport):
def __init__(self, socket_path: str) -> None:
self.socket_path = socket_path
super().__init__()
@override
def make_connection(
self, host: Union[tuple[str, dict[str, str]], str]
) -> UnixStreamHTTPConnection:
return UnixStreamHTTPConnection(self.socket_path)
def rpc() -> client.ServerProxy:
return client.ServerProxy(
"http://localhost", transport=UnixStreamTransport("/var/run/supervisor.sock")
)
def list_supervisor_processes(
filter_names: Optional[list[str]] = None, *, only_running: Optional[bool] = None
) -> list[str]:
results = []
processes = rpc().supervisor.getAllProcessInfo()
assert isinstance(processes, list)
for process in processes:
if process["group"] != process["name"]:
name = f"{process['group']}:{process['name']}"
else:
name = process["name"]
if filter_names:
match = False
for filter_name in filter_names:
# zulip-tornado:* matches zulip-tornado:9800 and zulip-tornado
if filter_name.endswith(":*") and (
name.startswith(filter_name[:-1]) or name == filter_name[:-2]
):
match = True
break
if name == filter_name:
match = True
break
if not match:
continue
if only_running is None:
results.append(name)
elif only_running == (process["statename"] in ("RUNNING", "STARTING")):
results.append(name)
return results