Stop tasks on nodes before deleting task queues (#801)

src/api-service/__app__/onefuzzlib/agent_events.py

@@ -238,6 +238,11 @@ def on_worker_event_done(machine_id: UUID, event: WorkerDoneEvent) -> Result[Non
             node.debug_keep_node = True
             node.save()
 
+    if not node.debug_keep_node:
+        node_task = NodeTasks.get(machine_id, event.task_id)
+        if node_task:
+            node_task.delete()
+
     event.stdout = event.stdout[-MAX_OUTPUT_SIZE:]
     event.stderr = event.stderr[-MAX_OUTPUT_SIZE:]
     task_event = TaskEvent(

src/api-service/__app__/onefuzzlib/tasks/main.py

@@ -26,7 +26,7 @@ from ..azure.storage import StorageType
 from ..events import send_event
 from ..orm import MappingIntStrAny, ORMMixin, QueryFilter
 from ..proxy_forward import ProxyForward
-from ..workers.nodes import Node
+from ..workers.nodes import Node, NodeTasks
 from ..workers.pools import Pool
 from ..workers.scalesets import Scaleset
 
@@ -124,14 +124,18 @@ class Task(BASE_TASK, ORMMixin):
         self.set_state(TaskState.waiting)
 
     def stopping(self) -> None:
-        # TODO: we need to 'unschedule' this task from the existing pools
-        from ..jobs import Job
-
         logging.info("stopping task: %s:%s", self.job_id, self.task_id)
         ProxyForward.remove_forward(self.task_id)
-        delete_queue(str(self.task_id), StorageType.corpus)
         Node.stop_task(self.task_id)
+        if not NodeTasks.get_nodes_by_task_id(self.task_id):
+            self.stopped()
+
+    def stopped(self) -> None:
         self.set_state(TaskState.stopped)
+        delete_queue(str(self.task_id), StorageType.corpus)
+
+        # TODO: we need to 'unschedule' this task from the existing pools
+        from ..jobs import Job
 
         job = Job.get(self.job_id)
         if job:

src/api-service/__app__/onefuzzlib/workers/nodes.py

@@ -239,6 +239,8 @@ class Node(BASE_NODE, ORMMixin):
             task = Task.get_by_task_id(entry.task_id)
             if isinstance(task, Task):
                 task.mark_failed(error)
+            if not self.debug_keep_node:
+                entry.delete()
 
     def could_shrink_scaleset(self) -> bool:
         from .scalesets import ScalesetShrinkQueue