From 6af7909378d2778a4e98fd398bd23eea62d31945 Mon Sep 17 00:00:00 2001
From: Giovanni Petrantoni <giovanni@fragcolor.xyz>
Date: Sat, 18 Jul 2020 21:00:44 +0900
Subject: [PATCH] fix internal gossip tests

---
 libp2p/protocols/pubsub/gossipsub.nim | 30 ++++++++++++++++++++-------
 tests/pubsub/testgossipinternal.nim   |  4 ++++
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/libp2p/protocols/pubsub/gossipsub.nim b/libp2p/protocols/pubsub/gossipsub.nim
index 60aff780f..ddafd13fa 100644
--- a/libp2p/protocols/pubsub/gossipsub.nim
+++ b/libp2p/protocols/pubsub/gossipsub.nim
@@ -103,6 +103,9 @@ type
     heartbeatLock: AsyncLock                   # heartbeat lock to prevent two consecutive concurrent heartbeats
     peerStats: Table[PubSubPeer, PeerStats]
 
+    when not defined(release):
+      prunedPeers: HashSet[PubSubPeer]
+
 declareGauge(libp2p_gossipsub_peers_per_topic_mesh,
   "gossipsub peers per topic in mesh",
   labels = ["topic"])
@@ -191,14 +194,27 @@ method handleConnect*(g: GossipSub, peer: PubSubPeer) =
     discard
 
 proc grafted(g: GossipSub, p: PubSubPeer, topic: string) =
-  var stats = g.peerStats[p]
-  var info = stats.topicInfos.mgetOrPut(topic, TopicInfo())
-  info.graftTime = Moment.now()
-  info.meshTime = 0.seconds
+  g.peerStats.withValue(p, stats) do:
+    var info = stats.topicInfos.mgetOrPut(topic, TopicInfo())
+    info.graftTime = Moment.now()
+    info.meshTime = 0.seconds
+  do:
+    raise newException(CatchableError, "TopicInfo key not found for " & $p)
 
-proc pruned(g: GossipSub, p: PubSubPeer, topic: string) =
-  var stats = g.peerStats[p]
-  var _ = stats.topicInfos[topic]
+proc pruned(g: GossipSub, p: PubSubPeer, topic: string) {.gcsafe.} =
+  g.peerStats.withValue(p, stats) do:
+    when not defined(release):
+      g.prunedPeers.incl(p)
+
+    var _ = stats.topicInfos[topic]
+  do:
+    when not defined(release):
+      if p in g.prunedPeers:
+        raise newException(CatchableError, "Dupe prune " & $p)
+      else:
+        raise newException(CatchableError, "TopicInfo key not found for " & $p)
+    else:
+      raise newException(CatchableError, "TopicInfo key not found for " & $p)
 
 proc replenishFanout(g: GossipSub, topic: string) =
   ## get fanout peers for a topic
diff --git a/tests/pubsub/testgossipinternal.nim b/tests/pubsub/testgossipinternal.nim
index 96d00f070..e5d6cd5d8 100644
--- a/tests/pubsub/testgossipinternal.nim
+++ b/tests/pubsub/testgossipinternal.nim
@@ -53,6 +53,8 @@ suite "GossipSub internal":
         let peer = newPubSubPeer(peerInfo, GossipSubCodec)
         peer.conn = conn
         gossipSub.peers[peerInfo.id] = peer
+        gossipSub.handleConnect(peer)
+        gossipSub.grafted(peer, topic)
         gossipSub.mesh[topic].incl(peer)
 
       check gossipSub.peers.len == 15
@@ -84,6 +86,8 @@ suite "GossipSub internal":
         let peer = newPubSubPeer(peerInfo, GossipSubCodec)
         peer.conn = conn
         gossipSub.peers[peerInfo.id] = peer
+        gossipSub.handleConnect(peer)
+        gossipSub.grafted(peer, topic)
         gossipSub.mesh[topic].incl(peer)
 
       check gossipSub.mesh[topic].len == 15