From 6f2162d5f34134c2342c6e05b5a616ea7ef6f758 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 11 Nov 2018 21:20:50 +0200
Subject: [PATCH] Add some backup checks when stopping maubot

---
 maubot/__main__.py           | 5 ++++-
 maubot/management/api/log.py | 7 +++++--
 maubot/server.py             | 1 +
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/maubot/__main__.py b/maubot/__main__.py
index ccfceb8..95118a7 100644
--- a/maubot/__main__.py
+++ b/maubot/__main__.py
@@ -90,7 +90,10 @@ except KeyboardInterrupt:
     log.debug("Closing websockets")
     loop.run_until_complete(stop_management_api())
     log.debug("Stopping server")
-    loop.run_until_complete(server.stop())
+    try:
+        loop.run_until_complete(asyncio.wait_for(server.stop(), 5, loop=loop))
+    except asyncio.TimeoutError:
+        log.warning("Stopping server timed out")
     log.debug("Closing event loop")
     loop.close()
     log.debug("Everything stopped, shutting down")
diff --git a/maubot/management/api/log.py b/maubot/management/api/log.py
index 8700629..467f16e 100644
--- a/maubot/management/api/log.py
+++ b/maubot/management/api/log.py
@@ -72,7 +72,7 @@ sockets = []
 async def stop_all() -> None:
     for socket in sockets:
         try:
-            await socket.close(1012)
+            await socket.close(code=1012)
         except Exception:
             pass
 
@@ -107,7 +107,10 @@ async def log_websocket(request: web.Request) -> web.WebSocketResponse:
             elif not authenticated:
                 await ws.send_json({"auth_success": False})
     except Exception:
-        pass
+        try:
+            await ws.close()
+        except Exception:
+            pass
     log_root.removeHandler(handler)
     log.debug(f"Connection from {request.remote} closed")
     sockets.remove(ws)
diff --git a/maubot/server.py b/maubot/server.py
index 827fd7e..deb17e6 100644
--- a/maubot/server.py
+++ b/maubot/server.py
@@ -99,6 +99,7 @@ class MaubotServer:
         self.log.info(f"Listening on {site.name}")
 
     async def stop(self) -> None:
+        await self.runner.shutdown()
         await self.runner.cleanup()
 
     @staticmethod