fix: Avoid removeAllListeners on sockets (#174)

packages/backend/src/Feed.ts

@@ -24,10 +24,10 @@ export default class Feed {
     this.id = nextId();
     this.socket = socket;
 
-    socket.on('message', (data) => this.handleCommand(data.toString()));
-    socket.on('error', () => this.disconnect());
-    socket.on('close', () => this.disconnect());
-    socket.on('pong', () => this.waitingForPong = false);
+    socket.on('message', this.handleCommand);
+    socket.on('error', this.disconnect);
+    socket.on('close', this.disconnect);
+    socket.on('pong', this.onPong);
   }
 
   public static feedVersion(): FeedMessage.Message {
@@ -230,8 +230,8 @@ export default class Feed {
     this.socket.send(data, this.handleError);
   }
 
-  private handleCommand(cmd: string) {
-    const [tag, payload] = cmd.split(':', 2) as [string, Maybe<string>];
+  private handleCommand = (data: WebSocket.Data) => {
+    const [tag, payload] = data.toString().split(':', 2) as [string, Maybe<string>];
 
     if (!payload) {
       return;
@@ -272,10 +272,17 @@ export default class Feed {
     }
   }
 
-  private disconnect() {
-    this.socket.removeAllListeners();
+  private disconnect = () => {
+    this.socket.removeListener('message', this.handleCommand);
+    this.socket.removeListener('error', this.disconnect);
+    this.socket.removeListener('close', this.disconnect);
+    this.socket.removeListener('pong', this.onPong);
     this.socket.terminate();
 
     this.events.emit('disconnect');
   }
+
+  private onPong = () => {
+    this.waitingForPong = false;
+  }
 }

packages/backend/src/Node.ts

@@ -99,32 +99,10 @@ export default class Node {
     this.lastMessage = timestamp();
     this.socket = socket;
 
-    socket.on('message', (data) => {
-      const message = parseMessage(data);
-
-      if (!message) {
-        return;
-      }
-
-      this.onMessage(message);
-    });
-
-    socket.on('close', () => {
-      console.log(`${this.name} has disconnected`);
-
-      this.disconnect();
-    });
-
-    socket.on('error', (error) => {
-      console.error(`${this.name} has errored`, error);
-
-      this.disconnect();
-    });
-
-    socket.on('pong', () => {
-      this.latency = (timestamp() - this.pingStart) as Types.Milliseconds;
-      this.pingStart = 0 as Types.Timestamp;
-    });
+    socket.on('message', this.onMessageData);
+    socket.on('close', this.disconnect);
+    socket.on('error', this.disconnect);
+    socket.on('pong', this.onPong);
 
     process.nextTick(() => {
       // Handle cached messages
@@ -237,14 +215,29 @@ export default class Node {
     return +(this.lastBlockAt || 0) as Types.Milliseconds;
   }
 
-  private disconnect() {
-    this.socket.removeAllListeners();
+  private disconnect = () => {
+    console.log(`${this.name} has disconnected`);
+
+    this.socket.removeListener('message', this.onMessageData);
+    this.socket.removeListener('close', this.disconnect);
+    this.socket.removeListener('error', this.disconnect);
+    this.socket.removeListener('pong', this.onPong);
     this.socket.close();
     this.socket.terminate();
 
     this.events.emit('disconnect');
   }
 
+  private onMessageData = (data: WebSocket.Data) => {
+    const message = parseMessage(data);
+
+    if (!message) {
+      return;
+    }
+
+    this.onMessage(message);
+  }
+
   private onMessage(message: Message) {
     this.lastMessage = timestamp();
 
@@ -422,4 +415,9 @@ export default class Node {
 
     return (+time - +this.lastBlockAt) as Types.Milliseconds;
   }
+
+  private onPong = () => {
+    this.latency = (timestamp() - this.pingStart) as Types.Milliseconds;
+    this.pingStart = 0 as Types.Timestamp;
+  }
 }