Properly NACK failed tasks

servicelayer/taskqueue.py

@@ -310,6 +310,19 @@ def mark_done(self, task: Task):
             # delete stages key
             pipe.delete(self.active_stages_key)
 
+    def mark_for_retry(self, task):
+        pipe = self.conn.pipeline()
+        stage_key = self.get_stage_key(stage)
+
+        log.info(f"Marking task {task.task_id} for retry after NACK")
+
+        pipe.sadd(make_key(stage_key, "pending"), task.task_id)
+        pipe.srem(make_key(stage_key, "running"), task.task_id)
+        pipe.delete(task.retry_key)
+        pipe.srem(stage_key, task.task_id)
+
+        pipe.set(self.last_update_key, pack_now())
+
     def is_done(self):
         status = self.get_status()
         return status["pending"] == 0 and status["running"] == 0
@@ -455,8 +468,10 @@ def process_blocking(self):
             try:
                 (task, channel, connection) = self.local_queue.get(timeout=TIMEOUT)
                 apply_task_context(task, v=self.version)
-                self.handle(task, channel)
-                cb = functools.partial(self.ack_message, task, channel)
+                if self.handle(task, channel) is True:
+                    cb = functools.partial(self.ack_message, task, channel)
+                else:
+                    cb = functools.partial(self.nack_message, task, channel)
                 connection.add_callback_threadsafe(cb)
             except Empty:
                 pass
@@ -480,16 +495,20 @@ def process_nonblocking(self):
                 else:
                     queue_active[queue] = True
                     task = get_task(body, method.delivery_tag)
-                    self.handle(task, channel)
+                    if self.handle(task, channel) is True:
+                        channel.basic_ack(task.delivery_tag)
+                    else:
+                        channel.basic_nack(task.delivery_tag)
 
     def process(self, blocking=True):
         if blocking:
             self.process_blocking()
         else:
             self.process_nonblocking()
 
-    def handle(self, task: Task, channel):
+    def handle(self, task: Task, channel) -> bool:
         """Execute a task."""
+        success = True
         try:
             dataset = Dataset(
                 conn=self.conn, name=dataset_from_collection_id(task.collection_id)
@@ -531,7 +550,7 @@ def handle(self, task: Task, channel):
                 log.info(
                     f"Sending a NACK for message {task.delivery_tag}"
                     f" for task_id {task.task_id}."
-                    f"Message will be requeued."
+                    f" Message will be requeued."
                 )
                 # In this case, a task ID was found neither in the
                 # list of Pending, nor the list of Running tasks
@@ -541,12 +560,13 @@ def handle(self, task: Task, channel):
                     retries=0,
                     failed_permanently=True,
                 ).inc()
-                if channel.is_open:
-                    channel.basic_nack(task.delivery_tag)
+                success = False
         except Exception:
             log.exception("Error in task handling")
+            success = False
         finally:
             self.after_task(task)
+        return success
 
     @abstractmethod
     def dispatch_task(self, task: Task) -> Task:
@@ -584,6 +604,17 @@ def ack_message(self, task, channel):
                 channel.basic_ack(task.delivery_tag)
         clear_contextvars()
 
+    def nack_message(self, task, channel):
+        """NACK task and update status."""
+        apply_task_context(task, v=self.version)
+        log.info(f"NACKing message {task.delivery_tag} for task_id {task.task_id}")
+        dataset = task.get_dataset(conn=self.conn)
+        # Sync state to redis
+        dataset.mark_for_retry(task)
+        if channel.is_open:
+            channel.basic_nack(task.delivery_tag)
+        clear_contextvars()
+
     def run(self):
         """Run a blocking worker instance"""
 

tests/test_taskqueue.py

@@ -142,6 +142,10 @@ def did_nack():
 
         dataset = Dataset(conn=conn, name=dataset_from_collection_id(collection_id))
         dataset.add_task(task_id, "test-op")
+        status = dataset.get_active_dataset_status(conn=conn)
+        stage = status["datasets"]["2"]["stages"][0]
+        assert stage["pending"] == 1
+        assert stage["running"] == 0
 
         worker = CountingWorker(queues=[test_queue_name], conn=conn, num_threads=1)
         assert not dataset.should_execute(task_id=task_id)
@@ -160,3 +164,8 @@ def did_nack():
                 dispatch_fn.assert_not_called()
 
         channel.close()
+
+        status = dataset.get_active_dataset_status(conn=conn)
+        stage = status["datasets"]["2"]["stages"][0]
+        assert stage["pending"] == 1
+        assert stage["running"] == 0