From c817a16d314e9a7189a322265e42979c28ef1f11 Mon Sep 17 00:00:00 2001
From: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com>
Date: Mon, 1 Jul 2019 19:16:42 +0200
Subject: [PATCH] feat: Don't count stale nodes (#160)

* feat: Don't count stale nodes

* fix: Don't decrement twice when removing a stale node

* chore: Remove unnecessary empty line
---
 packages/backend/src/Chain.ts | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/packages/backend/src/Chain.ts b/packages/backend/src/Chain.ts
index b44cb60..bd03763 100644
--- a/packages/backend/src/Chain.ts
+++ b/packages/backend/src/Chain.ts
@@ -10,6 +10,7 @@ const BLOCK_TIME_HISTORY = 10;
 export default class Chain {
   private nodes = new Set<Node>();
   private feeds = new FeedSet();
+  private count = 0;
 
   public readonly events = new EventEmitter();
   public readonly label: Types.ChainLabel;
@@ -28,12 +29,13 @@ export default class Chain {
   }
 
   public get nodeCount(): Types.NodeCount {
-    return this.nodes.size as Types.NodeCount;
+    return this.count as Types.NodeCount;
   }
 
   public addNode(node: Node) {
     console.log(`[${this.label}] new node: ${node.name}`);
 
+    this.count += 1;
     this.nodes.add(node);
     this.feeds.broadcast(Feed.addedNode(node));
 
@@ -81,7 +83,11 @@ export default class Chain {
     node.events.removeAllListeners();
 
     this.nodes.delete(node);
-    this.feeds.broadcast(Feed.removedNode(node));
+
+    if (!node.isStale) {
+      this.count -= 1;
+      this.feeds.broadcast(Feed.removedNode(node));
+    }
 
     this.events.emit('disconnect', this.nodeCount);
 
@@ -91,6 +97,7 @@ export default class Chain {
   }
 
   public staleNode(node: Node) {
+    this.count -= 1;
     this.feeds.broadcast(Feed.removedNode(node));
 
     if (this.height === node.best.number) {
@@ -168,6 +175,7 @@ export default class Chain {
     if (node.isStale) {
       node.isStale = false;
       this.feeds.broadcast(Feed.addedNode(node));
+      this.count += 1;
     } else {
       this.feeds.broadcast(Feed.imported(node));
     }