grace period disconnect should be per-client

lib/Client.js

@@ -1,21 +1,60 @@
 import http from 'http';
 import Debug from 'debug';
 import pump from 'pump';
+import EventEmitter from 'events';
 
 // A client encapsulates req/res handling using an agent
 //
 // If an agent is destroyed, the request handling will error
 // The caller is responsible for handling a failed request
-class Client {
+class Client extends EventEmitter {
     constructor(options) {
-        this.agent = options.agent;
+        super();
-        this.debug = Debug('lt:Client');
+
+        const agent = this.agent = options.agent;
+        const id = this.id = options.id;
+
+        this.debug = Debug(`lt:Client[${this.id}]`);
+
+        // client is given a grace period in which they can connect before they are _removed_
+        this.graceTimeout = setTimeout(() => {
+            this.close();
+        }, 1000).unref();
+
+        agent.on('online', () => {
+            this.debug('client online %s', id);
+            clearTimeout(this.graceTimeout);
+        });
+
+        agent.on('offline', () => {
+            this.debug('client offline %s', id);
+
+            // if there was a previous timeout set, we don't want to double trigger
+            clearTimeout(this.graceTimeout);
+
+            // client is given a grace period in which they can re-connect before they are _removed_
+            this.graceTimeout = setTimeout(() => {
+                this.close();
+            }, 1000).unref();
+        });
+
+        // TODO(roman): an agent error removes the client, the user needs to re-connect?
+        // how does a user realize they need to re-connect vs some random client being assigned same port?
+        agent.once('error', (err) => {
+            this.close();
+        });
     }
 
     stats() {
         return this.agent.stats();
     }
 
+    close() {
+        clearTimeout(this.graceTimeout);
+        this.agent.destroy();
+        this.emit('close');
+    }
+
     handleRequest(req, res) {
         this.debug('> %s', req.url);
         const opt = {

lib/Client.test.js

@@ -1,7 +1,6 @@
 import assert from 'assert';
 import http from 'http';
 import { Duplex } from 'stream';
-import EventEmitter from 'events';
 import WebSocket from 'ws';
 import net from 'net';
 

lib/ClientManager.js

@@ -21,6 +21,7 @@ class ClientManager {
 
         this.debug = Debug('lt:ClientManager');
 
+        // This is totally wrong :facepalm: this needs to be per-client...
         this.graceTimeout = null;
     }
 
@@ -42,35 +43,19 @@ class ClientManager {
             maxSockets: 10,
         });
 
-        agent.on('online', () => {
+        const client = new Client({
-            this.debug('client online %s', id);
+            id,
-            clearTimeout(this.graceTimeout);
+            agent,
         });
 
-        agent.on('offline', () => {
-            // TODO(roman): grace period for re-connecting
-            // this period is short as the client is expected to maintain connections actively
-            // if they client does not reconnect on a dropped connection they need to re-establish
-            this.debug('client offline %s', id);
-
-            // client is given a grace period in which they can re-connect before they are _removed_
-            this.graceTimeout = setTimeout(() => {
-                this.removeClient(id);
-            }, 1000);
-        });
-
-        // TODO(roman): an agent error removes the client, the user needs to re-connect?
-        // how does a user realize they need to re-connect vs some random client being assigned same port?
-        agent.once('error', (err) => {
-            this.removeClient(id);
-        });
-
-        const client = new Client({ agent });
-
         // add to clients map immediately
         // avoiding races with other clients requesting same id
         clients[id] = client;
 
+        client.once('close', () => {
+            this.removeClient(id);
+        });
+
         // try/catch used here to remove client id
         try {
             const info = await agent.listen();
@@ -96,11 +81,11 @@ class ClientManager {
         }
         --this.stats.tunnels;
         delete this.clients[id];
-        client.agent.destroy();
+        client.close();
     }
 
     hasClient(id) {
-        return this.clients[id];
+        return !!this.clients[id];
     }
 
     getClient(id) {

lib/ClientManager.test.js

@@ -54,4 +54,37 @@ describe('ClientManager', () => {
         await new Promise(resolve => setTimeout(resolve, 1500));
         assert(!manager.hasClient('foobar'));
     }).timeout(5000);
+
+    it('should remove correct client once it goes offline', async () => {
+        const manager = new ClientManager();
+        const clientFoo = await manager.newClient('foo');
+        const clientBar = await manager.newClient('bar');
+
+        const socket = await new Promise((resolve) => {
+            const netClient = net.createConnection({ port: clientFoo.port }, () => {
+                resolve(netClient);
+            });
+        });
+
+        await new Promise(resolve => setTimeout(resolve, 1500));
+
+        // foo should still be ok
+        assert(manager.hasClient('foo'));
+
+        // clientBar shound be removed - nothing connected to it
+        assert(!manager.hasClient('bar'));
+
+        manager.removeClient('foo');
+        socket.end();
+    }).timeout(5000);
+
+    it('should remove clients if they do not connect within 5 seconds', async () => {
+        const manager = new ClientManager();
+        const clientFoo = await manager.newClient('foo');
+        assert(manager.hasClient('foo'));
+
+        // wait past grace period (1s)
+        await new Promise(resolve => setTimeout(resolve, 1500));
+        assert(!manager.hasClient('foo'));
+    }).timeout(5000);
 });