Graceful shutdown of uvicorn starlette app with websockets
Question:
Given this sample Starlette app with an open websocket connection, how do you shut down the Starlette app? I am running on uvicorn. Whenever I press Ctrl+C
the output is Waiting for background tasks to complete.
which hangs forever.
from starlette.applications import Starlette
app = Starlette()
@app.websocket_route('/ws')
async def ws(websocket):
await websocket.accept()
while True:
# How to interrupt this while loop on the shutdown event?
await asyncio.sleep(0.1)
await websocket.close()
I tried switching a bool variable on the shutdown event but the variable never updates. It is always False
.
eg.
app.state.is_shutting_down = False
@app.on_event('shutdown')
async def shutdown():
app.state.is_shutting_down = True
@app.websocket_route('/ws')
async def ws(websocket):
await websocket.accept()
while app.state.is_shutting_down is False:
Answers:
I can’t really reproduce the problem, I get
INFO: Started server process [31]
INFO: Waiting for application startup.
INFO: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
^CINFO: Shutting down
INFO: Waiting for application shutdown.
INFO: Finished server process [31]
But what I do in my own async apps to gracefully shut down is:
import signal
async def shutdown(signal: signal):
"""
try to shut down gracefully
"""
logger.info("Received exit signal %s...", signal.name)
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
[task.cancel() for task in tasks]
logging.info("Canceling outstanding tasks")
await asyncio.gather(*tasks)
if __name__ == "__main__":
loop = asyncio.get_event_loop()
signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
for s in signals:
loop.add_signal_handler(
s, lambda s=s: asyncio.create_task(shutdown(s))
)
You probably have to move the await websocket.close()
into the shutdown
method.
Credits go to Lynn Root, writing about Graceful Shutdowns with asyncio here.
The reason your variable didn’t change is because the handlers for the "shutdown" event are executed after all the tasks have been executed (i.e. we end up with a deadlock as the shutdown handler waits for the async task and the async task waits for the shutdown handler).
Setting a signal handler on the asyncio event loop will probably not work as I believe only a single signal handler is allowed which uvicorn already sets for its own shutdown process.
Instead, you can Monkey Patch the uvicorn signal handler to detect the application shutdown and set your controlling variable in that new function.
import asyncio
from starlette.applications import Starlette
from uvicorn.main import Server
original_handler = Server.handle_exit
class AppStatus:
should_exit = False
@staticmethod
def handle_exit(*args, **kwargs):
AppStatus.should_exit = True
original_handler(*args, **kwargs)
Server.handle_exit = AppStatus.handle_exit
app = Starlette()
@app.websocket_route('/ws')
async def ws(websocket):
await websocket.accept()
while AppStatus.should_exit is False:
await asyncio.sleep(0.1)
await websocket.close()
print('Exited!')
Given this sample Starlette app with an open websocket connection, how do you shut down the Starlette app? I am running on uvicorn. Whenever I press Ctrl+C
the output is Waiting for background tasks to complete.
which hangs forever.
from starlette.applications import Starlette
app = Starlette()
@app.websocket_route('/ws')
async def ws(websocket):
await websocket.accept()
while True:
# How to interrupt this while loop on the shutdown event?
await asyncio.sleep(0.1)
await websocket.close()
I tried switching a bool variable on the shutdown event but the variable never updates. It is always False
.
eg.
app.state.is_shutting_down = False
@app.on_event('shutdown')
async def shutdown():
app.state.is_shutting_down = True
@app.websocket_route('/ws')
async def ws(websocket):
await websocket.accept()
while app.state.is_shutting_down is False:
I can’t really reproduce the problem, I get
INFO: Started server process [31]
INFO: Waiting for application startup.
INFO: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
^CINFO: Shutting down
INFO: Waiting for application shutdown.
INFO: Finished server process [31]
But what I do in my own async apps to gracefully shut down is:
import signal
async def shutdown(signal: signal):
"""
try to shut down gracefully
"""
logger.info("Received exit signal %s...", signal.name)
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
[task.cancel() for task in tasks]
logging.info("Canceling outstanding tasks")
await asyncio.gather(*tasks)
if __name__ == "__main__":
loop = asyncio.get_event_loop()
signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
for s in signals:
loop.add_signal_handler(
s, lambda s=s: asyncio.create_task(shutdown(s))
)
You probably have to move the await websocket.close()
into the shutdown
method.
Credits go to Lynn Root, writing about Graceful Shutdowns with asyncio here.
The reason your variable didn’t change is because the handlers for the "shutdown" event are executed after all the tasks have been executed (i.e. we end up with a deadlock as the shutdown handler waits for the async task and the async task waits for the shutdown handler).
Setting a signal handler on the asyncio event loop will probably not work as I believe only a single signal handler is allowed which uvicorn already sets for its own shutdown process.
Instead, you can Monkey Patch the uvicorn signal handler to detect the application shutdown and set your controlling variable in that new function.
import asyncio
from starlette.applications import Starlette
from uvicorn.main import Server
original_handler = Server.handle_exit
class AppStatus:
should_exit = False
@staticmethod
def handle_exit(*args, **kwargs):
AppStatus.should_exit = True
original_handler(*args, **kwargs)
Server.handle_exit = AppStatus.handle_exit
app = Starlette()
@app.websocket_route('/ws')
async def ws(websocket):
await websocket.accept()
while AppStatus.should_exit is False:
await asyncio.sleep(0.1)
await websocket.close()
print('Exited!')