Added timeouts and used one http session for SSE, added stopping manager tasks when destroy is called and removed references to tasks

Author: chillaq

splitio/client/factory.py

@@ -367,7 +367,7 @@ def __init__(  # pylint: disable=too-many-arguments
         self._manager_start_task = manager_start_task
         self._status = Status.NOT_INITIALIZED
         self._sdk_ready_flag = asyncio.Event()
-        asyncio.get_running_loop().create_task(self._update_status_when_ready_async())
+        self._ready_task = asyncio.get_running_loop().create_task(self._update_status_when_ready_async())
 
     async def _update_status_when_ready_async(self):
         """Wait until the sdk is ready and update the status for async mode."""
@@ -377,6 +377,7 @@ async def _update_status_when_ready_async(self):
 
         if self._manager_start_task is not None:
             await self._manager_start_task
+            self._manager_start_task = None
         await self._telemetry_init_producer.record_ready_time(get_current_epoch_time_ms() - self._ready_time)
         redundant_factory_count, active_factory_count = _get_active_and_redundant_count()
         await self._telemetry_init_producer.record_active_and_redundant_factories(active_factory_count, redundant_factory_count)
@@ -430,14 +431,25 @@ async def destroy(self, destroyed_event=None):
 
         try:
             _LOGGER.info('Factory destroy called, stopping tasks.')
+            if self._manager_start_task is not None and not self._manager_start_task.done():
+                self._manager_start_task.cancel()
+
             if self._sync_manager is not None:
                 await self._sync_manager.stop(True)
 
+                if not self._ready_task.done():
+                    self._ready_task.cancel()
+                    self._ready_task = None
+
                 if isinstance(self._storages['splits'], RedisSplitStorageAsync):
                     await self._get_storage('splits').redis.close()
 
                 if isinstance(self._sync_manager, ManagerAsync) and isinstance(self._telemetry_submitter, InMemoryTelemetrySubmitterAsync):
                     await self._telemetry_submitter._telemetry_api._client.close_session()
+
+                if isinstance(self._sync_manager, ManagerAsync) and self._sync_manager._streaming_enabled:
+                    await self._sync_manager._push._sse_client._client.close_session()
+
         except Exception as e:
             _LOGGER.error('Exception destroying factory.')
             _LOGGER.debug(str(e))

splitio/push/manager.py

@@ -350,9 +350,10 @@ async def stop(self, blocking=False):
         if self._token_task:
             self._token_task.cancel()
 
-        stop_task = asyncio.get_running_loop().create_task(self._stop_current_conn())
         if blocking:
-            await stop_task
+            await self._stop_current_conn()
+        else:
+            asyncio.get_running_loop().create_task(self._stop_current_conn())
 
     async def _event_handler(self, event):
         """
@@ -382,6 +383,7 @@ async def _token_refresh(self, current_token):
         :param current_token: token (parsed) JWT
         :type current_token: splitio.models.token.Token
         """
+        _LOGGER.debug("Next token refresh in " + str(self._get_time_period(current_token)) + " seconds")
         await asyncio.sleep(self._get_time_period(current_token))
         await self._stop_current_conn()
         self._running_task = asyncio.get_running_loop().create_task(self._trigger_connection_flow())
@@ -441,6 +443,7 @@ async def _trigger_connection_flow(self):
         finally:
             if self._token_task is not None:
                 self._token_task.cancel()
+                self._token_task = None
             self._running = False
             self._done.set()
 
@@ -529,4 +532,5 @@ async def _stop_current_conn(self):
         await self._sse_client.stop()
         self._running_task.cancel()
         await self._running_task
+        self._running_task = None
         _LOGGER.debug("SplitSSE tasks are stopped")

splitio/push/splitsse.py

@@ -219,7 +219,7 @@ async def start(self, token):
             _LOGGER.debug('stack trace: ', exc_info=True)
         finally:
             self.status = SplitSSEClient._Status.IDLE
-            _LOGGER.debug('sse connection ended.')
+            _LOGGER.debug('Split sse connection ended.')
             self._event_source_ended.set()
 
     async def stop(self):
@@ -230,7 +230,10 @@ async def stop(self):
             return
 
         await self._client.shutdown()
+# catching exception to avoid task hanging
         try:
             await self._event_source_ended.wait()
-        except asyncio.CancelledError:
+        except asyncio.CancelledError as e:
+            _LOGGER.error("Exception waiting for event source ended")
+            _LOGGER.debug('stack trace: ', exc_info=True)
             pass

splitio/push/sse.py

@@ -155,6 +155,8 @@ def __init__(self, socket_read_timeout=_DEFAULT_SOCKET_READ_TIMEOUT):
         self._socket_read_timeout = socket_read_timeout + socket_read_timeout * .3
         self._response = None
         self._done = asyncio.Event()
+        client_timeout = aiohttp.ClientTimeout(total=0, sock_read=self._socket_read_timeout)
+        self._sess = aiohttp.ClientSession(timeout=client_timeout)
 
     async def start(self, url, extra_headers=None):  # pylint:disable=protected-access
         """
@@ -168,46 +170,53 @@ async def start(self, url, extra_headers=None):  # pylint:disable=protected-acce
             raise RuntimeError('Client already started.')
 
         self._done.clear()
-        client_timeout = aiohttp.ClientTimeout(total=0, sock_read=self._socket_read_timeout)
-        async with aiohttp.ClientSession(timeout=client_timeout) as sess:
-            try:
-                async with sess.get(url, headers=get_headers(extra_headers)) as response:
-                    self._response = response
-                    event_builder = EventBuilder()
-                    async for line in response.content:
-                        if line.startswith(b':'):
-                            _LOGGER.debug("skipping emtpy line / comment")
-                            continue
-                        elif line in _EVENT_SEPARATORS:
-                            _LOGGER.debug("dispatching event: %s", event_builder.build())
-                            yield event_builder.build()
-                            event_builder = EventBuilder()
-                        else:
-                            event_builder.process_line(line)
-
-            except Exception as exc:  # pylint:disable=broad-except
-                if self._is_conn_closed_error(exc):
-                    _LOGGER.debug('sse connection ended.')
-                    return
-
-                _LOGGER.error('http client is throwing exceptions')
-                _LOGGER.error('stack trace: ', exc_info=True)
-
-            finally:
-                self._response = None
-                self._done.set()
+        try:
+            async with self._sess.get(url, headers=get_headers(extra_headers)) as response:
+                self._response = response
+                event_builder = EventBuilder()
+                async for line in response.content:
+                    if line.startswith(b':'):
+                        _LOGGER.debug("skipping emtpy line / comment")
+                        continue
+                    elif line in _EVENT_SEPARATORS:
+                        _LOGGER.debug("dispatching event: %s", event_builder.build())
+                        yield event_builder.build()
+                        event_builder = EventBuilder()
+                    else:
+                        event_builder.process_line(line)
+
+        except Exception as exc:  # pylint:disable=broad-except
+            if self._is_conn_closed_error(exc):
+                _LOGGER.debug('sse connection ended.')
+                return
+
+            _LOGGER.error('http client is throwing exceptions')
+            _LOGGER.error('stack trace: ', exc_info=True)
+
+        finally:
+            self._response = None
+            self._done.set()
 
     async def shutdown(self):
         """Close connection"""
         if self._response:
             self._response.close()
-        await self._done.wait()
+# catching exception to avoid task hanging
+        try:
+            await self._done.wait()
+        except asyncio.CancelledError:
+            _LOGGER.error("Exception waiting for event source ended")
+            _LOGGER.debug('stack trace: ', exc_info=True)
+            pass
 
     @staticmethod
     def _is_conn_closed_error(exc):
         """Check if the ReadError is caused by the connection being closed."""
         return isinstance(exc, aiohttp.ClientConnectionError) and str(exc) == "Connection closed"
 
+    async def close_session(self):
+        if not self._sess.closed:
+            await self._sess.close()
 
 def get_headers(extra=None):
     """

splitio/sync/manager.py

@@ -172,19 +172,20 @@ def __init__(self, synchronizer, auth_api, streaming_enabled, sdk_metadata, tele
             self._backoff = Backoff()
             self._queue = asyncio.Queue()
             self._push = PushManagerAsync(auth_api, synchronizer, self._queue, sdk_metadata, telemetry_runtime_producer, sse_url, client_key)
-            self._push_status_handler_task = None
+        self._stopped = False
 
     async def start(self, max_retry_attempts=_SYNC_ALL_NO_RETRIES):
         """Start the SDK synchronization tasks."""
+        self._stopped = False
         try:
             await self._synchronizer.sync_all(max_retry_attempts)
-            self._synchronizer.start_periodic_data_recording()
-            if self._streaming_enabled:
-                self._push_status_handler_task = asyncio.get_running_loop().create_task(self._streaming_feedback_handler())
-                self._push.start()
-            else:
-                self._synchronizer.start_periodic_fetching()
-
+            if not self._stopped:
+                self._synchronizer.start_periodic_data_recording()
+                if self._streaming_enabled:
+                    asyncio.get_running_loop().create_task(self._streaming_feedback_handler())
+                    self._push.start()
+                else:
+                    self._synchronizer.start_periodic_fetching()
         except (APIException, RuntimeError):
             _LOGGER.error('Exception raised starting Split Manager')
             _LOGGER.debug('Exception information: ', exc_info=True)
@@ -201,8 +202,9 @@ async def stop(self, blocking):
         if self._streaming_enabled:
             self._push_status_handler_active = False
             await self._queue.put(self._CENTINEL_EVENT)
-            await self._push.stop()
+            await self._push.stop(blocking)
         await self._synchronizer.shutdown(blocking)
+        self._stopped = True
 
     async def _streaming_feedback_handler(self):
         """

splitio/sync/synchronizer.py

@@ -520,7 +520,7 @@ def __init__(self, split_synchronizers, split_tasks):
         :type split_tasks: splitio.sync.synchronizer.SplitTasks
         """
         SynchronizerInMemoryBase.__init__(self, split_synchronizers, split_tasks)
-        self.stop_periodic_data_recording_task = None
+        self._shutdown = False
 
     async def _synchronize_segments(self):
         _LOGGER.debug('Starting segments synchronization')
@@ -551,6 +551,9 @@ async def synchronize_splits(self, till, sync_segments=True):
         :returns: whether the synchronization was successful or not.
         :rtype: bool
         """
+        if self._shutdown:
+            return
+
         _LOGGER.debug('Starting feature flags synchronization')
         try:
             new_segments = []
@@ -583,8 +586,9 @@ async def sync_all(self, max_retry_attempts=_SYNC_ALL_NO_RETRIES):
         :param max_retry_attempts: apply max attempts if it set to absilute integer.
         :type max_retry_attempts: int
         """
+        self._shutdown = False
         retry_attempts = 0
-        while True:
+        while not self._shutdown:
             try:
                 sync_result = await self.synchronize_splits(None, False)
                 if not sync_result.success and sync_result.error_code is not None and sync_result.error_code == 414:
@@ -609,7 +613,8 @@ async def sync_all(self, max_retry_attempts=_SYNC_ALL_NO_RETRIES):
                     if retry_attempts > max_retry_attempts:
                         break
                 how_long = self._backoff.get()
-                time.sleep(how_long)
+                if not self._shutdown:
+                    time.sleep(how_long)
 
         _LOGGER.error("Could not correctly synchronize feature flags and segments after %d attempts.", retry_attempts)
 
@@ -621,6 +626,7 @@ async def shutdown(self, blocking):
         :type blocking: bool
         """
         _LOGGER.debug('Shutting down tasks.')
+        self._shutdown = True
         await self._split_synchronizers.segment_sync.shutdown()
         await self.stop_periodic_fetching()
         await self.stop_periodic_data_recording(blocking)
@@ -639,10 +645,11 @@ async def stop_periodic_data_recording(self, blocking):
         :type blocking: bool
         """
         _LOGGER.debug('Stopping periodic data recording')
-        stop_periodic_data_recording_task = asyncio.get_running_loop().create_task(self._stop_periodic_data_recording())
         if blocking:
-            await stop_periodic_data_recording_task
+            await self._stop_periodic_data_recording()
             _LOGGER.debug('all tasks finished successfully.')
+        else:
+            asyncio.get_running_loop().create_task(self._stop_periodic_data_recording())
 
     async def _stop_periodic_data_recording(self):
         """
@@ -798,7 +805,6 @@ def __init__(self, split_synchronizers, split_tasks):
         :type split_tasks: splitio.sync.synchronizer.SplitTasks
         """
         RedisSynchronizerBase.__init__(self, split_synchronizers, split_tasks)
-        self.stop_periodic_data_recording_task = None
 
     async def shutdown(self, blocking):
         """
@@ -829,7 +835,7 @@ async def stop_periodic_data_recording(self, blocking):
             await self._stop_periodic_data_recording()
             _LOGGER.debug('all tasks finished successfully.')
         else:
-            self.stop_periodic_data_recording_task = asyncio.get_running_loop().create_task(self._stop_periodic_data_recording)
+            asyncio.get_running_loop().create_task(self._stop_periodic_data_recording)
 
 
 

splitio/tasks/util/asynctask.py

@@ -288,7 +288,7 @@ def start(self):
             return
         # Start execution
         self._completion_event.clear()
-        self._wrapper_task = asyncio.get_running_loop().create_task(self._execution_wrapper())
+        asyncio.get_running_loop().create_task(self._execution_wrapper())
 
     async def stop(self, wait_for_completion=False):
         """

splitio/tasks/util/workerpool.py

@@ -178,7 +178,7 @@ async def _do_work(self, message):
 
     def start(self):
         """Start the workers."""
-        self._task = asyncio.get_running_loop().create_task(self._schedule_work())
+        asyncio.get_running_loop().create_task(self._schedule_work())
 
     async def submit_work(self, jobs):
         """

tests/client/test_factory.py

@@ -699,9 +699,9 @@ class SplitFactoryAsyncTests(object):
     @pytest.mark.asyncio
     async def test_flag_sets_counts(self):
         factory = await get_factory_async("none", config={
-            'flagSetsFilter': ['set1', 'set2', 'set3']
+            'flagSetsFilter': ['set1', 'set2', 'set3'],
+            'streamEnabled': False
         })
-
         assert factory._telemetry_init_producer._telemetry_storage._tel_config._flag_sets == 3
         assert factory._telemetry_init_producer._telemetry_storage._tel_config._flag_sets_invalid == 0
         await factory.destroy()
@@ -741,7 +741,7 @@ async def synchronize_config(*_):
         mocker.patch('splitio.sync.telemetry.InMemoryTelemetrySubmitterAsync.synchronize_config', new=synchronize_config)
 
         # Start factory and make assertions
-        factory = await get_factory_async('some_api_key')
+        factory = await get_factory_async('some_api_key', config={'streamingEmabled': False})
         assert isinstance(factory, SplitFactoryAsync)
         assert isinstance(factory._storages['splits'], inmemmory.InMemorySplitStorageAsync)
         assert isinstance(factory._storages['segments'], inmemmory.InMemorySegmentStorageAsync)
@@ -859,6 +859,10 @@ async def stop(*_):
             pass
         factory._sync_manager.stop = stop
 
+        async def start(*_):
+            pass
+        factory._sync_manager.start = start
+
         try:
             await factory.block_until_ready(1)
         except:

tests/integration/test_client_e2e.py

@@ -2002,7 +2002,7 @@ async def _setup_method(self):
             await redis_client.set(split_storage._get_key(split['name']), json.dumps(split))
             if split.get('sets') is not None:
                 for flag_set in split.get('sets'):
-                    redis_client.sadd(split_storage._get_flag_set_key(flag_set), split['name'])
+                    await redis_client.sadd(split_storage._get_flag_set_key(flag_set), split['name'])
 
         await redis_client.set(split_storage._FEATURE_FLAG_TILL_KEY, data['till'])
 
@@ -2217,7 +2217,7 @@ async def _setup_method(self):
             await redis_client.set(split_storage._get_key(split['name']), json.dumps(split))
             if split.get('sets') is not None:
                 for flag_set in split.get('sets'):
-                    redis_client.sadd(split_storage._get_flag_set_key(flag_set), split['name'])
+                    await redis_client.sadd(split_storage._get_flag_set_key(flag_set), split['name'])
         await redis_client.set(split_storage._FEATURE_FLAG_TILL_KEY, data['till'])
 
         segment_fn = os.path.join(os.path.dirname(__file__), 'files', 'segmentEmployeesChanges.json')