fix: memory leak in Agent

Author: hexchain

lib/dispatcher/agent.js

@@ -45,22 +45,10 @@ class Agent extends DispatcherBase {
     }
 
     this[kOnConnect] = (origin, targets) => {
-      const result = this[kClients].get(origin)
-      if (result) {
-        result.count += 1
-      }
       this.emit('connect', origin, [this, ...targets])
     }
 
     this[kOnDisconnect] = (origin, targets, err) => {
-      const result = this[kClients].get(origin)
-      if (result) {
-        result.count -= 1
-        if (result.count <= 0) {
-          this[kClients].delete(origin)
-          result.dispatcher.destroy()
-        }
-      }
       this.emit('disconnect', origin, [this, ...targets], err)
     }
 
@@ -91,7 +79,23 @@ class Agent extends DispatcherBase {
     if (!dispatcher) {
       dispatcher = this[kFactory](opts.origin, this[kOptions])
         .on('drain', this[kOnDrain])
+        .on('connect', () => {
+          const result = this[kClients].get(key)
+          if (result) {
+            result.count += 1
+          }
+        })
         .on('connect', this[kOnConnect])
+        .on('disconnect', () => {
+          const result = this[kClients].get(key)
+          if (result) {
+            result.count -= 1
+            if (result.count <= 0) {
+              this[kClients].delete(key)
+              result.dispatcher.destroy()
+            }
+          }
+        })
         .on('disconnect', this[kOnDisconnect])
         .on('connectionError', this[kOnConnectionError])