Added clean_worker_registry()clean-worker-registry

5 files changed, 65 insertions, 6 deletions

diff --git a/rq/queue.py b/rq/queue.py
index 52cf275..e71f661 100644
--- a/rq/queue.py
+++ b/rq/queue.py
@@ -89,6 +89,17 @@ class Queue(object):
         """Returns the Redis key for this Queue."""
         return self._key
 
+    @property
+    def registry_cleaning_key(self):
+        """Redis key used to indicate this queue has been cleaned."""
+        return 'rq:clean_registries:%s' % self.name
+
+    def acquire_cleaning_lock(self):
+        """Returns a boolean indicating whether a lock to clean this queue
+        is acquired. A lock expires in 899 seconds (15 minutes - 1 second)
+        """
+        return self.connection.set(self.registry_cleaning_key, 1, nx=1, ex=899)
+
     def empty(self):
         """Removes all messages on the queue."""
         script = """
diff --git a/rq/worker.py b/rq/worker.py
index 97f80bd..4960d59 100644
--- a/rq/worker.py
+++ b/rq/worker.py
@@ -40,7 +40,7 @@ from .timeouts import JobTimeoutException, HorseMonitorTimeoutException, UnixSig
 from .utils import (backend_class, ensure_list, enum,
                     make_colorizer, utcformat, utcnow, utcparse)
 from .version import VERSION
-from .worker_registration import get_keys
+from .worker_registration import clean_worker_registry, get_keys
 
 try:
     from procname import setprocname
@@ -909,16 +909,20 @@ class Worker(object):
     def clean_registries(self):
         """Runs maintenance jobs on each Queue's registries."""
         for queue in self.queues:
-            self.log.info('Cleaning registries for queue: %s', queue.name)
-            clean_registries(queue)
+            # If there are multiple workers running, we only want 1 worker
+            # to run clean_registries().
+            if queue.acquire_cleaning_lock():
+                self.log.info('Cleaning registries for queue: %s', queue.name)
+                clean_registries(queue)
+                clean_worker_registry(queue)
         self.last_cleaned_at = utcnow()
 
     @property
     def should_run_maintenance_tasks(self):
-        """Maintenance tasks should run on first startup or every hour."""
+        """Maintenance tasks should run on first startup or 15 minutes."""
         if self.last_cleaned_at is None:
             return True
-        if (utcnow() - self.last_cleaned_at) > timedelta(hours=1):
+        if (utcnow() - self.last_cleaned_at) > timedelta(minutes=15):
             return True
         return False
 
diff --git a/rq/worker_registration.py b/rq/worker_registration.py
index d24fac3..3944bc7 100644
--- a/rq/worker_registration.py
+++ b/rq/worker_registration.py
@@ -43,3 +43,25 @@ def get_keys(queue=None, connection=None):
         redis_key = REDIS_WORKER_KEYS
 
     return {as_text(key) for key in redis.smembers(redis_key)}
+
+
+def clean_worker_registry(queue):
+    """Delete invalid worker keys in registry"""
+    keys = list(get_keys(queue))
+
+    with queue.connection.pipeline() as pipeline:
+
+        for key in keys:
+            pipeline.exists(key)
+        results = pipeline.execute()
+
+        invalid_keys = []
+
+        for i, key_exists in enumerate(results):
+            if not key_exists:
+                invalid_keys.append(keys[i])
+
+        if invalid_keys:
+            pipeline.srem(WORKERS_BY_QUEUE_KEY % queue.name, *invalid_keys)
+            pipeline.srem(REDIS_WORKER_KEYS, *invalid_keys)
+            pipeline.execute()
diff --git a/tests/test_worker.py b/tests/test_worker.py
index 3a42aee..765294a 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -743,6 +743,12 @@ class TestWorker(RQTestCase):
         self.assertEqual(self.testconn.zcard(foo_registry.key), 0)
         self.assertEqual(self.testconn.zcard(bar_registry.key), 0)
 
+        # worker.clean_registries() only runs once every 15 minutes
+        # If we add another key, calling clean_registries() should do nothing
+        self.testconn.zadd(bar_registry.key, {'bar': 1})
+        worker.clean_registries()
+        self.assertEqual(self.testconn.zcard(bar_registry.key), 1)
+
     def test_should_run_maintenance_tasks(self):
         """Workers should run maintenance tasks on startup and every hour."""
         queue = Queue(connection=self.testconn)
diff --git a/tests/test_worker_registration.py b/tests/test_worker_registration.py
index 1fe6612..0fe9ce8 100644
--- a/tests/test_worker_registration.py
+++ b/tests/test_worker_registration.py
@@ -1,7 +1,8 @@
 from tests import RQTestCase
 
 from rq import Queue, Worker
-from rq.worker_registration import (get_keys, register, unregister,
+from rq.worker_registration import (clean_worker_registry, get_keys, register,
+                                    unregister, REDIS_WORKER_KEYS,
                                     WORKERS_BY_QUEUE_KEY)
 
 
@@ -68,3 +69,18 @@ class TestWorkerRegistry(RQTestCase):
         unregister(worker1)
         unregister(worker2)
         unregister(worker3)
+
+    def test_clean_registry(self):
+        """clean_registry removes worker keys that don't exist in Redis"""
+        queue = Queue(name='foo')
+        worker = Worker([queue])
+
+        register(worker)
+        redis = worker.connection
+
+        self.assertTrue(redis.sismember(worker.redis_workers_keys, worker.key))
+        self.assertTrue(redis.sismember(REDIS_WORKER_KEYS, worker.key))
+
+        clean_worker_registry(queue)
+        self.assertFalse(redis.sismember(worker.redis_workers_keys, worker.key))
+        self.assertFalse(redis.sismember(REDIS_WORKER_KEYS, worker.key))