fix: ensure wantlist gauge gets decremented on disconnect

ipfs · Apr 3, 2020 · d310fe3 · d310fe3
1 parent d39c760
commit d310fe3
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/internal/peermanager/peerwantmanager.go b/internal/peermanager/peerwantmanager.go
@@ -50,6 +50,16 @@ func (pwm *peerWantManager) addPeer(p peer.ID) {
 
 // RemovePeer removes a peer and its associated wants from tracking
 func (pwm *peerWantManager) removePeer(p peer.ID) {
+	pws, ok := pwm.peerWants[p]
+	if !ok {
+		return
+	}
+
+	// Decrement the gauge by the number of pending want-blocks to the peer
+	for range pws.wantBlocks.Keys() {
+		pwm.wantBlockGauge.Dec()
+	}
+
 	delete(pwm.peerWants, p)
 }
 

diff --git a/internal/peermanager/peerwantmanager_test.go b/internal/peermanager/peerwantmanager_test.go
@@ -289,4 +289,10 @@ func TestStats(t *testing.T) {
 	if g.count != 3 {
 		t.Fatal("Expected 3 want-blocks", g.count)
 	}
+
+	pwm.removePeer(p0)
+
+	if g.count != 0 {
+		t.Fatal("Expected all want-blocks to be removed with peer", g.count)
+	}
 }