feat: add option for challenging/removing lock on task start

add challenge_lock_on_startup setting on Singleton
add celery_app setting as passing Celery application
is mandatory for inspection of celery workers

Author: tbrien

README.md

@@ -7,17 +7,19 @@ This is a baseclass for celery tasks that ensures only one instance of the task
 **Table of Contents**
 
 - [Celery-Singleton](#celery-singleton)
-    - [Prerequisites](#prerequisites)
-    - [Quick start](#quick-start)
-    - [How does it work?](#how-does-it-work)
-    - [Handling deadlocks](#handling-deadlocks)
-    - [Backends](#backends)
-    - [Task configuration](#task-configuration)
-        - [unique\_on](#uniqueon)
-        - [raise\_on\_duplicate](#raiseonduplicate)
-    - [App Configuration](#app-configuration)
-    - [Testing](#testing)
-    - [Contribute](#contribute)
+  - [Prerequisites](#prerequisites)
+  - [Quick start](#quick-start)
+  - [How does it work?](#how-does-it-work)
+  - [Handling deadlocks](#handling-deadlocks)
+  - [Backends](#backends)
+  - [Task configuration](#task-configuration)
+    - [unique\_on](#uniqueon)
+    - [raise\_on\_duplicate](#raiseonduplicate)
+    - [lock\_expiry](#lockexpiry)
+    - [callenge_lock_on_startup](#callengelockonstartup)
+  - [App Configuration](#app-configuration)
+  - [Testing](#testing)
+  - [Contribute](#contribute)
 
 <!-- markdown-toc end -->
 
@@ -182,6 +184,32 @@ assert task1 != task2  # These are two separate task instances
 
 This option can be applied globally in the [app config](#app-configuration) with `singleton_lock_expiry`. Task option supersedes the app config.
 
+### callenge_lock_on_startup 
+
+Some tasks may be stucked after unexpected shuttdown of a worker.
+If this situation should be avoided, the lock can be challenged on task startup.
+
+> **Use with caution** : this configuration may lead to unexpected behaviour. Try using classic [deadlocks handling](#markdown-header-handling-deadlocks) if possible.
+> this may lead to unexpected behaviour if celery is in a transition state at task startup
+> ex: Tasks beeing scheduled while no worker online will all start on worker startup
+
+```python
+from celery_singleton import Singleton
+from .local_path import app
+
+
+@app.task(
+    base=Singleton,
+    challenge_lock_on_startup=True,
+    celery_app=app,
+)
+def do_something(username):
+    time.sleep(5)
+
+task1 = do_something.delay('bob')
+# Worker cold shuttdown causes lock to remain
+task2 = do_something.delay('bob') # lock should be removed
+```
 
 ## App Configuration
 

celery_singleton/inspect_celery.py

@@ -0,0 +1,35 @@
+from celery import Celery
+
+class CeleryTask:
+    def __init__(self, name, args, kwargs):
+        self.name = name
+        self.args = args
+        self.kwargs = kwargs
+
+def list_celery_tasks_with_name(app: Celery, name: str) -> [CeleryTask]:
+    active_tasks = app.control.inspect().active()
+    if not active_tasks:
+        return []
+    global_task_list = []
+    try:
+        for worker, task_list in active_tasks.items():
+            for task in task_list:
+                if not task['name'] == name:
+                    continue
+                celery_task = CeleryTask(task["name"], task["args"], task["kwargs"])
+                global_task_list.append(celery_task)
+    except Exception:
+        pass
+    return global_task_list
+
+def are_worker_active(app: Celery) -> bool:
+    worke_pong = app.control.inspect().ping()
+    if not worke_pong:
+        return False
+    try:
+        for worker, response in worke_pong.items():
+            if response['ok']:
+                return True
+    except Exception:
+        pass
+    return False

celery_singleton/singleton.py

@@ -6,6 +6,7 @@
 from .config import Config
 from .exceptions import DuplicateTaskError
 from . import util
+from . import inspect_celery
 
 
 def clear_locks(app):
@@ -21,6 +22,8 @@ class Singleton(BaseTask):
     unique_on = None
     raise_on_duplicate = None
     lock_expiry = None
+    challenge_lock_on_startup = False
+    celery_app = None
 
     @property
     def _raise_on_duplicate(self):
@@ -54,6 +57,10 @@ def get_existing_task_id(self, lock):
         return self.singleton_backend.get(lock)
 
     def generate_lock(self, task_name, task_args=None, task_kwargs=None):
+        unique_on = self.unique_on
+        return self._generate_lock_filtered_by_unique_on(unique_on, task_name, task_args, task_kwargs)
+    
+    def _generate_lock_filtered_by_unique_on(self, unique_on, task_name, task_args=None, task_kwargs=None):
         unique_on = self.unique_on
         task_args = task_args or []
         task_kwargs = task_kwargs or {}
@@ -103,6 +110,9 @@ def apply_async(
             **options
         )
 
+        if self.challenge_lock_on_startup:
+            self.challenge_lock(lock, args, kwargs)
+
         task = self.lock_and_run(**run_args)
         if task:
             return task
@@ -126,6 +136,14 @@ def lock_and_run(self, lock, *args, task_id=None, **kwargs):
                 # Clear the lock if apply_async fails
                 self.unlock(lock)
                 raise
+    
+    def challenge_lock(self, lock, args, kwargs):
+        current_task_lock = self.generate_lock(self.name, args, kwargs)
+        active_tasks = inspect_celery.list_celery_tasks_with_name(self.celery_app, self.name)
+        active_tasks_lock = [self.generate_lock(task.name, task.args, task.kwargs) for task in active_tasks]
+        if lock in active_tasks_lock or not inspect_celery.are_worker_active(self.celery_app):
+            return
+        self.unlock(lock)    
 
     def release_lock(self, task_args=None, task_kwargs=None):
         lock = self.generate_lock(self.name, task_args, task_kwargs)

tests/test_inspect_celery.py

@@ -0,0 +1,121 @@
+import pytest
+import sys
+from unittest import mock
+from celery_singleton.inspect_celery import list_celery_tasks_with_name, are_worker_active, CeleryTask
+
+
+class TestInspectCelery:
+    def test__list_celery_tasks_with_name__should_return_matching_task(self):
+        # Given
+        active_tasks = {'celery@worker_host': [{'id': '262a1cf9-2c4f-4680-8261-7498fb39756c', 'name': 'simple_task', 'args': [1, 2, 3], 'kwargs': {}, 'type': 'simple_task', 'hostname': 'celery@worker_host', 'time_start': 1588508284.207397, 'acknowledged': True, 'delivery_info': {'exchange': '', 'routing_key': 'celery', 'priority': 0, 'redelivered': None}, 'worker_pid': 45895}]}
+        inspect_mock = mock.MagicMock()
+        inspect_mock.active.return_value = active_tasks
+        control_mock = mock.MagicMock()
+        control_mock.inspect.return_value = inspect_mock
+        mock_app = mock.MagicMock()
+        mock_app.control = control_mock
+
+        # When
+        task_list = list_celery_tasks_with_name(mock_app, "simple_task")
+
+        # Then
+        assert len(task_list) == 1
+        actual = task_list[0]
+        assert actual.name == 'simple_task'
+        assert actual.args == [1, 2, 3]
+        assert actual.kwargs == {}
+    
+    def test__list_celery_tasks_with_name__should_return_empty_list_when_no_matching_taskname(self):
+        # Given
+        active_tasks = {'celery@worker_host': [{'id': '262a1cf9-2c4f-4680-8261-7498fb39756c', 'name': 'simple_task', 'args': [1, 2, 3], 'kwargs': {}, 'type': 'simple_task', 'hostname': 'celery@worker_host', 'time_start': 1588508284.207397, 'acknowledged': True, 'delivery_info': {'exchange': '', 'routing_key': 'celery', 'priority': 0, 'redelivered': None}, 'worker_pid': 45895}]}
+        inspect_mock = mock.MagicMock()
+        inspect_mock.active.return_value = active_tasks
+        control_mock = mock.MagicMock()
+        control_mock.inspect.return_value = inspect_mock
+        mock_app = mock.MagicMock()
+        mock_app.control = control_mock
+
+        # When
+        task_list = list_celery_tasks_with_name(mock_app, "bad_task_name")
+
+        # Then
+        assert len(task_list) == 0
+    
+    def test__list_celery_tasks_with_name__should_return_empty_list_when_no_active_tasts(self):
+        # Given
+        inspect_mock = mock.MagicMock()
+        inspect_mock.active.return_value = None
+        control_mock = mock.MagicMock()
+        control_mock.inspect.return_value = inspect_mock
+        mock_app = mock.MagicMock()
+        mock_app.control = control_mock
+
+        # When
+        task_list = list_celery_tasks_with_name(mock_app, "any_name")
+
+        # Then
+        assert len(task_list) == 0
+    
+    def test__list_celery_tasks_with_name__should_return_empty_list_when_worker_answer_cannot_be_parsed(self):
+        # Given
+        inspect_mock = mock.MagicMock()
+        inspect_mock.active.return_value = {'worker': ['bad_task_definition']}
+        control_mock = mock.MagicMock()
+        control_mock.inspect.return_value = inspect_mock
+        mock_app = mock.MagicMock()
+        mock_app.control = control_mock
+
+        # When
+        task_list = list_celery_tasks_with_name(mock_app, "bad_task_name")
+
+        # Then
+        assert len(task_list) == 0
+    
+    def test__are_worker_active__should_return_true_if_worker_responds_to_ping(self):
+        # Given
+        active_workers = {u'celery@host': {u'ok': u'pong'}}
+        inspect_mock = mock.MagicMock()
+        inspect_mock.ping.return_value = active_workers
+        control_mock = mock.MagicMock()
+        control_mock.inspect.return_value = inspect_mock
+        mock_app = mock.MagicMock()
+        mock_app.control = control_mock
+
+        # When
+        active_workers_found = are_worker_active(mock_app)
+
+        # Then
+        assert active_workers_found
+    
+    def test__are_worker_active__should_return_false_if_worker_does_not_respond_to_ping(self):
+        # Given
+        active_workers = None
+        inspect_mock = mock.MagicMock()
+        inspect_mock.ping.return_value = active_workers
+        control_mock = mock.MagicMock()
+        control_mock.inspect.return_value = inspect_mock
+        mock_app = mock.MagicMock()
+        mock_app.control = control_mock
+
+        # When
+        active_workers_found = are_worker_active(mock_app)
+
+        # Then
+        assert not active_workers_found
+    
+    def test__are_worker_active__should_return_false_if_worker_respond_ko_to_ping(self):
+        # Given
+        active_workers = {u'celery@host': {u'not_ok': u'pong'}}
+        inspect_mock = mock.MagicMock()
+        inspect_mock.ping.return_value = active_workers
+        control_mock = mock.MagicMock()
+        control_mock.inspect.return_value = inspect_mock
+        mock_app = mock.MagicMock()
+        mock_app.control = control_mock
+
+        # When
+        active_workers_found = are_worker_active(mock_app)
+
+        # Then
+        assert not active_workers_found
+

tests/test_singleton.py

@@ -10,6 +10,7 @@
 from celery_singleton.backends.redis import RedisBackend
 from celery_singleton.backends import get_backend
 from celery_singleton.config import Config
+from celery_singleton import inspect_celery
 
 
 @pytest.fixture(scope="session")
@@ -313,3 +314,60 @@ def simple_task(*args):
             mock_lock.assert_called_once_with(
                 simple_task.singleton_backend, lock, result.task_id, expiry=60
             )
+
+class TestChallengeLockOnStartup:
+    @mock.patch.object(inspect_celery, "list_celery_tasks_with_name", return_value=[])
+    @mock.patch.object(inspect_celery, "are_worker_active", return_value=True)
+    def test__lock_invalidation__remove_lock_if_no_task_running(self, mock_celery_inspect, mock_celery_ping, scoped_app):
+        with scoped_app as app:
+            # Given
+            task_args=[1, 2, 3]
+            @app.task(name="simple_task", base=Singleton, challenge_lock_on_startup=True, celery_app=app)
+            def simple_task(*args):
+                return args
+            lock = simple_task.generate_lock(simple_task.name, task_args=task_args)
+            simple_task.unlock = mock.MagicMock()
+            simple_task.release_lock = mock.MagicMock() # avoid call to unlock on task completion
+
+            # When
+            simple_task.delay(1, 2, 3)
+
+            # Then
+            simple_task.unlock.assert_called_once_with(lock)
+    
+    @mock.patch.object(inspect_celery, "list_celery_tasks_with_name", return_value=[inspect_celery.CeleryTask('simple_task', args=[1, 2, 3], kwargs=None)])
+    @mock.patch.object(inspect_celery, "are_worker_active", return_value=True)
+    def test__lock_invalidation__do_not_remove_lock_if_task_running(self, mock_celery_inspect, mock_celery_ping, scoped_app):
+        with scoped_app as app:
+            # Given
+            task_name = "simple_task"
+            task_args=[1, 2, 3]
+            @app.task(name="simple_task", base=Singleton, challenge_lock_on_startup=True, celery_app=app)
+            def simple_task(*args):
+                return args
+            lock = simple_task.generate_lock(simple_task.name, task_args, None)
+            simple_task.unlock = mock.MagicMock()
+            simple_task.release_lock = mock.MagicMock() # avoid call to unlock on task completion
+
+            # When
+            simple_task.delay(1, 2, 3)
+
+            # Then
+            simple_task.unlock.assert_not_called()
+    
+    @mock.patch.object(inspect_celery, "list_celery_tasks_with_name", return_value=[])
+    @mock.patch.object(inspect_celery, "are_worker_active", return_value=False)
+    def test__lock_invalidation__do_not_remove_lock_if_no_worker_running(self, mock_celery_inspect, mock_celery_ping, scoped_app):
+        with scoped_app as app:
+            # Given
+            @app.task(name="simple_task", base=Singleton, challenge_lock_on_startup=True, celery_app=app)
+            def simple_task(*args):
+                return args
+            simple_task.unlock = mock.MagicMock()
+            simple_task.release_lock = mock.MagicMock() # avoid call to unlock on task completion
+
+            # When
+            simple_task.delay(1, 2, 3)
+
+            # Then
+            simple_task.unlock.assert_not_called()