Reorganize task revocation logic

Author: bjester

contentcuration/contentcuration/tests/helpers.py

@@ -2,6 +2,7 @@
 from importlib import import_module
 
 import mock
+from celery import states
 
 from contentcuration.models import TaskResult
 
@@ -20,6 +21,7 @@ def clear_tasks(except_task_id=None):
         qs = qs.exclude(task_id=except_task_id)
     for task_id in qs.values_list("task_id", flat=True):
         app.control.revoke(task_id, terminate=True)
+    qs.update(status=states.REVOKED)
 
 
 def mock_class_instance(target):

contentcuration/contentcuration/tests/test_asynctask.py

@@ -142,6 +142,9 @@ def _wait_for(self, async_result, timeout=30):
         with allow_join_result():
             return async_result.get(timeout=timeout)
 
+    def test_app_count_queued_tasks(self):
+        self.assertIsInstance(app.count_queued_tasks(), int)
+
     def test_asynctask_reports_success(self):
         """
         Tests that when an async task is created and completed, the Task object has a status of 'SUCCESS' and
@@ -245,3 +248,17 @@ def test_requeue_task(self):
         second_result = self._wait_for(second_async_result)
         self.assertIsNone(second_result)
         self.assertTrue(second_async_result.successful())
+
+    def test_revoke_task(self):
+        channel_id = uuid.uuid4()
+        async_result = test_task.enqueue(self.user, channel_id=channel_id)
+        test_task.revoke(channel_id=channel_id)
+
+        # this should raise an exception, even though revoked, because the task is in ready state but not success
+        with self.assertRaises(Exception):
+            self._wait_for(async_result)
+
+        try:
+            TaskResult.objects.get(task_id=async_result.task_id, status=states.REVOKED)
+        except TaskResult.DoesNotExist:
+            self.fail('Missing revoked task result')

contentcuration/contentcuration/utils/celery/app.py

@@ -54,37 +54,6 @@ def count_queued_tasks(self, queue_name="celery"):
             count = conn.default_channel.client.llen(queue_name)
         return count
 
-    def _revoke_matching(self, tasks, task_name, matching_kwargs=None):
-        count = 0
-        for task in tasks:
-            if task['name'] == task_name and (matching_kwargs is None or task['kwargs'] == matching_kwargs):
-                print('Revoking task {}'.format(task))
-                count += 1
-                self.control.revoke(task['id'], terminate=True)
-        return count
-
-    def revoke_reserved_tasks_by_name(self, task_name, matching_kwargs=None):
-        """
-        :param task_name: The string name of the task
-        :param matching_kwargs: A dict of kwargs to match to the task
-        :return: The count of revoked tasks
-        """
-        count = 0
-        for worker_name, tasks in self.control.inspect().reserved().items():
-            count += self._revoke_matching(tasks, task_name, matching_kwargs)
-        return count
-
-    def terminate_active_tasks_by_name(self, task_name, matching_kwargs=None):
-        """
-        :param task_name: The string name of the task
-        :param matching_kwargs: A dict of kwargs to match to the task
-        :return: The count of terminated tasks
-        """
-        count = 0
-        for worker_name, tasks in self.control.inspect().active().items():
-            count += self._revoke_matching(tasks, task_name, matching_kwargs)
-        return count
-
     def decode_result(self, result, status=None):
         """
         Decodes the celery result, like the raw result from the database, using celery tools

contentcuration/contentcuration/utils/celery/tasks.py

@@ -246,6 +246,25 @@ def requeue(self, **kwargs):
         logging.info(f"Re-queuing task {self.name} for user {task_result.user.pk} from {request.id} | {task_kwargs}")
         return self.enqueue(task_result.user, **task_kwargs)
 
+    def revoke(self, exclude_task_ids=None, **kwargs):
+        """
+        Revokes and terminates all unready tasks matching the kwargs
+        :param exclude_task_ids: Any task ids to exclude from this behavior
+        :param kwargs: Task keyword arguments that will be used to match against tasks
+        :return: The number of tasks revoked
+        """
+        task_ids = self.find_incomplete_ids(**self.backend.decode(self._prepare_kwargs(kwargs)))
+        if exclude_task_ids is not None:
+            task_ids = task_ids.exclude(task_id__in=task_ids)
+        count = 0
+        for task_id in task_ids:
+            logging.info(f"Revoking task {task_id}")
+            self.app.control.revoke(task_id, terminate=True)
+            count += 1
+        # be sure the database backend has these marked appropriately
+        self.TaskModel.objects.filter(task_id__in=task_ids).update(status=states.REVOKED)
+        return count
+
 
 class CeleryAsyncResult(AsyncResult):
     """