From a06398bade97d806a582e56d2ecf1a8fc0bac27f Mon Sep 17 00:00:00 2001 From: Marco Paolini Date: Wed, 16 Aug 2017 09:19:19 +0100 Subject: [PATCH] Update server-side websocket close FAQ Due to improvements done in #754, websocket `close()` can now be called from a separate task. --- docs/faq.rst | 36 +++++++++++++++--------------------- 1 file changed, 15 insertions(+), 21 deletions(-) diff --git a/docs/faq.rst b/docs/faq.rst index c82ae0005a4..d104cf7988d 100644 --- a/docs/faq.rst +++ b/docs/faq.rst @@ -162,49 +162,43 @@ For example we have an application with two endpoints: 2. ``/logout_user`` that when invoked needs to close all open websockets for that user. -Keep in mind that you can only ``.close()`` a websocket from inside -the handler task, and since the handler task is busy reading from the -websocket, it can't react to other events. - -One simple solution is keeping a shared registry of websocket handler -tasks for a user in the :class:`aiohttp.web.Application` instance and -``cancel()`` them in ``/logout_user`` handler:: +One simple solution is keeping a shared registry of websocket responses +for a user in the :class:`aiohttp.web.Application` instance and +call :meth:`aiohttp.web.WebSocketResponse.close` on all of them in ``/logout_user`` handler:: async def echo_handler(request): - ws = web.WebSocketResponse() + ws = aiohttp.web.WebSocketResponse() user_id = authenticate_user(request) await ws.prepare(request) - request.app['websockets'][user_id].add(asyncio.Task.current_task()) + request.app['websockets'][user_id].append(ws) try: async for msg in ws: - # handle incoming messages - ... - - except asyncio.CancelledError: - print('websocket cancelled') + ws.send_str(msg.data) finally: - request.app['websockets'][user_id].remove(asyncio.Task.current_task()) - await ws.close() + request.app['websockets'][user_id].remove(ws) + return ws + async def logout_handler(request): user_id = authenticate_user(request) - for task in request.app['websockets'][user_id]: - task.cancel() + ws_closers = [ws.close() for ws in request.app['websockets'][user_id] if not ws.closed] - # return response - ... + # Watch out, this will keep us from returing the response until all are closed + ws_closers and await asyncio.wait(ws_closers) + + return aiohttp.web.Response(text='OK') def main(): loop = asyncio.get_event_loop() app = aiohttp.web.Application(loop=loop) app.router.add_route('GET', '/echo', echo_handler) app.router.add_route('POST', '/logout', logout_handler) - app['websockets'] = defaultdict(set) + app['websockets'] = defaultdict(deque) aiohttp.web.run_app(app, host='localhost', port=8080)