diff --git a/src/api-service/__app__/onefuzzlib/jobs.py b/src/api-service/__app__/onefuzzlib/jobs.py index bdf1d2ee3..993c14e2e 100644 --- a/src/api-service/__app__/onefuzzlib/jobs.py +++ b/src/api-service/__app__/onefuzzlib/jobs.py @@ -85,6 +85,20 @@ class Job(BASE_JOB, ORMMixin): self.state = JobState.enabled self.save() + def stop_if_all_done(self) -> None: + not_stopped = [ + task + for task in Task.search(query={"job_id": [self.job_id]}) + if task.state != TaskState.stopped + ] + if not_stopped: + return + + logging.info( + JOB_LOG_PREFIX + "stopping job as all tasks are stopped: %s", self.job_id + ) + self.stopping() + def stopping(self) -> None: self.state = JobState.stopping logging.info(JOB_LOG_PREFIX + "stopping: %s", self.job_id) diff --git a/src/api-service/__app__/onefuzzlib/tasks/main.py b/src/api-service/__app__/onefuzzlib/tasks/main.py index 753117abc..2a04f3da2 100644 --- a/src/api-service/__app__/onefuzzlib/tasks/main.py +++ b/src/api-service/__app__/onefuzzlib/tasks/main.py @@ -126,6 +126,7 @@ class Task(BASE_TASK, ORMMixin): 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) @@ -133,6 +134,10 @@ class Task(BASE_TASK, ORMMixin): Node.stop_task(self.task_id) self.set_state(TaskState.stopped, send=False) + job = Job.get(self.job_id) + if job: + job.stop_if_all_done() + @classmethod def search_states( cls, *, job_id: Optional[UUID] = None, states: Optional[List[TaskState]] = None