Merge pull request #536 from splitio/async-pr-2

polishing

Author: chillaq

splitio/push/__init__.py

@@ -0,0 +1,13 @@
+class AuthException(Exception):
+    """Exception to raise when an API call fails."""
+
+    def __init__(self, custom_message, status_code=None):
+        """Constructor."""
+        Exception.__init__(self, custom_message)
+
+class SplitStorageException(Exception):
+    """Exception to raise when an API call fails."""
+
+    def __init__(self, custom_message, status_code=None):
+        """Constructor."""
+        Exception.__init__(self, custom_message)

splitio/push/manager.py

@@ -5,6 +5,7 @@
 from splitio.optional.loaders import asyncio, anext
 from splitio.api import APIException
 from splitio.util.time import get_current_epoch_time_ms
+from splitio.push import AuthException
 from splitio.push.splitsse import SplitSSEClient, SplitSSEClientAsync
 from splitio.push.sse import SSE_EVENT_ERROR
 from splitio.push.parser import parse_incoming_event, EventParsingException, EventType, \
@@ -315,7 +316,6 @@ def __init__(self, auth_api, synchronizer, feedback_loop, sdk_metadata, telemetr
         kwargs = {} if sse_url is None else {'base_url': sse_url}
         self._sse_client = SplitSSEClientAsync(sdk_metadata, client_key, **kwargs)
         self._running = False
-        self._done = asyncio.Event()
         self._telemetry_runtime_producer = telemetry_runtime_producer
         self._token_task = None
 
@@ -366,6 +366,7 @@ async def _event_handler(self, event):
         :param event: Incoming event
         :type event: splitio.push.sse.SSEEvent
         """
+        parsed = None
         try:
             parsed = parse_incoming_event(event)
             handle = self._event_handlers[parsed.event_type]
@@ -377,8 +378,8 @@ async def _event_handler(self, event):
         try:
             await handle(parsed)
         except Exception:  # pylint:disable=broad-except
-            _LOGGER.error('something went wrong when processing message of type %s',
-                          parsed.event_type)
+            event_type = "unknown" if parsed is None else parsed.event_type
+            _LOGGER.error('something went wrong when processing message of type %s', event_type)
             _LOGGER.debug(str(parsed), exc_info=True)
 
     async def _token_refresh(self, current_token):
@@ -396,15 +397,15 @@ async def _get_auth_token(self):
         """Get new auth token"""
         try:
             token = await self._auth_api.authenticate()
-        except APIException:
+        except APIException as e:
             _LOGGER.error('error performing sse auth request.')
             _LOGGER.debug('stack trace: ', exc_info=True)
             await self._feedback_loop.put(Status.PUSH_RETRYABLE_ERROR)
-            raise
+            raise AuthException(e)
 
         if token is not None and not token.push_enabled:
             await self._feedback_loop.put(Status.PUSH_NONRETRYABLE_ERROR)
-            raise Exception("Push is not enabled")
+            raise AuthException("Push is not enabled")
 
         await self._telemetry_runtime_producer.record_token_refreshes()
         await self._telemetry_runtime_producer.record_streaming_event((StreamingEventTypes.TOKEN_REFRESH, 1000 * token.exp,  get_current_epoch_time_ms()))
@@ -416,23 +417,11 @@ async def _trigger_connection_flow(self):
         self._status_tracker.reset()
 
         try:
-            try:
-                token = await self._get_auth_token()
-            except Exception as e:
-                _LOGGER.error("error getting auth token: " + str(e))
-                _LOGGER.debug("trace: ", exc_info=True)
-                return
-
+            token = await self._get_auth_token()
             events_source = self._sse_client.start(token)
-            self._done.clear()
             self._running = True
 
-            try:
-                first_event = await anext(events_source)
-            except StopAsyncIteration: # will enter here if there was an error
-                await self._feedback_loop.put(Status.PUSH_RETRYABLE_ERROR)
-                return
-
+            first_event = await anext(events_source)
             if first_event.data is not None:
                 await self._event_handler(first_event)
 
@@ -444,13 +433,17 @@ async def _trigger_connection_flow(self):
             async for event in events_source:
                 await self._event_handler(event)
             await self._handle_connection_end()  # TODO(mredolatti): this is not tested
+        except AuthException as e:
+            _LOGGER.error("error getting auth token: " + str(e))
+            _LOGGER.debug("trace: ", exc_info=True)
+        except StopAsyncIteration: # will enter here if there was an error
+            await self._feedback_loop.put(Status.PUSH_RETRYABLE_ERROR)
         finally:
             if self._token_task is not None:
                 self._token_task.cancel()
                 self._token_task = None
             self._running = False
             await self._processor.update_workers_status(False)
-            self._done.set()
 
     async def _handle_message(self, event):
         """

splitio/push/splitsse.py

@@ -22,6 +22,15 @@ class _Status(Enum):
         ERRORED = 2
         CONNECTED = 3
 
+    def __init__(self, base_url):
+        """
+        Construct a split sse client.
+
+        :param base_url: scheme + :// + host
+        :type base_url: str
+        """
+        self._base_url = base_url
+
     @staticmethod
     def _format_channels(channels):
         """
@@ -90,11 +99,11 @@ def __init__(self, event_callback, sdk_metadata, first_event_callback=None,
         :param client_key: client key.
         :type client_key: str
         """
+        SplitSSEClientBase.__init__(self, base_url)
         self._client = SSEClient(self._raw_event_handler)
         self._callback = event_callback
         self._on_connected = first_event_callback
         self._on_disconnected = connection_closed_callback
-        self._base_url = base_url
         self._status = SplitSSEClient._Status.IDLE
         self._sse_first_event = None
         self._sse_connection_closed = None
@@ -178,7 +187,7 @@ def __init__(self, sdk_metadata, client_key=None, base_url='https://streaming.sp
         :param base_url: scheme + :// + host
         :type base_url: str
         """
-        self._base_url = base_url
+        SplitSSEClientBase.__init__(self, base_url)
         self.status = SplitSSEClient._Status.IDLE
         self._metadata = headers_from_metadata(sdk_metadata, client_key)
         self._client = SSEClientAsync(self.KEEPALIVE_TIMEOUT)

splitio/push/workers.py

@@ -10,6 +10,7 @@
 
 from splitio.models.splits import from_raw
 from splitio.models.telemetry import UpdateFromSSE
+from splitio.push import SplitStorageException
 from splitio.push.parser import UpdateType
 from splitio.optional.loaders import asyncio
 from splitio.util.storage_helper import update_feature_flag_storage, update_feature_flag_storage_async
@@ -202,9 +203,28 @@ def is_running(self):
         """Return whether the working is running."""
         return self._running
 
+    def _apply_iff_if_needed(self, event):
+        if not self._check_instant_ff_update(event):
+            return False
+
+        try:
+            new_feature_flag = from_raw(json.loads(self._get_feature_flag_definition(event)))
+            segment_list = update_feature_flag_storage(self._feature_flag_storage, [new_feature_flag], event.change_number)
+            for segment_name in segment_list:
+                if self._segment_storage.get(segment_name) is None:
+                    _LOGGER.debug('Fetching new segment %s', segment_name)
+                    self._segment_handler(segment_name, event.change_number)
+
+            self._telemetry_runtime_producer.record_update_from_sse(UpdateFromSSE.SPLIT_UPDATE)
+            return True
+
+        except Exception as e:
+            raise SplitStorageException(e)
+
     def _check_instant_ff_update(self, event):
         if event.update_type == UpdateType.SPLIT_UPDATE and event.compression is not None and event.previous_change_number == self._feature_flag_storage.get_change_number():
             return True
+
         return False
 
     def _run(self):
@@ -217,28 +237,19 @@ def _run(self):
                 continue
             _LOGGER.debug('Processing feature flag update %d', event.change_number)
             try:
-                if self._check_instant_ff_update(event):
-                    try:
-                        new_feature_flag = from_raw(json.loads(self._get_feature_flag_definition(event)))
-                        segment_list = update_feature_flag_storage(self._feature_flag_storage, [new_feature_flag], event.change_number)
-                        for segment_name in segment_list:
-                            if self._segment_storage.get(segment_name) is None:
-                                _LOGGER.debug('Fetching new segment %s', segment_name)
-                                self._segment_handler(segment_name, event.change_number)
-
-                        self._telemetry_runtime_producer.record_update_from_sse(UpdateFromSSE.SPLIT_UPDATE)
-                        continue
-                    except Exception as e:
-                        _LOGGER.error('Exception raised in updating feature flag')
-                        _LOGGER.debug('Exception information: ', exc_info=True)
-                        pass
+                if self._apply_iff_if_needed(event):
+                    continue
+
                 sync_result = self._handler(event.change_number)
                 if not sync_result.success and sync_result.error_code is not None and sync_result.error_code == 414:
                     _LOGGER.error("URI too long exception caught, sync failed")
 
                 if not sync_result.success:
                     _LOGGER.error("feature flags sync failed")
 
+            except SplitStorageException as e:  # pylint: disable=broad-except
+                _LOGGER.error('Exception Updating Feature Flag')
+                _LOGGER.debug('Exception information: ', exc_info=True)
             except Exception as e:  # pylint: disable=broad-except
                 _LOGGER.error('Exception raised in feature flag synchronization')
                 _LOGGER.debug('Exception information: ', exc_info=True)
@@ -297,6 +308,24 @@ def is_running(self):
         """Return whether the working is running."""
         return self._running
 
+    async def _apply_iff_if_needed(self, event):
+        if not await self._check_instant_ff_update(event):
+            return False
+        try:
+            new_feature_flag = from_raw(json.loads(self._get_feature_flag_definition(event)))
+            segment_list = await update_feature_flag_storage_async(self._feature_flag_storage, [new_feature_flag], event.change_number)
+            for segment_name in segment_list:
+                if await self._segment_storage.get(segment_name) is None:
+                    _LOGGER.debug('Fetching new segment %s', segment_name)
+                    await self._segment_handler(segment_name, event.change_number)
+
+            await self._telemetry_runtime_producer.record_update_from_sse(UpdateFromSSE.SPLIT_UPDATE)
+            return True
+
+        except Exception as e:
+            raise SplitStorageException(e)
+
+
     async def _check_instant_ff_update(self, event):
         if event.update_type == UpdateType.SPLIT_UPDATE and event.compression is not None and event.previous_change_number == await self._feature_flag_storage.get_change_number():
             return True
@@ -312,22 +341,12 @@ async def _run(self):
                 continue
             _LOGGER.debug('Processing split_update %d', event.change_number)
             try:
-                if await self._check_instant_ff_update(event):
-                    try:
-                        new_feature_flag = from_raw(json.loads(self._get_feature_flag_definition(event)))
-                        segment_list = await update_feature_flag_storage_async(self._feature_flag_storage, [new_feature_flag], event.change_number)
-                        for segment_name in segment_list:
-                            if await self._segment_storage.get(segment_name) is None:
-                                _LOGGER.debug('Fetching new segment %s', segment_name)
-                                await self._segment_handler(segment_name, event.change_number)
-
-                        await self._telemetry_runtime_producer.record_update_from_sse(UpdateFromSSE.SPLIT_UPDATE)
-                        continue
-                    except Exception as e:
-                        _LOGGER.error('Exception raised in updating feature flag')
-                        _LOGGER.debug('Exception information: ', exc_info=True)
-                        pass
+                if await self._apply_iff_if_needed(event):
+                    continue
                 await self._handler(event.change_number)
+            except SplitStorageException as e:  # pylint: disable=broad-except
+                _LOGGER.error('Exception Updating Feature Flag')
+                _LOGGER.debug('Exception information: ', exc_info=True)
             except Exception as e:  # pylint: disable=broad-except
                 _LOGGER.error('Exception raised in split synchronization')
                 _LOGGER.debug('Exception information: ', exc_info=True)

splitio/storage/redis.py

@@ -157,11 +157,6 @@ def kill_locally(self, feature_flag_name, default_treatment, change_number):
 class RedisSplitStorage(RedisSplitStorageBase):
     """Redis-based storage for feature flags."""
 
-    _FEATURE_FLAG_KEY = 'SPLITIO.split.{feature_flag_name}'
-    _FEATURE_FLAG_TILL_KEY = 'SPLITIO.splits.till'
-    _TRAFFIC_TYPE_KEY = 'SPLITIO.trafficType.{traffic_type_name}'
-    _FLAG_SET_KEY = 'SPLITIO.flagSet.{flag_set}'
-
     def __init__(self, redis_client, enable_caching=False, max_age=DEFAULT_MAX_AGE, config_flag_sets=[]):
         """
         Class constructor.
@@ -213,7 +208,8 @@ def get_feature_flags_by_sets(self, flag_sets):
 
             keys = [self._get_flag_set_key(flag_set) for flag_set in sets_to_fetch]
             pipe = self._pipe()
-            [pipe.smembers(key) for key in keys]
+            for key in keys:
+                pipe.smembers(key)
             result_sets = pipe.execute()
             _LOGGER.debug("Fetchting Feature flags by set [%s] from redis" % (keys))
             _LOGGER.debug(result_sets)
@@ -342,7 +338,9 @@ def __init__(self, redis_client, enable_caching=False, max_age=DEFAULT_MAX_AGE,
         self.flag_set_filter = FlagSetsFilter(config_flag_sets)
         self._pipe = self.redis.pipeline
         if enable_caching:
-            self._cache = LocalMemoryCacheAsync(None, None, max_age)
+            self._feature_flag_cache = LocalMemoryCacheAsync(None, None, max_age)
+            self._traffic_type_cache = LocalMemoryCacheAsync(None, None, max_age)
+
 
     async def get(self, feature_flag_name):  # pylint: disable=method-hidden
         """
@@ -359,15 +357,16 @@ async def get(self, feature_flag_name):  # pylint: disable=method-hidden
         :type change_number: int
         """
         try:
-            if self._enable_caching and await self._cache.get_key(feature_flag_name) is not None:
-                raw = await self._cache.get_key(feature_flag_name)
-            else:
-                raw = await self.redis.get(self._get_key(feature_flag_name))
+            raw_feature_flags = None
+            if self._enable_caching:
+                raw_feature_flags = await self._feature_flag_cache.get_key(feature_flag_name)
+            if raw_feature_flags is None:
+                raw_feature_flags = await self.redis.get(self._get_key(feature_flag_name))
                 if self._enable_caching:
-                    await self._cache.add_key(feature_flag_name, raw)
+                    await self._feature_flag_cache.add_key(feature_flag_name, raw_feature_flags)
                 _LOGGER.debug("Fetchting feature flag [%s] from redis" % feature_flag_name)
-                _LOGGER.debug(raw)
-            return splits.from_raw(json.loads(raw)) if raw is not None else None
+                _LOGGER.debug(raw_feature_flags)
+            return splits.from_raw(json.loads(raw_feature_flags)) if raw_feature_flags is not None else None
 
         except RedisAdapterException:
             _LOGGER.error('Error fetching feature flag from storage')
@@ -410,13 +409,13 @@ async def fetch_many(self, feature_flag_names):
         """
         to_return = dict()
         try:
-            if self._enable_caching and await self._cache.get_key(frozenset(feature_flag_names)) is not None:
-                raw_feature_flags = await self._cache.get_key(frozenset(feature_flag_names))
-            else:
-                keys = [self._get_key(feature_flag_name) for feature_flag_name in feature_flag_names]
-                raw_feature_flags = await self.redis.mget(keys)
+            raw_feature_flags = None
+            if self._enable_caching:
+                raw_feature_flags = await self._feature_flag_cache.get_key(frozenset(feature_flag_names))
+            if raw_feature_flags is None:
+                raw_feature_flags = await self.redis.mget([self._get_key(feature_flag_name) for feature_flag_name in feature_flag_names])
                 if self._enable_caching:
-                    await self._cache.add_key(frozenset(feature_flag_names), raw_feature_flags)
+                    await self._feature_flag_cache.add_key(frozenset(feature_flag_names), raw_feature_flags)
             for i in range(len(feature_flag_names)):
                 feature_flag = None
                 try:
@@ -439,13 +438,14 @@ async def is_valid_traffic_type(self, traffic_type_name):  # pylint: disable=met
         :rtype: bool
         """
         try:
-            if self._enable_caching and await self._cache.get_key(traffic_type_name) is not None:
-                raw = await self._cache.get_key(traffic_type_name)
-            else:
-                raw = await self.redis.get(self._get_traffic_type_key(traffic_type_name))
+            raw_traffic_type = None
+            if self._enable_caching:
+                raw_traffic_type = await self._traffic_type_cache.get_key(traffic_type_name)
+            if raw_traffic_type is None:
+                raw_traffic_type = await self.redis.get(self._get_traffic_type_key(traffic_type_name))
                 if self._enable_caching:
-                    await self._cache.add_key(traffic_type_name, raw)
-            count = json.loads(raw) if raw else 0
+                    await self._traffic_type_cache.add_key(traffic_type_name, raw_traffic_type)
+            count = json.loads(raw_traffic_type) if raw_traffic_type else 0
             return count > 0
 
         except RedisAdapterException:

splitio/sync/split.py

@@ -112,8 +112,7 @@ def _fetch_until(self, fetch_options, till=None):
                 _LOGGER.error('Exception raised while fetching feature flags')
                 _LOGGER.debug('Exception information: ', exc_info=True)
                 raise exc
-            fetched_feature_flags = []
-            [fetched_feature_flags.append(splits.from_raw(feature_flag)) for feature_flag in feature_flag_changes.get('splits', [])]
+            fetched_feature_flags = [(splits.from_raw(feature_flag)) for feature_flag in feature_flag_changes.get('splits', [])]
             segment_list = update_feature_flag_storage(self._feature_flag_storage, fetched_feature_flags, feature_flag_changes['till'])
             if feature_flag_changes['till'] == feature_flag_changes['since']:
                 return feature_flag_changes['till'], segment_list