From df5dc2f6911c572463706648f3b97e43a40d6c23 Mon Sep 17 00:00:00 2001
From: Tom Moor <tom.moor@gmail.com>
Date: Mon, 4 Oct 2021 18:20:42 -0700
Subject: [PATCH] fix: Improve graceful shutdown

---
 package.json                     |  2 +-
 server/collaboration/tracing.js  | 15 ++++++++++-----
 server/index.js                  |  2 ++
 server/services/collaboration.js |  4 ++++
 server/services/websockets.js    |  4 ++++
 5 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/package.json b/package.json
index 996709abb..8beabfab2 100644
--- a/package.json
+++ b/package.json
@@ -10,7 +10,7 @@
     "build:webpack": "webpack --config webpack.config.prod.js",
     "build": "yarn clean && yarn build:webpack && yarn build:i18n && yarn build:server",
     "start": "node ./build/server/index.js",
-    "dev": "yarn concurrently --kill-others -n api,collaboration  -c \"blue,magenta\" \"node --inspect=0.0.0.0 build/server/index.js --services=websockets,admin,web,worker\" \"node build/server/index.js --services=collaboration --port=4000\"",
+    "dev": "yarn concurrently -n api,collaboration  -c \"blue,magenta\" \"node --inspect=0.0.0.0 build/server/index.js --services=websockets,admin,web,worker\" \"node build/server/index.js --services=collaboration --port=4000\"",
     "dev:watch": "nodemon --exec \"yarn build:server && yarn build:i18n && yarn dev\" -e js --ignore build/ --ignore app/ --ignore flow-typed/",
     "lint": "eslint app server shared",
     "deploy": "git push heroku master",
diff --git a/server/collaboration/tracing.js b/server/collaboration/tracing.js
index d8b9e62d9..51bcfde55 100644
--- a/server/collaboration/tracing.js
+++ b/server/collaboration/tracing.js
@@ -2,7 +2,7 @@
 import Metrics from "../logging/metrics";
 
 export default class Tracing {
-  async onCreateDocument({
+  onCreateDocument({
     documentName,
     instance,
   }: {
@@ -17,11 +17,11 @@ export default class Tracing {
     );
   }
 
-  async onAuthenticationFailed({ documentName }: { documentName: string }) {
+  onAuthenticationFailed({ documentName }: { documentName: string }) {
     Metrics.increment("collaboration.authentication_failed", { documentName });
   }
 
-  async onConnect({
+  onConnect({
     documentName,
     instance,
   }: {
@@ -35,7 +35,7 @@ export default class Tracing {
     );
   }
 
-  async onDisconnect({
+  onDisconnect({
     documentName,
     instance,
   }: {
@@ -54,7 +54,12 @@ export default class Tracing {
     );
   }
 
-  async onChange({ documentName }: { documentName: string }) {
+  onChange({ documentName }: { documentName: string }) {
     Metrics.increment("collaboration.change", { documentName });
   }
+
+  onDestroy() {
+    Metrics.gaugePerInstance("collaboration.connections_count", 0);
+    Metrics.gaugePerInstance("collaboration.documents_count", 0);
+  }
 }
diff --git a/server/index.js b/server/index.js
index e12122f09..b3e915039 100644
--- a/server/index.js
+++ b/server/index.js
@@ -108,6 +108,8 @@ async function start(id: string, disconnect: () => void) {
 
   function shutdown() {
     Logger.info("lifecycle", "Stopping server");
+
+    server.emit("shutdown");
     server.stop(disconnect);
   }
 }
diff --git a/server/services/collaboration.js b/server/services/collaboration.js
index fe34c2322..2bf0a1671 100644
--- a/server/services/collaboration.js
+++ b/server/services/collaboration.js
@@ -36,4 +36,8 @@ export default function init(app: Koa, server: http.Server) {
   app.use(websocket());
   app.use(router.routes());
   app.use(router.allowedMethods());
+
+  server.on("shutdown", () => {
+    hocuspocus.destroy();
+  });
 }
diff --git a/server/services/websockets.js b/server/services/websockets.js
index 456423851..7bb813182 100644
--- a/server/services/websockets.js
+++ b/server/services/websockets.js
@@ -23,6 +23,10 @@ export default function init(app: Koa, server: http.Server) {
     cookie: false,
   });
 
+  server.on("shutdown", () => {
+    Metrics.gaugePerInstance("websockets.count", 0);
+  });
+
   io.adapter(
     socketRedisAdapter({
       pubClient: client,