From 4219e3eae4998ccf1a4cdfa1f6773ae1d33b0048 Mon Sep 17 00:00:00 2001
From: gammazero <11790789+gammazero@users.noreply.github.com>
Date: Mon, 8 Jul 2024 15:34:02 -0700
Subject: [PATCH] Fix memory leak on BlockPresenceManager

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
---
 .../blockpresencemanager.go                   | 34 +++++++++++++++++--
 .../blockpresencemanager_test.go              | 21 ++++++++++++
 .../internal/session/sessionwantsender.go     |  1 +
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/bitswap/client/internal/blockpresencemanager/blockpresencemanager.go b/bitswap/client/internal/blockpresencemanager/blockpresencemanager.go
index 1b76acc5b8..8595c600f3 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 b977c28ffa..0a1ba7d80f 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 390fdf29dd..1beefeb94b 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)
 			}