From d310fe30d4fe9bf889900b32bb8a91393f0d8a0f Mon Sep 17 00:00:00 2001
From: Dirk McCormick <dirkmdev@gmail.com>
Date: Fri, 3 Apr 2020 14:36:15 -0400
Subject: [PATCH] fix: ensure wantlist gauge gets decremented on disconnect

---
 internal/peermanager/peerwantmanager.go      | 10 ++++++++++
 internal/peermanager/peerwantmanager_test.go |  6 ++++++
 2 files changed, 16 insertions(+)

diff --git a/internal/peermanager/peerwantmanager.go b/internal/peermanager/peerwantmanager.go
index b0c843a..08914bb 100644
--- 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
index 9cfa941..a56df16 100644
--- 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)
+	}
 }
-- 
GitLab