Merge pull request #427 from splitio/async-tasks-asynctask

added asynctask async class

Author: chillaq

splitio/tasks/util/asynctask.py

@@ -2,14 +2,13 @@
 import threading
 import logging
 import queue
-
+from splitio.optional.loaders import asyncio
 
 __TASK_STOP__ = 0
 __TASK_FORCE_RUN__ = 1
 
 _LOGGER = logging.getLogger(__name__)
 
-
 def _safe_run(func):
     """
     Execute a function wrapped in a try-except block.
@@ -30,6 +29,26 @@ def _safe_run(func):
         _LOGGER.debug('Original traceback:', exc_info=True)
         return False
 
+async def _safe_run_async(func):
+    """
+    Execute a function wrapped in a try-except block.
+
+    If anything goes wrong returns false instead of propagating the exception.
+
+    :param func: Function to be executed, receives no arguments and it's return
+        value is ignored.
+    """
+    try:
+        await func()
+        return True
+    except Exception:  # pylint: disable=broad-except
+        # Catch any exception that might happen to avoid the periodic task
+        # from ending and allowing for a recovery, as well as preventing
+        # an exception from propagating and breaking the main thread
+        _LOGGER.error('Something went wrong when running passed function.')
+        _LOGGER.debug('Original traceback:', exc_info=True)
+        return False
+
 
 class AsyncTask(object):  # pylint: disable=too-many-instance-attributes
     """
@@ -166,3 +185,133 @@ def force_execution(self):
     def running(self):
         """Return whether the task is running or not."""
         return self._running
+
+
+class AsyncTaskAsync(object):  # pylint: disable=too-many-instance-attributes
+    """
+    Asyncrhonous controllable task async class.
+
+    This class creates is used to wrap around a function to treat it as a
+    periodic task. This task can be stopped, it's execution can be forced, and
+    it's status (whether it's running or not) can be obtained from the task
+    object.
+    It also allows for "on init" and "on stop" functions to be passed.
+    """
+
+
+    def __init__(self, main, period, on_init=None, on_stop=None):
+        """
+        Class constructor.
+
+        :param main: Main function to be executed periodically
+        :type main: callable
+        :param period: How many seconds to wait between executions
+        :type period: int
+        :param on_init: Function to be executed ONCE before the main one
+        :type on_init: callable
+        :param on_stop: Function to be executed ONCE after the task has finished
+        :type on_stop: callable
+        """
+        self._on_init = on_init
+        self._main = main
+        self._on_stop = on_stop
+        self._period = period
+        self._messages = asyncio.Queue()
+        self._running = False
+        self._completion_event = None
+
+    async def _execution_wrapper(self):
+        """
+        Execute user defined function in separate thread.
+
+        It will execute the "on init" hook is available. If an exception is
+        raised it will abort execution, otherwise it will enter an infinite
+        loop in which the main function is executed every <period> seconds.
+        After stop has been called the "on stop" hook will be invoked if
+        available.
+
+        All custom functions are run within a _safe_run() function which
+        prevents exceptions from being propagated.
+        """
+        try:
+            if self._on_init is not None:
+                if not await _safe_run_async(self._on_init):
+                    _LOGGER.error("Error running task initialization function, aborting execution")
+                    self._running = False
+                    return
+            self._running = True
+
+            while self._running:
+                try:
+                    msg = self._messages.get_nowait()
+                    if msg == __TASK_STOP__:
+                        _LOGGER.debug("Stop signal received. finishing task execution")
+                        break
+                    elif msg == __TASK_FORCE_RUN__:
+                        _LOGGER.debug("Force execution signal received. Running now")
+                        if not await _safe_run_async(self._main):
+                            _LOGGER.error("An error occurred when executing the task. "
+                                          "Retrying after perio expires")
+                        continue
+                except asyncio.QueueEmpty:
+                    # If no message was received, the timeout has expired
+                    # and we're ready for a new execution
+                    pass
+                except asyncio.CancelledError:
+                    break
+
+                await asyncio.sleep(self._period)
+                if not await _safe_run_async(self._main):
+                    _LOGGER.error(
+                        "An error occurred when executing the task. "
+                        "Retrying after period expires"
+                    )
+        finally:
+            await self._cleanup()
+
+    async def _cleanup(self):
+        """Execute on_stop callback, set event if needed, update status."""
+        if self._on_stop is not None:
+            if not await _safe_run_async(self._on_stop):
+                _LOGGER.error("An error occurred when executing the task's OnStop hook. ")
+
+        self._running = False
+        self._completion_event.set()
+
+    def start(self):
+        """Start the async task."""
+        if self._running:
+            _LOGGER.warning("Task is already running. Ignoring .start() call")
+            return
+        # Start execution
+        self._completion_event = asyncio.Event()
+        asyncio.get_running_loop().create_task(self._execution_wrapper())
+
+    async def stop(self, wait_for_completion=False):
+        """
+        Send a signal to the thread in order to stop it. If the task is not running do nothing.
+
+        Optionally accept an event to be set upon task completion.
+
+        :param event: Event to set when the task completes.
+        :type event: threading.Event
+        """
+        if not self._running:
+            return
+
+        # Queue is of infinite size, should not raise an exception
+        self._messages.put_nowait(__TASK_STOP__)
+
+        if wait_for_completion:
+            await self._completion_event.wait()
+
+    def force_execution(self):
+        """Force an execution of the task without waiting for the period to end."""
+        if not self._running:
+            return
+        # Queue is of infinite size, should not raise an exception
+        self._messages.put_nowait(__TASK_FORCE_RUN__)
+
+    def running(self):
+        """Return whether the task is running or not."""
+        return self._running

tests/tasks/util/test_asynctask.py

@@ -2,8 +2,10 @@
 
 import time
 import threading
-from splitio.tasks.util import asynctask
+import pytest
 
+from splitio.tasks.util import asynctask
+from splitio.optional.loaders import asyncio
 
 class AsyncTaskTests(object):
     """AsyncTask test cases."""
@@ -116,3 +118,140 @@ def test_force_run(self, mocker):
         assert on_stop.mock_calls == [mocker.call()]
         assert len(main_func.mock_calls) == 2
         assert not task.running()
+
+
+class AsyncTaskAsyncTests(object):
+    """AsyncTask test cases."""
+
+    @pytest.mark.asyncio
+    async def test_default_task_flow(self, mocker):
+        """Test the default execution flow of an asynctask."""
+        self.main_called = 0
+        async def main_func():
+            self.main_called += 1
+
+        self.init_called = 0
+        async def on_init():
+            self.init_called += 1
+
+        self.stop_called = 0
+        async def on_stop():
+            self.stop_called += 1
+
+        task = asynctask.AsyncTaskAsync(main_func, 0.5, on_init, on_stop)
+        task.start()
+        await asyncio.sleep(1)
+        assert task.running()
+        await task.stop(True)
+
+        assert 0 < self.main_called <= 2
+        assert self.init_called == 1
+        assert self.stop_called == 1
+        assert not task.running()
+
+    @pytest.mark.asyncio
+    async def test_main_exception_skips_iteration(self, mocker):
+        """Test that an exception in the main func only skips current iteration."""
+        self.main_called = 0
+        async def raise_exception():
+            self.main_called += 1
+            raise Exception('something')
+        main_func = raise_exception
+
+        self.init_called = 0
+        async def on_init():
+            self.init_called += 1
+
+        self.stop_called = 0
+        async def on_stop():
+            self.stop_called += 1
+
+        task = asynctask.AsyncTaskAsync(main_func, 0.1, on_init, on_stop)
+        task.start()
+        await asyncio.sleep(1)
+        assert task.running()
+        await task.stop(True)
+
+        assert 9 <= self.main_called <= 10
+        assert self.init_called == 1
+        assert self.stop_called == 1
+        assert not task.running()
+
+    @pytest.mark.asyncio
+    async def test_on_init_failure_aborts_task(self, mocker):
+        """Test that if the on_init callback fails, the task never runs."""
+        self.main_called = 0
+        async def main_func():
+            self.main_called += 1
+
+        self.init_called = 0
+        async def on_init():
+            self.init_called += 1
+            raise Exception('something')
+
+        self.stop_called = 0
+        async def on_stop():
+            self.stop_called += 1
+
+        task = asynctask.AsyncTaskAsync(main_func, 0.1, on_init, on_stop)
+        task.start()
+        await asyncio.sleep(0.5)
+        assert not task.running() # Since on_init fails, task never starts
+        await task.stop(True)
+
+        assert self.init_called == 1
+        assert self.stop_called == 1
+        assert self.main_called == 0
+        assert not task.running()
+
+    @pytest.mark.asyncio
+    async def test_on_stop_failure_ends_gacefully(self, mocker):
+        """Test that if the on_init callback fails, the task never runs."""
+        self.main_called = 0
+        async def main_func():
+            self.main_called += 1
+
+        self.init_called = 0
+        async def on_init():
+            self.init_called += 1
+
+        self.stop_called = 0
+        async def on_stop():
+            self.stop_called += 1
+            raise Exception('something')
+
+        task = asynctask.AsyncTaskAsync(main_func, 0.1, on_init, on_stop)
+        task.start()
+        await asyncio.sleep(1)
+        await task.stop(True)
+        assert 9 <= self.main_called <= 10
+        assert self.init_called == 1
+        assert self.stop_called == 1
+
+    @pytest.mark.asyncio
+    async def test_force_run(self, mocker):
+        """Test that if the on_init callback fails, the task never runs."""
+        self.main_called = 0
+        async def main_func():
+            self.main_called += 1
+
+        self.init_called = 0
+        async def on_init():
+            self.init_called += 1
+
+        self.stop_called = 0
+        async def on_stop():
+            self.stop_called += 1
+
+        task = asynctask.AsyncTaskAsync(main_func, 5, on_init, on_stop)
+        task.start()
+        await asyncio.sleep(1)
+        assert task.running()
+        task.force_execution()
+        task.force_execution()
+        await task.stop(True)
+
+        assert self.main_called == 3
+        assert self.init_called == 1
+        assert self.stop_called == 1
+        assert not task.running()