Reimplemented tasks to use greenlets

Author: Joel Collins

examples/simple_thing.py

@@ -1,3 +1,11 @@
+#!/usr/bin/env python
+from gevent import monkey
+
+# Patch most system modules. Leave threads untouched so we can still use them normally if needed.
+print("Monkey patching with Gevenet")
+monkey.patch_all(thread=False)
+print("Monkey patching successful")
+
 import random
 import math
 import time
@@ -13,7 +21,7 @@
 from labthings.server.view import View
 from labthings.server.find import find_component
 from labthings.server import fields
-from labthings.core.tasks import taskify
+from labthings.core.tasks import taskify, update_task_data
 
 
 """
@@ -22,6 +30,14 @@
 """
 
 
+from gevent.monkey import get_original
+
+get_ident = get_original("_thread", "get_ident")
+
+print(f"ROOT IDENT")
+print(get_ident())
+
+
 class MyComponent:
     def __init__(self):
         self.x_range = range(-100, 100)
@@ -50,6 +66,7 @@ def average_data(self, n: int):
 
         for i in range(n):
             summed_data = [summed_data[i] + el for i, el in enumerate(self.data)]
+            update_task_data({"data": summed_data})
             time.sleep(0.25)
 
         summed_data = [i / n for i in summed_data]

labthings/core/tasks/pool.py

@@ -1,6 +1,7 @@
 import threading
 import logging
 from functools import wraps
+from gevent import getcurrent
 
 from .thread import TaskThread
 
@@ -114,7 +115,7 @@ def current_task():
     Returns:
         TaskThread -- Currently running Task thread.
     """
-    current_task_thread = threading.current_thread()
+    current_task_thread = getcurrent()
     if not isinstance(current_task_thread, TaskThread):
         return None
     return current_task_thread

labthings/core/tasks/thread.py

@@ -1,34 +1,23 @@
-import ctypes
+from gevent import Greenlet
 import datetime
 import logging
 import traceback
 import uuid
 
-from gevent.monkey import get_original
-
-# Guarantee that Task threads will always be proper system threads, regardless of Gevent patches
-Thread = get_original("threading", "Thread")
-Event = get_original("threading", "Event")
-Lock = get_original("threading", "Lock")
-
 _LOG = logging.getLogger(__name__)
 
 
 class ThreadTerminationError(SystemExit):
     """Sibling of SystemExit, but specific to thread termination."""
 
 
-class TaskThread(Thread):
-    def __init__(self, target=None, name=None, args=None, kwargs=None, daemon=True):
-        Thread.__init__(
-            self,
-            group=None,
-            target=target,
-            name=name,
-            args=args,
-            kwargs=kwargs,
-            daemon=daemon,
-        )
+class TaskKillException(Exception):
+    """Sibling of SystemExit, but specific to thread termination."""
+
+
+class TaskThread(Greenlet):
+    def __init__(self, target=None, args=None, kwargs=None):
+        Greenlet.__init__(self)
         # Handle arguments
         if args is None:
             args = ()
@@ -56,10 +45,6 @@ def __init__(self, target=None, name=None, args=None, kwargs=None, daemon=True):
         self.progress: int = None  # Percent progress of the task
         self.data = {}  # Dictionary of custom data added during the task
 
-        # Stuff for handling termination
-        self._running_lock = Lock()  # Lock obtained while self._target is running
-        self._killed = Event()  # Event triggered when thread is manually terminated
-
     @property
     def id(self):
         """Return ID of current TaskThread"""
@@ -86,6 +71,9 @@ def update_data(self, data: dict):
         # Store data to be used before task finishes (eg for real-time plotting)
         self.data.update(data)
 
+    def _run(self):
+        return self._thread_proc(self._target)(*self._args, **self._kwargs)
+
     def _thread_proc(self, f):
         """
         Wraps the target function to handle recording `status` and `return` to `state`.
@@ -110,95 +98,12 @@ def wrapped(*args, **kwargs):
 
         return wrapped
 
-    def run(self):
-        """Overrides default threading.Thread run() method"""
-        logging.debug((self._args, self._kwargs))
-        try:
-            with self._running_lock:
-                if self._killed.is_set():
-                    raise ThreadTerminationError()
-                if self._target:
-                    self._thread_proc(self._target)(*self._args, **self._kwargs)
-        finally:
-            # Avoid a refcycle if the thread is running a function with
-            # an argument that has a member that points to the thread.
-            del self._target, self._args, self._kwargs
-
-    def wait(self):
-        """Start waiting for the task to finish before returning"""
-        print("Joining thread {}".format(self))
-        self.join()
-        return self._return_value
-
-    def async_raise(self, exc_type):
-        """Raise an exception in this thread."""
-        # Should only be called on a started thread, so raise otherwise.
-        if self.ident is None:
-            raise RuntimeError(
-                "Cannot halt a thread that hasn't started. "
-                "No valid running thread identifier."
-            )
-
-        # If the thread has died we don't want to raise an exception so log.
-        if not self.is_alive():
-            _LOG.debug(
-                "Not raising %s because thread %s (%s) is not alive",
-                exc_type,
-                self.name,
-                self.ident,
-            )
-            return
-
-        result = ctypes.pythonapi.PyThreadState_SetAsyncExc(
-            ctypes.c_long(self.ident), ctypes.py_object(exc_type)
-        )
-        if result == 0 and self.is_alive():
-            # Don't raise an exception an error unnecessarily if the thread is dead.
-            raise ValueError("Thread ID was invalid.", self.ident)
-        elif result > 1:
-            # Something bad happened, call with a NULL exception to undo.
-            ctypes.pythonapi.PyThreadState_SetAsyncExc(self.ident, None)
-            raise RuntimeError(
-                "Error: PyThreadState_SetAsyncExc %s %s (%s) %s"
-                % (exc_type, self.name, self.ident, result)
-            )
-
-    def _is_thread_proc_running(self):
-        """
-        Test if thread funtion (_thread_proc) is running,
-        by attemtping to acquire the lock _thread_proc acquires at runtime.
-        Returns:
-            bool: If _thread_proc is currently running
-        """
-        could_acquire = self._running_lock.acquire(False)  # skipcq: PYL-E1111
-        if could_acquire:
-            self._running_lock.release()
-            return False
-        return True
-
-    def terminate(self):
-        """
-        Raise ThreadTerminatedException in the context of the given thread,
-        which should cause the thread to exit silently.
-        """
-        _LOG.warning(f"Terminating thread {self}")
-        self._killed.set()
-        if not self.is_alive():
-            logging.debug("Cannot kill thread that is no longer running.")
-            return
-        if not self._is_thread_proc_running():
-            logging.debug(
-                "Thread's _thread_proc function is no longer running, "
-                "will not kill; letting thread exit gracefully."
-            )
-            return
-        self.async_raise(ThreadTerminationError)
-
-        # Wait for the thread for finish closing. If the threaded function has cleanup code in a try-except,
-        # this pause allows it to finish running before the main process can continue.
-        while self._is_thread_proc_running():
-            pass
-
+    def kill(self, exception=TaskKillException, block=True, timeout=None):
+        # Kill the greenlet
+        Greenlet.kill(self, exception=exception, block=block, timeout=timeout)
         # Set state to terminated
         self._status = "terminated"
         self.progress = None
+
+    def terminate(self):
+        return self.kill()

labthings/server/views/tasks.py

@@ -31,11 +31,13 @@ def get(self, task_id):
 
         Includes progress and intermediate data.
         """
+        task_dict = tasks.dictionary()
 
-        task = tasks.dictionary().get(task_id)
-        if not task:
+        if not task_id in task_dict:
             return abort(404)  # 404 Not Found
 
+        task = task_dict.get(task_id)
+
         return task
 
     @marshal_with(TaskSchema())
@@ -45,11 +47,13 @@ def delete(self, task_id):
 
         If the task is finished, deletes its entry.
         """
+        task_dict = tasks.dictionary()
 
-        task = tasks.dictionary().get(task_id)
-        if not task:
+        if not task_id in task_dict:
             return abort(404)  # 404 Not Found
 
-        task.terminate()
+        task = task_dict.get(task_id)
+
+        task.kill(block=True, timeout=3)
 
         return task