asyncio Event/Condition not working with aiohttp? - python

The following is a reduced version of server that periodically serves any connected clients with telemetry in the form of json strings. This was my initial attempt, in which the main loop pushes data to all connected clients. However, I cannot simply let the handler terminate after "registering" the client. The connection will be closed. So I need to block the handler until the main loop determines the client has disconnected. Signalling the handler through an Event simply does nothing.
#routes.get('/telemetry/json')
async def handler(request: Request):
global CLIENT
CLIENT = await StreamResponse().prepare(request)
log.debug(f"Wait for {EVENT}")
await EVENT.wait() # This never wakes up!
log.debug(f"Client {request.remote} disconnected")
async def main():
global EVENT
EVENT = Event()
app = web.Application()
app.add_routes(routes)
runner = web.AppRunner(app)
await runner.setup()
await web.TCPSite(runner, port=8080).start()
while True:
await sleep(1)
if CLIENT is None:
continue
try:
await CLIENT.write('FLUSH\n'.encode('utf-8'))
await CLIENT.drain()
except ConnectionResetError:
log.debug(f"Notify {EVENT}")
EVENT.set()
log.addHandler(logging.StreamHandler())
log.setLevel(10)
asyncio.run(main())
To clarify: The use of global CLIENT and EVENT is not how it is intended. The handling of multipple clients was removed to make the example code as short as possible.

Related

Python Websocket and Async using

I try to create a websocket server, I wanna make a client to exchange data from server to client, but now my data from other process, I need to make a queue accept data from other process, it makes my main websocket function blocked, the final result is that could not reconnect after client connection break, I think it blocked in the code of queue.
Here is my part of my code:
class RecorderEventHook(object):
def __init__(self, high_event_mq):
self.high_event_mq = high_event_mq
self.msg = None
self.loop = None
# #wrap_keep_alive
async def on_msg_event(self, websocket):
try:
# async for message in websocket:
while True:
msg = self.high_event_mq.get()
await websocket.send(json.dumps(msg))
# msg
except Exception as error:
print(error)
async def event_controller(self):
await websockets.serve(self.on_msg_event, 'localhost', 8888)
def start(self):
loop = asyncio.new_event_loop()
loop.create_task(self.event_controller())
loop.run_forever()
I try to save connected websocket object and using in other thread(in same process), but it failed and mentions
"xxxx" function never waited
I want to be able to receive data from other processes without affecting the normal reconnection of the client.
Anybody help and big appreciate.

Python asyncio: Cancel streams server and all clients

I have a streams server that handles multiple independent clients. When I shut it down I want to notify all clients that the server has shut down.
I figured out how to close the server to new connections, but not how to cancel the specific handlers waiting for client data.
So far the only solution I found is to cancel all tasks in the loop, but this doesn't work for me as I have other tasks that must finish their jobs first.
Does asyncio provide some interface for this or do I have to keep track of all connections myself and cancel them once the server shuts down? I would prefer if the connection handler catches an exception when it calls await reader.readuntil() and not in the middle of execution, but this is not required.
Right now the client looses connection without warning. With this it cannot tell if it was a network issue or if the server shut down.
import asyncio
import signal
server = None
shutdown = False
async def important_task():
while not shutdown:
await asyncio.sleep(10)
print("I am important")
async def handle_conn(reader,writer):
print("Got connection")
try:
while True:
text = await reader.readuntil(b'\n')
# Do stuff
writer.write( text ) # Echo example
await writer.drain()
except serverShutdownException: # How do I cause something like this?
writer.write(b"Goodbye")
await writer.drain()
finally:
writer.close()
await writer.wait_closed()
def handle_sig(num,frame):
global shutdown
print(f"Caught {num}")
server.close()
shutdown = True
async def serve():
global server
server = await asyncio.start_server(handle_conn,"127.0.0.1",8080)
try:
await server.serve_forever()
except asyncio.CancelledError:
pass
await server.wait_closed()
# wait for all handlers to be done?
def main():
signal.signal(signal.SIGINT, handle_sig)
loop = asyncio.get_event_loop()
t1 = loop.create_task(serve())
t2 = loop.create_task(important_task())
loop.run_until_complete(asyncio.gather(t1,t2))
main()

How can I stop my websocket connection from closing?

I'm new to websockets and asyncio and I'm trying to get a simple example working. I would like to create a server that accepts connections from multiple clients and concurrently runs a loop that sends a message to every connection once a second. I'm also trying to use asyncio.run() which I believe is preferred to the get_event_loop() code on which many examples are based.
Here's my code so far:
import asyncio
import websockets
USERS = set()
async def register(websocket, path):
USERS.add(websocket)
await websocket.send("Successfully registered")
async def count():
count = 0
while True:
print(f"Iteration: {count}")
if USERS:
for user in USERS:
await user.send("Sending message back to client")
await asyncio.sleep(1)
count +=1
async def serve():
server = await websockets.serve(register, 'localhost', 8765)
await server.wait_closed()
print("Server closed")
async def main():
await asyncio.gather(count(), serve())
asyncio.run(main())
When I run this the count coroutine works until I make a connection from a client. At this point the connection is successfully registered but when I try to send a message back to the client in count() I get an error because the connection is already closed. How should I change my code to stop this from happening?
How should I change my code to stop this from happening?
The problem might be that your handler, the register coroutine, is returning immediately, which prompts websockets to close the connection. Try to change it like this:
async def register(websocket, path):
USERS.add(websocket)
await websocket.send("Successfully registered")
await asyncio.Event().wait()
If that helps, you can put the event in USERS along with websocket, so that you have a way to terminate the connection to the client when you want to.

listen to multiple socket with websockets and asyncio

I am trying to create a script in python that listens to multiple sockets using websockets and asyncio, the problem is that no matter what I do it only listen to the first socket I call.
I think its the infinite loop, what are my option to solve this? using threads for each sockets?
async def start_socket(self, event):
payload = json.dumps(event)
loop = asyncio.get_event_loop()
self.tasks.append(loop.create_task(
self.subscribe(event)))
# this should not block the rest of the code
await asyncio.gather(*tasks)
def test(self):
# I want to be able to add corotines at a different time
self.start_socket(event1)
# some code
self.start_socket(event2)
this is what I did eventually, that way its not blocking the main thread and all subscriptions are working in parallel.
def subscribe(self, payload):
ws = websocket.WebSocket(sslopt={"cert_reqs": ssl.CERT_NONE})
ws.connect(url)
ws.send(payload)
while True:
result = ws.recv()
print("Received '%s'" % result)
def start_thread(self, loop):
asyncio.set_event_loop(loop)
loop.run_forever()
def start_socket(self, **kwargs):
worker_loop = asyncio.new_event_loop()
worker = Thread(target=self.start_thread, args=(worker_loop,))
worker.start()
worker_loop.call_soon_threadsafe(self.subscribe, payload)
def listen(self):
self.start_socket(payload1)
# code
self.start_socket(payload2)
# code
self.start_socket(payload3)
Your code appears incomplete, but what you've shown has two issues. One is that run_until_complete accepts a coroutine object (or other kind of future), not a coroutine function. So it should be:
# note parentheses after your_async_function()
asyncio.get_event_loop().run_until_complete(your_async_function())
the problem is that no matter what I do it only listen to the first socket I call. I think its the infinite loop, what are my option to solve this? using threads for each sockets?
The infinite loop is not the problem, asyncio is designed to support such "infinite loops". The problem is that you are trying to do everything in one coroutine, whereas you should be creating one coroutine per websocket. This is not a problem, as coroutines are very lightweight.
For example (untested):
async def subscribe_all(self, payload):
loop = asyncio.get_event_loop()
# create a task for each URL
for url in url_list:
tasks.append(loop.create_task(self.subscribe_one(url, payload)))
# run all tasks in parallel
await asyncio.gather(*tasks)
async def subsribe_one(self, url, payload):
async with websockets.connect(url) as websocket:
await websocket.send(payload)
while True:
msg = await websocket.recv()
print(msg)
One way to efficiently listen to multiple websocket connections from a websocket server is to keep a list of connected clients and essentially juggle multiple conversations in parallel.
E.g. A simple server that sends random # to each connected client every few secs:
import os
import asyncio
import websockets
import random
websocket_clients = set()
async def handle_socket_connection(websocket, path):
"""Handles the whole lifecycle of each client's websocket connection."""
websocket_clients.add(websocket)
print(f'New connection from: {websocket.remote_address} ({len(websocket_clients)} total)')
try:
# This loop will keep listening on the socket until its closed.
async for raw_message in websocket:
print(f'Got: [{raw_message}] from socket [{id(websocket)}]')
except websockets.exceptions.ConnectionClosedError as cce:
pass
finally:
print(f'Disconnected from socket [{id(websocket)}]...')
websocket_clients.remove(websocket)
async def broadcast_random_number(loop):
"""Keeps sending a random # to each connected websocket client"""
while True:
for c in websocket_clients:
num = str(random.randint(10, 99))
print(f'Sending [{num}] to socket [{id(c)}]')
await c.send(num)
await asyncio.sleep(2)
if __name__ == "__main__":
loop = asyncio.get_event_loop()
try:
socket_server = websockets.serve(handle_socket_connection, 'localhost', 6789)
print(f'Started socket server: {socket_server} ...')
loop.run_until_complete(socket_server)
loop.run_until_complete(broadcast_random_number(loop))
loop.run_forever()
finally:
loop.close()
print(f"Successfully shutdown [{loop}].")
A simple client that connects to the server and listens for the numbers:
import asyncio
import random
import websockets
async def handle_message():
uri = "ws://localhost:6789"
async with websockets.connect(uri) as websocket:
msg = 'Please send me a number...'
print(f'Sending [{msg}] to [{websocket}]')
await websocket.send(msg)
while True:
got_back = await websocket.recv()
print(f"Got: {got_back}")
asyncio.get_event_loop().run_until_complete(handle_message())
Mixing up threads and asyncio is more trouble than its worth and you still have code that will block on the most wasteful steps like network IO (which is the essential benefit of using asyncio).
You need to run each coroutine asynchronously in an event loop, call any blocking calls with await and define each method that interacts with any awaitable interactions with an async
See a working e.g.: https://github.com/adnantium/websocket_client_server

Discord.py background task breaks but task is still pending (using websocket)

I'm very inexperienced with asyncio and asynchronous programming in general, so i've been having a hard time trying to use a synchronous websocket module with the async discord.py module. I am trying to setup a bot which constantly listens for websocket messages, and if it receives data, it performs some calculations and sends a message in discord. This should run indefinitely, and the websocket changes its origin routinely. Here is some commented code of what i'm trying to accomplish:
import requests
import websocket
import discord
import asyncio
from time import sleep
client = discord.Client() # Initialize the discord client
class Wrapper: # This just holds some variables so that I can use them without worrying about scope
data = None
first_time = True
is_running = False
socket_url = None
async def init_bot(): # The problem begins here
def on_message(ws, message):
if is_valid(message['data']): # is_valid is just a name for a special comparison I do with the message
Wrapper.data = message['data']
print('Received valid data')
def on_error(ws, error):
print('There was an error connecting to the websocket.')
ws = websocket.WebSocketApp(Wrapper.socket_url,
on_message=on_message,
on_error=on_error)
while True:
ws.run_forever() # I believe this is connected to the discord event loop?
# Using ws.close() here breaks the program when it receives data for the second time
async def start():
await client.wait_until_ready()
def get_response(hq_id, headers):
response = requests.get('my_data_source')
try:
return json.loads(response.text)
except:
return None
print('Starting websocket')
await init_bot()
while True: # This should run forever, but the outer coroutine gets task pending errors
print('Running main loop')
if Wrapper.first_time:
Wrapper.first_time = False
on_login = await client.send_message(Config.public_channel, embed=Config.waiting_embed()) # Sends an embed to a channel
while Wrapper.socket_url is None:
response = get_response(ID, HEADERS)
try:
Wrapper.socket_url = response['socket_url'] # Assume this sets the socket url to a websocket in the form ws://anyhost.com
await client.edit_message(on_login, embed=Config.connect_embed())
Wrapper.is_running = True
await asyncio.sleep(3)
except:
await asyncio.sleep(60) # The response will only sometimes include a proper socket_url, so we wait for one
if Wrapper.is_running:
while Wrapper.data is None: # Is this blocking? I essentially want this while loop to end when we have data from line 18
await asyncio.sleep(1)
if Wrapper.data is not None:
data_message = await client.send_message(Config.public_channel, embed=Wrapper.data[0])
await client.add_reaction(ans_message, '👍')
await client.add_reaction(ans_message, '👎')
if Wrapper.data[1] == True: # Marks the last message, so we want to reset the bot to its initial state
Wrapper.is_running = False
Wrapper.first_time = True
Wrapper.socket_url = None
await asyncio.sleep(100) # Sleep for a period of time in order to make sure the socket url is closed
Wrapper.data = None
await asyncio.sleep(3)
print('While loop ended?')
#client.event
async def on_ready():
print(f'My Bot\n')
client.loop.create_task(start())
client.run('<TOKEN>')
I've tried several variations of the above, but the error I typically get is something along these lines:
File "mybot.py", line 247, in <module>
client.loop.create_task(start())
task: <Task pending coro=<start() running at mybot.py> wait_for=<Future pending cb=[BaseSelectorEventLoop._sock_connect_done(1016)(), <TaskWakeupMethWrapper object at 0x000002545FFFC8B8>()]
You cannot just mix asyncio-aware code, such as discord, with synchronous websockets code. Since nothing is awaited in init_bot, calling await init_bot() completely stops the event loop.
Instead, you need to run the websocket code (init_bot function in your case) in a separate thread and await an appropriate event. For example:
def init_bot(loop, w):
def on_message(ws, message):
w.data = message['data']
loop.call_soon_threadsafe(w.event.set)
# ...
async def start():
# ...
loop = asyncio.get_event_loop()
w = Wrapper()
w.event = asyncio.Event()
threading.Thread(target=lambda: init_bot(loop, w)).start()
# ...
# instead of while Wrapper.data is None ...
await w.event.wait()
# ... process data
w.event.clear()

Categories