From 00b091a1b1dfb16d9ac22bab718bf80bc4ebf299 Mon Sep 17 00:00:00 2001
From: cjihrig <cjihrig@gmail.com>
Date: Fri, 21 Jun 2019 15:05:58 -0400
Subject: [PATCH] grpc-js: shutdown improvements

This commit maintains a Set of all active sessions. This allows
tryShutdown() to gracefully stop the server properly (as
recommended in the Node HTTP2 documentation). The same Set of
sessions also allows forceShutdown() to be implemented.
---
 packages/grpc-js/src/server.ts              | 51 +++++++++++++++++----
 packages/grpc-js/test/test-server-errors.ts |  4 +-
 packages/grpc-js/test/test-server.ts        | 15 ------
 3 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/packages/grpc-js/src/server.ts b/packages/grpc-js/src/server.ts
index 6342b5e1..0d92881b 100644
--- a/packages/grpc-js/src/server.ts
+++ b/packages/grpc-js/src/server.ts
@@ -91,6 +91,7 @@ export class Server {
     string,
     UntypedHandler
   >();
+  private sessions = new Set<http2.ServerHttp2Session>();
   private started = false;
 
   constructor(options?: object) {}
@@ -223,7 +224,22 @@ export class Server {
   }
 
   forceShutdown(): void {
-    throw new Error('Not yet implemented');
+    // Close the server if it is still running.
+    if (this.http2Server && this.http2Server.listening) {
+      this.http2Server.close();
+    }
+
+    this.started = false;
+
+    // Always destroy any available sessions. It's possible that one or more
+    // tryShutdown() calls are in progress. Don't wait on them to finish.
+    this.sessions.forEach(session => {
+      // Cast NGHTTP2_CANCEL to any because TypeScript doesn't seem to
+      // recognize destroy(code) as a valid signature.
+      // tslint:disable-next-line:no-any
+      session.destroy(http2.constants.NGHTTP2_CANCEL as any);
+    });
+    this.sessions.clear();
   }
 
   register<RequestType, ResponseType>(
@@ -259,17 +275,34 @@ export class Server {
   }
 
   tryShutdown(callback: (error?: Error) => void): void {
-    callback = typeof callback === 'function' ? callback : noop;
+    let pendingChecks = 0;
 
-    if (this.http2Server === null) {
-      callback(new Error('server is not running'));
-      return;
+    function maybeCallback(): void {
+      pendingChecks--;
+
+      if (pendingChecks === 0) {
+        callback();
+      }
     }
 
-    this.http2Server.close((err?: Error) => {
-      this.started = false;
-      callback(err);
+    // Close the server if necessary.
+    this.started = false;
+
+    if (this.http2Server && this.http2Server.listening) {
+      pendingChecks++;
+      this.http2Server.close(maybeCallback);
+    }
+
+    // If any sessions are active, close them gracefully.
+    pendingChecks += this.sessions.size;
+    this.sessions.forEach(session => {
+      session.close(maybeCallback);
     });
+
+    // If the server is closed and there are no active sessions, just call back.
+    if (pendingChecks === 0) {
+      callback();
+    }
   }
 
   addHttp2Port(): void {
@@ -352,6 +385,8 @@ export class Server {
         session.destroy();
         return;
       }
+
+      this.sessions.add(session);
     });
   }
 }
diff --git a/packages/grpc-js/test/test-server-errors.ts b/packages/grpc-js/test/test-server-errors.ts
index d429330d..ffabaac5 100644
--- a/packages/grpc-js/test/test-server-errors.ts
+++ b/packages/grpc-js/test/test-server-errors.ts
@@ -386,9 +386,9 @@ describe('Other conditions', () => {
     });
   });
 
-  after(done => {
+  after(() => {
     client.close();
-    server.tryShutdown(done);
+    server.forceShutdown();
   });
 
   describe('Server receiving bad input', () => {
diff --git a/packages/grpc-js/test/test-server.ts b/packages/grpc-js/test/test-server.ts
index 09259481..3e02c3d5 100644
--- a/packages/grpc-js/test/test-server.ts
+++ b/packages/grpc-js/test/test-server.ts
@@ -127,17 +127,6 @@ describe('Server', () => {
     });
   });
 
-  describe('tryShutdown', () => {
-    it('calls back with an error if the server is not running', done => {
-      const server = new Server();
-
-      server.tryShutdown(err => {
-        assert(err !== undefined && err.message === 'server is not running');
-        done();
-      });
-    });
-  });
-
   describe('start', () => {
     let server: Server;
 
@@ -237,10 +226,6 @@ describe('Server', () => {
       server.addProtoService();
     }, /Not implemented. Use addService\(\) instead/);
 
-    assert.throws(() => {
-      server.forceShutdown();
-    }, /Not yet implemented/);
-
     assert.throws(() => {
       server.addHttp2Port();
     }, /Not yet implemented/);