From 690a6451dbda140a9a2061a1507ebfccbfc887f6 Mon Sep 17 00:00:00 2001
From: Ethan Wolinsky <58117461+eswolinsky3241@users.noreply.github.com>
Date: Fri, 23 Feb 2024 21:44:44 -0500
Subject: [PATCH] Delete maintenance lock after registries cleaned (#2024)

* Delete maintenance lock after registries cleaned

* Fix test to assert that lock is released

* Remove deprecated test case
---
 rq/queue.py          |  4 ++++
 rq/worker.py         |  1 +
 tests/test_worker.py | 12 +++---------
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/rq/queue.py b/rq/queue.py
index 2c572c93..813a9d39 100644
--- a/rq/queue.py
+++ b/rq/queue.py
@@ -253,6 +253,10 @@ class Queue:
             return False
         return lock_acquired
 
+    def release_maintenance_lock(self):
+        """Deletes the maintenance lock after registries have been cleaned"""
+        self.connection.delete(self.registry_cleaning_key)
+
     def empty(self):
         """Removes all messages on the queue.
         This is currently being done using a Lua script,
diff --git a/rq/worker.py b/rq/worker.py
index ac8091c2..a5083f2e 100644
--- a/rq/worker.py
+++ b/rq/worker.py
@@ -450,6 +450,7 @@ class BaseWorker:
                 clean_registries(queue)
                 worker_registration.clean_worker_registry(queue)
                 clean_intermediate_queue(self, queue)
+                queue.release_maintenance_lock()
         self.last_cleaned_at = utcnow()
 
     def get_redis_server_version(self):
diff --git a/tests/test_worker.py b/tests/test_worker.py
index dab3bf36..6a4b32cc 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -1043,14 +1043,8 @@ class TestWorker(RQTestCase):
         self.assertEqual(worker.last_cleaned_at, None)
         worker.clean_registries()
         self.assertNotEqual(worker.last_cleaned_at, None)
-        self.assertEqual(self.connection.zcard(foo_registry.key), 0)
-        self.assertEqual(self.connection.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.connection.zadd(bar_registry.key, {'bar': 1})
-        worker.clean_registries()
-        self.assertEqual(self.connection.zcard(bar_registry.key), 1)
+        self.assertEqual(len(foo_registry), 0)
+        self.assertEqual(len(bar_registry), 0)
 
     def test_should_run_maintenance_tasks(self):
         """Workers should run maintenance tasks on startup and every hour."""
@@ -1079,7 +1073,7 @@ class TestWorker(RQTestCase):
 
         worker = Worker(queue, connection=self.connection)
         worker.work(burst=True)
-        self.assertEqual(self.connection.zcard(registry.key), 0)
+        self.assertEqual(len(registry), 0)
 
     def test_job_dependency_race_condition(self):
         """Dependencies added while the job gets finished shouldn't get lost."""