I am trying to create a tcp server that handles small amount of connections (about a 100, but could grow to more), but each tcp connection will be persistent and stream GBs of data off which I will be uploading to Elasticsearch database via REST API.
I have tried two independent ways (asyncio and multiprocessing), and I see MP has the data uploaded quicker, but at the cost of a ton of processes, and asyncio only uses 1 process, but the data takes about 5 mins to get into the DB. I was wondering if there is a hybrid way such that have a batching mechanism to back 5 connections into a child process and run an asyncio loop in there to handle the connections.
Kinda what I am thinking:
async def dial_out_server(args):
conn_handler = ClientConnection(args.elastic_server)
server = await asyncio.start_server(
conn_handler.handle_connection, args.host, args.port)
addr = server.sockets[0].getsockname()
print(f'Serving on {addr}')
async with server:
#batch the connections and spawn a new process with a loop here.
#Not sure what is best way or if that is possible
await server.serve_forever()
if __name__ == '__main__':
parser = ArgumentParser()
parser.add_argument("-a", "--host", dest="host", help="host", required=True)
parser.add_argument("-r", "--port", dest="port", help="port", required=True)
parser.add_argument("-e", "--elastic_server", dest="elastic_server", help="Elastic Server", required=True)
args = parser.parse_args()
loop = uvloop.new_event_loop()
asyncio.set_event_loop(loop)
asyncio.run(dial_out_server(args))
Thanks
Full Code:
from utils import process_cisco_encoding
from telemetry_pb2 import Telemetry
from argparse import ArgumentParser
from struct import Struct, unpack
from aiohttp import ClientSession
from logging.handlers import RotatingFileHandler, QueueHandler
import grpc
import logging
import asyncio
import json
import logging
import traceback
import uvloop
class Error(Exception):
pass
class GetIndexListError(Error):
def __init__(self, traceback, response_json, message, e):
self.traceback = traceback
self.response = response_json
self.message = message
self.exception = e
class PostDataError(Error):
def __init__(self, traceback, response_json, data, message, e):
self.traceback = traceback
self.response = response_json
self.data = data
self.message = message
self.exception = e
class PutIndexError(Error):
def __init__(self, traceback, response_json, message, e):
self.traceback = traceback
self.response = response_json
self.message = message
self.exception = e
class ElasticSearchError(Error):
def __init__(self, response_json, message):
self.response = response_json
self.message = message
class ClientConnection(object):
def __init__(self, elastic_server):
self.elastic_server = elastic_server
self.lock = asyncio.Lock()
self.log = None
async def get_index_list(self, url):
indices = []
try:
async with ClientSession() as session:
async with session.get(url) as response:
response = await response.read()
response = json.loads(response.decode())
for key in response:
if not key.startswith('.'):
indices.append(key)
return indices
except Exception as e:
raise GetIndexListError(traceback.print_exc(), response, "Got Exception while trying to get index list", e)
async def post_data(self, data_to_post):
headers = {'Content-Type': "application/x-ndjson"}
url = f"http://{self.elastic_server}:9200/_bulk"
try:
async with ClientSession() as session:
response = await session.post(url, data=data_to_post, headers=headers)
return response
except Exception as e:
raise PostDataError(traceback.print_exc(), data_to_post, "Got Exception while trying to post data", e)
async def put_index(self, index):
url = f"http://{self.elastic_server}:9200/{index}"
headers = {'Content-Type': "application/json"}
mapping = {"mappings": {"properties": {"#timestamp": {"type": "date"}}}}
try:
async with ClientSession() as session:
response = await session.put(url, json=mapping, headers=headers)
return response
except Exception as e:
raise PutDataError(traceback.print_exc(), response, f"Got Exception while trying to put index {index}", e)
async def init_logger(self, address):
log_name = "dial-out.log"
log = logging.getLogger(log_name)
log.setLevel(logging.INFO)
file_handler = RotatingFileHandler(log_name, maxBytes=536870912, backupCount=2)
screen_handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s %(processName)-10s %(name)s %(levelname)-8s %(message)s')
file_handler.setFormatter(formatter)
screen_handler.setFormatter(formatter)
log.addHandler(file_handler)
log.addHandler(screen_handler)
return log
async def handle_connection(self, reader, writer):
try:
HEADER_SIZE = 12
header_struct = Struct('>hhhhi')
_UNPACK_HEADER = header_struct.unpack
address = writer.get_extra_info('peername')
if self.log is None:
self.log = await self.init_logger(address)
self.log.info(f"Got Connection from {address[0]}:{address[1]}")
while True:
header_data = await reader.read(HEADER_SIZE)
msg_type, encode_type, msg_version, flags, msg_length = _UNPACK_HEADER(header_data)
encoding = {1:'gpb', 2:'json'}[encode_type]
msg_data = b''
if encode_type == 1:
while len(msg_data) < msg_length:
packet = await reader.read(msg_length - len(msg_data))
msg_data += packet
sorted_by_index = {}
converted_decode_segments = process_cisco_encoding([msg_data])
for converted_decode_segment in converted_decode_segments:
if not converted_decode_segment["_index"] in sorted_by_index.keys():
sorted_by_index[converted_decode_segment["_index"]] = [converted_decode_segment]
else:
sorted_by_index[converted_decode_segment["_index"]].append(converted_decode_segment)
index_list = await self.get_index_list(f"http://{self.elastic_server}:9200/*")
for index in sorted_by_index.keys():
if index not in index_list:
async with self.lock:
index_list = await self.get_index_list(f"http://{self.elastic_server}:9200/*")
if index not in index_list:
self.log.info("Acciqured lock to put index in elasticsearch")
response = await self.put_index(index)
if response.status is not 200:
raise ElasticSearchError(await response.json(), "Unable to put index into Elasticsearch")
else:
index_list.append(index)
else:
segment_list = sorted_by_index[index]
elastic_index = {'index': {'_index': f'{index}'}}
payload_list = [elastic_index]
for segment in segment_list:
segment.pop('_index', None)
payload_list.append(segment)
payload_list.append(elastic_index)
payload_list.pop()
data_to_post = '\n'.join(json.dumps(d) for d in payload_list)
data_to_post += '\n'
response = await self.post_data(data_to_post)
if response.status is not 200:
raise ElasticSearchError(await response.json(), "Unable to put data into Elasticsearch")
except GetIndexListError as e:
self.log.error(e.message)
self.log.error(e.traceback)
self.log.error(e.response)
self.log.error(e.exception)
await writer.drain()
self.log.error(f"Closing connection from {address[0]}")
writer.close()
except PostDataError as e:
self.log.error(e.message)
self.log.error(e.traceback)
self.log.error(e.response)
self.log.error(e.exception)
self.log.error(e.data)
await writer.drain()
self.log.error(f"Closing connection from {address[0]}")
writer.close()
except PutIndexError as e:
self.log.error(e.message)
self.log.error(e.traceback)
self.log.error(e.response)
await writer.drain()
self.log.error(f"Closing connection from {address[0]}")
writer.close()
except ElasticSearchError as e:
self.log.error(e.message)
self.log.error(e.response)
await writer.drain()
self.log.error(f"Closing connection from {address[0]}")
writer.close()
except Exception as e:
self.log.error(e)
self.log.error(traceback.print_exc())
await writer.drain()
self.log.error(f"Closing connection from {address[0]}")
writer.close()
async def dial_out_server(args):
conn_handler = ClientConnection(args.elastic_server)
server = await asyncio.start_server(
conn_handler.handle_connection, args.host, args.port)
addr = server.sockets[0].getsockname()
print(f'Serving on {addr}')
async with server:
await server.serve_forever()
if __name__ == '__main__':
parser = ArgumentParser()
parser.add_argument("-a", "--host", dest="host", help="host", required=True)
parser.add_argument("-r", "--port", dest="port", help="port", required=True)
parser.add_argument("-e", "--elastic_server", dest="elastic_server", help="Elastic Server", required=True)
args = parser.parse_args()
loop = uvloop.new_event_loop()
asyncio.set_event_loop(loop)
asyncio.run(dial_out_server(args))
Related
I'm trying to post a dataframe to a Discord channel. However, I am having issue getting Discord.py to close the connection and move on the next task. I've tried using the event loop as suggested in this thread (How to run async function in Airflow?) as well as asyncio.run() function. Not really familiar with the async and hoping to get some pointers here. Below is my code in Python that I've tried importing in DAG and Task without success. Thanks in advance!
Airflow: 2.5.1
Python: 3.7
import discord
from tabulate import tabulate
import asyncio
import pandas as pd
async def post_to_discord(df, channel_id, bot_token, as_message=True, num_rows=5):
intents = discord.Intents.default()
intents.members = True
client = discord.Client(intents=intents)
try:
#client.event
async def on_ready():
channel = client.get_channel(channel_id)
if as_message:
# Post the dataframe as a message, num_rows rows at a time
for i in range(0, len(df), num_rows):
message = tabulate(df.iloc[i:i+num_rows,:], headers='keys', tablefmt='pipe', showindex=False)
await channel.send(message)
else:
# Send the dataframe as a CSV file
df.to_csv("dataframe.csv", index=False)
with open("dataframe.csv", "rb") as f:
await channel.send(file=discord.File(f))
# client.run(bot_token)
await client.start(bot_token)
await client.wait_until_ready()
finally:
await client.close()
async def main(df, channel_id, bot_token, as_message=True, num_rows=5):
# loop = asyncio.get_event_loop()
# result = loop.run_until_complete(post_to_discord(df, channel_id, bot_token, as_message, num_rows))
result = asyncio.run(post_to_discord(df, channel_id, bot_token, as_message, num_rows))
await result
return result
if __name__ =='__main__':
main()
It seems like your script works but the server is blocking the open socket (and kudos - the discord server is good at that). So we will work by creating a ping function (adopted from another answer).
def ping(ip, port):
try:
s = socket.socket() # TCP - standard values are `socket.AF_INET, socket.SOCK_STREAM` so you don't have to write them
s.settimeout(2)
print('[DEBUG] connect')
s.connect((ip, int(port)))
#result = s.connect_ex((ip, int(port)))
#print('result:', result)
return True
except socket.timeout as ex:
print('[DEBUG] timeout')
return True
except Exception as ex:
print('[Exception]', ex)
return False
finally:
print('[DEBUG] close')
s.close()
Feel free to test your ID
id = ...channel number...
print(id, type(id))
and you should see
<built-in function id> <class 'builtin_function_or_method'>
Then let us move on to improving your code:
import discord
import asyncio
import time # you are not using this module
import socket
import os
from tabulate import tabulate
import pandas as pd # as pd is not required
def ping(ip, port):
try:
s = socket.socket() # TCP - standard values are `socket.AF_INET, socket.SOCK_STREAM` so you don't have to write them
s.settimeout(2)
print('[DEBUG] connect')
s.connect((ip, int(port)))
#result = s.connect_ex((ip, int(port)))
#print('result:', result)
return True
except socket.timeout as ex:
print('[DEBUG] timeout')
return True
except Exception as ex:
print('[Exception]', ex)
return False
finally:
print('[DEBUG] close')
s.close()
TOKEN = os.getenv('DISCORD_TOKEN')
client = discord.Client()
async def post_to_discord(df, channel_id, bot_token, as_message=True, num_rows=5):
intents = discord.Intents.default()
intents.members = True
client = discord.Client(intents=intents)
try:
#client.event
async def on_ready():
channel = client.get_channel(channel_id)
if as_message:
# Post the dataframe as a message, num_rows rows at a time
for i in range(0, len(df), num_rows):
message = tabulate(df.iloc[i:i+num_rows,:], headers='keys', tablefmt='pipe', showindex=False)
await channel.send(message)
else:
# Send the dataframe as a CSV file
df.to_csv("dataframe.csv", index=False)
with open("dataframe.csv", "rb") as f:
await channel.send(file=discord.File(f))
# client.run(bot_token)
await client.start(bot_token)
await client.wait_until_ready()
while True:
online = ping("26.51.174.109", "25565") #modify it as you see fit
#online = ping("192.168.1.101", "8081") #same as above
if online:
print("server online")
#await channel.edit(name="Server Status - Online")
else:
print("server offline")
#await channel.edit(name="Server Status - Offline")
await asyncio.sleep(5)
# optional - client.run(TOKEN)
finally:
await client.close()
async def main(df, channel_id, bot_token, as_message=True, num_rows=5):
# loop = asyncio.get_event_loop()
# result = loop.run_until_complete(post_to_discord(df, channel_id, bot_token, as_message, num_rows))
result = asyncio.run(post_to_discord(df, channel_id, bot_token, as_message, num_rows))
await result
return result
if __name__ =='__main__':
main()
I want to do the following thing:
I want to make a websocket, which prints me all events about esports and I want to use https://sofascore.com
I've inspected the network requests as usual and it seems that I need to send a Auth WebSocket Content first, then one for subscribing the right sport and then I will receive my events I need.
I've wrote the following code:
import websockets
import asyncio
from websockets.extensions import permessage_deflate
async def esports():
async with websockets.connect('wss://ws.sofascore.com:9222/', compression='deflate') as websocket:
msg = await websocket.recv()
print(f"From Server: {msg}")
t = await websocket.send(
'CONNECT {"no_responders":true,"protocol":1,"verbose":false,"pedantic":false,"user":"none","pass":"none","lang":"nats.ws","version":"1.8.1","headers":true}')
await websocket.send("PING")
pong = await websocket.recv()
print(f"From Server: {pong}")
await websocket.send(
'SUB sport.esports 6')
while (True):
msg = await websocket.recv()
print(f"From Server: {msg}")
asyncio.get_event_loop().run_until_complete(esports())
I know that the websocket is compressed as permessage_deflate, when I saw into the request headers of the websocket.
But I still get an error:
Traceback (most recent call last):
File "C:\Users\Coding\Desktop\websockett.py", line 23, in <module>
asyncio.get_event_loop().run_until_complete(esports())
File "C:\Users\Coding\AppData\Local\Programs\Python\Python39-32\lib\asyncio\base_events.py", line 642, in run_until_complete
return future.result()
File "C:\Users\Coding\Desktop\websockett.py", line 15, in esports
await websocket.send(
File "C:\Users\Coding\AppData\Roaming\Python\Python39\site-packages\websockets\legacy\protocol.py", line 620, in send
await self.ensure_open()
File "C:\Users\Coding\AppData\Roaming\Python\Python39\site-packages\websockets\legacy\protocol.py", line 921, in ensure_open
raise self.connection_closed_exc()
websockets.exceptions.ConnectionClosedError: received 1008 (policy violation) Authentication Timeout; then sent 1008 (policy violation) Authentication Timeout
Process finished with exit code 1
EDIT:
I have now found out that the whole thing works with the Nats network. Is there any way to use Nats with a Libary that also supports the websockets?
Haven't found one on github or pypi unfortunately....
Ideally you would be able to use the nats-py library:
import asyncio
import nats
async def handler(msg):
print(f"From server: {msg}")
async def main():
nc = await nats.connect("wss://ws.sofascore.com:9222")
await nc.subscribe("sport.esports", cb=handler)
if __name__ == "__main__":
loop = asyncio.get_event_loop()
loop.run_until_complete(main())
try:
loop.run_forever()
finally:
loop.close()
However, this library does not currently support connecting with WebSockets, so
the above doesn't work (yet - it looks like it's being worked on right
now).
For your code, the only reason it fails is that the messages you're sending
don't end with \r\n, which the NATS protocol requires. The code works as
expected with this change:
import asyncio
import websockets
async def esports():
async with websockets.connect('wss://ws.sofascore.com:9222') as websocket:
msg = await websocket.recv()
print(f"From Server: {msg}")
await websocket.send(
'CONNECT {"no_responders":true,"protocol":1,"verbose":false,"pedantic":false,"user":"none","pass":"none","lang":"nats.ws","version":"1.8.1","headers":true}\r\n'
)
await websocket.send("SUB sport.esports 1\r\n")
async for msg in websocket:
print(f"From Server: {msg}")
asyncio.run(esports())
Of course this will eventually get disconnected because it doesn't respond to
PING messages. Here's a little more fleshed out script which implements enough
of the NATS protocol to log the sport.esports messages:
import asyncio
import json
import textwrap
from dataclasses import dataclass
import websockets
class SofaError(Exception):
pass
def message_string(message, data=None, pretty=False):
s = message
if data is not None:
if pretty:
s += json.dumps(data, indent=2)
else:
s += json.dumps(data, separators=(",", ":"))
return s
def log(pre, message, data=None):
print(textwrap.indent(message_string(message, data, True), pre))
def recv_log(message, data=None):
log("< ", message, data)
async def send(websocket, message, data=None):
log("> ", message, data)
data = (message_string(message, data, False) + "\r\n").encode()
await websocket.send(data)
async def connect_and_subscribe(websocket):
connect_options = {
"no_responders": True,
"protocol": 1,
"verbose": False,
"pedantic": False,
"user": "none",
"pass": "none",
"lang": "nats.ws",
"version": "1.8.1",
"headers": True,
}
await send(websocket, "CONNECT ", connect_options)
await send(websocket, "SUB sport.esports 1")
#dataclass
class NatsMsg:
subject: str
sid: str
reply_to: str
size: int
payload: bytes
def parse_msg(info_line, pending_data):
if not info_line:
raise SofaError("No payload information received")
info = [b.decode(errors="replace") for b in info_line.split(b" ")]
if len(info) == 3:
subject, sid, size = info
reply_to = None
elif len(info) == 4:
subject, sid, reply_to, size = info
else:
raise SofaError("Unrecognized info format")
try:
size = int(size)
except ValueError:
raise SofaError("Bad payload size")
if len(pending_data) < size:
raise SofaError("Incomplete payload")
payload = pending_data[:size]
pending_data = pending_data[size:]
return NatsMsg(subject, sid, reply_to, size, payload), pending_data
async def handler(websocket, ws_message, connected):
while len(ws_message):
nats_message, _, ws_message = ws_message.partition(b"\r\n")
if not nats_message:
continue
op, _, rest = nats_message.partition(b" ")
if op == b"-ERR":
recv_log(nats_message.decode(errors="replace"))
err = rest.strip(b"'").decode(errors="replace") if rest else "(No message received)"
raise SofaError(f"Server error: {err}")
elif op == b"INFO":
info_options = json.loads(rest) if rest else None
recv_log("INFO ", info_options)
if not connected:
await connect_and_subscribe(websocket)
connected = True
elif op == b"PING":
recv_log("PING")
await send(websocket, "PONG")
elif op == b"MSG":
try:
msg, ws_message = parse_msg(rest, ws_message)
except SofaError as e:
recv_log(f"MSG (Error: {e}) {rest}")
continue
msg_info = (
f"MSG subject={msg.subject} sid={msg.sid} "
f"reply-to={msg.reply_to} nbytes={msg.size}:\n"
)
try:
decoded = msg.payload.decode()
data = json.loads(decoded)
except UnicodeError:
recv_log(f"{msg_info}{msg.payload}")
except json.JSONDecodeError:
recv_log(f"{msg_info}{decoded}")
else:
recv_log(msg_info, data)
else:
recv_log(f"(Unhandled op) {nats_message.decode(errors='replace')}")
return connected
async def main():
async with websockets.connect("wss://ws.sofascore.com:9222") as websocket:
connected = False
async for message in websocket:
connected = await handler(websocket, message, connected)
if __name__ == "__main__":
asyncio.run(main())
I have started to work with Asyncio recently and I was wondering if that is possible thing to do.
I made a TCP Echo server that get "commands" from client in order to execute methods with binance.client (which is stored on external file called GetData), I've also made a client that connects to binance websocket to receive live updates.
The problem is I want to combine the TCP Server and the Client that connects to binance websocket so when the server is running it would be able to handle the streaming data and still let the server handle new client connections.
Server Code:
import asyncio
from GetData import *
class EchoServerProtocol(asyncio.Protocol):
def connection_made(self, transport):
peername = transport.get_extra_info('peername')
print('Connection from {}'.format(peername))
self.transport = transport
def data_received(self, data):
recivedData = data.decode()
print('Data received: {!r}'.format(recivedData))
commands = recivedData.split(' ')
if commands[0] == 'ping':
result = 'pong!'
print(result)
elif commands[0] == 'buy' and commands[1].upper() in get_all_assets():
symbol = commands[1].upper()
result = f'Buying request for {symbol}'
try:
current_status = buy_asset(symbol, commands[2])
result = f"You bought: {current_status['asset_amount']} {symbol}, Timestamp: {current_status['time']}"
print(result)
# Save on firebase >>
except IndexError:
print("Please mention usdt amount.")
else:
result = 'invalid command!'
print(result)
self.transport.write(data)
print('Close the client socket')
self.transport.close()
def main():
loop = asyncio.get_event_loop()
coro = loop.create_server(EchoServerProtocol, '127.0.0.1', 8070)
print("Server is now running")
server = loop.run_until_complete(coro)
try:
loop.run_forever()
except KeyboardInterrupt:
print("exit")
finally:
server.close()
loop.close()
if __name__ == "__main__":
asyncio.run(main())
Client code:
from autobahn.asyncio.websocket import WebSocketClientProtocol, WebSocketClientFactory
from GetData import *
import os
import json
class BinanceMarketStream(WebSocketClientProtocol):
def onConnect(self, response):
print("Server connected: {0}".format(response.peer))
self.df_historicalData = get_historical_kline('BTCUSDT', '1m')
def onOpen(self):
print("WebSocket connection open.")
def onMessage(self, payload, isBinary):
if isBinary:
print("Binary message received: {0} bytes".format(len(payload)))
else:
data = payload.decode('utf8')
os.system('cls' if os.name == 'nt' else 'clear')
json_message = json.loads(data)
candle = json_message['k']
is_candle_closed = candle['x']
currentCandle = set_to_df(candle)
self.df_historicalData = update_currentCandle(
self.df_historicalData, currentCandle)
if is_candle_closed:
newCandle = set_to_df(candle)
self.df_historicalData = append_newcandle(
self.df_historicalData, newCandle)
# Trading Algorithms goes to here...
def onClose(self, wasClean, code, reason):
print("WebSocket connection closed: {0}".format(reason))
if __name__ == '__main__':
import asyncio
factory = WebSocketClientFactory(
"wss://stream.binance.com:9443/ws/btcusdt#kline_1m")
factory.protocol = BinanceMarketStream
loop = asyncio.get_event_loop()
coro = loop.create_connection(
factory, "stream.binance.com", 9443, ssl=True)
loop.run_until_complete(coro)
loop.run_forever()
loop.close()
I can also add the client that connects to the TCP server if needed.
I am trying to receive messages via websocket-client module and be able to use received messages for other purposes (e.g. execute buy/sell orders based on incoming messages).
Here is what I have so far:
import websocket
import time
import json
def on_message(ws, message):
try:
current_price = json.loads(message)
print(current_price["price"]) # data type is dict.. only showing values for the key 'price'
except:
print("Please wait..")
time.sleep(1)
def on_error(ws, error):
print(error)
def on_close(ws):
print("### closed ###")
def on_open(ws):
sub_params = {'type': 'subscribe', 'product_ids': ['BTC-USD'], 'channels': ['ticker']}
ws.send(json.dumps(sub_params))
if __name__ == "__main__":
websocket.enableTrace(False)
ws = websocket.WebSocketApp("wss://ws-feed.pro.coinbase.com/",
on_open = on_open,
on_message = on_message,
on_error = on_error,
on_close = on_close)
ws.run_forever()
Running this code will print the current Bitcoin price (current_price) as they come in through its websocket feed.
What I want to do next is to be able to access that variable current_price outside of the websocket function, and I am having a difficulty here. Writing anything beyond ws.run_forever() will be ignored because the websocket event loop will never end.
So I tried running the websocket on a separate thread with 'threading' mordule:
import websocket
import json
import threading
current_price = 0
def on_message(ws, message):
global current_price
current_price = message
def on_error(ws, error):
print(error)
def on_close(ws):
print("### closed ###")
def on_open(ws):
sub_params = {'type': 'subscribe', 'product_ids': ['BTC-USD'], 'channels': ['ticker']}
ws.send(json.dumps(sub_params))
if __name__ == "__main__":
websocket.enableTrace(False)
ws = websocket.WebSocketApp("wss://ws-feed.pro.coinbase.com/",
on_open = on_open,
on_message = on_message,
on_error = on_error,
on_close = on_close)
ws_thread = threading.Thread(target = ws.run_forever)
ws_thread.start()
print(current_price)
and this returns 0. What can I do to make this work?
Not sure if this is the most appropriate answer, but found a way to make this work.
import queue
.
.
.
.
def on_message(ws, message):
current_price = message
q.put(current_price)
.
.
.
ws_thread.start()
while True:
print(q.get())
The key to is is to use functools partials to wrap the run_forever thread's callback functions. I've built a simple multithreaded demonstrator:
from websocket import WebSocketApp
import time
from threading import Thread
from queue import Queue
from functools import partial
class websocket_client:
def __init__(self):
self.responseQ = Queue() #our message response Q
#add the response Qs on by wrapping the callbacks in functools partials
self.websocket = WebSocketApp("ws://echo.websocket.events",
on_open = partial(self.on_open, responseQ=self.responseQ),
on_message = partial(self.on_message, responseQ=self.responseQ),
on_error = partial(self.on_error, responseQ=self.responseQ),
on_close = partial(self.on_close, responseQ=self.responseQ))
# start run_forever as thread. request keep alive pings every 10 seconds with 5 second timouts
self.controller = Thread(target=self.websocket.run_forever, args=(None, None, 10, 5))
self.controller.start()
def on_open(self, websocket, responseQ):
#print("Opening Connection (run_forever)")
responseQ.put("Connected") #transmit status information
def on_error(self, websocket, error, responseQ):
#print("Connection Error (run_forever):", error)
responseQ.put(error) # or error messages
def on_message(self, websocket, message, responseQ):
#print("Message Received (run_forever):", message)
responseQ.put(message)
def on_close(self, websocket, status, message, responseQ):
#print("Closing Connection (run_forever):", status, message)
responseQ.put({'status': message, 'message': message}) #send close status and message in a dictionary
if __name__ == '__main__':
client = websocket_client() # make the client
while client.responseQ.get() != "Connected": time.sleep(0.1) #block until connected
print("\nSERVER GREETING:", client.responseQ.get()) #greeting from server
while True:
message_to_send = time.strftime("%b %d %Y %H:%M:%S")
client.websocket.send(message_to_send)
message_received = client.responseQ.get()
print("Main Thread: SENT --> '%s' RECEIVED --> '%s'" %(message_to_send, message_received))
time.sleep(1)
I am trying to implement an application specific reverse-proxy for jupyter notebooks using aiohttp. It works fine for http requests, but the websocket forwarding does not work. Requests from the browser arrive and get forwarded, but there are no responses from jupyter forthcoming. I assume my websocket client code somehow does not react to incoming messages from jupyter.
The only indication on the jupyter side that something is amiss are messages like this:
WebSocket ping timeout after 90009 ms.
so here is my attempt at writing the proxy
from aiohttp import web
from aiohttp import client
import aiohttp
import logging
import pprint
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
baseUrl = 'http://0.0.0.0:8888'
mountPoint = '/fakeUuid'
async def handler(req):
proxyPath = req.match_info.get('proxyPath','no proxyPath placeholder defined')
reqH = req.headers.copy()
if reqH['connection'] == 'Upgrade' and reqH['upgrade'] == 'websocket' and req.method == 'GET':
ws_server = web.WebSocketResponse()
await ws_server.prepare(req)
logger.info('##### WS_SERVER %s' % pprint.pformat(ws_server))
client_session = aiohttp.ClientSession()
async with client_session.ws_connect(baseUrl+req.path_qs,
headers = { 'cookie': reqH['cookie'] },
) as ws_client:
logger.info('##### WS_CLIENT %s' % pprint.pformat(ws_client))
async for server_msg in ws_server:
logger.info('>>> msg from browser: %s',pprint.pformat(server_msg))
if server_msg.type == aiohttp.WSMsgType.TEXT:
await ws_client.send_str(server_msg.data)
else:
await ws_client.send_bytes(server_msg.data)
async for client_msg in ws_client:
logger.info('>>> msg from jupyter: %s',pprint.pformat(client_msg))
if client_msg.tp == aiohttp.WSMsgType.TEXT:
await ws_server.send_str(client_msg.data)
else:
await ws_server.send_bytes(client_msg.data)
return ws_server
else:
async with client.request(
req.method,baseUrl+mountPoint+proxyPath,
headers = reqH,
allow_redirects=False,
data = await req.read()
) as res:
headers = res.headers.copy()
body = await res.read()
return web.Response(
headers = headers,
status = res.status,
body = body
)
return ws_server
app = web.Application()
app.router.add_route('*',mountPoint + '{proxyPath:.*}', handler)
web.run_app(app,port=3984)
Lesson learned: the two async for are blocking in the flow of the current function. By running them with asyncio.wait I can get them to run at the same time. The resulting program looks like this:
from aiohttp import web
from aiohttp import client
import aiohttp
import asyncio
import logging
import pprint
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
baseUrl = 'http://0.0.0.0:8888'
mountPoint = '/fakeUuid'
async def handler(req):
proxyPath = req.match_info.get('proxyPath','no proxyPath placeholder defined')
reqH = req.headers.copy()
if reqH['connection'] == 'Upgrade' and reqH['upgrade'] == 'websocket' and req.method == 'GET':
ws_server = web.WebSocketResponse()
await ws_server.prepare(req)
logger.info('##### WS_SERVER %s' % pprint.pformat(ws_server))
client_session = aiohttp.ClientSession(cookies=req.cookies)
async with client_session.ws_connect(
baseUrl+req.path_qs,
},
) as ws_client:
logger.info('##### WS_CLIENT %s' % pprint.pformat(ws_client))
async def wsforward(ws_from,ws_to):
async for msg in ws_from:
logger.info('>>> msg: %s',pprint.pformat(msg))
mt = msg.type
md = msg.data
if mt == aiohttp.WSMsgType.TEXT:
await ws_to.send_str(md)
elif mt == aiohttp.WSMsgType.BINARY:
await ws_to.send_bytes(md)
elif mt == aiohttp.WSMsgType.PING:
await ws_to.ping()
elif mt == aiohttp.WSMsgType.PONG:
await ws_to.pong()
elif ws_to.closed:
await ws_to.close(code=ws_to.close_code,message=msg.extra)
else:
raise ValueError('unexpecte message type: %s',pprint.pformat(msg))
finished,unfinished = await asyncio.wait([wsforward(ws_server,ws_client),wsforward(ws_client,ws_server)],return_when=asyncio.FIRST_COMPLETED)
return ws_server
else:
async with client.request(
req.method,baseUrl+mountPoint+proxyPath,
headers = reqH,
allow_redirects=False,
data = await req.read()
) as res:
headers = res.headers.copy()
body = await res.read()
return web.Response(
headers = headers,
status = res.status,
body = body
)
return ws_server
app = web.Application()
app.router.add_route('*',mountPoint + '{proxyPath:.*}', handler)
web.run_app(app,port=3984)