Cleanup everything on API shutdown

via app.stop()

jupyterhub/apihandlers/hub.py

@@ -47,9 +47,8 @@ class ShutdownAPIHandler(APIHandler):
         self.set_status(202)
         self.finish(json.dumps({"message": "Shutting down Hub"}))
 
-        # stop the eventloop, which will trigger cleanup
-        loop = IOLoop.current()
-        loop.add_callback(loop.stop)
+        # instruct the app to stop, which will trigger cleanup
+        app.stop()
 
 
 class RootAPIHandler(APIHandler):

jupyterhub/app.py

@@ -3234,9 +3234,15 @@ class JupyterHub(Application):
         loop.make_current()
         loop.run_sync(self.cleanup)
 
-    async def shutdown_cancel_tasks(self, sig):
+    async def shutdown_cancel_tasks(self, sig=None):
         """Cancel all other tasks of the event loop and initiate cleanup"""
+        if sig is None:
+            self.log.critical("Initiating shutdown...")
+        else:
             self.log.critical("Received signal %s, initiating shutdown...", sig.name)
+
+        await self.cleanup()
+
         tasks = [t for t in asyncio_all_tasks() if t is not asyncio_current_task()]
 
         if tasks:
@@ -3253,7 +3259,6 @@ class JupyterHub(Application):
             tasks = [t for t in asyncio_all_tasks()]
             for t in tasks:
                 self.log.debug("Task status: %s", t)
-        await self.cleanup()
         asyncio.get_event_loop().stop()
 
     def stop(self):
@@ -3261,7 +3266,7 @@ class JupyterHub(Application):
             return
         if self.http_server:
             self.http_server.stop()
-        self.io_loop.add_callback(self.io_loop.stop)
+        self.io_loop.add_callback(self.shutdown_cancel_tasks)
 
     async def start_show_config(self):
         """Async wrapper around base start_show_config method"""

jupyterhub/tests/conftest.py

@@ -188,6 +188,8 @@ def cleanup_after(request, io_loop):
         if not MockHub.initialized():
             return
         app = MockHub.instance()
+        if app.db_file.closed:
+            return
         for uid, user in list(app.users.items()):
             for name, spawner in list(user.spawners.items()):
                 if spawner.active:

jupyterhub/tests/mocking.py

@@ -325,26 +325,28 @@ class MockHub(JupyterHub):
         roles.assign_default_roles(self.db, entity=user)
         self.db.commit()
 
-    def stop(self):
-        super().stop()
+    _stop_called = False
 
+    def stop(self):
+        if self._stop_called:
+            return
+        self._stop_called = True
         # run cleanup in a background thread
         # to avoid multiple eventloops in the same thread errors from asyncio
 
         def cleanup():
-            asyncio.set_event_loop(asyncio.new_event_loop())
-            loop = IOLoop.current()
-            loop.run_sync(self.cleanup)
+            loop = asyncio.new_event_loop()
+            loop.run_until_complete(self.cleanup())
             loop.close()
 
-        pool = ThreadPoolExecutor(1)
+        with ThreadPoolExecutor(1) as pool:
             f = pool.submit(cleanup)
             # wait for cleanup to finish
             f.result()
-        pool.shutdown()
 
-        # ignore the call that will fire in atexit
-        self.cleanup = lambda: None
+        # prevent redundant atexit from running
+        self._atexit_ran = True
+        super().stop()
         self.db_file.close()
 
     async def login_user(self, name):

jupyterhub/tests/test_api.py

@@ -2095,14 +2095,23 @@ def test_shutdown(app):
         )
         return r
 
-    real_stop = loop.stop
+    real_stop = loop.asyncio_loop.stop
 
     def stop():
         stop.called = True
         loop.call_later(1, real_stop)
 
-    with mock.patch.object(loop, 'stop', stop):
+    real_cleanup = app.cleanup
+
+    def cleanup():
+        cleanup.called = True
+        return real_cleanup()
+
+    app.cleanup = cleanup
+
+    with mock.patch.object(loop.asyncio_loop, 'stop', stop):
         r = loop.run_sync(shutdown, timeout=5)
     r.raise_for_status()
     reply = r.json()
+    assert cleanup.called
     assert stop.called