Added refactored workerpool and updated sync.segment

Author: chillaq

splitio/sync/segment.py

@@ -357,12 +357,11 @@ async def synchronize_segments(self, segment_names = None, dont_wait = False):
         if segment_names is None:
             segment_names = await self._feature_flag_storage.get_segment_names()
 
-        for segment_name in segment_names:
-            await self._worker_pool.submit_work(segment_name)
+        jobs = await self._worker_pool.submit_work(segment_names)
         if (dont_wait):
             return True
-        await asyncio.sleep(.5)
-        return not await self._worker_pool.wait_for_completion()
+        await jobs.await_completion()
+        return not self._worker_pool.pop_failed()
 
     async def segment_exist_in_storage(self, segment_name):
         """

splitio/tasks/util/workerpool.py

@@ -7,8 +7,6 @@
 from splitio.optional.loaders import asyncio
 
 _LOGGER = logging.getLogger(__name__)
-_ASYNC_SLEEP_SECONDS = 0.3
-
 
 class WorkerPool(object):
     """Worker pool class to implement single producer/multiple consumer."""
@@ -141,121 +139,94 @@ def _wait_workers_shutdown(self, event):
 class WorkerPoolAsync(object):
     """Worker pool async class to implement single producer/multiple consumer."""
 
+    _abort = object()
+
     def __init__(self, worker_count, worker_func):
         """
         Class constructor.
 
         :param worker_count: Number of workers for the pool.
         :type worker_func: Function to be executed by the workers whenever a messages is fetched.
         """
+        self._semaphore = asyncio.Semaphore(worker_count)
+        self._queue = asyncio.Queue()
+        self._handler = worker_func
+        self._aborted = False
         self._failed = False
-        self._running = False
-        self._incoming = asyncio.Queue()
-        self._worker_count = worker_count
-        self._worker_func = worker_func
-        self.current_workers = []
 
+    async def _schedule_work(self):
+        """wrap the message handler execution."""
+        while True:
+            message = await self._queue.get()
+            if message == self._abort:
+                self._aborted = True
+                return
+            asyncio.get_running_loop().create_task(self._do_work(message))
+
+    async def _do_work(self, message):
+        """process a single message."""
+        try:
+            await self._semaphore.acquire() # wait until "there's a free worker"
+            if self._aborted: # check in case the pool was shutdown while we were waiting for a worker
+                return
+            await self._handler(message._message)
+        except Exception:
+            _LOGGER.error("Something went wrong when processing message %s", message)
+            _LOGGER.debug('Original traceback: ', exc_info=True)
+            self._failed = True
+        message._complete.set()
+        self._semaphore.release() # signal worker is idle
 
     def start(self):
         """Start the workers."""
-        self._running = True
-        self._worker_pool_task = asyncio.get_running_loop().create_task(self._wrapper())
+        self._task = asyncio.get_running_loop().create_task(self._schedule_work())
 
-    async def _safe_run(self, message):
+    async def submit_work(self, jobs):
         """
-        Execute the user funcion for a given message without raising exceptions.
-
-        :param func: User defined function.
-        :type func: callable
-        :param message: Message fetched from the queue.
-        :param message: object
+        Add a new message to the work-queue.
 
-        :return True if no everything goes well. False otherwise.
-        :rtype bool
+        :param message: New message to add.
+        :type message: object.
         """
-        try:
-            await self._worker_func(message)
-            return True
-        except Exception:  # pylint: disable=broad-except
-            _LOGGER.error("Something went wrong when processing message %s", message)
-            _LOGGER.error('Original traceback: ', exc_info=True)
-            return False
+        self.jobs = jobs
+        if len(jobs) == 1:
+            wrapped = TaskCompletionWraper(jobs[0])
+            await self._queue.put(wrapped)
+            return wrapped
 
-    async def _wrapper(self):
-        """
-        Fetch message, execute tasks, and acknowledge results.
+        tasks = [TaskCompletionWraper(job) for job in jobs]
+        for w in tasks:
+            await self._queue.put(w)
 
-        :param worker_number: # (id) of worker whose function will be executed.
-        :type worker_number: int
-        :param func: User defined function.
-        :type func: callable.
-        """
-        self.current_workers = []
-        while self._running:
-            try:
-                if len(self.current_workers) == self._worker_count or self._incoming.qsize() == 0:
-                    await asyncio.sleep(_ASYNC_SLEEP_SECONDS)
-                    self._check_and_clean_workers()
-                    continue
-                message = await self._incoming.get()
-                # For some reason message can be None in python2 implementation of queue.
-                # This method must be both ignored and acknowledged with .task_done()
-                # otherwise .join() will halt.
-                if message is None:
-                    _LOGGER.debug('spurious message received. acking and ignoring.')
-                    continue
+        return BatchCompletionWrapper(tasks)
 
-                # If the task is successfully executed, the ack is done AFTERWARDS,
-                # to avoid race conditions on SDK initialization.
-                _LOGGER.debug("processing message '%s'", message)
-                self.current_workers.append([asyncio.get_running_loop().create_task(self._safe_run(message)), message])
+    async def stop(self, event=None):
+        """abort all execution (except currently running handlers)."""
+        await self._queue.put(self._abort)
 
-                # check tasks status
-                self._check_and_clean_workers()
-            except queue.Empty:
-                # No message was fetched, just keep waiting.
-                pass
+    def pop_failed(self):
+        old = self._failed
+        self._failed = False
+        return old
 
-    def _check_and_clean_workers(self):
-        found_running = False
-        for task in self.current_workers:
-            if task[0].done():
-                self.current_workers.remove(task)
-                if not task[0].result():
-                    self._failed = True
-                    _LOGGER.error(
-                        ("Something went wrong during the execution, "
-                        "removing message \"%s\" from queue.",
-                        task[1])
-                    )
-            else:
-                found_running = True
-        return found_running
 
-    async def submit_work(self, message):
-        """
-        Add a new message to the work-queue.
+class TaskCompletionWraper:
+    """Task completion class"""
+    def __init__(self, message):
+        self._message = message
+        self._complete = asyncio.Event()
 
-        :param message: New message to add.
-        :type message: object.
-        """
-        await self._incoming.put(message)
-        _LOGGER.debug('queued message %s for processing.', message)
+    async def await_completion(self):
+        await self._complete.wait()
 
-    async def wait_for_completion(self):
-        """Block until the work queue is empty."""
-        _LOGGER.debug('waiting for all messages to be processed.')
-        if self._incoming.qsize() > 0:
-            await self._incoming.join()
-        _LOGGER.debug('all messages processed.')
-        old = self._failed
-        self._failed = False
-        self._running = False
-        return old
+    def _mark_as_complete(self):
+        self._complete.set()
 
-    async def stop(self, event=None):
-        """Stop all worker nodes."""
-        await self.wait_for_completion()
-        while self._check_and_clean_workers():
-            await asyncio.sleep(_ASYNC_SLEEP_SECONDS)
-        self._worker_pool_task.cancel()
+
+class BatchCompletionWrapper:
+    """Batch completion class"""
+    def __init__(self, tasks):
+        self._tasks = tasks
+
+    async def await_completion(self):
+        await asyncio.gather(*[task.await_completion() for task in self._tasks])

tests/sync/test_segments_synchronizer.py

@@ -202,17 +202,22 @@ async def get_segment_names():
         split_storage.get_segment_names = get_segment_names
 
         storage = mocker.Mock(spec=SegmentStorage)
-        async def get_change_number():
+        async def get_change_number(*args):
             return -1
         storage.get_change_number = get_change_number
 
+        async def put(*args):
+            pass
+        storage.put = put
+
         api = mocker.Mock()
-        async def run(x):
+        async def run(*args):
             raise APIException("something broke")
         api.fetch_segment = run
 
         segments_synchronizer = SegmentSynchronizerAsync(api, split_storage, storage)
         assert not await segments_synchronizer.synchronize_segments()
+        await segments_synchronizer.shutdown()
 
     @pytest.mark.asyncio
     async def test_synchronize_segments(self, mocker):
@@ -295,6 +300,8 @@ async def fetch_segment_mock(segment_name, change_number, fetch_options):
             assert segment.name in segments_to_validate
             segments_to_validate.remove(segment.name)
 
+        await segments_synchronizer.shutdown()
+
     @pytest.mark.asyncio
     async def test_synchronize_segment(self, mocker):
         """Test particular segment update."""
@@ -339,6 +346,8 @@ async def fetch_segment_mock(segment_name, change_number, fetch_options):
         assert (self.segment[0], self.change[0], self.options[0]) == ('segmentA', -1, FetchOptions(True))
         assert (self.segment[1], self.change[1], self.options[1]) == ('segmentA', 123, FetchOptions(True))
 
+        await segments_synchronizer.shutdown()
+
     @pytest.mark.asyncio
     async def test_synchronize_segment_cdn(self, mocker):
         """Test particular segment update cdn bypass."""
@@ -401,14 +410,18 @@ async def fetch_segment_mock(segment_name, change_number, fetch_options):
         await segments_synchronizer.synchronize_segment('segmentA', 12345)
         assert (self.segment[7], self.change[7], self.options[7]) == ('segmentA', 12345, FetchOptions(True, 1234))
         assert len(self.segment) == 8 # 2 ok + BACKOFF(2 since==till + 2 re-attempts) + CDN(2 since==till)
+        await segments_synchronizer.shutdown()
 
     @pytest.mark.asyncio
     async def test_recreate(self, mocker):
         """Test recreate logic."""
         segments_synchronizer = SegmentSynchronizerAsync(mocker.Mock(), mocker.Mock(), mocker.Mock())
         current_pool = segments_synchronizer._worker_pool
+        await segments_synchronizer.shutdown()
         segments_synchronizer.recreate()
+
         assert segments_synchronizer._worker_pool != current_pool
+        await segments_synchronizer.shutdown()
 
 
 class LocalSegmentsSynchronizerTests(object):

tests/tasks/util/test_workerpool.py

@@ -89,14 +89,16 @@ async def worker_func(num):
 
         wpool = workerpool.WorkerPoolAsync(10, worker_func)
         wpool.start()
+        jobs = []
         for num in range(0, 11):
-            await wpool.submit_work(str(num))
+            jobs.append(str(num))
 
-        await asyncio.sleep(1)
+        task = await wpool.submit_work(jobs)
+        await task.await_completion()
         await wpool.stop()
-        assert wpool._running == False
         for num in range(0, 11):
             assert str(num) in calls
+        assert not wpool.pop_failed()
 
     @pytest.mark.asyncio
     async def test_fail_in_msg_doesnt_break(self):
@@ -114,9 +116,10 @@ async def do_work(self, work):
         wpool = workerpool.WorkerPoolAsync(50, worker.do_work)
         wpool.start()
         for num in range(0, 100):
-            await wpool.submit_work(str(num))
+            await wpool.submit_work([str(num)])
         await asyncio.sleep(1)
         await wpool.stop()
+        assert wpool.pop_failed()
 
         for num in range(0, 100):
             if num != 55:
@@ -138,9 +141,10 @@ async def do_work(self, work):
         worker = Worker()
         wpool = workerpool.WorkerPoolAsync(50, worker.do_work)
         wpool.start()
+        jobs = []
         for num in range(0, 100):
-            await wpool.submit_work(str(num))
-
-        await asyncio.sleep(1)
-        await wpool.wait_for_completion()
+            jobs.append(str(num))
+        task = await wpool.submit_work(jobs)
+        await task.await_completion()
+        await wpool.stop()
         assert len(worker.worked) == 100