Marmelator/pyinfra
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 36 Wiki Activity

feat: ensure shutdown of queue manager

Browse Source
This commit is contained in:
Jonathan Kössler 2024-11-13 16:30:18 +01:00
parent 3d3c76b466
commit 1540c2894e
3 changed files with 12 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
pyinfra/examples.py
View File

@ -26,16 +26,20 @@ from pyinfra.webserver.utils import (
shutdown_flag = False shutdown_flag = False
async def graceful_shutdown(manager, queue_task, webserver_task): async def graceful_shutdown(manager: AsyncQueueManager, queue_task, webserver_task):
global shutdown_flag global shutdown_flag
shutdown_flag = True shutdown_flag = True
logger.info("SIGTERM received, shutting down gracefully...") logger.info("SIGTERM received, shutting down gracefully...")
if queue_task and not queue_task.done(): if queue_task and not queue_task.done():
queue_task.cancel() queue_task.cancel()
await manager.shutdown()
if webserver_task and not webserver_task.done(): if webserver_task and not webserver_task.done():
webserver_task.cancel() webserver_task.cancel()
await asyncio.gather(queue_task, webserver_task, return_exceptions=True)
# explicitly shutdown manager and webserver
await asyncio.gather(queue_task, manager.shutdown(), webserver_task, return_exceptions=True)
logger.info("Shutdown complete.") logger.info("Shutdown complete.")
@ -72,19 +76,18 @@ async def run_async_queues(manager: AsyncQueueManager, app, port, host):
sys.exit(1) sys.exit(1)
finally: finally:
if shutdown_flag: if shutdown_flag:
logger.info("Graceful shutdown already in progress.") logger.debug("Graceful shutdown already in progress.")
else: else:
logger.warning("Initiating shutdown due to error or manual interruption.") logger.warning("Initiating shutdown due to error or manual interruption.")
if not tenant_api_available: if not tenant_api_available:
sys.exit(0) sys.exit(0)
if queue_task and not queue_task.done(): if queue_task and not queue_task.done():
queue_task.cancel() queue_task.cancel()
await manager.shutdown()
if webserver_task and not webserver_task.done(): if webserver_task and not webserver_task.done():
webserver_task.cancel() webserver_task.cancel()
await asyncio.gather(queue_task, webserver_task, return_exceptions=True) await asyncio.gather(queue_task, manager.shutdown(), webserver_task, return_exceptions=True)
logger.info("Shutdown complete.") logger.info("Shutdown complete.")

4
pyinfra/queue/async_manager.py
View File

@ -290,7 +290,7 @@ class AsyncQueueManager:
while self.message_count != 0: while self.message_count != 0:
logger.debug(f"Messages are still being processed: {self.message_count=} ") logger.debug(f"Messages are still being processed: {self.message_count=} ")
await asyncio.sleep(2) await asyncio.sleep(2)
await self.channel.close() await self.channel.close(exc=asyncio.CancelledError)
logger.debug("Channel closed.") logger.debug("Channel closed.")
else: else:
logger.debug("No channel to close.") logger.debug("No channel to close.")
@ -304,7 +304,7 @@ class AsyncQueueManager:
async def close_connection(self) -> None: async def close_connection(self) -> None:
try: try:
if self.connection and not self.connection.is_closed: if self.connection and not self.connection.is_closed:
await self.connection.close() await self.connection.close(exc=asyncio.CancelledError)
logger.debug("Connection closed.") logger.debug("Connection closed.")
else: else:
logger.debug("No connection to close.") logger.debug("No connection to close.")

2
pyinfra/webserver/utils.py
View File

@ -65,7 +65,7 @@ async def run_async_webserver(app: FastAPI, port: int, host: str):
try: try:
await server.serve() await server.serve()
except asyncio.CancelledError: except asyncio.CancelledError:
logger.info("Webserver was cancelled.") logger.debug("Webserver was cancelled.")
server.should_exit = True server.should_exit = True
await server.shutdown() await server.shutdown()
except Exception as e: except Exception as e: