feat: revoke etcd lease when server shutdown

edvardchen · Jul 29, 2019 · b639483 · b639483
1 parent cf56c1b
commit b639483
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 13 deletions.
diff --git a/__tests__/client.test.ts b/__tests__/client.test.ts
@@ -29,25 +29,23 @@ describe('grpc proxy', () => {
   });
 
   let proxy: GrpcObject;
-  let revokers: (() => void)[];
+  let revokers: Server[];
   beforeAll(async () => {
     revokers = await Promise.all(
-      [50051, 50052].map(port => {
+      [50051, 50052].map(async port => {
         // const port = Math.floor(Math.random() * 1e4);
         // start server
         const server = startGreeterServer(port) as Server;
 
-        return register({
+        await register({
+          server,
           etcdKV: {
             key: `test-services:helloworld.Greeter:localhost:${port}`,
           },
           etcdHosts: hosts,
-        }).then(revoke => {
-          return () => {
-            server.forceShutdown();
-            revoke();
-          };
         });
+
+        return server;
       })
     );
 
@@ -59,7 +57,7 @@ describe('grpc proxy', () => {
   });
 
   afterAll(() => {
-    revokers.map(item => item());
+    revokers.map(item => item.forceShutdown());
   });
 
   it('round=robin', async () => {
@@ -74,7 +72,7 @@ describe('grpc proxy', () => {
 
     // revoke
     const revoker = revokers.shift();
-    revoker && revoker();
+    revoker && revoker.forceShutdown();
     await sleep(200);
     const fourth = proxy.Greeter.getChannel().getTarget();
     const fifth = proxy.Greeter.getChannel().getTarget();

diff --git a/src/register.ts b/src/register.ts
@@ -1,11 +1,14 @@
 import { Etcd3, Lease } from 'etcd3';
 import Debug from 'debug';
+import { Server } from 'grpc';
 
 const debug = Debug('grpclb:server');
 
 type Str = string | (() => string);
 
 type RegisterOptions = {
+  server?: Server;
+
   /** etcd lease ttl in seconds */
   ttl?: number;
 
@@ -17,6 +20,7 @@ type RegisterOptions = {
 };
 
 export default async function register({
+  server,
   ttl = 10, // 10 seconds,
   etcdKV: { key: _key, value: _value = '' },
   etcdHosts = process.env.ETCD_HOSTS,
@@ -48,7 +52,23 @@ export default async function register({
 
   await grantLease();
 
-  return () => {
-    lease.revoke();
-  };
+  function revoke(): Promise<void> {
+    return lease.revoke();
+  }
+
+  if (server) {
+    const old = server.tryShutdown;
+    server.tryShutdown = callback => {
+      revoke();
+      old.call(server, callback);
+    };
+
+    const original = server.forceShutdown;
+    server.forceShutdown = () => {
+      revoke();
+      original.call(server);
+    };
+  }
+
+  return revoke;
 }