fix: race in SessionInterestManager (#324)

Author: dirkmc

Diff for: internal/sessioninterestmanager/sessioninterestmanager.go

@@ -1,13 +1,17 @@
 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
 }
@@ -20,29 +24,52 @@ func New() *SessionInterestManager {
 	}
 }
 
+// When the client asks the session for blocks, the session calls
+// RecordSessionInterest() with those cids.
 func (sim *SessionInterestManager) RecordSessionInterest(ses uint64, ks []cid.Cid) {
+	sim.lk.Lock()
+	defer sim.lk.Unlock()
+
 	sim.interested.Add(ks, ses)
 	sim.wanted.Add(ks, ses)
 }
 
+// When the session shuts down it calls RemoveSessionInterest().
 func (sim *SessionInterestManager) RemoveSessionInterest(ses uint64) []cid.Cid {
+	sim.lk.Lock()
+	defer sim.lk.Unlock()
+
 	sim.wanted.RemoveSession(ses)
 	return sim.interested.RemoveSession(ses)
 }
 
+// When the session receives blocks, it calls RemoveSessionWants().
 func (sim *SessionInterestManager) RemoveSessionWants(ses uint64, wants []cid.Cid) {
+	sim.lk.Lock()
+	defer sim.lk.Unlock()
+
 	sim.wanted.RemoveSessionKeys(ses, wants)
 }
 
+// The session calls FilterSessionInterested() to filter the sets of keys for
+// those that the session is interested in
 func (sim *SessionInterestManager) FilterSessionInterested(ses uint64, ksets ...[]cid.Cid) [][]cid.Cid {
+	sim.lk.RLock()
+	defer sim.lk.RUnlock()
+
 	kres := make([][]cid.Cid, len(ksets))
 	for i, ks := range ksets {
 		kres[i] = sim.interested.SessionHas(ses, ks).Keys()
 	}
 	return kres
 }
 
+// When bitswap receives blocks it calls SplitWantedUnwanted() to discard
+// unwanted blocks
 func (sim *SessionInterestManager) SplitWantedUnwanted(blks []blocks.Block) ([]blocks.Block, []blocks.Block) {
+	sim.lk.RLock()
+	defer sim.lk.RUnlock()
+
 	// Get the wanted block keys
 	ks := make([]cid.Cid, len(blks))
 	for _, b := range blks {
@@ -63,7 +90,12 @@ func (sim *SessionInterestManager) SplitWantedUnwanted(blks []blocks.Block) ([]b
 	return wantedBlks, notWantedBlks
 }
 
+// When the WantManager receives a message is calls InterestedSessions() to
+// find out which sessions are interested in the message.
 func (sim *SessionInterestManager) InterestedSessions(blks []cid.Cid, haves []cid.Cid, dontHaves []cid.Cid) []uint64 {
+	sim.lk.RLock()
+	defer sim.lk.RUnlock()
+
 	ks := make([]cid.Cid, 0, len(blks)+len(haves)+len(dontHaves))
 	ks = append(ks, blks...)
 	ks = append(ks, haves...)