diff --git a/internal/peermanager/peerwantmanager.go b/internal/peermanager/peerwantmanager.go index b0c843a2ea9e60c86cf7af1a4d18258355bedec0..08914bbcae7e1a359b8f9ff207c26d04fed8ab25 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 9cfa9410fdacb0f657dab248ea7716106e6a45fa..a56df168a2e75ab0e041cc52b5d138d4d36a1624 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) + } }