diff --git a/labthings/core/tasks/pool.py b/labthings/core/tasks/pool.py index 3fb8e411..04aeaa81 100644 --- a/labthings/core/tasks/pool.py +++ b/labthings/core/tasks/pool.py @@ -1,11 +1,10 @@ -import threading import logging from functools import wraps from gevent import getcurrent from .thread import TaskThread -from flask import copy_current_request_context +from flask import copy_current_request_context, has_request_context class TaskMaster: @@ -38,21 +37,26 @@ def states(self): def new(self, f, *args, **kwargs): # copy_current_request_context allows threads to access flask current_app - task = TaskThread( - target=copy_current_request_context(f), args=args, kwargs=kwargs - ) + if has_request_context(): + target = copy_current_request_context(f) + else: + target = f + task = TaskThread(target=target, args=args, kwargs=kwargs) self._tasks.append(task) return task def remove(self, task_id): for task in self._tasks: - if (task.id == task_id) and not task.isAlive(): - del task + if (str(task.id) == str(task_id)) and task.dead: + self._tasks.remove(task) def cleanup(self): - for task in self._tasks: - if not task.isAlive(): - del task + for i, task in enumerate(self._tasks): + if task.dead: + # Mark for delection + self._tasks[i] = None + # Remove items marked for deletion + self._tasks = [t for t in self._tasks if t] # Task management