diff --git a/CHANGELOG.md b/CHANGELOG.md index 8cf7a40ae..15d6aa0fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,8 @@ The following emojis are used to highlight certain changes: ### Fixed +- `bitswap/client` fix memory leak in BlockPresenceManager due to unlimited map growth. + ### Security ## [v0.21.0] diff --git a/bitswap/client/internal/blockpresencemanager/blockpresencemanager.go b/bitswap/client/internal/blockpresencemanager/blockpresencemanager.go index 1b76acc5b..8595c600f 100644 --- a/bitswap/client/internal/blockpresencemanager/blockpresencemanager.go +++ b/bitswap/client/internal/blockpresencemanager/blockpresencemanager.go @@ -15,9 +15,7 @@ type BlockPresenceManager struct { } func New() *BlockPresenceManager { - return &BlockPresenceManager{ - presence: make(map[cid.Cid]map[peer.ID]bool), - } + return &BlockPresenceManager{} } // ReceiveFrom is called when a peer sends us information about which blocks @@ -26,6 +24,10 @@ func (bpm *BlockPresenceManager) ReceiveFrom(p peer.ID, haves []cid.Cid, dontHav bpm.Lock() defer bpm.Unlock() + if bpm.presence == nil { + bpm.presence = make(map[cid.Cid]map[peer.ID]bool) + } + for _, c := range haves { bpm.updateBlockPresence(p, c, true) } @@ -75,6 +77,10 @@ func (bpm *BlockPresenceManager) AllPeersDoNotHaveBlock(peers []peer.ID, ks []ci bpm.RLock() defer bpm.RUnlock() + if len(bpm.presence) == 0 { + return nil + } + var res []cid.Cid for _, c := range ks { if bpm.allDontHave(peers, c) { @@ -90,6 +96,9 @@ func (bpm *BlockPresenceManager) allDontHave(peers []peer.ID, c cid.Cid) bool { if !cok { return false } + if len(ps) == 0 { + return false + } // Check if we explicitly know that all the given peers do not have the cid for _, p := range peers { @@ -108,6 +117,25 @@ func (bpm *BlockPresenceManager) RemoveKeys(ks []cid.Cid) { for _, c := range ks { delete(bpm.presence, c) } + if len(bpm.presence) == 0 { + bpm.presence = nil + } +} + +// RemovePeers removes the given peer from every cid key in the presence map. +func (bpm *BlockPresenceManager) RemovePeer(p peer.ID) { + bpm.Lock() + defer bpm.Unlock() + + for c, pm := range bpm.presence { + delete(pm, p) + if len(pm) == 0 { + delete(bpm.presence, c) + } + } + if len(bpm.presence) == 0 { + bpm.presence = nil + } } // HasKey indicates whether the BlockPresenceManager is tracking the given key diff --git a/bitswap/client/internal/blockpresencemanager/blockpresencemanager_test.go b/bitswap/client/internal/blockpresencemanager/blockpresencemanager_test.go index b977c28ff..0a1ba7d80 100644 --- a/bitswap/client/internal/blockpresencemanager/blockpresencemanager_test.go +++ b/bitswap/client/internal/blockpresencemanager/blockpresencemanager_test.go @@ -93,6 +93,27 @@ func TestBlockPresenceManager(t *testing.T) { if bpm.PeerDoesNotHaveBlock(p, c1) { t.Fatal(expDoesNotHaveFalseMsg) } + + bpm.ReceiveFrom(p, []cid.Cid{c0}, []cid.Cid{c1}) + if !bpm.PeerHasBlock(p, c0) { + t.Fatal(expHasTrueMsg) + } + if !bpm.PeerDoesNotHaveBlock(p, c1) { + t.Fatal(expDoesNotHaveTrueMsg) + } + bpm.RemovePeer(p) + if bpm.PeerHasBlock(p, c0) { + t.Fatal(expHasFalseMsg) + } + if bpm.PeerDoesNotHaveBlock(p, c0) { + t.Fatal(expDoesNotHaveFalseMsg) + } + if bpm.PeerHasBlock(p, c1) { + t.Fatal(expHasFalseMsg) + } + if bpm.PeerDoesNotHaveBlock(p, c1) { + t.Fatal(expDoesNotHaveFalseMsg) + } } func TestAddRemoveMulti(t *testing.T) { diff --git a/bitswap/client/internal/session/sessionwantsender.go b/bitswap/client/internal/session/sessionwantsender.go index 390fdf29d..1beefeb94 100644 --- a/bitswap/client/internal/session/sessionwantsender.go +++ b/bitswap/client/internal/session/sessionwantsender.go @@ -455,6 +455,7 @@ func (sws *sessionWantSender) processUpdates(updates []update) []cid.Cid { go func() { for p := range prunePeers { // Peer doesn't have anything we want, so remove it + sws.bpm.RemovePeer(p) log.Infof("peer %s sent too many dont haves, removing from session %d", p, sws.ID()) sws.SignalAvailability(p, false) }