From e41fd53742ddbb4160412404b4bbd3c7696651de Mon Sep 17 00:00:00 2001 From: guillaumemichel Date: Wed, 29 Jan 2025 16:54:51 +0100 Subject: [PATCH 1/2] fix(bitswap): simplify sessioninterestmanager --- .../sessioninterestmanager.go | 37 ++++--------------- .../sessioninterestmanager_test.go | 4 +- .../internal/sessionmanager/sessionmanager.go | 2 +- 3 files changed, 10 insertions(+), 33 deletions(-) diff --git a/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager.go b/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager.go index 3756005d8..e8c27937c 100644 --- a/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager.go +++ b/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager.go @@ -11,21 +11,14 @@ import ( // SessionInterestManager records the CIDs that each session is interested in. type SessionInterestManager struct { lk sync.RWMutex - wants map[cid.Cid]map[uint64]bool + wants map[cid.Cid]map[uint64]struct{} } // New initializes a new SessionInterestManager. func New() *SessionInterestManager { return &SessionInterestManager{ - // Map of cids -> sessions -> bool - // - // The boolean indicates whether the session still wants the block - // or is just interested in receiving messages about it. - // - // Note that once the block is received the session no longer wants - // the block, but still wants to receive messages from peers who have - // the block as they may have other blocks the session is interested in. - wants: make(map[cid.Cid]map[uint64]bool), + // Map of cids -> set of sessions that want this cid + wants: make(map[cid.Cid]map[uint64]struct{}), } } @@ -39,9 +32,9 @@ func (sim *SessionInterestManager) RecordSessionInterest(ses uint64, ks []cid.Ci for _, c := range ks { // Record that the session wants the blocks if want, ok := sim.wants[c]; ok { - want[ses] = true + want[ses] = struct{}{} } else { - sim.wants[c] = map[uint64]bool{ses: true} + sim.wants[c] = map[uint64]struct{}{ses: {}} } } } @@ -73,23 +66,7 @@ func (sim *SessionInterestManager) RemoveSession(ses uint64) []cid.Cid { } // When the session receives blocks, it calls RemoveSessionWants(). -func (sim *SessionInterestManager) RemoveSessionWants(ses uint64, ks []cid.Cid) { - sim.lk.Lock() - defer sim.lk.Unlock() - - // For each key - for _, c := range ks { - // If the session wanted the block - if wanted, ok := sim.wants[c][ses]; ok && wanted { - // Mark the block as unwanted - sim.wants[c][ses] = false - } - } -} - -// When a request is cancelled, the session calls RemoveSessionInterested(). -// Returns the keys that no session is interested in any more. -func (sim *SessionInterestManager) RemoveSessionInterested(ses uint64, ks []cid.Cid) []cid.Cid { +func (sim *SessionInterestManager) RemoveSessionWants(ses uint64, ks []cid.Cid) []cid.Cid { sim.lk.Lock() defer sim.lk.Unlock() @@ -153,7 +130,7 @@ func (sim *SessionInterestManager) SplitWantedUnwanted(blks []blocks.Block) ([]b // For each session that is interested in the key for ses := range sim.wants[c] { // If the session wants the key (rather than just being interested) - if wanted, ok := sim.wants[c][ses]; ok && wanted { + if _, ok := sim.wants[c][ses]; ok { // Add the key to the set wantedKs.Add(c) } diff --git a/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager_test.go b/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager_test.go index a9779e297..242630273 100644 --- a/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager_test.go +++ b/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager_test.go @@ -121,7 +121,7 @@ func TestRemoveSessionInterested(t *testing.T) { sim.RecordSessionInterest(ses1, cids1) sim.RecordSessionInterest(ses2, cids2) - res := sim.RemoveSessionInterested(ses1, []cid.Cid{cids1[0]}) + res := sim.RemoveSessionWants(ses1, []cid.Cid{cids1[0]}) if len(res) != 1 { t.Fatal("Expected no interested sessions left") } @@ -131,7 +131,7 @@ func TestRemoveSessionInterested(t *testing.T) { t.Fatal("Expected ses1 still interested in one cid") } - res = sim.RemoveSessionInterested(ses1, cids1) + res = sim.RemoveSessionWants(ses1, cids1) if len(res) != 0 { t.Fatal("Expected ses2 to be interested in one cid") } diff --git a/bitswap/client/internal/sessionmanager/sessionmanager.go b/bitswap/client/internal/sessionmanager/sessionmanager.go index 0d2b24330..80eac9f2f 100644 --- a/bitswap/client/internal/sessionmanager/sessionmanager.go +++ b/bitswap/client/internal/sessionmanager/sessionmanager.go @@ -181,7 +181,7 @@ func (sm *SessionManager) ReceiveFrom(ctx context.Context, p peer.ID, blks []cid func (sm *SessionManager) CancelSessionWants(sesid uint64, wants []cid.Cid) { // Remove session's interest in the given blocks - returns the keys that no // session is interested in anymore. - cancelKs := sm.sessionInterestManager.RemoveSessionInterested(sesid, wants) + cancelKs := sm.sessionInterestManager.RemoveSessionWants(sesid, wants) sm.cancelWants(cancelKs) } From 37cb7ea2970a5930e05ea94383411f89d4a76ebd Mon Sep 17 00:00:00 2001 From: Guillaume Michel Date: Fri, 31 Jan 2025 16:37:28 +0100 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: Andrew Gillis <11790789+gammazero@users.noreply.github.com> --- .../internal/sessioninterestmanager/sessioninterestmanager.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager.go b/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager.go index e8c27937c..0504a63bd 100644 --- a/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager.go +++ b/bitswap/client/internal/sessioninterestmanager/sessioninterestmanager.go @@ -127,9 +127,8 @@ func (sim *SessionInterestManager) SplitWantedUnwanted(blks []blocks.Block) ([]b wantedKs := cid.NewSet() for _, b := range blks { c := b.Cid() - // For each session that is interested in the key + // For each session that wants the key. for ses := range sim.wants[c] { - // If the session wants the key (rather than just being interested) if _, ok := sim.wants[c][ses]; ok { // Add the key to the set wantedKs.Add(c)