Prefilter polled queues

tasktiger/worker.py

@@ -3,6 +3,7 @@
 import errno
 import fcntl
 import random
+import re
 import select
 import signal
 import socket
@@ -23,6 +24,7 @@
 from .timeouts import JobTimeoutException
 
 LOCK_REDIS_KEY = "qslock"
+REDIS_GLOB_CHARACTER_PATTERN = re.compile(r"([\\?*\[\]])")
 
 __all__ = ["Worker"]
 
@@ -1156,10 +1158,24 @@ def _queue_periodic_tasks(self):
             )
 
     def _refresh_queue_set(self):
-        self._queue_set = set(
-            self._filter_queues(self.connection.smembers(self._key(QUEUED)))
+        self._queue_set = set(self._filter_queues(self._retrieve_queues()))
+
+    def _retrieve_queues(self):
+        key = self._key(QUEUED)
+
+        if len(self.only_queues) != 1:
+            return self.connection.smembers(key)
+
+        # Escape special Redis glob characters in the queue name
+        match = (
+            REDIS_GLOB_CHARACTER_PATTERN.sub(
+                r"\\\1", list(self.only_queues)[0]
+            )
+            + "*"
         )
 
+        return set(self.connection.sscan_iter(key, match=match, count=100000))
+
     def run(self, once=False, force_once=False):
         """
         Main loop of the worker.

tests/test_base.py

@@ -739,19 +739,38 @@ def test_batch_lock_key(self):
         Worker(self.tiger).run(once=True)
         self._ensure_queues(queued={"batch": 0})
 
-    def test_only_queues(self):
+    @pytest.mark.parametrize(
+        "only_queues,expected_unprocessed",
+        [
+            (["a"], {"b": 1, "b.a": 1, "[ab\\c]*?": 2}),
+            (["a", "b"], {"[ab\\c]*?": 2}),
+            (["a", "[ab\\c]*?"], {"b": 1, "b.a": 1}),
+            # tests that queue prefiltering works correctly with queues
+            # containing special characters used in Redis glob patterns
+            (["[ab\\c]*?"], {"a": 1, "a.a": 1, "b": 1, "b.a": 1}),
+            # all queues selected, all queues processed
+            (["a", "b", "[ab\\c]*?"], {}),
+            # no queue restriction, all queues processed
+            ([], {}),
+        ],
+    )
+    def test_only_queues(self, only_queues, expected_unprocessed):
         self.tiger.delay(simple_task, queue="a")
         self.tiger.delay(simple_task, queue="a.a")
         self.tiger.delay(simple_task, queue="b")
         self.tiger.delay(simple_task, queue="b.a")
+        self.tiger.delay(simple_task, queue="[ab\\c]*?")
+        self.tiger.delay(simple_task, queue="[ab\\c]*?")
 
-        self._ensure_queues(queued={"a": 1, "a.a": 1, "b": 1, "b.a": 1})
+        self._ensure_queues(
+            queued={"a": 1, "a.a": 1, "b": 1, "b.a": 1, "[ab\\c]*?": 2}
+        )
 
-        self.tiger.config["ONLY_QUEUES"] = ["a"]
+        self.tiger.config["ONLY_QUEUES"] = only_queues
 
         Worker(self.tiger).run(once=True)
 
-        self._ensure_queues(queued={"b": 1, "b.a": 1})
+        self._ensure_queues(queued=expected_unprocessed)
 
     def test_exclude_queues(self):
         """