Below is a simple echo server. But if the client does not send anything for 10 seconds, I want to close the connection.
import asyncio
async def process(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
print("awaiting for data")
line = await reader.readline()
print(f"received {line}")
writer.write(line)
print(f"sent {line}")
await writer.drain()
print(f"Drained")
async def timeout(task: asyncio.Task, duration):
print("timeout started")
await asyncio.sleep(duration)
print("client unresponsive, cancelling")
task.cancel()
print("task cancelled")
async def new_session(reader, writer):
print("new session started")
task = asyncio.create_task(process(reader, writer))
timer = asyncio.create_task(timeout(task, 10))
await task
print("task complete")
timer.cancel()
print("timer cancelled")
writer.close()
print("writer closed")
async def a_main():
server = await asyncio.start_server(new_session, port=8088)
await server.serve_forever()
if __name__ == '__main__':
asyncio.run(a_main())
If the client sends a message, it works fine. But the other case, when client is silent, it does not work
When client sends message:
new session started
awaiting for data
timeout started
received b'slkdfjsdlkfj\r\n'
sent b'slkdfjsdlkfj\r\n'
Drained
task complete
timer cancelled
writer closed
When client is silent after opening connection
new session started
awaiting for data
timeout started
client unresponsive, cancelling
task cancelled
There is no task complete, timer cancelled, writer closed.
What is the issue with above code?
Is there a better way to implement timeouts?
Update
Figured out the problem, Looks like the task was actually cancelled, but the exception got silently ignored, Fixed the problem by catching CancelledError
async def new_session(reader, writer):
print("new session started")
task = asyncio.create_task(process(reader, writer))
timer = asyncio.create_task(timeout(task, 10))
try:
await task
except asyncio.CancelledError:
print(f"Task took too long and was cancelled by timer")
print("task complete")
timer.cancel()
print("timer cancelled")
writer.close()
print("writer closed")
Second part still remains. Is there a better way to implement timeouts?
Update2
Complete code using wait_for. The timeout code is no longer needed. Check accepted solution below:
async def new_session(reader, writer):
print("new session started")
try:
await asyncio.wait_for(process(reader, writer), timeout=5)
except asyncio.TimeoutError as te:
print(f'time is up!{te}')
finally:
writer.close()
print("writer closed")
I use the following code when making a connection. I'd suggest using wait_for similarly for your code.
fut = asyncio.open_connection( self.host, self.port, loop=self.loop )
try:
r, w = await asyncio.wait_for(fut, timeout=self.connection_timeout)
except asyncio.TimeoutError:
pass
Is there a better way to implement timeouts?
You can use asyncio.wait_for instead of timeout. It has similar semantics, but already comes with asyncio. Also, you can await the future it returns to detect if the timeout has occurred.
Related
I'm trying to understand how to use asyncio streams for multiple connections that will keep sending messages until a predefined condition or a socket timeout. Looking at Python docs, they provide the following example for a TCP server based on asyncio streams:
import asyncio
async def handle_echo(reader, writer):
data = await reader.read(100)
message = data.decode()
addr = writer.get_extra_info('peername')
print(f"Received {message!r} from {addr!r}")
print(f"Send: {message!r}")
writer.write(data)
await writer.drain()
print("Close the connection")
writer.close()
async def main():
server = await asyncio.start_server(
handle_echo, '127.0.0.1', 8888)
addrs = ', '.join(str(sock.getsockname()) for sock in server.sockets)
print(f'Serving on {addrs}')
async with server:
await server.serve_forever()
asyncio.run(main())
What I'm trying to do is more complex and it looks more like so (a lot of it is pseudocode, written in capital letters or with implementation omitted):
import asyncio
async def io_control(queue):
while true:
...
# do I/O control in this function ...
async def data_processing(queue):
while true:
...
# perform data handling
async def handle_data(reader, writer):
data = await reader.read()
message = data.decode()
addr = writer.get_extra_info('peername')
print(f"Received {message!r} from {addr!r}")
#do stuff with a queue - pass messages to other two async functions as needed
#keep open until something happens
if(ERROR or SOCKET_TIMEOUT):
writer.close()
async def server(queue):
server = await asyncio.start_server(
handle_data, '127.0.0.1', 8888)
addrs = ', '.join(str(sock.getsockname()) for sock in server.sockets)
print(f'Serving on {addrs}')
async with server:
await server.serve_forever()
async def main():
queue_io = asyncio.Queue()
queue_data = asyncio.Queue()
asyncio.run(server(queue_data))
asyncio.run(data_handling(queue_data))
asyncio.run(io_control(queue_io))
asyncio.run(main())
Does this look feasible? I'm not used to working with co-routines (I'm coming from more of a multi-threading paradigm), so I'm not sure if what I'm doing is right or if I have to explicitly include yields or do any extra stuff.
If I understand correctly, you just need the TCP server to be able to handle multiple concurrent connections. The start_server function should already give you everything you need.
The first parameter client_connected_cb is a coroutine function called whenever a client establishes a connection. If you introduce a loop into that function (in your example code handle_data), you can keep the connection open until some criterion is met. What conditions exactly should lead to closing the connection is up to you, and the implementation details will obviously depend on that. The simplest approach I can imagine is something like this:
import asyncio
import logging
log = logging.getLogger(__name__)
async def handle_data(reader, writer):
while True:
data = (await reader.readline()).decode().strip()
if not data:
log.debug("client disconnected")
break
response = await your_data_processing_function(data)
writer.write(response.encode())
await writer.drain()
...
async def main():
server = await asyncio.start_server(handle_data, '127.0.0.1', 8888)
async with server:
await server.serve_forever()
if __name__ == '__main__':
asyncio.run(main())
There is theoretically no limit for the number of concurrent connections.
If your client_connected_cb is a coroutine function, each new connection will schedule a new task for the event loop. That is where the concurrency comes from. The magic then happens at the point of awaiting new data from the client; that is where the event loop can switch execution to another coroutine. All this happens behind the scenes, so to speak.
If you want to introduce a timeout, you could wrap the awaitable readline coroutine in a wait_for for example and then catch the TimeoutError exiting the loop.
Hope this helps.
Today I found very strange problem with asyncio or aiohttp.
I wrote very simple server and client which use Websockets. When server gets connection from client, it creates two tasks, one task listens to data from client, another one send data to client.
If client decides to finish session, it sends close to server, listen_on_socket (server) Task finishes fine, but send_to_socket (server) Task became frozen if it contains asyncio.sleep inside of the Task. I can not even cancel the frozen task.
What's the reason of the problem and how can I handle it?
I have the following aiohttp server code as example:
from aiohttp import web, WSMsgType
import asyncio
async def send_to_socket(ws: web.WebSocketResponse):
"""helper func which send messages to socket"""
for i in range(10):
try:
if ws.closed:
break
else:
await ws.send_str(f"I am super socket server-{i} !!!")
except Exception as ex:
print(ex)
break
# remove await asyncio.sleep(0.5) and it works !
print("| send_to_socket | St sleeping")
await asyncio.sleep(0.5)
print("| send_to_socket | Stopped sleeping") # you will not get the message
if not ws.closed:
await ws.send_str("close")
print("| send_to_socket | Finished sending")
async def listen_on_socket(ws: web.WebSocketResponse, send_task: asyncio.Task):
"""helper func which Listen messages to socket"""
async for msg in ws:
if msg.type == WSMsgType.TEXT:
if msg.data == "close":
await ws.close()
send_task.cancel()
print(send_task.cancelled(), send_task.done(), send_task)
break
elif msg.type == WSMsgType.ERROR:
print(f'ws connection closed with exception {ws.exception()}')
print("* listen_on_socket * Finished listening")
async def websocket_handler(req: web.Request) -> web.WebSocketResponse:
"""Socket aiohttp handler"""
ws = web.WebSocketResponse()
print(f"Handler | Started websocket: {id(ws)}")
await ws.prepare(req)
t = asyncio.create_task(send_to_socket(ws))
await asyncio.gather(listen_on_socket(ws, t), t)
print("Handler | websocket connection closed")
return ws
if __name__ == '__main__':
app = web.Application()
app.router.add_get("/socket", websocket_handler)
web.run_app(app, host="0.0.0.0", port=9999)
I have the following aiohttp client code as example:
from aiohttp import ClientSession
import aiohttp
import asyncio
async def client():
n = 3
async with ClientSession() as session:
async with session.ws_connect('http://localhost:9999/socket') as ws:
async for msg in ws:
if n == 0:
await ws.send_str("close")
break
if msg.type == aiohttp.WSMsgType.TEXT:
if msg.data == "close":
await ws.close()
break
else:
print(msg.data)
n -= 1
elif msg.type == aiohttp.WSMsgType.ERROR:
break
print("Client stopped")
if __name__ == '__main__':
asyncio.run(client())
It isn't freezes, just your cancellation and logging a bit incorrect, you should await for cancelled task
async def listen_on_socket(ws: web.WebSocketResponse, send_task: asyncio.Task):
"""helper func which Listen messages to socket"""
async for msg in ws:
if msg.type == WSMsgType.TEXT:
if msg.data == "close":
await ws.close()
send_task.cancel()
try:
await send_task
except asyncio.CancelledError:
print("send task cancelled")
print(send_task.cancelled(), send_task.done(), send_task)
break
elif msg.type == WSMsgType.ERROR:
print(f'ws connection closed with exception {ws.exception()}')
print("* listen_on_socket * Finished listening")
Also there should be set return_exceptions=True in the gather call inside the websocket_handler to prevent exception propagation.
You could just wrap all the function body with try-finally block and ensure it finishes fine (sure just for debugging, not in final implementation).
From aiohttp documentation: Reading from the WebSocket (await ws.receive()) must only be done inside the request handler task; however, writing (ws.send_str(...)) to the WebSocket, closing (await ws.close()) and canceling the handler task may be delegated to other tasks.
Hereby the mistake was that I created reading from ws task in listen_on_socket.
Solution. Changes only in server, client is the same:
from aiohttp import web, WSMsgType
import asyncio
async def send_to_socket(ws: web.WebSocketResponse):
"""helper func which send messages to socket"""
for i in range(4):
try:
if ws.closed:
break
else:
await ws.send_str(f"I am super socket server-{i} !!!")
except Exception as ex:
print(ex)
break
await asyncio.sleep(1.5)
if not ws.closed:
await ws.send_str("close")
print(f"| send_to_socket | Finished sending {id(ws)}")
async def websocket_handler(req: web.Request) -> web.WebSocketResponse:
"""Socket aiohttp handler"""
ws = web.WebSocketResponse()
print(f"Handler | Started websocket: {id(ws)}")
await ws.prepare(req)
# create another task for writing
asyncio.create_task(send_to_socket(ws))
async for msg in ws:
if msg.type == WSMsgType.TEXT:
if msg.data == "close":
await ws.close()
break
elif msg.type == WSMsgType.ERROR:
print(f'ws connection closed with exception {ws.exception()}')
print(f"Connection {id(ws)} is finished")
return ws
if __name__ == '__main__':
app = web.Application()
app.router.add_get("/socket", websocket_handler)
web.run_app(app, host="0.0.0.0", port=9999)
I am currently implementing the TCP socket protocol. The protocol requires sending heartbeat messages every five minutes. I am implementing a protocol using asyncio in Python. The source code below is a program that connects to localhost:8889, sends hello, and disconnects the socket after 1 second. In this case, the connection is disconnected after one second (if this actually happens, the network is down or the server is disconnected). The problem is that the send_heartbeat function waits 5 minutes without knowing that the socket is down. I would like to cancel the coroutine immediately instead of waiting 5 minutes when the socket is disconnected. What's the best way to do it?
import asyncio
async def run(host: str, port: int):
while True:
try:
reader, writer = await asyncio.open_connection(host, port)
except OSError as e:
print('connection failed:', e)
await asyncio.sleep(0.5)
continue
await asyncio.wait([
handle_stream(reader, writer),
send_heartbeat(reader, writer),
], return_when=asyncio.FIRST_COMPLETED) # will stop after 1 second
writer.close() # close socket after 1 second
await writer.wait_closed()
async def handle_stream(reader, writer):
writer.write(b'hello\n') # will success because socket is alive
await writer.drain()
await asyncio.sleep(1)
async def send_heartbeat(reader, writer):
while True:
await asyncio.sleep(300)
heartbeat_message = b'heartbeat\n'
writer.write(heartbeat_message) # will fail because socket is already closed after 1 second
await writer.drain()
if __name__ == '__main__':
asyncio.run(run('127.0.0.1', 8889))
You can cancel the sleep by canceling a task that executes it. Creating send_heartbeat as a separate task ensures that it runs in parallel to handle_stream while you await the latter:
async def run(host: str, port: int):
while True:
...
heartbeat = asyncio.create_task(send_heartbeat(reader, writer))
try:
await handle_stream(reader, writer)
finally:
heartbeat.cancel()
writer.close()
await writer.wait_closed()
BTW, since you're awaiting writer.drain() inside handle_stream, there is no guarantee that handle_stream will always complete in 1 second. This might be a place where you might want to avoid the drain, or you can use asyncio.wait_for when awaiting handle_stream(...).
I need to write some async code which runs a subprocess as part of its tasks. Even though I am using asyncio.subprocess my code is still blocking. My server looks like this:
import asyncio
import asyncio.subprocess
import websockets
async def handler(websocket, path):
while True:
data = await websocket.recv()
print('I received a message')
player = await asyncio.create_subprocess_exec(
'sleep', '5',
stdin=asyncio.subprocess.DEVNULL,
stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL)
await player.wait()
print('Finished waiting')
server = websockets.serve(handler, '0.0.0.0', '8000')
asyncio.get_event_loop().run_until_complete(server)
asyncio.get_event_loop().run_forever()
And a very basic client:
import asyncio
import websockets
async def client():
async with websockets.connect('ws://localhost:8000') as websocket:
for i in range(5):
await websocket.send('message')
await asyncio.sleep(0.5)
asyncio.get_event_loop().run_until_complete(client())
I would expect the output to look like this:
I received a message
I received a message
I received a message
I received a message
I received a message
Finished waiting
Finished waiting
Finished waiting
Finished waiting
Finished waiting
But instead I get this:
I received a message
Finished waiting
I received a message
Finished waiting
I received a message
Finished waiting
I received a message
Finished waiting
I received a message
Finished waiting
With a 5 second wait after each "I received a message" line.
The line await player.wait() does not block other async operations, but waits for 5 seconds!
If you don't want to wait for the response, try using ensure_future() instead:
# add:
async def wait_for_player(player, path):
print("Waiting...", path)
await player.wait()
print("Done", path)
# and replace await player.wait() with:
asyncio.ensure_future(wait_for_player(player, path))
You can actually also move create_subprocess_exec() to wait_for_player().
To see your code is not blocking see try these:
Client:
import asyncio
import websockets
async def client(n):
async with websockets.connect('ws://localhost:8000/{}/'.format(n)) as websocket:
print(n, "start")
for i in range(5):
print(n, i)
await websocket.send('message')
await asyncio.sleep(0.5)
print(n, "done")
tasks = [client(i) for i in range(5)]
asyncio.get_event_loop().run_until_complete(asyncio.wait(tasks))
Server:
import asyncio
import asyncio.subprocess
import random
import websockets
async def handler(websocket, path):
try:
while True:
data = await websocket.recv()
pause = random.randint(1, 5)
print('I received a message', path, "Pausing:", pause)
player = await asyncio.create_subprocess_exec(
'sleep', str(pause),
stdin=asyncio.subprocess.DEVNULL,
stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL)
await player.wait()
print('Finished waiting', path)
except websockets.ConnectionClosed:
print("Connection closed!", path)
server = websockets.serve(handler, '0.0.0.0', '8000')
asyncio.get_event_loop().run_until_complete(server)
asyncio.get_event_loop().run_forever()
Your ws server seems ok. Actually it is your client that is blocking. If you want to test the async behavior of your server, You need to make asynchronous requests. The for loop in your client blocks the thread. So remove it and instead, use asyncio.gather to run your client() method 5 times asynchronously
import asyncio
import websockets
async def client():
async with websockets.connect('ws://localhost:8000') as websocket:
await websocket.send('message')
await asyncio.sleep(0.5)
tasks = asyncio.gather(*[client() for i in range(5)])
asyncio.get_event_loop().run_until_complete(tasks)
I am trying to implement a basic websocket client using asyncio and websockets with Python 3.5.2.
Basically, I want connect_to_dealer to be a blocking call, but wait for the websocket message on a different thread.
After reading some docs (I have very little exp with Python), I concluded that asyncio.ensure_future() passing a coroutine (listen_for_message) was the way to go.
Now, I get to run listen_for_message on a different thread, but from within the coroutine I can't seem to use await or any other mechanism to make the calls synchronous. If I do it, the execution waits forever (it hangs) even for a simple sleep.
I'd like to know what I'm doing wrong.
async def listen_for_message(self, future, websocket):
while (True):
try:
await asyncio.sleep(1) # It hangs here
print('Listening for a message...')
message = await websocket.recv() # If I remove the sleep, hangs here
print("< {}".format(message))
future.set_result(message)
future.done()
except websockets.ConnectionClosed as cc:
print('Connection closed')
except Exception as e:
print('Something happened')
def handle_connect_message(self, future):
# We must first remove the websocket-specific payload because we're only interested in the connect protocol msg
print(future.result)
async def connect_to_dealer(self):
print('connect to dealer')
websocket = await websockets.connect('wss://mywebsocket'))
hello_message = await websocket.recv()
print("< {}".format(hello_message))
# We need to parse the connection ID out of the message
connection_id = hello_message['connectionId']
print('Got connection id {}'.format(connection_id))
sub_response = requests.put('https://subscribetotraffic{user_id}?connection={connection_id}'.format(user_id='username', connection_id=connection_id), headers=headers)
if sub_response.status_code == 200:
print('Now we\'re observing traffic')
else:
print('Oops request failed with {code}'.format(code=sub_response.status_code))
# Now we need to handle messages but continue with the regular execution
try:
future = asyncio.get_event_loop().create_future()
future.add_done_callback(self.handle_connect_message)
asyncio.ensure_future(self.listen_for_message(future, websocket))
except Exception as e:
print(e)
Is there a specific reason you need to work with explicit futures?
With asyncio you can use a combination of coroutines and Tasks to achieve most purposes. Tasks are essentially wrapped coroutines that go about cranking themselves over in the background, independently of other async code, so you don't have to explicitly manage their flow or juggle them with other bits of code.
I am not entirely sure of your end goal, but perhaps the approach elaborated below gives you something to work with:
import asyncio
async def listen_for_message(websocket):
while True:
await asyncio.sleep(0)
try:
print('Listening for a message...')
message = await websocket.recv()
print("< {}".format(message))
except websockets.ConnectionClosed as cc:
print('Connection closed')
except Exception as e:
print('Something happened')
async def connect_to_dealer():
print('connect to dealer')
websocket = await websockets.connect('wss://mywebsocket')
hello_message = await websocket.recv()
print("< {}".format(hello_message))
# We need to parse the connection ID out of the message
connection_id = hello_message['connectionId']
print('Got connection id {}'.format(connection_id))
sub_response = requests.put('https://subscribetotraffic{user_id}?connection={connection_id}'.format(
user_id='username', connection_id=connection_id), headers=headers)
if sub_response.status_code == 200:
print('Now we\'re observing traffic')
else:
print('Oops request failed with {code}'.format(code=sub_response.status_code))
async def my_app():
# this will block until connect_to_dealer() returns
websocket = await connect_to_dealer()
# start listen_for_message() in its own task wrapper, so doing it continues in the background
asyncio.ensure_future(listen_for_message(websocket))
# you can continue with other code here that can now coexist with listen_for_message()
if __name__ == '__main__':
loop = asyncio.get_event_loop()
loop.run_until_complete(my_app())
loop.run_forever()