Add query graceful shutdown for rolling updates (#10106)

* Merge pull request #35 from populationgenomics/add-query-graceful-shutdown

Add query graceful shutdown

* Remove unused argument from query:on_shutdown

query/deployment.yaml

@@ -16,6 +16,7 @@ spec:
         app: query
         hail.is/sha: "{{ code.sha }}"
     spec:
+      terminationGracePeriodSeconds: 28800  # 8 hours
       serviceAccountName: query
 {% if deploy %}
       priorityClassName: production

query/query/query.py

@@ -220,6 +220,15 @@ async def on_cleanup(app):
     await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task()))
 
 
+async def on_shutdown(_):
+    # Filter the asyncio.current_task(), because if we await
+    # the current task we'll end up in a deadlock
+    remaining_tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
+    log.info(f"On shutdown request received, with {len(remaining_tasks)} remaining tasks")
+    await asyncio.wait(*remaining_tasks)
+    log.info("All tasks on shutdown have completed")
+
+
 def run():
     app = web.Application()
 
@@ -229,6 +238,7 @@ def run():
 
     app.on_startup.append(on_startup)
     app.on_cleanup.append(on_cleanup)
+    app.on_shutdown.append(on_shutdown)
 
     deploy_config = get_deploy_config()
     web.run_app(