chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,285 @@
|
||||
# This module provides broker clients for querying queue lengths from message brokers.
|
||||
# Adapted from Flower's broker.py (https://github.com/mher/flower/blob/master/flower/utils/broker.py)
|
||||
# with the following modification:
|
||||
# - Added close() method to BrokerBase and RedisBase for resource cleanup
|
||||
|
||||
import json
|
||||
import logging
|
||||
import numbers
|
||||
import socket
|
||||
from urllib.parse import quote, unquote, urljoin, urlparse
|
||||
|
||||
from tornado import httpclient, ioloop
|
||||
|
||||
from ray.serve._private.constants import SERVE_LOGGER_NAME
|
||||
|
||||
try:
|
||||
import redis
|
||||
except ImportError:
|
||||
redis = None
|
||||
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
|
||||
class BrokerBase:
|
||||
def __init__(self, broker_url, *_, **__):
|
||||
purl = urlparse(broker_url)
|
||||
self.host = purl.hostname
|
||||
self.port = purl.port
|
||||
self.vhost = purl.path[1:]
|
||||
|
||||
username = purl.username
|
||||
password = purl.password
|
||||
|
||||
self.username = unquote(username) if username else username
|
||||
self.password = unquote(password) if password else password
|
||||
|
||||
async def queues(self, names):
|
||||
raise NotImplementedError
|
||||
|
||||
def close(self):
|
||||
"""Close any open connections. Override in subclasses as needed."""
|
||||
pass
|
||||
|
||||
|
||||
class RabbitMQ(BrokerBase):
|
||||
def __init__(self, broker_url, http_api, io_loop=None, **__):
|
||||
super().__init__(broker_url)
|
||||
self.io_loop = io_loop or ioloop.IOLoop.instance()
|
||||
|
||||
self.host = self.host or "localhost"
|
||||
self.port = self.port or 15672
|
||||
self.vhost = quote(self.vhost, "") or "/" if self.vhost != "/" else self.vhost
|
||||
self.username = self.username or "guest"
|
||||
self.password = self.password or "guest"
|
||||
|
||||
if not http_api:
|
||||
http_api = f"http://{self.username}:{self.password}@{self.host}:{self.port}/api/{self.vhost}"
|
||||
|
||||
try:
|
||||
self.validate_http_api(http_api)
|
||||
except ValueError:
|
||||
logger.error("Invalid broker api url: %s", http_api)
|
||||
|
||||
self.http_api = http_api
|
||||
|
||||
async def queues(self, names):
|
||||
url = urljoin(self.http_api, "queues/" + self.vhost)
|
||||
api_url = urlparse(self.http_api)
|
||||
username = unquote(api_url.username or "") or self.username
|
||||
password = unquote(api_url.password or "") or self.password
|
||||
|
||||
http_client = httpclient.AsyncHTTPClient()
|
||||
try:
|
||||
response = await http_client.fetch(
|
||||
url,
|
||||
auth_username=username,
|
||||
auth_password=password,
|
||||
connect_timeout=1.0,
|
||||
request_timeout=2.0,
|
||||
validate_cert=False,
|
||||
)
|
||||
except (socket.error, httpclient.HTTPError) as e:
|
||||
logger.error("RabbitMQ management API call failed: %s", e)
|
||||
return []
|
||||
finally:
|
||||
http_client.close()
|
||||
|
||||
if response.code == 200:
|
||||
info = json.loads(response.body.decode())
|
||||
return [x for x in info if x["name"] in names]
|
||||
response.rethrow()
|
||||
|
||||
@classmethod
|
||||
def validate_http_api(cls, http_api):
|
||||
url = urlparse(http_api)
|
||||
if url.scheme not in ("http", "https"):
|
||||
raise ValueError(f"Invalid http api schema: {url.scheme}")
|
||||
|
||||
|
||||
class RedisBase(BrokerBase):
|
||||
DEFAULT_SEP = "\x06\x16"
|
||||
DEFAULT_PRIORITY_STEPS = [0, 3, 6, 9]
|
||||
|
||||
def __init__(self, broker_url, *_, **kwargs):
|
||||
super().__init__(broker_url)
|
||||
self.redis = None
|
||||
|
||||
if not redis:
|
||||
raise ImportError("redis library is required")
|
||||
|
||||
broker_options = kwargs.get("broker_options", {})
|
||||
self.priority_steps = broker_options.get(
|
||||
"priority_steps", self.DEFAULT_PRIORITY_STEPS
|
||||
)
|
||||
self.sep = broker_options.get("sep", self.DEFAULT_SEP)
|
||||
self.broker_prefix = broker_options.get("global_keyprefix", "")
|
||||
|
||||
def _q_for_pri(self, queue, pri):
|
||||
if pri not in self.priority_steps:
|
||||
raise ValueError("Priority not in priority steps")
|
||||
# pylint: disable=consider-using-f-string
|
||||
return "{0}{1}{2}".format(*((queue, self.sep, pri) if pri else (queue, "", "")))
|
||||
|
||||
async def queues(self, names):
|
||||
queue_stats = []
|
||||
for name in names:
|
||||
priority_names = [
|
||||
self.broker_prefix + self._q_for_pri(name, pri)
|
||||
for pri in self.priority_steps
|
||||
]
|
||||
queue_stats.append(
|
||||
{
|
||||
"name": name,
|
||||
"messages": sum((self.redis.llen(x) for x in priority_names)),
|
||||
}
|
||||
)
|
||||
return queue_stats
|
||||
|
||||
def close(self):
|
||||
"""Close the Redis connection."""
|
||||
if self.redis is not None:
|
||||
self.redis.close()
|
||||
self.redis = None
|
||||
|
||||
|
||||
class Redis(RedisBase):
|
||||
def __init__(self, broker_url, *args, **kwargs):
|
||||
super().__init__(broker_url, *args, **kwargs)
|
||||
self.host = self.host or "localhost"
|
||||
self.port = self.port or 6379
|
||||
self.vhost = self._prepare_virtual_host(self.vhost)
|
||||
self.redis = self._get_redis_client()
|
||||
|
||||
def _prepare_virtual_host(self, vhost):
|
||||
if not isinstance(vhost, numbers.Integral):
|
||||
if not vhost or vhost == "/":
|
||||
vhost = 0
|
||||
elif vhost.startswith("/"):
|
||||
vhost = vhost[1:]
|
||||
try:
|
||||
vhost = int(vhost)
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"Database is int between 0 and limit - 1, not {vhost}"
|
||||
) from exc
|
||||
return vhost
|
||||
|
||||
def _get_redis_client_args(self):
|
||||
return {
|
||||
"host": self.host,
|
||||
"port": self.port,
|
||||
"db": self.vhost,
|
||||
"username": self.username,
|
||||
"password": self.password,
|
||||
}
|
||||
|
||||
def _get_redis_client(self):
|
||||
return redis.Redis(**self._get_redis_client_args())
|
||||
|
||||
|
||||
class RedisSentinel(RedisBase):
|
||||
def __init__(self, broker_url, *args, **kwargs):
|
||||
super().__init__(broker_url, *args, **kwargs)
|
||||
broker_options = kwargs.get("broker_options", {})
|
||||
broker_use_ssl = kwargs.get("broker_use_ssl", None)
|
||||
self.host = self.host or "localhost"
|
||||
self.port = self.port or 26379
|
||||
self.vhost = self._prepare_virtual_host(self.vhost)
|
||||
self.master_name = self._prepare_master_name(broker_options)
|
||||
self.redis = self._get_redis_client(broker_options, broker_use_ssl)
|
||||
|
||||
def _prepare_virtual_host(self, vhost):
|
||||
if not isinstance(vhost, numbers.Integral):
|
||||
if not vhost or vhost == "/":
|
||||
vhost = 0
|
||||
elif vhost.startswith("/"):
|
||||
vhost = vhost[1:]
|
||||
try:
|
||||
vhost = int(vhost)
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"Database is int between 0 and limit - 1, not {vhost}"
|
||||
) from exc
|
||||
return vhost
|
||||
|
||||
def _prepare_master_name(self, broker_options):
|
||||
try:
|
||||
master_name = broker_options["master_name"]
|
||||
except KeyError as exc:
|
||||
raise ValueError("master_name is required for Sentinel broker") from exc
|
||||
return master_name
|
||||
|
||||
def _get_redis_client(self, broker_options, broker_use_ssl):
|
||||
connection_kwargs = {
|
||||
"password": self.password,
|
||||
"sentinel_kwargs": broker_options.get("sentinel_kwargs"),
|
||||
}
|
||||
if isinstance(broker_use_ssl, dict):
|
||||
connection_kwargs["ssl"] = True
|
||||
connection_kwargs.update(broker_use_ssl)
|
||||
# get all sentinel hosts from Celery App config and use them to initialize Sentinel
|
||||
sentinel = redis.sentinel.Sentinel(
|
||||
[(self.host, self.port)], **connection_kwargs
|
||||
)
|
||||
redis_client = sentinel.master_for(self.master_name)
|
||||
return redis_client
|
||||
|
||||
|
||||
class RedisSocket(RedisBase):
|
||||
def __init__(self, broker_url, *args, **kwargs):
|
||||
super().__init__(broker_url, *args, **kwargs)
|
||||
self.redis = redis.Redis(
|
||||
unix_socket_path="/" + self.vhost, password=self.password
|
||||
)
|
||||
|
||||
|
||||
class RedisSsl(Redis):
|
||||
"""
|
||||
Redis SSL class offering connection to the broker over SSL.
|
||||
This does not currently support SSL settings through the url, only through
|
||||
the broker_use_ssl celery configuration.
|
||||
"""
|
||||
|
||||
def __init__(self, broker_url, *args, **kwargs):
|
||||
if "broker_use_ssl" not in kwargs:
|
||||
raise ValueError("rediss broker requires broker_use_ssl")
|
||||
self.broker_use_ssl = kwargs.get("broker_use_ssl", {})
|
||||
super().__init__(broker_url, *args, **kwargs)
|
||||
|
||||
def _get_redis_client_args(self):
|
||||
client_args = super()._get_redis_client_args()
|
||||
client_args["ssl"] = True
|
||||
if isinstance(self.broker_use_ssl, dict):
|
||||
client_args.update(self.broker_use_ssl)
|
||||
return client_args
|
||||
|
||||
|
||||
class Broker:
|
||||
"""Factory returning the appropriate broker client based on URL scheme.
|
||||
|
||||
Supported schemes:
|
||||
``amqp`` or ``amqps`` -> :class:`RabbitMQ`
|
||||
``redis`` -> :class:`Redis`
|
||||
``rediss`` -> :class:`RedisSsl`
|
||||
``redis+socket`` -> :class:`RedisSocket`
|
||||
``sentinel`` -> :class:`RedisSentinel`
|
||||
"""
|
||||
|
||||
def __new__(cls, broker_url, *args, **kwargs):
|
||||
scheme = urlparse(broker_url).scheme
|
||||
if scheme in ("amqp", "amqps"):
|
||||
return RabbitMQ(broker_url, *args, **kwargs)
|
||||
if scheme == "redis":
|
||||
return Redis(broker_url, *args, **kwargs)
|
||||
if scheme == "rediss":
|
||||
return RedisSsl(broker_url, *args, **kwargs)
|
||||
if scheme == "redis+socket":
|
||||
return RedisSocket(broker_url, *args, **kwargs)
|
||||
if scheme == "sentinel":
|
||||
return RedisSentinel(broker_url, *args, **kwargs)
|
||||
raise NotImplementedError
|
||||
|
||||
async def queues(self, names):
|
||||
raise NotImplementedError
|
||||
Reference in New Issue
Block a user