Merge pull request #386 from splitio/async-sse-class

Async sse class

Author: chillaq

splitio/push/sse.py

@@ -1,23 +1,46 @@
 """Low-level SSE Client."""
 import logging
 import socket
+import abc
+import urllib
 from collections import namedtuple
 from http.client import HTTPConnection, HTTPSConnection
 from urllib.parse import urlparse
+import pytest
 
+from splitio.optional.loaders import asyncio, aiohttp
+from splitio.api.client import HttpClientException
 
 _LOGGER = logging.getLogger(__name__)
 
-
 SSE_EVENT_ERROR = 'error'
 SSE_EVENT_MESSAGE = 'message'
-
+_DEFAULT_HEADERS = {'accept': 'text/event-stream'}
+_EVENT_SEPARATORS = set([b'\n', b'\r\n'])
+_DEFAULT_ASYNC_TIMEOUT = 300
 
 SSEEvent = namedtuple('SSEEvent', ['event_id', 'event', 'retry', 'data'])
 
 
 __ENDING_CHARS = set(['\n', ''])
 
+def _get_request_parameters(url, extra_headers):
+    """
+    Parse URL and headers
+
+    :param url: url to connect to
+    :type url: str
+
+    :param extra_headers: additional headers
+    :type extra_headers: dict[str, str]
+
+    :returns: processed URL and Headers
+    :rtype: str, dict
+    """
+    url = urlparse(url)
+    headers = _DEFAULT_HEADERS.copy()
+    headers.update(extra_headers if extra_headers is not None else {})
+    return url, headers
 
 class EventBuilder(object):
     """Event builder class."""
@@ -46,12 +69,19 @@ def build(self):
         return SSEEvent(self._lines.get('id'), self._lines.get('event'),
                         self._lines.get('retry'), self._lines.get('data'))
 
+class SSEClientBase(object, metaclass=abc.ABCMeta):
+    """Worker template."""
 
-class SSEClient(object):
-    """SSE Client implementation."""
+    @abc.abstractmethod
+    def start(self, url, extra_headers, timeout):  # pylint:disable=protected-access
+        """Connect and start listening for events."""
 
-    _DEFAULT_HEADERS = {'accept': 'text/event-stream'}
-    _EVENT_SEPARATORS = set([b'\n', b'\r\n'])
+    @abc.abstractmethod
+    def shutdown(self):
+        """Shutdown the current connection."""
+
+class SSEClient(SSEClientBase):
+    """SSE Client implementation."""
 
     def __init__(self, callback):
         """
@@ -81,7 +111,7 @@ def _read_events(self):
                 elif line.startswith(b':'):  # comment. Skip
                     _LOGGER.debug("skipping sse comment")
                     continue
-                elif line in self._EVENT_SEPARATORS:
+                elif line in _EVENT_SEPARATORS:
                     event = event_builder.build()
                     _LOGGER.debug("dispatching event: %s", event)
                     self._event_callback(event)
@@ -117,9 +147,7 @@ def start(self, url, extra_headers=None, timeout=socket._GLOBAL_DEFAULT_TIMEOUT)
             raise RuntimeError('Client already started.')
 
         self._shutdown_requested = False
-        url = urlparse(url)
-        headers = self._DEFAULT_HEADERS.copy()
-        headers.update(extra_headers if extra_headers is not None else {})
+        url, headers = _get_request_parameters(url, extra_headers)
         self._conn = (HTTPSConnection(url.hostname, url.port, timeout=timeout)
                       if url.scheme == 'https'
                       else HTTPConnection(url.hostname, port=url.port, timeout=timeout))
@@ -139,3 +167,111 @@ def shutdown(self):
 
         self._shutdown_requested = True
         self._conn.sock.shutdown(socket.SHUT_RDWR)
+
+class SSEClientAsync(SSEClientBase):
+    """SSE Client implementation."""
+
+    def __init__(self, callback):
+        """
+        Construct an SSE client.
+
+        :param callback: function to call when an event is received
+        :type callback: callable
+        """
+        self._conn = None
+        self._event_callback = callback
+        self._shutdown_requested = False
+
+    async def _read_events(self, response):
+        """
+        Read events from the supplied connection.
+
+        :returns: True if the connection was ended by us. False if it was closed by the serve.
+        :rtype: bool
+        """
+        try:
+            event_builder = EventBuilder()
+            while not self._shutdown_requested:
+                line = await response.readline()
+                if line is None or len(line) <= 0:  # connection ended
+                    break
+                elif line.startswith(b':'):  # comment. Skip
+                    _LOGGER.debug("skipping sse comment")
+                    continue
+                elif line in _EVENT_SEPARATORS:
+                    event = event_builder.build()
+                    _LOGGER.debug("dispatching event: %s", event)
+                    await self._event_callback(event)
+                    event_builder = EventBuilder()
+                else:
+                    event_builder.process_line(line)
+        except asyncio.CancelledError:
+           _LOGGER.debug("Cancellation request, proceeding to cancel.")
+           raise
+        except Exception:  # pylint:disable=broad-except
+            _LOGGER.debug('sse connection ended.')
+            _LOGGER.debug('stack trace: ', exc_info=True)
+        finally:
+            await self._conn.close()
+            self._conn = None  # clear so it can be started again
+
+        return self._shutdown_requested
+
+    async def start(self, url, extra_headers=None, timeout=_DEFAULT_ASYNC_TIMEOUT):  # pylint:disable=protected-access
+        """
+        Connect and start listening for events.
+
+        :param url: url to connect to
+        :type url: str
+
+        :param extra_headers: additional headers
+        :type extra_headers: dict[str, str]
+
+        :param timeout: connection & read timeout
+        :type timeout: float
+
+        :returns: True if the connection was ended by us. False if it was closed by the serve.
+        :rtype: bool
+        """
+        _LOGGER.debug("Async SSEClient Started")
+        if self._conn is not None:
+            raise RuntimeError('Client already started.')
+
+        self._shutdown_requested = False
+        url = urlparse(url)
+        headers = _DEFAULT_HEADERS.copy()
+        headers.update(extra_headers if extra_headers is not None else {})
+        parsed_url =  urllib.parse.urljoin(url[0] + "://" + url[1], url[2])
+        params=url[4]
+        try:
+            self._conn = aiohttp.connector.TCPConnector()
+            async with aiohttp.client.ClientSession(
+                connector=self._conn,
+                headers=headers,
+                timeout=aiohttp.ClientTimeout(timeout)
+                ) as self._session:
+                reader = await self._session.request(
+                    "GET",
+                    parsed_url,
+                    params=params
+                )
+                return await self._read_events(reader.content)
+        except aiohttp.ClientError as exc:  # pylint: disable=broad-except
+            _LOGGER.error(str(exc))
+            raise HttpClientException('http client is throwing exceptions') from exc
+
+    async def shutdown(self):
+        """Shutdown the current connection."""
+        _LOGGER.debug("Async SSEClient Shutdown")
+        if self._conn is None:
+            _LOGGER.warning("no sse connection has been started on this SSEClient instance. Ignoring")
+            return
+
+        if self._shutdown_requested:
+            _LOGGER.warning("shutdown already requested")
+            return
+
+        self._shutdown_requested = True
+        sock = self._session.connector._loop._ssock
+        sock.shutdown(socket.SHUT_RDWR)
+        await self._conn.close()

tests/push/test_sse.py

@@ -3,9 +3,10 @@
 import time
 import threading
 import pytest
-from splitio.push.sse import SSEClient, SSEEvent
-from tests.helpers.mockserver import SSEMockServer
 
+from splitio.push.sse import SSEClient, SSEEvent, SSEClientAsync
+from splitio.optional.loaders import asyncio
+from tests.helpers.mockserver import SSEMockServer
 
 class SSEClientTests(object):
     """SSEClient test cases."""
@@ -123,3 +124,113 @@ def runner():
         ]
 
         assert client._conn is None
+
+class SSEClientAsyncTests(object):
+    """SSEClient test cases."""
+
+#    @pytest.mark.asyncio
+    async def test_sse_client_disconnects(self):
+        """Test correct initialization. Client ends the connection."""
+        server = SSEMockServer()
+        server.start()
+
+        events = []
+        async def callback(event):
+            """Callback."""
+            events.append(event)
+
+        client = SSEClientAsync(callback)
+
+        async def connect_split_sse_client():
+            await client.start('http://127.0.0.1:' + str(server.port()))
+
+        self._client_task = asyncio.gather(connect_split_sse_client())
+        server.publish({'id': '1'})
+        server.publish({'id': '2', 'event': 'message', 'data': 'abc'})
+        server.publish({'id': '3', 'event': 'message', 'data': 'def'})
+        server.publish({'id': '4', 'event': 'message', 'data': 'ghi'})
+        await asyncio.sleep(1)
+        await client.shutdown()
+        self._client_task.cancel()
+        await asyncio.sleep(1)
+
+        assert events == [
+            SSEEvent('1', None, None, None),
+            SSEEvent('2', 'message', None, 'abc'),
+            SSEEvent('3', 'message', None, 'def'),
+            SSEEvent('4', 'message', None, 'ghi')
+        ]
+        assert client._conn is None
+        server.publish(server.GRACEFUL_REQUEST_END)
+        server.stop()
+
+    async def test_sse_server_disconnects(self):
+        """Test correct initialization. Server ends connection."""
+        server = SSEMockServer()
+        server.start()
+
+        events = []
+        async def callback(event):
+            """Callback."""
+            events.append(event)
+
+        client = SSEClientAsync(callback)
+
+        async def start_client():
+            await client.start('http://127.0.0.1:' + str(server.port()))
+
+        asyncio.gather(start_client())
+        server.publish({'id': '1'})
+        server.publish({'id': '2', 'event': 'message', 'data': 'abc'})
+        server.publish({'id': '3', 'event': 'message', 'data': 'def'})
+        server.publish({'id': '4', 'event': 'message', 'data': 'ghi'})
+        server.publish(server.GRACEFUL_REQUEST_END)
+
+        await asyncio.sleep(1)
+        server.stop()
+        await asyncio.sleep(1)
+
+        assert events == [
+            SSEEvent('1', None, None, None),
+            SSEEvent('2', 'message', None, 'abc'),
+            SSEEvent('3', 'message', None, 'def'),
+            SSEEvent('4', 'message', None, 'ghi')
+        ]
+
+        assert client._conn is None
+
+    async def test_sse_server_disconnects_abruptly(self):
+        """Test correct initialization. Server ends connection."""
+        server = SSEMockServer()
+        server.start()
+
+        events = []
+        async def callback(event):
+            """Callback."""
+            events.append(event)
+
+        client = SSEClientAsync(callback)
+
+        async def runner():
+            """SSE client runner thread."""
+            await client.start('http://127.0.0.1:' + str(server.port()))
+
+        client_task = asyncio.gather(runner())
+
+        server.publish({'id': '1'})
+        server.publish({'id': '2', 'event': 'message', 'data': 'abc'})
+        server.publish({'id': '3', 'event': 'message', 'data': 'def'})
+        server.publish({'id': '4', 'event': 'message', 'data': 'ghi'})
+        await asyncio.sleep(1)
+        server.publish(server.VIOLENT_REQUEST_END)
+        server.stop()
+        await asyncio.sleep(1)
+
+        assert events == [
+            SSEEvent('1', None, None, None),
+            SSEEvent('2', 'message', None, 'abc'),
+            SSEEvent('3', 'message', None, 'def'),
+            SSEEvent('4', 'message', None, 'ghi')
+        ]
+
+        assert client._conn is None