Refactor SessionInterestManager

internal/sessioninterestmanager/sessioninterestmanager.go

@@ -3,24 +3,29 @@ package sessioninterestmanager
 import (
 	"sync"
 
-	bsswl "github.com/ipfs/go-bitswap/internal/sessionwantlist"
 	blocks "github.com/ipfs/go-block-format"
 
 	cid "github.com/ipfs/go-cid"
 )
 
 // SessionInterestManager records the CIDs that each session is interested in.
 type SessionInterestManager struct {
-	lk         sync.RWMutex
-	interested *bsswl.SessionWantlist
-	wanted     *bsswl.SessionWantlist
+	lk    sync.RWMutex
+	wants map[cid.Cid]map[uint64]bool
 }
 
 // New initializes a new SessionInterestManager.
 func New() *SessionInterestManager {
 	return &SessionInterestManager{
-		interested: bsswl.NewSessionWantlist(),
-		wanted:     bsswl.NewSessionWantlist(),
+		// 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),
 	}
 }
 
@@ -30,8 +35,15 @@ func (sim *SessionInterestManager) RecordSessionInterest(ses uint64, ks []cid.Ci
 	sim.lk.Lock()
 	defer sim.lk.Unlock()
 
-	sim.interested.Add(ks, ses)
-	sim.wanted.Add(ks, ses)
+	// For each key
+	for _, c := range ks {
+		// Record that the session wants the blocks
+		if want, ok := sim.wants[c]; ok {
+			want[ses] = true
+		} else {
+			sim.wants[c] = map[uint64]bool{ses: true}
+		}
+	}
 }
 
 // When the session shuts down it calls RemoveSessionInterest().
@@ -40,26 +52,68 @@ func (sim *SessionInterestManager) RemoveSession(ses uint64) []cid.Cid {
 	sim.lk.Lock()
 	defer sim.lk.Unlock()
 
-	sim.wanted.RemoveSession(ses)
-	return sim.interested.RemoveSession(ses)
+	// The keys that no session is interested in
+	deletedKs := make([]cid.Cid, 0)
+
+	// For each known key
+	for c := range sim.wants {
+		// Remove the session from the list of sessions that want the key
+		delete(sim.wants[c], ses)
+
+		// If there are no more sessions that want the key
+		if len(sim.wants[c]) == 0 {
+			// Clean up the list memory
+			delete(sim.wants, c)
+			// Add the key to the list of keys that no session is interested in
+			deletedKs = append(deletedKs, c)
+		}
+	}
+
+	return deletedKs
 }
 
 // When the session receives blocks, it calls RemoveSessionWants().
-func (sim *SessionInterestManager) RemoveSessionWants(ses uint64, wants []cid.Cid) {
+func (sim *SessionInterestManager) RemoveSessionWants(ses uint64, ks []cid.Cid) {
 	sim.lk.Lock()
 	defer sim.lk.Unlock()
 
-	sim.wanted.RemoveSessionKeys(ses, wants)
+	// 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, wants []cid.Cid) []cid.Cid {
+func (sim *SessionInterestManager) RemoveSessionInterested(ses uint64, ks []cid.Cid) []cid.Cid {
 	sim.lk.Lock()
 	defer sim.lk.Unlock()
 
-	sim.wanted.RemoveSessionKeys(ses, wants)
-	return sim.interested.RemoveSessionKeys(ses, wants)
+	// The keys that no session is interested in
+	deletedKs := make([]cid.Cid, 0, len(ks))
+
+	// For each key
+	for _, c := range ks {
+		// If there is a list of sessions that want the key
+		if _, ok := sim.wants[c]; ok {
+			// Remove the session from the list of sessions that want the key
+			delete(sim.wants[c], ses)
+
+			// If there are no more sessions that want the key
+			if len(sim.wants[c]) == 0 {
+				// Clean up the list memory
+				delete(sim.wants, c)
+				// Add the key to the list of keys that no session is interested in
+				deletedKs = append(deletedKs, c)
+			}
+		}
+	}
+
+	return deletedKs
 }
 
 // The session calls FilterSessionInterested() to filter the sets of keys for
@@ -68,9 +122,20 @@ func (sim *SessionInterestManager) FilterSessionInterested(ses uint64, ksets ...
 	sim.lk.RLock()
 	defer sim.lk.RUnlock()
 
+	// For each set of keys
 	kres := make([][]cid.Cid, len(ksets))
 	for i, ks := range ksets {
-		kres[i] = sim.interested.SessionHas(ses, ks).Keys()
+		// The set of keys that at least one session is interested in
+		has := make([]cid.Cid, 0, len(ks))
+
+		// For each key in the list
+		for _, c := range ks {
+			// If there is a session that's interested, add the key to the set
+			if _, ok := sim.wants[c][ses]; ok {
+				has = append(has, c)
+			}
+		}
+		kres[i] = has
 	}
 	return kres
 }
@@ -81,12 +146,19 @@ func (sim *SessionInterestManager) SplitWantedUnwanted(blks []blocks.Block) ([]b
 	sim.lk.RLock()
 	defer sim.lk.RUnlock()
 
-	// Get the wanted block keys
-	ks := make([]cid.Cid, len(blks))
+	// Get the wanted block keys as a set
+	wantedKs := cid.NewSet()
 	for _, b := range blks {
-		ks = append(ks, b.Cid())
+		c := b.Cid()
+		// 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 {
+				// Add the key to the set
+				wantedKs.Add(c)
+			}
+		}
 	}
-	wantedKs := sim.wanted.Has(ks)
 
 	// Separate the blocks into wanted and unwanted
 	wantedBlks := make([]blocks.Block, 0, len(blks))
@@ -112,5 +184,18 @@ func (sim *SessionInterestManager) InterestedSessions(blks []cid.Cid, haves []ci
 	ks = append(ks, haves...)
 	ks = append(ks, dontHaves...)
 
-	return sim.interested.SessionsFor(ks)
+	// Create a set of sessions that are interested in the keys
+	sesSet := make(map[uint64]struct{})
+	for _, c := range ks {
+		for s := range sim.wants[c] {
+			sesSet[s] = struct{}{}
+		}
+	}
+
+	// Convert the set into a list
+	ses := make([]uint64, 0, len(sesSet))
+	for s := range sesSet {
+		ses = append(ses, s)
+	}
+	return ses
 }