Fixes to polygon WS connection retries, C-c, error messages.

xuperbo · Jul 12, 2019 · 21b9fec · 21b9fec
1 parent 5ca968b
commit 21b9fec
Showing 1 changed file with 17 additions and 13 deletions.
diff --git a/alpaca_trade_api/polygon/stream2.py b/alpaca_trade_api/polygon/stream2.py
@@ -63,13 +63,12 @@ async def _consume_msg(self):
                     stream = update.get('ev')
                     if stream is not None:
                         await self._dispatch(stream, update)
-        except websockets.exceptions.ConnectionClosedError:
+        except websockets.exceptions.ConnectionClosedError as e:
             await self._dispatch('status',
                                  {'ev': 'status',
                                   'status': 'disconnected',
                                   'message':
-                                  'Polygon Disconnected Unexpectedly'})
-        finally:
+                                  f'Polygon Disconnected Unexpectedly ({e})'})
             if self._ws is not None:
                 await self._ws.close()
             self._ws = None
@@ -78,16 +77,21 @@ async def _consume_msg(self):
     async def _ensure_ws(self):
         if self._ws is not None:
             return
-        try:
-            await self.connect()
-        except Exception:
-            self._ws = None
-            self._retries += 1
-            time.sleep(self._retry_wait)
-            if self._retries <= self._retry:
-                asyncio.ensure_future(self._ensure_ws())
-            else:
-                raise ConnectionError("Max Retries Exceeded")
+
+        while self._retries <= self._retry:
+            try:
+                await self.connect()
+            except (ConnectionRefusedError, ConnectionError) as e:
+                await self._dispatch('status',
+                                     {'ev': 'status',
+                                      'status': 'connect failed',
+                                      'message':
+                                      f'Connection Failed ({e})'})
+                self._ws = None
+                self._retries += 1
+                time.sleep(self._retry_wait * self._retry)
+        else:
+            raise ConnectionError("Max Retries Exceeded")
 
     async def subscribe(self, channels):
         '''Start subscribing channels.