add websocket multiplexer

bfxapi/websockets/GenericWebsocket.py

@@ -6,6 +6,8 @@ import asyncio
 import websockets
 import socket
 import json
+import time
+from threading import Thread
 
 from pyee import EventEmitter
 from ..utils.CustomLogger import CustomLogger
@@ -26,55 +28,126 @@ def is_json(myjson):
         return False
     return True
 
+class Socket():
+    def __init__(self, sId):
+        self.ws = None
+        self.isConnected = False
+        self.isAuthenticated = False
+        self.id = sId
+
+    def set_connected(self):
+        self.isConnected = True
+
+    def set_disconnected(self):
+        self.isConnected = False
+
+    def set_authenticated(self):
+        self.isAuthenticated = True
+
+    def set_websocket(self, ws):
+        self.ws = ws
+
+def _start_event_worker():
+    async def event_sleep_process():
+        """
+        sleeping process for event emitter to schedule on
+        """
+        while True:
+            await asyncio.sleep(0)
+    def start_loop(loop):
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(event_sleep_process())
+    event_loop = asyncio.new_event_loop()
+    worker = Thread(target=start_loop, args=(event_loop,))
+    worker.start()
+    return event_loop
+
 class GenericWebsocket:
     """
     Websocket object used to contain the base functionality of a websocket.
     Inlcudes an event emitter and a standard websocket client.
     """
 
-    def __init__(self, host, logLevel='INFO', loop=None, max_retries=5):
+    def __init__(self, host, logLevel='INFO', loop=None, max_retries=5,
+                 create_event_emitter=_start_event_worker):
         self.host = host
         self.logger = CustomLogger('BfxWebsocket', logLevel=logLevel)
         self.loop = loop or asyncio.get_event_loop()
-        self.events = EventEmitter(
-            scheduler=asyncio.ensure_future, loop=self.loop)
         # overide 'error' event to stop it raising an exception
         # self.events.on('error', self.on_error)
         self.ws = None
         self.max_retries = max_retries
         self.attempt_retry = True
+        self.sockets = {}
+        # start seperate process for the even emitter
+        eventLoop = create_event_emitter()
+        self.events = EventEmitter(scheduler=asyncio.ensure_future, loop=eventLoop)
 
     def run(self):
         """
-        Run the websocket connection indefinitely
+        Starte the websocket connection. This functions spawns the initial socket
+        thread and connection.
         """
-        self.loop.run_until_complete(self._main(self.host))
+        self._start_new_socket()
 
     def get_task_executable(self):
         """
         Get the run indefinitely asyncio task
         """
-        return self._main(self.host)
+        return self._run_socket()
 
-    async def _connect(self, host):
-        async with websockets.connect(host) as websocket:
-            self.ws = websocket
-            self.logger.info("Wesocket connected to {}".format(host))
+    def _start_new_socket(self, socketId=None):
+        if not socketId:
+            socketId = len(self.sockets)
+        def start_loop(loop):
+            asyncio.set_event_loop(loop)
+            loop.run_until_complete(self._run_socket())
+        worker_loop = asyncio.new_event_loop()
+        worker = Thread(target=start_loop, args=(worker_loop,))
+        worker.start()
+        return socketId
+
+    def _wait_for_socket(self, socket_id):
+        """
+        Block until the given socket connection is open
+        """
+        while True:
+            socket = self.sockets.get(socket_id, False)
+            if socket:
+                if socket.isConnected and socket.ws:
+                    return
+            time.sleep(0.01)
+
+    async def _connect(self, socket):
+        async with websockets.connect(self.host) as websocket:
+            self.sockets[socket.id].set_websocket(websocket)
+            self.sockets[socket.id].set_connected()
+            self.logger.info("Wesocket connected to {}".format(self.host))
             while True:
                 await asyncio.sleep(0)
                 message = await websocket.recv()
-                await self.on_message(message)
+                await self.on_message(socket.id, message)
 
-    def get_ws(self):
-        return self.ws
+    def get_socket(self, socketId):
+        return self.sockets[socketId]
 
-    async def _main(self, host):
+    def get_authenticated_socket(self):
+        for socketId in self.sockets:
+            if self.sockets[socketId].isAuthenticated:
+                return self.sockets[socketId]
+        return None
+
+    async def _run_socket(self):
         retries = 0
+        sId =  len(self.sockets)
+        s = Socket(sId)
+        self.sockets[sId] = s
         while retries < self.max_retries and self.attempt_retry:
             try:
-                await self._connect(host)
+                await self._connect(s)
                 retries = 0
             except (ConnectionClosed, socket.error) as e:
+                self.sockets[sId].set_disconnected()
                 self._emit('disconnected')
                 if (not self.attempt_retry):
                     return