diff --git a/celery_exporter/__main__.py b/celery_exporter/__main__.py index 2737dcd..74d8367 100644 --- a/celery_exporter/__main__.py +++ b/celery_exporter/__main__.py @@ -11,7 +11,7 @@ from .utils import generate_broker_use_ssl, get_transport_scheme LOG_FORMAT = "[%(asctime)s] %(name)s:%(levelname)s: %(message)s" - +CELERY_DEFAULT_QUEUE = "celery" @click.command(context_settings={"auto_envvar_prefix": "CELERY_EXPORTER"}) @click.option( @@ -102,6 +102,15 @@ @click.option( "--verbose", is_flag=True, allow_from_autoenv=False, help="Enable verbose logging." ) +@click.option( + "--queue", + "-q", + type=str, + show_default=True, + show_envvar=True, + default=CELERY_DEFAULT_QUEUE, + help="Queue name for metrics.", +) def main( broker_url, listen_address, @@ -116,6 +125,7 @@ def main( ssl_keyfile, tz, verbose, + queue, ): # pragma: no cover if verbose: @@ -155,6 +165,7 @@ def main( transport_options, enable_events, broker_use_ssl, + queue, ) celery_exporter.start() diff --git a/celery_exporter/core.py b/celery_exporter/core.py index f0612ee..c464c5c 100644 --- a/celery_exporter/core.py +++ b/celery_exporter/core.py @@ -23,6 +23,7 @@ def __init__( transport_options=None, enable_events=False, broker_use_ssl=None, + queue='celery', ): self._listen_address = listen_address self._max_tasks = max_tasks @@ -31,10 +32,11 @@ def __init__( self._app = celery.Celery(broker=broker_url, broker_use_ssl=broker_use_ssl) self._app.conf.broker_transport_options = transport_options or {} + self._queue = queue def start(self): - setup_metrics(self._app, self._namespace) + setup_metrics(self._app, self._namespace, self._queue) self._start_httpd() @@ -42,6 +44,7 @@ def start(self): app=self._app, namespace=self._namespace, max_tasks_in_memory=self._max_tasks, + queue=self._queue ) t.daemon = True t.start() diff --git a/celery_exporter/monitor.py b/celery_exporter/monitor.py index 4f7125e..7d964eb 100644 --- a/celery_exporter/monitor.py +++ b/celery_exporter/monitor.py @@ -18,9 +18,10 @@ class TaskThread(threading.Thread): exposed from Celery using its eventing system. """ - def __init__(self, app, namespace, max_tasks_in_memory, *args, **kwargs): + def __init__(self, app, namespace, max_tasks_in_memory, queue, *args, **kwargs): self._app = app self._namespace = namespace + self._queue = queue self.log = logging.getLogger("task-thread") self._state = CeleryState(max_tasks_in_memory=max_tasks_in_memory) self._known_states = set() @@ -56,12 +57,12 @@ def _monitor(self): # pragma: no cover recv = self._app.events.Receiver( conn, handlers={"*": self._process_event} ) - setup_metrics(self._app, self._namespace) + setup_metrics(self._app, self._namespace, self._queue) self.log.info("Start capturing events...") recv.capture(limit=None, timeout=None, wakeup=True) except Exception: self.log.exception("Connection failed") - setup_metrics(self._app, self._namespace) + setup_metrics(self._app, self._namespace, self._queue) time.sleep(5) @@ -109,13 +110,13 @@ def enable_events(self): self._app.control.enable_events() -def setup_metrics(app, namespace): +def setup_metrics(app, namespace, queue='celery'): """ This initializes the available metrics with default values so that even before the first event is received, data can be exposed. """ WORKERS.labels(namespace=namespace) - config = get_config(app) + config = get_config(app, queue=queue) if not config: # pragma: no cover for metric in TASKS.collect(): diff --git a/celery_exporter/utils.py b/celery_exporter/utils.py index 7c2317c..3cdb63a 100644 --- a/celery_exporter/utils.py +++ b/celery_exporter/utils.py @@ -2,7 +2,6 @@ from itertools import chain from urllib.parse import urlparse -CELERY_DEFAULT_QUEUE = "celery" CELERY_MISSING_DATA = "undefined" @@ -15,7 +14,7 @@ def _gen_wildcards(name): return res -def get_config(app): +def get_config(app, queue): res = dict() try: registered_tasks = app.control.inspect().registered_tasks().values() @@ -26,7 +25,7 @@ def get_config(app): default_queues = [] for task_name in set(chain.from_iterable(registered_tasks)): for conf in confs.values(): - default = conf.get("task_default_queue", CELERY_DEFAULT_QUEUE) + default = conf.get("task_default_queue", queue) default_queues.append(default) if task_name in res and res[task_name] not in default_queues: break diff --git a/test/test_func.py b/test/test_func.py index 9b0faf3..4214e8d 100644 --- a/test/test_func.py +++ b/test/test_func.py @@ -28,6 +28,7 @@ def setUp(self): max_tasks=TestCeleryExporter.max_tasks, namespace=TestCeleryExporter.namespace, enable_events=True, + queue='celery', ) def test_setup_metrics(self): diff --git a/test/test_unit.py b/test/test_unit.py index 1f6afc8..75099a0 100644 --- a/test/test_unit.py +++ b/test/test_unit.py @@ -36,7 +36,7 @@ def setUp(self): "celery@12311847jsa2": {}, } registered.return_value = {"celery@d6f95e9e24fc": [self.task, "trial"]} - setup_metrics(self.app, self.namespace) # reset metrics + setup_metrics(self.app, self.namespace, queue=self.queue) # reset metrics def test_initial_metric_values(self): self._assert_task_states(celery.states.ALL_STATES, 0) @@ -70,7 +70,7 @@ def test_workers_count(self): ) with patch.object(self.app.control, "ping") as mock_ping: - w = WorkerMonitoringThread(app=self.app, namespace=self.namespace) + w = WorkerMonitoringThread(app=self.app, namespace=self.namespace, queue=self.queue) mock_ping.return_value = [] w.update_workers_count() @@ -116,7 +116,7 @@ def test_tasks_events(self): runtime = 234.5 m = TaskThread( - app=self.app, namespace=self.namespace, max_tasks_in_memory=self.max_tasks + app=self.app, namespace=self.namespace, max_tasks_in_memory=self.max_tasks, queue=self.queue ) self._assert_task_states(celery.states.ALL_STATES, 0)