Skip to content

Commit

Permalink
Fix memory leak on BlockPresenceManager
Browse files Browse the repository at this point in the history
Limit BlockPresenceManager map growth be doing the following:

- Use nil map when BlockPresenceManager CID map is empty.
- Delete peer map when from CID map when peer map is empty.
- Remove peers from BlockPresenceManager when peers are pruned from session.

This allows GC to free memory when maps in BlockPresenceManager become empty.

Additional improvement: Do not look for lists of keys or peers in empty maps; return early instead.

Fix for issue #574
  • Loading branch information
gammazero committed Jul 8, 2024
1 parent 733fa55 commit 4219e3e
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
}
Expand Down Expand Up @@ -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) {
Expand All @@ -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 {
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
1 change: 1 addition & 0 deletions bitswap/client/internal/session/sessionwantsender.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down

0 comments on commit 4219e3e

Please sign in to comment.