Skip to content

Commit

Permalink
Added test and implemented lifespan.shutdown.failed
Browse files Browse the repository at this point in the history
  • Loading branch information
euri10 committed Oct 6, 2020
1 parent a504c56 commit e5f60d5
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 2 deletions.
30 changes: 30 additions & 0 deletions tests/test_lifespan.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,36 @@ async def test():
loop.run_until_complete(test())


@pytest.mark.parametrize("mode", ("auto", "on"))
@pytest.mark.parametrize("raise_exception", (True, False))
def test_lifespan_with_failed_shutdown(mode, raise_exception):
async def app(scope, receive, send):
message = await receive()
assert message["type"] == "lifespan.startup"
await send({"type": "lifespan.startup.complete"})
message = await receive()
assert message["type"] == "lifespan.shutdown"
await send({"type": "lifespan.shutdown.failed"})

if raise_exception:
# App should be able to re-raise an exception if startup failed.
raise RuntimeError()

async def test():
config = Config(app=app, lifespan=mode)
lifespan = LifespanOn(config)

await lifespan.startup()
assert not lifespan.startup_failed
await lifespan.shutdown()
assert lifespan.shutdown_failed
assert lifespan.error_occured is raise_exception
assert lifespan.should_exit

loop = asyncio.new_event_loop()
loop.run_until_complete(test())


@pytest.mark.parametrize("mode", ("auto", "on"))
def test_lifespan_scope_asgi3app(mode):
async def asgi3app(scope, receive, send):
Expand Down
21 changes: 19 additions & 2 deletions uvicorn/lifespan/on.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ def __init__(self, config: Config) -> None:
self.receive_queue: "Queue[LifespanReceiveMessage]" = asyncio.Queue()
self.error_occured = False
self.startup_failed = False
self.shutdown_failed = False
self.should_exit = False

async def startup(self) -> None:
Expand All @@ -43,7 +44,14 @@ async def shutdown(self) -> None:
self.logger.info("Waiting for application shutdown.")
await self.receive_queue.put({"type": "lifespan.shutdown"})
await self.shutdown_event.wait()
self.logger.info("Application shutdown complete.")

if self.shutdown_failed or (
self.error_occured and self.config.lifespan == "on"
):
self.logger.error("Application shutdown failed. Exiting.")
self.should_exit = True
else:
self.logger.info("Application shutdown complete.")

async def main(self) -> None:
try:
Expand All @@ -56,7 +64,7 @@ async def main(self) -> None:
except BaseException as exc:
self.asgi = None
self.error_occured = True
if self.startup_failed:
if self.startup_failed or self.shutdown_failed:
return
if self.config.lifespan == "auto":
msg = "ASGI 'lifespan' protocol appears unsupported."
Expand All @@ -73,6 +81,7 @@ async def send(self, message: LifespanSendMessage) -> None:
"lifespan.startup.complete",
"lifespan.startup.failed",
"lifespan.shutdown.complete",
"lifespan.shutdown.failed",
)

if message["type"] == "lifespan.startup.complete":
Expand All @@ -93,5 +102,13 @@ async def send(self, message: LifespanSendMessage) -> None:
assert not self.shutdown_event.is_set(), STATE_TRANSITION_ERROR
self.shutdown_event.set()

elif message["type"] == "lifespan.shutdown.failed":
assert self.startup_event.is_set(), STATE_TRANSITION_ERROR
assert not self.shutdown_event.is_set(), STATE_TRANSITION_ERROR
self.shutdown_event.set()
self.shutdown_failed = True
if message.get("message"):
self.logger.error(message["message"])

async def receive(self) -> LifespanReceiveMessage:
return await self.receive_queue.get()

0 comments on commit e5f60d5

Please sign in to comment.