Fixed cleanup and context bugs

Author: Joel Collins

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