From 799b1a7bb0a78d0dea519b272e6a3f4c6ee60f9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C8=98tefan=20Talpalaru?= <stefantalpalaru@yahoo.com>
Date: Wed, 26 Jun 2019 14:48:04 +0200
Subject: [PATCH] initial statistics collection setup

---
 eth/common/eth_types.nim | 3 +++
 eth/common/utils.nim     | 2 ++
 eth/p2p/peer_pool.nim    | 3 ++-
 eth/p2p/rlpx.nim         | 3 ++-
 4 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/eth/common/eth_types.nim b/eth/common/eth_types.nim
index f66c82a..2415a0e 100644
--- a/eth/common/eth_types.nim
+++ b/eth/common/eth_types.nim
@@ -168,6 +168,9 @@ type
     OK
     Error
 
+  NimbusStats* = object
+    num_peers*: int
+
 when BlockNumber is int64:
   ## The goal of these templates is to make it easier to switch
   ## the block number type to a different representation
diff --git a/eth/common/utils.nim b/eth/common/utils.nim
index 215af36..ff89a97 100644
--- a/eth/common/utils.nim
+++ b/eth/common/utils.nim
@@ -8,3 +8,5 @@ proc parseAddress*(hexString: string): EthAddress =
 proc `$`*(a: EthAddress): string =
   a.toHex()
 
+var nimbusStats*: NimbusStats
+
diff --git a/eth/p2p/peer_pool.nim b/eth/p2p/peer_pool.nim
index e73339c..490f4eb 100644
--- a/eth/p2p/peer_pool.nim
+++ b/eth/p2p/peer_pool.nim
@@ -3,7 +3,7 @@
 
 import
   os, tables, times, random, sequtils, options,
-  chronos, chronicles, eth/[rlp, keys],
+  chronos, chronicles, eth/[rlp, keys, common],
   private/p2p_types, discovery, kademlia, rlpx
 
 const
@@ -110,6 +110,7 @@ proc getRandomBootnode(p: PeerPool): Option[Node] =
 proc addPeer*(pool: PeerPool, peer: Peer) =
   doAssert(peer.remote notin pool.connectedNodes)
   pool.connectedNodes[peer.remote] = peer
+  inc(nimbusStats.num_peers)
   for o in pool.observers.values:
     if not o.onPeerConnected.isNil:
       if o.protocol.isNil or peer.supports(o.protocol):
diff --git a/eth/p2p/rlpx.nim b/eth/p2p/rlpx.nim
index ee974e2..40c5967 100644
--- a/eth/p2p/rlpx.nim
+++ b/eth/p2p/rlpx.nim
@@ -787,8 +787,9 @@ proc removePeer(network: EthereumNode, peer: Peer) =
   # have been dropped already from the peers side.
   # E.g. when receiving a p2p.disconnect message from a peer, a race will happen
   # between which side disconnects first.
-  if network.peerPool != nil and not peer.remote.isNil:
+  if network.peerPool != nil and not peer.remote.isNil and peer.remote in network.peerPool.connectedNodes:
     network.peerPool.connectedNodes.del(peer.remote)
+    dec(nimbusStats.num_peers)
 
     # Note: we need to do this check as disconnect (and thus removePeer)
     # currently can get called before the dispatcher is initialized.