Merge pull request #424 from splitio/async-sync-segment-workerpool

Added workerpool and sync.segment async classes

Author: chillaq

splitio/sync/segment.py

@@ -10,34 +10,36 @@
 from splitio.util.backoff import Backoff
 from splitio.optional.loaders import asyncio, aiofiles
 from splitio.sync import util
+from splitio.optional.loaders import asyncio
 
 _LOGGER = logging.getLogger(__name__)
 
 
 _ON_DEMAND_FETCH_BACKOFF_BASE = 10  # backoff base starting at 10 seconds
 _ON_DEMAND_FETCH_BACKOFF_MAX_WAIT = 60  # don't sleep for more than 1 minute
 _ON_DEMAND_FETCH_BACKOFF_MAX_RETRIES = 10
+_MAX_WORKERS = 10
 
 
 class SegmentSynchronizer(object):
-    def __init__(self, segment_api, split_storage, segment_storage):
+    def __init__(self, segment_api, feature_flag_storage, segment_storage):
         """
         Class constructor.
 
         :param segment_api: API to retrieve segments from backend.
         :type segment_api: splitio.api.SegmentApi
 
-        :param split_storage: Feature Flag Storage.
-        :type split_storage: splitio.storage.InMemorySplitStorage
+        :param feature_flag_storage: Feature Flag Storage.
+        :type feature_flag_storage: splitio.storage.InMemorySplitStorage
 
         :param segment_storage: Segment storage reference.
         :type segment_storage: splitio.storage.SegmentStorage
 
         """
         self._api = segment_api
-        self._split_storage = split_storage
+        self._feature_flag_storage = feature_flag_storage
         self._segment_storage = segment_storage
-        self._worker_pool = workerpool.WorkerPool(10, self.synchronize_segment)
+        self._worker_pool = workerpool.WorkerPool(_MAX_WORKERS, self.synchronize_segment)
         self._worker_pool.start()
         self._backoff = Backoff(
                                 _ON_DEMAND_FETCH_BACKOFF_BASE,
@@ -48,7 +50,7 @@ def recreate(self):
         Create worker_pool on forked processes.
 
         """
-        self._worker_pool = workerpool.WorkerPool(10, self.synchronize_segment)
+        self._worker_pool = workerpool.WorkerPool(_MAX_WORKERS, self.synchronize_segment)
         self._worker_pool.start()
 
     def shutdown(self):
@@ -176,7 +178,7 @@ def synchronize_segments(self, segment_names = None, dont_wait = False):
         :rtype: bool
         """
         if segment_names is None:
-            segment_names = self._split_storage.get_segment_names()
+            segment_names = self._feature_flag_storage.get_segment_names()
 
         for segment_name in segment_names:
             self._worker_pool.submit_work(segment_name)
@@ -196,6 +198,184 @@ def segment_exist_in_storage(self, segment_name):
         """
         return self._segment_storage.get(segment_name) != None
 
+
+class SegmentSynchronizerAsync(object):
+    def __init__(self, segment_api, feature_flag_storage, segment_storage):
+        """
+        Class constructor.
+
+        :param segment_api: API to retrieve segments from backend.
+        :type segment_api: splitio.api.SegmentApi
+
+        :param feature_flag_storage: Feature Flag Storage.
+        :type feature_flag_storage: splitio.storage.InMemorySplitStorage
+
+        :param segment_storage: Segment storage reference.
+        :type segment_storage: splitio.storage.SegmentStorage
+
+        """
+        self._api = segment_api
+        self._feature_flag_storage = feature_flag_storage
+        self._segment_storage = segment_storage
+        self._worker_pool = workerpool.WorkerPoolAsync(_MAX_WORKERS, self.synchronize_segment)
+        self._worker_pool.start()
+        self._backoff = Backoff(
+                                _ON_DEMAND_FETCH_BACKOFF_BASE,
+                                _ON_DEMAND_FETCH_BACKOFF_MAX_WAIT)
+
+    def recreate(self):
+        """
+        Create worker_pool on forked processes.
+
+        """
+        self._worker_pool = workerpool.WorkerPoolAsync(_MAX_WORKERS, self.synchronize_segment)
+        self._worker_pool.start()
+
+    async def shutdown(self):
+        """
+        Shutdown worker_pool
+
+        """
+        await self._worker_pool.stop()
+
+    async def _fetch_until(self, segment_name, fetch_options, till=None):
+        """
+        Hit endpoint, update storage and return when since==till.
+
+        :param segment_name: Name of the segment to update.
+        :type segment_name: str
+
+        :param fetch_options Fetch options for getting segment definitions.
+        :type fetch_options splitio.api.FetchOptions
+
+        :param till: Passed till from Streaming.
+        :type till: int
+
+        :return: last change number
+        :rtype: int
+        """
+        while True:  # Fetch until since==till
+            change_number = await self._segment_storage.get_change_number(segment_name)
+            if change_number is None:
+                change_number = -1
+            if till is not None and till < change_number:
+                # the passed till is less than change_number, no need to perform updates
+                return change_number
+
+            try:
+                segment_changes = await self._api.fetch_segment(segment_name, change_number,
+                                                          fetch_options)
+            except APIException as exc:
+                _LOGGER.error('Exception raised while fetching segment %s', segment_name)
+                _LOGGER.debug('Exception information: ', exc_info=True)
+                raise exc
+
+            if change_number == -1:  # first time fetching the segment
+                new_segment = segments.from_raw(segment_changes)
+                await self._segment_storage.put(new_segment)
+            else:
+                await self._segment_storage.update(
+                    segment_name,
+                    segment_changes['added'],
+                    segment_changes['removed'],
+                    segment_changes['till']
+                )
+
+            if segment_changes['till'] == segment_changes['since']:
+                return segment_changes['till']
+
+    async def _attempt_segment_sync(self, segment_name, fetch_options, till=None):
+        """
+        Hit endpoint, update storage and return True if sync is complete.
+
+        :param segment_name: Name of the segment to update.
+        :type segment_name: str
+
+        :param fetch_options Fetch options for getting feature flag definitions.
+        :type fetch_options splitio.api.FetchOptions
+
+        :param till: Passed till from Streaming.
+        :type till: int
+
+        :return: Flags to check if it should perform bypass or operation ended
+        :rtype: bool, int, int
+        """
+        self._backoff.reset()
+        remaining_attempts = _ON_DEMAND_FETCH_BACKOFF_MAX_RETRIES
+        while True:
+            remaining_attempts -= 1
+            change_number = await self._fetch_until(segment_name, fetch_options, till)
+            if till is None or till <= change_number:
+                return True, remaining_attempts, change_number
+            elif remaining_attempts <= 0:
+                return False, remaining_attempts, change_number
+            how_long = self._backoff.get()
+            await asyncio.sleep(how_long)
+
+    async def synchronize_segment(self, segment_name, till=None):
+        """
+        Update a segment from queue
+
+        :param segment_name: Name of the segment to update.
+        :type segment_name: str
+
+        :param till: ChangeNumber received.
+        :type till: int
+
+        :return: True if no error occurs. False otherwise.
+        :rtype: bool
+        """
+        fetch_options = FetchOptions(True)  # Set Cache-Control to no-cache
+        successful_sync, remaining_attempts, change_number = await self._attempt_segment_sync(segment_name, fetch_options, till)
+        attempts = _ON_DEMAND_FETCH_BACKOFF_MAX_RETRIES - remaining_attempts
+        if successful_sync:  # succedeed sync
+            _LOGGER.debug('Refresh completed in %d attempts.', attempts)
+            return True
+        with_cdn_bypass = FetchOptions(True, change_number)  # Set flag for bypassing CDN
+        without_cdn_successful_sync, remaining_attempts, change_number = await self._attempt_segment_sync(segment_name, with_cdn_bypass, till)
+        without_cdn_attempts = _ON_DEMAND_FETCH_BACKOFF_MAX_RETRIES - remaining_attempts
+        if without_cdn_successful_sync:
+            _LOGGER.debug('Refresh completed bypassing the CDN in %d attempts.',
+                          without_cdn_attempts)
+            return True
+        _LOGGER.debug('No changes fetched after %d attempts with CDN bypassed.',
+                        without_cdn_attempts)
+        return False
+
+    async def synchronize_segments(self, segment_names = None, dont_wait = False):
+        """
+        Submit all current segments and wait for them to finish depend on dont_wait flag, then set the ready flag.
+
+        :param segment_names: Optional, array of segment names to update.
+        :type segment_name: {str}
+
+        :param dont_wait: Optional, instruct the function to not wait for task completion
+        :type segment_name: boolean
+
+        :return: True if no error occurs or dont_wait flag is True. False otherwise.
+        :rtype: bool
+        """
+        if segment_names is None:
+            segment_names = await self._feature_flag_storage.get_segment_names()
+
+        jobs = await self._worker_pool.submit_work(segment_names)
+        if (dont_wait):
+            return True
+        return await jobs.await_completion()
+
+    async def segment_exist_in_storage(self, segment_name):
+        """
+        Check if a segment exists in the storage
+
+        :param segment_name: Name of the segment
+        :type segment_name: str
+
+        :return: True if segment exist. False otherwise.
+        :rtype: bool
+        """
+        return await self._segment_storage.get(segment_name) != None
+
+
 class LocalSegmentSynchronizerBase(object):
     """Localhost mode segment base synchronizer."""
 

splitio/tasks/util/workerpool.py

@@ -4,10 +4,10 @@
 from threading import Thread, Event
 import queue
 
+from splitio.optional.loaders import asyncio
 
 _LOGGER = logging.getLogger(__name__)
 
-
 class WorkerPool(object):
     """Worker pool class to implement single producer/multiple consumer."""
 
@@ -134,3 +134,95 @@ def _wait_workers_shutdown(self, event):
             for worker_event in self._worker_events:
                 worker_event.wait()
             event.set()
+
+
+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
+
+    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)
+            message._failed = True
+        message._complete.set()
+        self._semaphore.release() # signal worker is idle
+
+    def start(self):
+        """Start the workers."""
+        self._task = asyncio.get_running_loop().create_task(self._schedule_work())
+
+    async def submit_work(self, jobs):
+        """
+        Add a new message to the work-queue.
+
+        :param message: New message to add.
+        :type message: object.
+        """
+        self.jobs = jobs
+        if len(jobs) == 1:
+            wrapped = TaskCompletionWraper(jobs[0])
+            await self._queue.put(wrapped)
+            return wrapped
+
+        tasks = [TaskCompletionWraper(job) for job in jobs]
+        for w in tasks:
+            await self._queue.put(w)
+
+        return BatchCompletionWrapper(tasks)
+
+    async def stop(self, event=None):
+        """abort all execution (except currently running handlers)."""
+        await self._queue.put(self._abort)
+
+
+class TaskCompletionWraper:
+    """Task completion class"""
+    def __init__(self, message):
+        self._message = message
+        self._complete = asyncio.Event()
+        self._failed = False
+
+    async def await_completion(self):
+        await self._complete.wait()
+
+    def _mark_as_complete(self):
+        self._complete.set()
+
+
+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])
+        return not any(task._failed for task in self._tasks)