This code tries to connect for 20 seconds. If the connection is established and lost then it tries to reconnect for 10 seconds. If the client can't connect or reconnect then it prints "Failure" and exits. If the server responds with "exit" then the client exits.
import asyncio
import websockets
class Client:
async def connect_with_retries(self, uri: str):
while True:
try:
return await websockets.connect(uri)
except (ConnectionError, websockets.exceptions.WebSocketException):
await asyncio.sleep(5)
async def loop(self, uri: str):
connection_timeout = 20
try:
while True:
websocket = await asyncio.wait_for(self.connect_with_retries(uri), connection_timeout)
connection_timeout = 10
try:
while True:
await websocket.send("req")
resp = await websocket.recv()
if resp == "exit":
return
await asyncio.sleep(1)
except (ConnectionError, websockets.exceptions.WebSocketException) as exc:
pass
finally:
await websocket.close()
except asyncio.TimeoutError:
print("Failure")
client = Client()
asyncio.run(client.loop("ws://localhost:8000"))
I don't like the explicit websocket.close(). How to use asyncio.wait_for with the context manager of the websocket?
Related
I'm trying to do a client/server with websockets and I want that when the client sends a message, it needs to wait to receive a response from the server. So the client can never send two messages in a row without a previous response.
I tried using Lock from asyncio but it doesn't work. Here's server.py:
async def runServer(self):
server = await websockets.serve(
self.onConnect,
"localhost",
port=8765
)
print("Server started listening to new connections...")
await server.wait_closed()
async def onConnect(self, ws):
try:
while True:
message = await ws.recv()
if message == "...":
...
And here is client.py
async def connect(self):
try:
lock = asyncio.Lock()
async with websockets.connect(
"ws://localhost:8765",
) as ws:
print("Connected to the switch.")
async with lock:
await ws.send("First message")
await ws.send("Second message")
while True:
message = await ws.recv()
print(f"Message received : {message}")
except Exception as e:
print(e)
And I would want the client to wait for a response from the first message and after that to send the second message.
Try this example:
server.py (server waits for msg1 and msg2 and responds accordingly):
async def runServer():
server = await websockets.serve(onConnect, "localhost", port=8765)
print("Server started listening to new connections...")
await server.wait_closed()
async def onConnect(ws):
try:
while True:
message = await ws.recv()
if message == "msg1":
await ws.send("response from server 1")
elif message == "msg2":
await ws.send("response from server 2")
except websockets.exceptions.ConnectionClosedOK:
print("Client closed...")
if __name__ == "__main__":
asyncio.run(runServer())
client.py (first sends msg1, prints response and afterwards sends msg2 and prints response):
async def connect():
async with websockets.connect(
"ws://localhost:8765",
) as ws:
print("Connected to the switch.")
await ws.send("msg1")
response = await ws.recv()
print("Response from the server:", response)
await ws.send("msg2")
response = await ws.recv()
print("Response from the server:", response)
if __name__ == "__main__":
asyncio.run(connect())
The clients prints:
Connected to the switch.
Response from the server: response from server 1
Response from the server: response from server 2
currently I'm using websockets to pass through data that I receive from a Redis queue (pub/sub). But for some reason the websocket doesn't send messages when using this redis queue.
What my code looks like
My code works as folllow:
I accept the socket connection
I connect to the redis queue
For each message that I receive from the subscription, i sent a message through the socket. (at the moment only text for testing)
#check_route.websocket_route("/check")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
redis = Redis(host='::1', port=6379, db=1)
subscribe = redis.pubsub()
subscribe.subscribe('websocket_queue')
try:
for result in subscribe.listen():
await websocket.send_text('test')
print('test send')
except Exception as e:
await websocket.close()
raise e
The issue with the code
When I'm using this code it's just not sending the message through the socket. But when I accept the websocket within the subscribe.listen() loop it does work but it reconnects every time (see code below).
#check_route.websocket_route("/check")
async def websocket_endpoint(websocket: WebSocket):
redis = Redis(host='::1', port=6379, db=1)
subscribe = redis.pubsub()
subscribe.subscribe('websocket_queue')
try:
for result in subscribe.listen():
await websocket.accept()
await websocket.send_text('test')
print('test send')
except Exception as e:
await websocket.close()
raise e
I think that the subscribe.listen() causes some problems that make the websocket do nothing when websocket.accept() is outside the for loop.
I hope someone knows whats wrong with this.
I'm not sure if this will work, but you could try this:
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
redis = Redis(host='::1', port=6379, db=1)
subscribe = redis.pubsub()
subscribe.subscribe('websocket_queue')
try:
results = await subscribe.listen()
for result in results:
await websocket.send_text('test')
print('test send')
except Exception as e:
await websocket.close()
raise e
After a few days more research I found a solution for this issue. I solved it by using aioredis. This solution is based on the following GitHub Gist.
import json
import aioredis
from fastapi import APIRouter, WebSocket
from app.service.config_service import load_config
check_route = APIRouter()
#check_route.websocket("/check")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
# ---------------------------- REDIS REQUIREMENTS ---------------------------- #
config = load_config()
redis_uri: str = f"redis://{config.redis.host}:{config.redis.port}"
redis_channel = config.redis.redis_socket_queue.channel
redis = await aioredis.create_redis_pool(redis_uri)
# ------------------ SEND SUBSCRIBE RESULT THROUGH WEBSOCKET ----------------- #
(channel,) = await redis.subscribe(redis_channel)
assert isinstance(channel, aioredis.Channel)
try:
while True:
response_raw = await channel.get()
response_str = response_raw.decode("utf-8")
response = json.loads(response_str)
if response:
await websocket.send_json({
"event": 'NEW_CHECK_RESULT',
"data": response
})
except Exception as e:
raise e
I try to build a base http server with the following code.
async def handle_client(client, address):
print('connection start')
data = await loop.sock_recv(client, 1024)
resp = b'HTTP/1.1 404 NOT FOUND\r\n\r\n<h1>404 NOT FOUND</h1>'
await loop.sock_sendall(client, resp)
client.close()
async def run_server():
while True:
client, address = await loop.sock_accept(server)
print('start')
loop.create_task(handle_client(client,address))
print(client)
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.bind(('localhost', 3006))
server.listen(8)
print(1)
loop = asyncio.get_event_loop()
loop.run_until_complete(run_server())
The output I expect to get is
1
start
connection start
But the actual result of running is
1
start
start
start
It seems that the function in loop.create_task() is not being run, so now I got confuesed., what is the correct way to use loop.create_task()?
You need to await the task that is created via loop.create_task(), otherwise run_server() will schedule the task and then just exit before the result has been returned.
Try changing run_server() to the following:
async def run_server():
while True:
client, address = await loop.sock_accept(server)
print('start')
await loop.create_task(handle_client(client,address))
print(client)
I'm trying to make the Postgress LISTEN async with FastAPI so that the WebSocket connection is not blocking while waiting for Postgres table updates.
What I got so far:
router = APIRouter()
#router.websocket("/pg_notify")
async def get_notifications(websocket: WebSocket):
await websocket.accept()
conn = psycopg2.connect("*****")
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
curs = conn.cursor()
curs.execute("LISTEN channel;")
while True:
try:
conn.poll()
while conn.notifies:
print("Waiting for notification...")
notify = await conn.notifies.pop(0)
print(notify.payload)
except Exception as e:
print("exception triggered: ", str(e))
await websocket.close()
This way an exception is raised on the conn.notifies,pop(0):
object psycopg2.extensions.Notify can't be used in the 'await' expression
Finally the below works for me.
while True:
try:
async with aiopg.create_pool(dsn) as pool:
async with pool.acquire() as conn:
listener = listen(conn, "channel")
task_pg_listen = asyncio.create_task(listener)
result = await task_pg_listen
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)