Fix: Task is removed from group on retry

Author: Schweigi

eb_sqs/redis/redis_group_client.py

@@ -17,11 +17,20 @@ def _key_name(self, group_id):
         # type: (unicode) -> None
         return '{}{}'.format(settings.REDIS_KEY_PREFIX, group_id)
 
+    def _task_identifier(self, worker_task):
+        # type: (WorkerTask) -> unicode
+        if worker_task.retry_id:
+            return '{}-{}'.format(worker_task.id, worker_task.retry_id)
+        else:
+            return worker_task.id
+
     def add(self, worker_task):
         # type: (WorkerTask) -> None
         name = self._key_name(worker_task.group_id)
+        value = self._task_identifier(worker_task)
+
         pipe = self._redis_client.pipeline()
-        pipe.sadd(name, worker_task.id)\
+        pipe.sadd(name, value)\
             .expire(name, settings.REDIS_EXPIRY)\
             .execute()
 
@@ -31,7 +40,9 @@ def remove(self, worker_task):
         :return: True if last task in group
         """
         name = self._key_name(worker_task.group_id)
-        if self._redis_client.srem(name, worker_task.id) > 0:
+        value = self._task_identifier(worker_task)
+
+        if self._redis_client.srem(name, value) > 0:
             return self._redis_client.scard(name) == 0
         else:
             return False

eb_sqs/tests/worker/tests_worker.py

@@ -113,11 +113,13 @@ def test_group(self):
         settings.GROUP_CALLBACK_TASK = Mock()
 
         group_set = set()
-        self.group_mock.add.side_effect = lambda task: group_set.add(task.id)
-        self.group_mock.remove.side_effect = lambda task: len(group_set) == 0 if group_set.discard(task.id) is None else False
+        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
 
         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()
 
         settings.GROUP_CALLBACK_TASK = None
@@ -126,30 +128,34 @@ def test_group_match_retries_reached(self):
         settings.GROUP_CALLBACK_TASK = Mock()
 
         group_set = set()
-        self.group_mock.add.side_effect = lambda task: group_set.add(task.id)
+        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(
-            task.id) is None else False
+            '{}-{}'.format(task.id, task.retry_id)) is None else False
 
         with self.assertRaises(MaxRetriesReachedException):
             max_retries_group_task.delay(group_id='group-id', execute_inline=True)
 
+        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, False)
+        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
 
         with self.assertRaises(MaxRetriesReachedException):
-            task = WorkerTask('id', None, 'queue', dummy_task, [], {'msg': 'Hello World!'}, 2, 0, False)
+            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)
@@ -163,7 +169,7 @@ def test_retry_max_reached_execution(self):
     def test_retry_no_limit(self):
         dummy_task.retry_num = 0
 
-        task = WorkerTask('id', None, 'queue', dummy_task, [], {'msg': 'Hello World!'}, 2, 0, False)
+        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, False)

eb_sqs/tests/worker/tests_worker_task.py

@@ -18,10 +18,10 @@ def dummy_function():
 
 class WorkerTaskTest(TestCase):
     def setUp(self):
-        self.dummy_msg = '{"queue": "default", "retry": 0, "func": "eb_sqs.tests.worker.tests_worker_task.dummy_function", "kwargs": {}, "maxRetries": 5, "args": [], "pickle": false, "id": "id-1", "groupId": "group-5"}'
+        self.dummy_msg = '{"queue": "default", "retryId": "retry-uuid", "retry": 0, "func": "eb_sqs.tests.worker.tests_worker_task.dummy_function", "kwargs": {}, "maxRetries": 5, "args": [], "pickle": false, "id": "id-1", "groupId": "group-5"}'
 
     def test_serialize_worker_task(self):
-        worker_task = WorkerTask('id-1', 'group-5', 'default', dummy_function, [], {}, 5, 0, False)
+        worker_task = WorkerTask('id-1', 'group-5', 'default', dummy_function, [], {}, 5, 0, 'retry-uuid', False)
         msg = worker_task.serialize()
 
         self.assertEqual(msg, self.dummy_msg)
@@ -37,9 +37,10 @@ def test_deserialize_worker_task(self):
         self.assertEqual(worker_task.kwargs, {})
         self.assertEqual(worker_task.max_retries, 5)
         self.assertEqual(worker_task.retry, 0)
+        self.assertEqual(worker_task.retry_id, 'retry-uuid')
 
     def test_serialize_pickle(self):
-        worker_task1 = WorkerTask('id-1', None, 'default', dummy_function, [], {'object': TestObject()}, 5, 0, True)
+        worker_task1 = WorkerTask('id-1', None, 'default', dummy_function, [], {'object': TestObject()}, 5, 0, None, True)
         msg = worker_task1.serialize()
 
         worker_task2 = WorkerTask.deserialize(msg)

eb_sqs/worker/worker.py

@@ -73,7 +73,7 @@ def execute(self, msg):
     def delay(self, group_id, queue_name, func, args, kwargs, max_retries, use_pickle, delay, execute_inline):
         # type: (unicode, unicode, Any, tuple, dict, int, bool, int, bool) -> Any
         id = unicode(uuid.uuid4())
-        worker_task = WorkerTask(id, group_id, queue_name, func, args, kwargs, max_retries, 0, use_pickle)
+        worker_task = WorkerTask(id, group_id, queue_name, func, args, kwargs, max_retries, 0, None, use_pickle)
         return self._enqueue_task(worker_task, delay, execute_inline, False, True)
 
     def retry(self, worker_task, delay, execute_inline, count_retries):
@@ -87,9 +87,9 @@ def _enqueue_task(self, worker_task, delay, execute_inline, is_retry, count_retr
                 if count_retries:
                     worker_task.retry += 1
                     if worker_task.retry > worker_task.max_retries:
-                        self._remove_from_group(worker_task)
                         raise MaxRetriesReachedException(worker_task.retry)
                 worker_task.retry_scheduled = True
+                worker_task.retry_id = unicode(uuid.uuid4())
 
             self._add_to_group(worker_task)
 
@@ -111,6 +111,7 @@ def _enqueue_task(self, worker_task, delay, execute_inline, is_retry, count_retr
                 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)
@@ -140,9 +141,10 @@ def _add_to_group(self, worker_task):
         # type: (WorkerTask) -> None
         if worker_task.group_id and not worker_task.retry_scheduled:
             logger.debug(
-                'Add task %s (%s) to group %s',
+                'Add task %s (%s, retry-id: %s) to group %s',
                 worker_task.abs_func_name,
                 worker_task.id,
+                worker_task.retry_id,
                 worker_task.group_id,
             )
 
@@ -152,9 +154,10 @@ def _remove_from_group(self, worker_task):
         # type: (WorkerTask) -> None
         if worker_task.group_id and not worker_task.retry_scheduled:
             logger.debug(
-                'Remove task %s (%s) from group %s',
+                'Remove task %s (%s, retry-id: %s) from group %s',
                 worker_task.abs_func_name,
                 worker_task.id,
+                worker_task.retry_id,
                 worker_task.group_id,
             )
 

eb_sqs/worker/worker_task.py

@@ -11,8 +11,8 @@
 
 
 class WorkerTask(object):
-    def __init__(self, id, group_id, queue, func, args, kwargs, max_retries, retry, use_pickle):
-        # type: (unicode, unicode, unicode, Any, tuple, dict, int, int, bool) -> None
+    def __init__(self, id, group_id, queue, func, args, kwargs, max_retries, retry, retry_id, use_pickle):
+        # type: (unicode, unicode, unicode, Any, tuple, dict, int, int, unicode, bool) -> None
         super(WorkerTask, self).__init__()
         self.id = id
         self.group_id = group_id
@@ -22,6 +22,7 @@ def __init__(self, id, group_id, queue, func, args, kwargs, max_retries, retry,
         self.kwargs = kwargs
         self.max_retries = max_retries
         self.retry = retry
+        self.retry_id = retry_id
         self.use_pickle = use_pickle
 
         self.abs_func_name = '{}.{}'.format(self.func.__module__, self.func.func_name)
@@ -48,6 +49,7 @@ def serialize(self):
                 'kwargs': kwargs,
                 'maxRetries': self.max_retries,
                 'retry': self.retry,
+                'retryId': self.retry_id,
                 'pickle': self.use_pickle,
             }
 
@@ -79,8 +81,9 @@ def deserialize(msg):
         kwargs = WorkerTask._unpickle_args(task['kwargs']) if use_pickle else task['kwargs']
         max_retries = task['maxRetries']
         retry = task['retry']
+        retry_id = task.get('retryId')
 
-        return WorkerTask(id, group_id, queue, func, args, kwargs, max_retries, retry, use_pickle)
+        return WorkerTask(id, group_id, queue, func, args, kwargs, max_retries, retry, retry_id, use_pickle)
 
     @staticmethod
     def _unpickle_args(args):