132 lines
3.3 KiB
Python
132 lines
3.3 KiB
Python
import asyncio
|
|
import websockets
|
|
from threading import Thread
|
|
from queue import Queue
|
|
from time import sleep
|
|
import json
|
|
from ws.connection import IConnection
|
|
|
|
|
|
class WsClient:
|
|
def __init__(self, listener: IConnection):
|
|
self.listener = listener
|
|
self.loop = asyncio.new_event_loop()
|
|
self.stop = None
|
|
self.bg_thread = None
|
|
|
|
def bg_run(self, uri):
|
|
print("bg_run: started")
|
|
asyncio.set_event_loop(self.loop)
|
|
fut = asyncio.ensure_future(self.run_client(uri), loop=self.loop)
|
|
self.loop.run_until_complete(fut)
|
|
self.stop = None
|
|
print("bg_run: terminated")
|
|
|
|
def connect(self, uri):
|
|
if self.bg_thread is not None and self.bg_thread.is_alive():
|
|
return
|
|
|
|
self.stop: asyncio.Future[None] = self.loop.create_future()
|
|
self.bg_thread = Thread(target=self.bg_run, args=(uri,))
|
|
self.bg_thread.setDaemon(True)
|
|
self.bg_thread.start()
|
|
|
|
def disconnect(self):
|
|
if self.stop is not None:
|
|
self.loop.call_soon_threadsafe(self.stop.set_result, None)
|
|
self.bg_thread.join()
|
|
|
|
async def run_client(self, uri):
|
|
try:
|
|
websocket = await websockets.connect(uri, loop=self.loop)
|
|
except Exception as exc:
|
|
print(f"Failed to connect to {uri}: {exc}.")
|
|
return
|
|
|
|
print("handler: got connection to {}".format(websocket.remote_address))
|
|
self.listener.on_connect(self.loop)
|
|
try:
|
|
path = "/"
|
|
consumer_task = asyncio.ensure_future(self.handler_recv(websocket, path), loop=self.loop)
|
|
producer_task = asyncio.ensure_future(self.handler_send(websocket, path), loop=self.loop)
|
|
done, pending = await asyncio.wait([consumer_task, producer_task, self.stop], return_when=asyncio.FIRST_COMPLETED,)
|
|
for task in pending:
|
|
task.cancel()
|
|
finally:
|
|
print("handler: lost connection from {}".format(websocket.remote_address))
|
|
await websocket.close()
|
|
self.listener.on_disconnect()
|
|
|
|
async def handler_recv(self, websocket, path):
|
|
while True:
|
|
try:
|
|
data = await websocket.recv()
|
|
await self.listener.on_recv(data)
|
|
except:
|
|
print("handler_recv: lost connection from {}".format(websocket.remote_address))
|
|
return
|
|
|
|
async def handler_send(self, websocket, path):
|
|
while True:
|
|
data = await self.listener.on_send()
|
|
if data is not None:
|
|
try:
|
|
await websocket.send(json.dumps(data))
|
|
except:
|
|
print("handler_send: lost connection from {}".format(websocket.remote_address))
|
|
break
|
|
|
|
|
|
class PeriodicUpdater(IConnection):
|
|
def __init__(self):
|
|
self.ws_client = WsClient(self)
|
|
|
|
self.state = Queue()
|
|
self.thread_cancel = True
|
|
self.thread = None
|
|
|
|
def connect(self, uri):
|
|
self.ws_client.connect(uri)
|
|
|
|
def disconnect(self):
|
|
self.ws_client.disconnect()
|
|
|
|
def on_connect(self):
|
|
self.thread_cancel = False
|
|
self.thread = Thread(target=self.run)
|
|
self.thread.setDaemon(True)
|
|
self.thread.start()
|
|
|
|
def on_disconnect(self):
|
|
self.thread_cancel = True
|
|
self.thread.join(1)
|
|
|
|
async def on_send(self):
|
|
return await asyncio.get_event_loop().run_in_executor(None, self.state.get)
|
|
|
|
async def on_recv(self, data):
|
|
print("R: {}".format(data))
|
|
|
|
def send(self, data):
|
|
self.state.put_nowait(data)
|
|
|
|
def run(self):
|
|
count = 0
|
|
while not self.thread_cancel:
|
|
self.send({"Value": count})
|
|
count += 1
|
|
sleep(0.5)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
ws = PeriodicUpdater()
|
|
ws.connect(uri="ws://localhost:8765")
|
|
|
|
sleep(5)
|
|
ws.disconnect()
|
|
sleep(1)
|
|
ws.connect(uri="ws://localhost:8765")
|
|
|
|
while True:
|
|
sleep(1) |