Merge pull request #111 from alpacahq/feature/alpaca-stream-reconnect

Add automatic reconnection for Alpaca stream

Author: ttt733

alpaca_trade_api/polygon/streamconn.py

@@ -88,10 +88,7 @@ async def _recv(self):
                 msg = json.loads(r)
                 for update in msg:
                     yield update
-        except websockets.exceptions.ConnectionClosed:
-            # Ignore, occurs on self.close() such as after KeyboardInterrupt
-            pass
-        except websockets.exceptions.ConnectionClosedError as e:
+        except Exception as e:
             await self._dispatch({'ev': 'status',
                                   'status': 'disconnected',
                                   'message':
@@ -121,13 +118,12 @@ async def _ensure_ws(self):
                 await self.connect()
                 if self._streams:
                     await self.subscribe(self._streams)
-
                 break
-            except (ConnectionRefusedError, ConnectionError) as e:
+            except Exception as e:
                 await self._dispatch({'ev': 'status',
                                       'status': 'connect failed',
                                       'message':
-                                      f'Connection Failed ({e})'})
+                                      f'Polygon Connection Failed ({e})'})
                 self._ws = None
                 self._retries += 1
                 time.sleep(self._retry_wait * self._retry)

alpaca_trade_api/stream2.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import os
 import re
 import websockets
 from .common import get_base_url, get_credentials
@@ -16,11 +17,15 @@ def __init__(self, key_id=None, secret_key=None, base_url=None):
         self._handlers = {}
         self._handler_symbols = {}
         self._base_url = base_url
+        self._streams = set([])
         self._ws = None
+        self._retry = int(os.environ.get('APCA_RETRY_MAX', 3))
+        self._retry_wait = int(os.environ.get('APCA_RETRY_WAIT', 3))
+        self._retries = 0
         self.polygon = None
         try:
             self.loop = asyncio.get_event_loop()
-        except:
+        except Exception:
             self.loop = asyncio.new_event_loop()
             asyncio.set_event_loop(self.loop)
 
@@ -43,12 +48,13 @@ async def _connect(self):
                 ("Invalid Alpaca API credentials, Failed to authenticate: {}"
                     .format(msg))
             )
+        else:
+            self._retries = 0
 
         self._ws = ws
         await self._dispatch('authorized', msg)
 
         asyncio.ensure_future(self._consume_msg())
-        return ws
 
     async def _consume_msg(self):
         ws = self._ws
@@ -61,9 +67,9 @@ async def _consume_msg(self):
                 stream = msg.get('stream')
                 if stream is not None:
                     await self._dispatch(stream, msg)
-        finally:
-            await ws.close()
-            self._ws = None
+        except Exception:
+            await self.close()
+            asyncio.ensure_future(self._ensure_ws())
 
     async def _ensure_polygon(self):
         if self.polygon is not None:
@@ -79,10 +85,22 @@ async def _ensure_polygon(self):
     async def _ensure_ws(self):
         if self._ws is not None:
             return
-        self._ws = await self._connect()
+
+        while self._retries <= self._retry:
+            try:
+                await self._connect()
+                if self._streams:
+                    await self.subscribe(self._streams)
+                break
+            except Exception:
+                self._ws = None
+                self._retries += 1
+                await asyncio.sleep(self._retry_wait * self._retry)
+        else:
+            raise ConnectionError("Max Retries Exceeded")
 
     async def subscribe(self, channels):
-        '''Start subscribing channels.
+        '''Start subscribing to channels.
         If the necessary connection isn't open yet, it opens now.
         '''
         ws_channels = []
@@ -94,6 +112,7 @@ async def subscribe(self, channels):
                 ws_channels.append(c)
 
         if len(ws_channels) > 0:
+            self._streams |= set(ws_channels)
             await self._ensure_ws()
             await self._ws.send(json.dumps({
                 'action': 'listen',
@@ -129,7 +148,7 @@ async def unsubscribe(self, channels):
             await self.polygon.unsubscribe(polygon_channels)
 
     def run(self, initial_channels=[]):
-        '''Run forever and block until exception is rasised.
+        '''Run forever and block until exception is raised.
         initial_channels is the channels to start with.
         '''
         loop = self.loop
@@ -146,8 +165,10 @@ async def close(self):
         '''Close any of open connections'''
         if self._ws is not None:
             await self._ws.close()
+            self._ws = None
         if self.polygon is not None:
             await self.polygon.close()
+            self.polygon = None
 
     def _cast(self, channel, msg):
         if channel == 'account_updates':

tests/test_stream2.py

@@ -77,10 +77,6 @@ class TestException(Exception):
     async def on_raise(conn, stream, msg):
         raise TestException()
 
-    with pytest.raises(TestException):
-        _run(conn._consume_msg())
-    assert ws.close.mock.called
-
     # _ensure_polygon
     conn = StreamConn('key-id', 'secret-key')
     with mock.patch('alpaca_trade_api.stream2.polygon') as polygon:
@@ -94,7 +90,6 @@ async def on_raise(conn, stream, msg):
     conn._connect = AsyncMock()
     _run(conn._ensure_ws())
     assert conn._connect.mock.called
-    assert conn._ws is not None
 
     # subscribe
     conn = StreamConn('key-id', 'secret-key')
@@ -116,8 +111,8 @@ async def on_raise(conn, stream, msg):
     conn.polygon = mock.Mock()
     conn.polygon.close = AsyncMock()
     _run(conn.close())
-    assert conn._ws.close.mock.called
-    assert conn.polygon.close.mock.called
+    assert conn._ws is None
+    assert conn.polygon is None
 
     # _cast
     conn = StreamConn('key-id', 'secret-key')