Fix failing task group removal

Author: Schweigi

eb_sqs/tests/worker/tests_worker.py

@@ -11,28 +11,47 @@
 from eb_sqs.worker.worker import Worker
 from eb_sqs.worker.worker_exceptions import MaxRetriesReachedException
 from eb_sqs.worker.worker_factory import WorkerFactory
-from eb_sqs.worker.worker_task import WorkerTask
 
+class TestException(Exception):
+    pass
 
 @task()
 def dummy_task(msg):
     return msg
 
+@task(max_retries=100)
+def retries_task(num_of_retries):
+    if retries_task.retry_num < num_of_retries:
+        retries_task.retry(execute_inline=True)
+
+@task(max_retries=5)
+def max_retries_task():
+    max_retries_task.retry(execute_inline=True)
+
+@task(max_retries=100)
+def repeating_group_task(num_of_retries):
+    if repeating_group_task.retry_num < num_of_retries:
+        repeating_group_task.retry(execute_inline=True)
+
 @task()
-def repeating_group_task(count):
-    if count > 0:
-        repeating_group_task.delay(count - 1, group_id='group-id', execute_inline=True)
+def exception_group_task():
+    raise TestException()
+
+@task(max_retries=100)
+def exception_repeating_group_task(num_of_retries):
+    if exception_repeating_group_task.retry_num == num_of_retries:
+        raise TestException()
+    else:
+        exception_repeating_group_task.retry(execute_inline=True)
 
 @task(max_retries=5)
 def max_retries_group_task():
-    repeating_group_task.delay(3, group_id='group-id', execute_inline=True)
     max_retries_group_task.retry(execute_inline=True)
 
 global_group_mock = Mock()
 
 class WorkerTest(TestCase):
     def setUp(self):
-        settings.FORCE_SERIALIZATION = True
         settings.DEAD_LETTER_MODE = False
 
         self.queue_mock = Mock(autospec=QueueClient)
@@ -44,6 +63,12 @@ def setUp(self):
         factory_mock.create.return_value = self.worker
         settings.WORKER_FACTORY = factory_mock
 
+    def setUpGroupsHandling(self):
+        self.group_set = set()
+        self.group_mock.add.side_effect = lambda task: self.group_set.add('{}-{}'.format(task.id, task.retry_id))
+        self.group_mock.remove.side_effect = lambda task: len(self.group_set) == 0 if self.group_set.discard(
+            '{}-{}'.format(task.id, task.retry_id)) is None else False
+
     def test_worker_execution_no_group(self):
         msg = '{"id": "id-1", "retry": 0, "queue": "default", "maxRetries": 5, "args": [], "func": "eb_sqs.tests.worker.tests_worker.dummy_task", "kwargs": {"msg": "Hello World!"}}'
 
@@ -81,6 +106,7 @@ def test_delay(self):
     def test_delay_inline(self):
         result = self.worker.delay(None, 'queue', dummy_task, [], {'msg': 'Hello World!'}, 5, False, 0, True)
 
+        self.group_mock.add.assert_not_called()
         self.queue_mock.add_message.assert_not_called()
         self.assertEqual(result, 'Hello World!')
 
@@ -89,94 +115,82 @@ def test_delay_with_group(self):
 
         self.group_mock.add.assert_called_once()
 
-    def test_group_callback(self):
+    def test_retry_max_reached_execution(self):
+        with self.assertRaises(MaxRetriesReachedException):
+            max_retries_task.delay(execute_inline=True)
+
+    def test_retry_no_limit(self):
+        retries_task.delay(10, execute_inline=True)
+
+        self.assertEqual(retries_task.retry_num, 10)
+
+    def test_group(self):
         settings.GROUP_CALLBACK_TASK = Mock()
 
         self.worker.delay('group-id', 'queue', dummy_task, [], {'msg': 'Hello World!'}, 5, False, 0, True)
 
-        self.group_mock.remove.assert_called_once()
         settings.GROUP_CALLBACK_TASK.delay.assert_called_once()
-
         settings.GROUP_CALLBACK_TASK = None
 
-    def test_group_callback_string(self):
-        settings.GROUP_CALLBACK_TASK = 'eb_sqs.tests.worker.tests_worker.global_group_mock'
+    def test_group_with_exception(self):
+        settings.GROUP_CALLBACK_TASK = Mock()
+        self.setUpGroupsHandling()
 
-        self.worker.delay('group-id', 'queue', dummy_task, [], {'msg': 'Hello World!'}, 5, False, 0, True)
+        with self.assertRaises(TestException):
+            exception_group_task.delay(group_id='group-id', execute_inline=True)
 
-        self.group_mock.remove.assert_called_once()
-        global_group_mock.delay.assert_called_once()
+        self.assertEqual(len(self.group_set), 0)
+        self.assertEqual(self.group_mock.add.call_count, 1)
+        self.assertEqual(self.group_mock.remove.call_count, 1)
 
+        settings.GROUP_CALLBACK_TASK.delay.assert_called_once()
         settings.GROUP_CALLBACK_TASK = None
 
-    def test_group(self):
+    def test_group_retries(self):
         settings.GROUP_CALLBACK_TASK = Mock()
-
-        group_set = set()
-        self.group_mock.add.side_effect = lambda task: group_set.add('{}-{}'.format(task.id, task.retry_id))
-        self.group_mock.remove.side_effect = lambda task: len(group_set) == 0 if group_set.discard(
-            '{}-{}'.format(task.id, task.retry_id)) is None else False
+        self.setUpGroupsHandling()
 
         repeating_group_task.delay(3, group_id='group-id', execute_inline=True)
 
-        self.assertEqual(len(group_set), 0)
-        settings.GROUP_CALLBACK_TASK.delay.assert_called_once()
+        self.assertEqual(len(self.group_set), 0)
+        self.assertEqual(self.group_mock.add.call_count, 4)
+        self.assertEqual(self.group_mock.remove.call_count, 4)
 
+        settings.GROUP_CALLBACK_TASK.delay.assert_called_once()
         settings.GROUP_CALLBACK_TASK = None
 
-    def test_group_match_retries_reached(self):
+    def test_group_exception_in_retries(self):
         settings.GROUP_CALLBACK_TASK = Mock()
+        self.setUpGroupsHandling()
 
-        group_set = set()
-        self.group_mock.add.side_effect = lambda task: group_set.add('{}-{}'.format(task.id, task.retry_id))
-        self.group_mock.remove.side_effect = lambda task: len(group_set) == 0 if group_set.discard(
-            '{}-{}'.format(task.id, task.retry_id)) is None else False
+        with self.assertRaises(TestException):
+            exception_repeating_group_task.delay(2, group_id='group-id', execute_inline=True)
 
-        with self.assertRaises(MaxRetriesReachedException):
-            max_retries_group_task.delay(group_id='group-id', execute_inline=True)
+        self.assertEqual(len(self.group_set), 0)
+        self.assertEqual(self.group_mock.add.call_count, 3)
+        self.assertEqual(self.group_mock.remove.call_count, 3)
 
-        self.assertEqual(len(group_set), 0)
         settings.GROUP_CALLBACK_TASK.delay.assert_called_once()
-
         settings.GROUP_CALLBACK_TASK = None
 
-    def test_retry_execution(self):
-        task = WorkerTask('id', None, 'queue', dummy_task, [], {'msg': 'Hello World!'}, 5, 0, None, False)
-        self.assertEqual(dummy_task.retry_num, 0)
-
-        self.worker.retry(task, 0, False, True)
-
-        self.queue_mock.add_message.assert_called_once()
-        self.assertEqual(task.id, 'id')
-        self.assertEqual(task.retry, 1)
-        self.assertIsNotNone(task.retry_id)
-
-    def test_retry_max_reached_execution(self):
-        dummy_task.retry_num = 0
+    def test_group_match_retries_reached(self):
+        settings.GROUP_CALLBACK_TASK = Mock()
+        self.setUpGroupsHandling()
 
         with self.assertRaises(MaxRetriesReachedException):
-            task = WorkerTask('id', None, 'queue', dummy_task, [], {'msg': 'Hello World!'}, 2, 0, None, False)
-            self.assertEqual(dummy_task.retry_num, 0)
-
-            self.worker.retry(task, 0, True, True)
-            self.assertEqual(dummy_task.retry_num, 1)
-
-            self.worker.retry(task, 0, True, True)
-            self.assertEqual(dummy_task.retry_num, 2)
-
-            self.worker.retry(task, 0, True, True)
+            max_retries_group_task.delay(group_id='group-id', execute_inline=True)
 
-    def test_retry_no_limit(self):
-        dummy_task.retry_num = 0
+        self.assertEqual(len(self.group_set), 0)
+        self.assertEqual(self.group_mock.add.call_count, 5)
+        self.assertEqual(self.group_mock.remove.call_count, 5)
 
-        task = WorkerTask('id', None, 'queue', dummy_task, [], {'msg': 'Hello World!'}, 2, 0, None, False)
-        self.assertEqual(dummy_task.retry_num, 0)
+        settings.GROUP_CALLBACK_TASK.delay.assert_called_once()
+        settings.GROUP_CALLBACK_TASK = None
 
-        self.worker.retry(task, 0, True, False)
-        self.assertEqual(dummy_task.retry_num, 0)
+    def test_group_callback_string(self):
+        settings.GROUP_CALLBACK_TASK = 'eb_sqs.tests.worker.tests_worker.global_group_mock'
 
-        self.worker.retry(task, 0, True, False)
-        self.assertEqual(dummy_task.retry_num, 0)
+        self.worker.delay('group-id', 'queue', dummy_task, [], {'msg': 'Hello World!'}, 5, False, 0, True)
 
-        self.worker.retry(task, 0, True, False)
-        self.assertEqual(dummy_task.retry_num, 0)
+        global_group_mock.delay.assert_called_once()
+        settings.GROUP_CALLBACK_TASK = None

eb_sqs/worker/worker.py

@@ -81,18 +81,17 @@ def delay(self, group_id, queue_name, func, args, kwargs, max_retries, use_pickl
 
     def retry(self, worker_task, delay, execute_inline, count_retries):
         # type: (WorkerTask, int, bool, bool) -> Any
+        worker_task = worker_task.copy(settings.FORCE_SERIALIZATION)
+        worker_task.retry_id = unicode(uuid.uuid4())
         return self._enqueue_task(worker_task, delay, execute_inline, True, count_retries)
 
     def _enqueue_task(self, worker_task, delay, execute_inline, is_retry, count_retries):
         # type: (WorkerTask, int, bool, bool, bool) -> Any
         try:
-            if is_retry:
-                if count_retries:
-                    worker_task.retry += 1
-                    if worker_task.retry > worker_task.max_retries:
-                        raise MaxRetriesReachedException(worker_task.retry)
-                worker_task.retry_scheduled = True
-                worker_task.retry_id = unicode(uuid.uuid4())
+            if is_retry and count_retries:
+                worker_task.retry += 1
+                if worker_task.retry >= worker_task.max_retries:
+                    raise MaxRetriesReachedException(worker_task.retry)
 
             self._add_to_group(worker_task)
 
@@ -107,22 +106,15 @@ def _enqueue_task(self, worker_task, delay, execute_inline, is_retry, count_retr
                         ', inline' if execute_inline else '')
 
             if execute_inline:
-                if settings.FORCE_SERIALIZATION:
-                    return self._execute_task(WorkerTask.deserialize(worker_task.serialize()))
-                else:
-                    return self._execute_task(worker_task)
+                return self._execute_task(worker_task)
             else:
                 self.queue_client.add_message(worker_task.queue, worker_task.serialize(), delay)
                 return None
-        except MaxRetriesReachedException:
-            self._remove_from_group(worker_task)
-            raise
         except QueueDoesNotExistException as ex:
             self._remove_from_group(worker_task)
             raise InvalidQueueException(ex.queue_name)
         except QueueClientException as ex:
             self._remove_from_group(worker_task)
-
             logger.exception('Task %s (%s, retry-id: %s) failed to enqueue to %s: %s',
                         worker_task.abs_func_name,
                         worker_task.id,
@@ -131,20 +123,18 @@ def _enqueue_task(self, worker_task, delay, execute_inline, is_retry, count_retr
                         ex)
 
             raise QueueException()
-        except Exception:
-            self._remove_from_group(worker_task)
-            raise
 
     def _execute_task(self, worker_task):
         # type: (WorkerTask) -> Any
-        worker_task.retry_scheduled = False
-        result = worker_task.execute()
-        self._remove_from_group(worker_task)
-        return result
+        try:
+            result = worker_task.execute()
+            return result
+        finally:
+            self._remove_from_group(worker_task)
 
     def _add_to_group(self, worker_task):
         # type: (WorkerTask) -> None
-        if worker_task.group_id and not worker_task.retry_scheduled:
+        if worker_task.group_id:
             logger.debug(
                 'Add task %s (%s, retry-id: %s) to group %s',
                 worker_task.abs_func_name,
@@ -157,7 +147,7 @@ def _add_to_group(self, worker_task):
 
     def _remove_from_group(self, worker_task):
         # type: (WorkerTask) -> None
-        if worker_task.group_id and not worker_task.retry_scheduled:
+        if worker_task.group_id:
             logger.debug(
                 'Remove task %s (%s, retry-id: %s) from group %s',
                 worker_task.abs_func_name,

eb_sqs/worker/worker_task.py

@@ -26,7 +26,6 @@ def __init__(self, id, group_id, queue, func, args, kwargs, max_retries, retry,
         self.use_pickle = use_pickle
 
         self.abs_func_name = '{}.{}'.format(self.func.__module__, self.func.func_name)
-        self.retry_scheduled = False
 
     def execute(self):
         # type: () -> Any
@@ -55,6 +54,24 @@ def serialize(self):
 
         return json.dumps(task)
 
+    def copy(self, use_serialization):
+        # type: (bool) -> WorkerTask
+        if use_serialization:
+            return WorkerTask.deserialize(self.serialize())
+        else:
+            return WorkerTask(
+                self.id,
+                self.group_id,
+                self.queue,
+                self.func,
+                self.args,
+                self.kwargs,
+                self.max_retries,
+                self.retry,
+                self.retry_id,
+                self.use_pickle,
+            )
+
     @staticmethod
     def _pickle_args(args):
         # type: (dict) -> unicode