fix: revoke leave correctly

__tests__/client-static.test.ts

@@ -44,7 +44,13 @@ describe('client static - grpc client pool', () => {
     });
 
     afterAll(async () => {
-      revokers.map(item => item.forceShutdown());
+      await Promise.all(
+        revokers.map(item => {
+          return new Promise(resolve => {
+            item.tryShutdown(resolve);
+          });
+        })
+      );
       await destroyGlobalPool();
     });
 
@@ -73,7 +79,7 @@ describe('client static - grpc client pool', () => {
 
       // revoke
       const revoker = revokers.shift();
-      revoker && revoker.forceShutdown();
+      revoker && (await revoker.forceShutdown());
       await sleep(200);
       const fourth = pool
         .get()

__tests__/client.test.ts

@@ -7,6 +7,7 @@ import { register, createGrpcProxy } from '../src';
 import { hosts, parseKV } from './helper/etcd';
 import sleep from './helper/sleep';
 import { destroyGlobalPool } from '../src/client/globalClientPool';
+// import log from 'why-is-node-running';
 
 describe('client dynamic - grpc proxy', () => {
   let pkgDef: GrpcObject;
@@ -51,8 +52,14 @@ describe('client dynamic - grpc proxy', () => {
     );
   });
 
-  afterAll(() => {
-    revokers.map(item => item.forceShutdown());
+  afterAll(async () => {
+    return Promise.all(
+      revokers.map(item => {
+        return new Promise(resolve => {
+          item.tryShutdown(resolve);
+        });
+      })
+    );
   });
 
   let proxy: GrpcObject;
@@ -75,7 +82,7 @@ describe('client dynamic - grpc proxy', () => {
 
     // revoke
     const revoker = revokers.shift();
-    revoker && revoker.forceShutdown();
+    revoker && (await revoker.forceShutdown());
     await sleep(200);
     const fourth = proxy.Greeter.getChannel().getTarget();
     const fifth = proxy.Greeter.getChannel().getTarget();

package.json

@@ -59,7 +59,8 @@
     "lint-staged": "^8.1.0",
     "prettier": "^1.18.2",
     "ts-jest": "^24.0.2",
-    "typescript": "^3.5.1"
+    "typescript": "^3.5.1",
+    "why-is-node-running": "^2.1.0"
   },
   "peerDependencies": {
     "grpc": "^1.22.2"

src/register.ts

@@ -52,21 +52,23 @@ export default async function register({
 
   await grantLease();
 
-  function revoke(): Promise<void> {
-    return lease.revoke();
+  async function revoke(): Promise<void> {
+    await lease.revoke();
   }
 
   if (server) {
     const old = server.tryShutdown;
     server.tryShutdown = callback => {
-      revoke();
-      old.call(server, callback);
+      revoke().finally(() => {
+        old.call(server, callback);
+      });
     };
 
     const original = server.forceShutdown;
     server.forceShutdown = () => {
-      revoke();
-      original.call(server);
+      return revoke().finally(() => {
+        original.call(server);
+      });
     };
   }