fix: Improve automatic queue finding for Celery (#85)

judoscale/celery/collector.py

@@ -73,12 +73,17 @@ def __init__(self, config: Config, broker: Celery):
         logger.debug(f"Redis is at {self.redis.connection_pool}")
 
         system_queues = {"unacked", "unacked_index"}
-        user_queues = {
-            q.decode("utf-8") if isinstance(q, bytes) else q
-            for q in self.redis.scan_iter(match="[^_]*", _type="list")
-        }
-        logger.debug(f"Found initial queues: {list(user_queues)}")
-        self._celery_queues = user_queues - system_queues
+        for q in self.redis.scan_iter(_type="list"):
+            queue_name = q.decode() if isinstance(q, bytes) else q
+            if (
+                queue_name in system_queues
+                or queue_name.startswith("_kombu")
+                or queue_name.endswith("celery.pidbox")
+            ):
+                continue
+            self._celery_queues.add(queue_name)
+
+        logger.debug(f"Found initial queues: {list(self._celery_queues)}")
         self.task_sent_handler.start()
 
     @property

tests/test_collectors.py

@@ -144,6 +144,17 @@ def test_queues_empty(self, heroku_worker_1, celery):
         collector = CeleryMetricsCollector(heroku_worker_1, celery)
         assert collector.queues == set()
 
+    def test_removing_various_celery_queues(self, heroku_worker_1, celery):
+        celery.connection_for_read().channel().client.scan_iter.return_value = [
+            b"unacked",
+            b"unacked_index",
+            b"_kombu.binding.celeryev",
+            b"e752fa70-f772-3c04-b05d-79b2a79ce766.reply.celery.pidbox",
+            b"user_queue",
+        ]
+        collector = CeleryMetricsCollector(heroku_worker_1, celery)
+        assert collector.queues == {"user_queue"}
+
     def test_collect_empty_queue(self, worker_1, celery):
         celery.connection_for_read().channel().client.scan_iter.return_value = [b"foo"]
         celery.connection_for_read().channel().client.lindex.return_value = None