Repository: ipfs/go-bitswap
Pull request #324 (2020-03-24 17:37:25)

fix: race in SessionInterestManager
Fixes https://github.com/ipfs/go-bitswap/issues/313

`SessionWantList` has its own internal locking, so it would be more efficient to use a custom data structure for `SessionInterestManager` rather than using two `SessionWantList`s. However I would prefer to go with this solution for now as it's a simpler change, and revisit if there's a performance issue.


Patch:
diff --git a/internal/peermanager/peermanager.go b/internal/peermanager/peermanager.go
index 726d4be..c2159b1 100644
--- a/internal/peermanager/peermanager.go
+++ b/internal/peermanager/peermanager.go
@@ -4,12 +4,15 @@ import (
 	"context"
 	"sync"
 
+	logging "github.com/ipfs/go-log"
 	"github.com/ipfs/go-metrics-interface"
 
 	cid "github.com/ipfs/go-cid"
 	peer "github.com/libp2p/go-libp2p-core/peer"
 )
 
+var log = logging.Logger("bs:peermgr")
+
 // PeerQueue provides a queue of messages to be sent for a single peer.
 type PeerQueue interface {
 	AddBroadcastWantHaves([]cid.Cid)
@@ -94,11 +97,11 @@ func (pm *PeerManager) Connected(p peer.ID, initialWantHaves []cid.Cid) {
 	// If this is the first connection to the peer
 	if pq.refcnt == 1 {
 		// Inform the peer want manager that there's a new peer
-		pm.pwm.AddPeer(p)
+		pm.pwm.addPeer(p)
 		// Record that the want-haves are being sent to the peer
-		pm.pwm.PrepareSendWants(p, nil, initialWantHaves)
+		_, wantHaves := pm.pwm.prepareSendWants(p, nil, initialWantHaves)
 		// Broadcast any live want-haves to the newly connected peers
-		pq.pq.AddBroadcastWantHaves(initialWantHaves)
+		pq.pq.AddBroadcastWantHaves(wantHaves)
 		// Inform the sessions that the peer has connected
 		pm.signalAvailability(p, true)
 	}
@@ -126,7 +129,7 @@ func (pm *PeerManager) Disconnected(p peer.ID) {
 	// Clean up the peer
 	delete(pm.peerQueues, p)
 	pq.pq.Shutdown()
-	pm.pwm.RemovePeer(p)
+	pm.pwm.removePeer(p)
 }
 
 // BroadcastWantHaves broadcasts want-haves to all peers (used by the session
@@ -137,7 +140,7 @@ func (pm *PeerManager) BroadcastWantHaves(ctx context.Context, wantHaves []cid.C
 	pm.pqLk.Lock()
 	defer pm.pqLk.Unlock()
 
-	for p, ks := range pm.pwm.PrepareBroadcastWantHaves(wantHaves) {
+	for p, ks := range pm.pwm.prepareBroadcastWantHaves(wantHaves) {
 		if pqi, ok := pm.peerQueues[p]; ok {
 			pqi.pq.AddBroadcastWantHaves(ks)
 		}
@@ -151,7 +154,7 @@ func (pm *PeerManager) SendWants(ctx context.Context, p peer.ID, wantBlocks []ci
 	defer pm.pqLk.Unlock()
 
 	if pqi, ok := pm.peerQueues[p]; ok {
-		wblks, whvs := pm.pwm.PrepareSendWants(p, wantBlocks, wantHaves)
+		wblks, whvs := pm.pwm.prepareSendWants(p, wantBlocks, wantHaves)
 		pqi.pq.AddWants(wblks, whvs)
 	}
 }
@@ -163,7 +166,7 @@ func (pm *PeerManager) SendCancels(ctx context.Context, cancelKs []cid.Cid) {
 	defer pm.pqLk.Unlock()
 
 	// Send a CANCEL to each peer that has been sent a want-block or want-have
-	for p, ks := range pm.pwm.PrepareSendCancels(cancelKs) {
+	for p, ks := range pm.pwm.prepareSendCancels(cancelKs) {
 		if pqi, ok := pm.peerQueues[p]; ok {
 			pqi.pq.AddCancels(ks)
 		}
@@ -175,7 +178,7 @@ func (pm *PeerManager) CurrentWants() []cid.Cid {
 	pm.pqLk.RLock()
 	defer pm.pqLk.RUnlock()
 
-	return pm.pwm.GetWants()
+	return pm.pwm.getWants()
 }
 
 // CurrentWantBlocks returns the list of pending want-blocks
@@ -183,7 +186,7 @@ func (pm *PeerManager) CurrentWantBlocks() []cid.Cid {
 	pm.pqLk.RLock()
 	defer pm.pqLk.RUnlock()
 
-	return pm.pwm.GetWantBlocks()
+	return pm.pwm.getWantBlocks()
 }
 
 // CurrentWantHaves returns the list of pending want-haves
@@ -191,7 +194,7 @@ func (pm *PeerManager) CurrentWantHaves() []cid.Cid {
 	pm.pqLk.RLock()
 	defer pm.pqLk.RUnlock()
 
-	return pm.pwm.GetWantHaves()
+	return pm.pwm.getWantHaves()
 }
 
 func (pm *PeerManager) getOrCreate(p peer.ID) *peerQueueInstance {
diff --git a/internal/peermanager/peerwantmanager.go b/internal/peermanager/peerwantmanager.go
index 27e37cc..b0c843a 100644
--- a/internal/peermanager/peerwantmanager.go
+++ b/internal/peermanager/peerwantmanager.go
@@ -39,7 +39,7 @@ func newPeerWantManager(wantBlockGauge Gauge) *peerWantManager {
 }
 
 // AddPeer adds a peer whose wants we need to keep track of
-func (pwm *peerWantManager) AddPeer(p peer.ID) {
+func (pwm *peerWantManager) addPeer(p peer.ID) {
 	if _, ok := pwm.peerWants[p]; !ok {
 		pwm.peerWants[p] = &peerWant{
 			wantBlocks: cid.NewSet(),
@@ -49,13 +49,13 @@ func (pwm *peerWantManager) AddPeer(p peer.ID) {
 }
 
 // RemovePeer removes a peer and its associated wants from tracking
-func (pwm *peerWantManager) RemovePeer(p peer.ID) {
+func (pwm *peerWantManager) removePeer(p peer.ID) {
 	delete(pwm.peerWants, p)
 }
 
 // PrepareBroadcastWantHaves filters the list of want-haves for each peer,
 // returning a map of peers to the want-haves they have not yet been sent.
-func (pwm *peerWantManager) PrepareBroadcastWantHaves(wantHaves []cid.Cid) map[peer.ID][]cid.Cid {
+func (pwm *peerWantManager) prepareBroadcastWantHaves(wantHaves []cid.Cid) map[peer.ID][]cid.Cid {
 	res := make(map[peer.ID][]cid.Cid)
 
 	// Iterate over all known peers
@@ -81,40 +81,49 @@ func (pwm *peerWantManager) PrepareBroadcastWantHaves(wantHaves []cid.Cid) map[p
 
 // PrepareSendWants filters the list of want-blocks and want-haves such that
 // it only contains wants that have not already been sent to the peer.
-func (pwm *peerWantManager) PrepareSendWants(p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) ([]cid.Cid, []cid.Cid) {
+func (pwm *peerWantManager) prepareSendWants(p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) ([]cid.Cid, []cid.Cid) {
 	resWantBlks := make([]cid.Cid, 0)
 	resWantHvs := make([]cid.Cid, 0)
 
 	// Get the existing want-blocks and want-haves for the peer
-	if pws, ok := pwm.peerWants[p]; ok {
-		// Iterate over the requested want-blocks
-		for _, c := range wantBlocks {
-			// If the want-block hasn't been sent to the peer
-			if !pws.wantBlocks.Has(c) {
-				// Record that the CID was sent as a want-block
-				pws.wantBlocks.Add(c)
+	pws, ok := pwm.peerWants[p]
+
+	if !ok {
+		// In practice this should never happen:
+		// - PeerManager calls addPeer() as soon as the peer connects
+		// - PeerManager calls removePeer() as soon as the peer disconnects
+		// - All calls to PeerWantManager are locked
+		log.Errorf("prepareSendWants() called with peer %s but peer not found in peerWantManager", string(p))
+		return resWantBlks, resWantHvs
+	}
 
-				// Add the CID to the results
-				resWantBlks = append(resWantBlks, c)
+	// Iterate over the requested want-blocks
+	for _, c := range wantBlocks {
+		// If the want-block hasn't been sent to the peer
+		if !pws.wantBlocks.Has(c) {
+			// Record that the CID was sent as a want-block
+			pws.wantBlocks.Add(c)
 
-				// Make sure the CID is no longer recorded as a want-have
-				pws.wantHaves.Remove(c)
+			// Add the CID to the results
+			resWantBlks = append(resWantBlks, c)
 
-				// Increment the count of want-blocks
-				pwm.wantBlockGauge.Inc()
-			}
+			// Make sure the CID is no longer recorded as a want-have
+			pws.wantHaves.Remove(c)
+
+			// Increment the count of want-blocks
+			pwm.wantBlockGauge.Inc()
 		}
+	}
 
-		// Iterate over the requested want-haves
-		for _, c := range wantHaves {
-			// If the CID has not been sent as a want-block or want-have
-			if !pws.wantBlocks.Has(c) && !pws.wantHaves.Has(c) {
-				// Record that the CID was sent as a want-have
-				pws.wantHaves.Add(c)
+	// Iterate over the requested want-haves
+	for _, c := range wantHaves {
+		// If the CID has not been sent as a want-block or want-have
+		if !pws.wantBlocks.Has(c) && !pws.wantHaves.Has(c) {
+			// Record that the CID was sent as a want-have
+			pws.wantHaves.Add(c)
 
-				// Add the CID to the results
-				resWantHvs = append(resWantHvs, c)
-			}
+			// Add the CID to the results
+			resWantHvs = append(resWantHvs, c)
 		}
 	}
 
@@ -124,7 +133,7 @@ func (pwm *peerWantManager) PrepareSendWants(p peer.ID, wantBlocks []cid.Cid, wa
 // PrepareSendCancels filters the list of cancels for each peer,
 // returning a map of peers which only contains cancels for wants that have
 // been sent to the peer.
-func (pwm *peerWantManager) PrepareSendCancels(cancelKs []cid.Cid) map[peer.ID][]cid.Cid {
+func (pwm *peerWantManager) prepareSendCancels(cancelKs []cid.Cid) map[peer.ID][]cid.Cid {
 	res := make(map[peer.ID][]cid.Cid)
 
 	// Iterate over all known peers
@@ -158,7 +167,7 @@ func (pwm *peerWantManager) PrepareSendCancels(cancelKs []cid.Cid) map[peer.ID][
 }
 
 // GetWantBlocks returns the set of all want-blocks sent to all peers
-func (pwm *peerWantManager) GetWantBlocks() []cid.Cid {
+func (pwm *peerWantManager) getWantBlocks() []cid.Cid {
 	res := cid.NewSet()
 
 	// Iterate over all known peers
@@ -174,7 +183,7 @@ func (pwm *peerWantManager) GetWantBlocks() []cid.Cid {
 }
 
 // GetWantHaves returns the set of all want-haves sent to all peers
-func (pwm *peerWantManager) GetWantHaves() []cid.Cid {
+func (pwm *peerWantManager) getWantHaves() []cid.Cid {
 	res := cid.NewSet()
 
 	// Iterate over all known peers
@@ -190,7 +199,7 @@ func (pwm *peerWantManager) GetWantHaves() []cid.Cid {
 }
 
 // GetWants returns the set of all wants (both want-blocks and want-haves).
-func (pwm *peerWantManager) GetWants() []cid.Cid {
+func (pwm *peerWantManager) getWants() []cid.Cid {
 	res := cid.NewSet()
 
 	// Iterate over all known peers
diff --git a/internal/session/session.go b/internal/session/session.go
index 412faba..34a7375 100644
--- a/internal/session/session.go
+++ b/internal/session/session.go
@@ -159,7 +159,7 @@ func New(ctx context.Context,
 		periodicSearchDelay: periodicSearchDelay,
 		self:                self,
 	}
-	s.sws = newSessionWantSender(ctx, id, pm, sprm, bpm, s.onWantsSent, s.onPeersExhausted)
+	s.sws = newSessionWantSender(id, pm, sprm, bpm, s.onWantsSent, s.onPeersExhausted)
 
 	go s.run(ctx)
 
@@ -387,6 +387,9 @@ func (s *Session) handleShutdown() {
 	s.idleTick.Stop()
 	// Shut down the session peer manager
 	s.sprm.Shutdown()
+	// Shut down the sessionWantSender (blocks until sessionWantSender stops
+	// sending)
+	s.sws.Shutdown()
 	// Remove the session from the want manager
 	s.wm.RemoveSession(s.ctx, s.id)
 }
diff --git a/internal/session/sessionwantsender.go b/internal/session/sessionwantsender.go
index 7af7b32..8ccba8f 100644
--- a/internal/session/sessionwantsender.go
+++ b/internal/session/sessionwantsender.go
@@ -71,8 +71,13 @@ type onPeersExhaustedFn func([]cid.Cid)
 // consults the peer response tracker (records which peers sent us blocks).
 //
 type sessionWantSender struct {
-	// When the context is cancelled, sessionWantSender shuts down
+	// The context is used when sending wants
 	ctx context.Context
+	// Called to shutdown the sessionWantSender
+	shutdown func()
+	// The sessionWantSender uses the closed channel to signal when it's
+	// finished shutting down
+	closed chan struct{}
 	// The session ID
 	sessionID uint64
 	// A channel that collects incoming changes (events)
@@ -97,11 +102,14 @@ type sessionWantSender struct {
 	onPeersExhausted onPeersExhaustedFn
 }
 
-func newSessionWantSender(ctx context.Context, sid uint64, pm PeerManager, spm SessionPeerManager,
+func newSessionWantSender(sid uint64, pm PeerManager, spm SessionPeerManager,
 	bpm *bsbpm.BlockPresenceManager, onSend onSendFn, onPeersExhausted onPeersExhaustedFn) sessionWantSender {
 
+	ctx, cancel := context.WithCancel(context.Background())
 	sws := sessionWantSender{
 		ctx:                      ctx,
+		shutdown:                 cancel,
+		closed:                   make(chan struct{}),
 		sessionID:                sid,
 		changes:                  make(chan change, changesBufferSize),
 		wants:                    make(map[cid.Cid]*wantInfo),
@@ -158,12 +166,25 @@ func (sws *sessionWantSender) Run() {
 		case ch := <-sws.changes:
 			sws.onChange([]change{ch})
 		case <-sws.ctx.Done():
-			sws.shutdown()
+			// Unregister the session with the PeerManager
+			sws.pm.UnregisterSession(sws.sessionID)
+
+			// Close the 'closed' channel to signal to Shutdown() that the run
+			// loop has exited
+			close(sws.closed)
 			return
 		}
 	}
 }
 
+// Shutdown the sessionWantSender
+func (sws *sessionWantSender) Shutdown() {
+	// Signal to the run loop to stop processing
+	sws.shutdown()
+	// Wait for run loop to complete
+	<-sws.closed
+}
+
 // addChange adds a new change to the queue
 func (sws *sessionWantSender) addChange(c change) {
 	select {
@@ -172,11 +193,6 @@ func (sws *sessionWantSender) addChange(c change) {
 	}
 }
 
-// shutdown unregisters the session with the PeerManager
-func (sws *sessionWantSender) shutdown() {
-	sws.pm.UnregisterSession(sws.sessionID)
-}
-
 // collectChanges collects all the changes that have occurred since the last
 // invocation of onChange
 func (sws *sessionWantSender) collectChanges(changes []change) []change {
diff --git a/internal/sessioninterestmanager/sessioninterestmanager.go b/internal/sessioninterestmanager/sessioninterestmanager.go
index e85a645..46888c9 100644
--- a/internal/sessioninterestmanager/sessioninterestmanager.go
+++ b/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,21 +24,39 @@ 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()
@@ -42,7 +64,12 @@ func (sim *SessionInterestManager) FilterSessionInterested(ses uint64, ksets ...
 	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...)
diff --git a/internal/wantmanager/wantmanager.go b/internal/wantmanager/wantmanager.go
index b34056b..908f9dc 100644
--- a/internal/wantmanager/wantmanager.go
+++ b/internal/wantmanager/wantmanager.go
@@ -89,7 +89,7 @@ func (wm *WantManager) BroadcastWantHaves(ctx context.Context, ses uint64, wantH
 
 // RemoveSession is called when the session is shut down
 func (wm *WantManager) RemoveSession(ctx context.Context, ses uint64) {
-	// Remove session's interest in the given blocks
+	// Remove session's interest in the given blocks.
 	cancelKs := wm.sim.RemoveSessionInterest(ses)
 
 	// Remove broadcast want-haves for session


Test patch:
diff --git a/internal/peermanager/peerwantmanager_test.go b/internal/peermanager/peerwantmanager_test.go
index 0172a68..9cfa941 100644
--- a/internal/peermanager/peerwantmanager_test.go
+++ b/internal/peermanager/peerwantmanager_test.go
@@ -22,10 +22,10 @@ func (g *gauge) Dec() {
 func TestEmpty(t *testing.T) {
 	pwm := newPeerWantManager(&gauge{})
 
-	if len(pwm.GetWantBlocks()) > 0 {
+	if len(pwm.getWantBlocks()) > 0 {
 		t.Fatal("Expected GetWantBlocks() to have length 0")
 	}
-	if len(pwm.GetWantHaves()) > 0 {
+	if len(pwm.getWantHaves()) > 0 {
 		t.Fatal("Expected GetWantHaves() to have length 0")
 	}
 }
@@ -38,11 +38,11 @@ func TestPrepareBroadcastWantHaves(t *testing.T) {
 	cids2 := testutil.GenerateCids(2)
 	cids3 := testutil.GenerateCids(2)
 
-	pwm.AddPeer(peers[0])
-	pwm.AddPeer(peers[1])
+	pwm.addPeer(peers[0])
+	pwm.addPeer(peers[1])
 
 	// Broadcast 2 cids to 2 peers
-	bcst := pwm.PrepareBroadcastWantHaves(cids)
+	bcst := pwm.prepareBroadcastWantHaves(cids)
 	if len(bcst) != 2 {
 		t.Fatal("Expected 2 peers")
 	}
@@ -53,13 +53,13 @@ func TestPrepareBroadcastWantHaves(t *testing.T) {
 	}
 
 	// Broadcasting same cids should have no effect
-	bcst2 := pwm.PrepareBroadcastWantHaves(cids)
+	bcst2 := pwm.prepareBroadcastWantHaves(cids)
 	if len(bcst2) != 0 {
 		t.Fatal("Expected 0 peers")
 	}
 
 	// Broadcast 2 other cids
-	bcst3 := pwm.PrepareBroadcastWantHaves(cids2)
+	bcst3 := pwm.prepareBroadcastWantHaves(cids2)
 	if len(bcst3) != 2 {
 		t.Fatal("Expected 2 peers")
 	}
@@ -70,7 +70,7 @@ func TestPrepareBroadcastWantHaves(t *testing.T) {
 	}
 
 	// Broadcast mix of old and new cids
-	bcst4 := pwm.PrepareBroadcastWantHaves(append(cids, cids3...))
+	bcst4 := pwm.prepareBroadcastWantHaves(append(cids, cids3...))
 	if len(bcst4) != 2 {
 		t.Fatal("Expected 2 peers")
 	}
@@ -84,9 +84,9 @@ func TestPrepareBroadcastWantHaves(t *testing.T) {
 	// Sending want-block for a cid should prevent broadcast to that peer
 	cids4 := testutil.GenerateCids(4)
 	wantBlocks := []cid.Cid{cids4[0], cids4[2]}
-	pwm.PrepareSendWants(peers[0], wantBlocks, []cid.Cid{})
+	pwm.prepareSendWants(peers[0], wantBlocks, []cid.Cid{})
 
-	bcst5 := pwm.PrepareBroadcastWantHaves(cids4)
+	bcst5 := pwm.prepareBroadcastWantHaves(cids4)
 	if len(bcst4) != 2 {
 		t.Fatal("Expected 2 peers")
 	}
@@ -105,8 +105,8 @@ func TestPrepareBroadcastWantHaves(t *testing.T) {
 	}
 
 	// Add another peer
-	pwm.AddPeer(peers[2])
-	bcst6 := pwm.PrepareBroadcastWantHaves(cids)
+	pwm.addPeer(peers[2])
+	bcst6 := pwm.prepareBroadcastWantHaves(cids)
 	if len(bcst6) != 1 {
 		t.Fatal("Expected 1 peer")
 	}
@@ -126,11 +126,11 @@ func TestPrepareSendWants(t *testing.T) {
 	cids := testutil.GenerateCids(2)
 	cids2 := testutil.GenerateCids(2)
 
-	pwm.AddPeer(p0)
-	pwm.AddPeer(p1)
+	pwm.addPeer(p0)
+	pwm.addPeer(p1)
 
 	// Send 2 want-blocks and 2 want-haves to p0
-	wb, wh := pwm.PrepareSendWants(p0, cids, cids2)
+	wb, wh := pwm.prepareSendWants(p0, cids, cids2)
 	if !testutil.MatchKeysIgnoreOrder(wb, cids) {
 		t.Fatal("Expected 2 want-blocks")
 	}
@@ -143,7 +143,7 @@ func TestPrepareSendWants(t *testing.T) {
 	// - 1 old want-have  and 2 new want-haves
 	cids3 := testutil.GenerateCids(2)
 	cids4 := testutil.GenerateCids(2)
-	wb2, wh2 := pwm.PrepareSendWants(p0, append(cids3, cids[0]), append(cids4, cids2[0]))
+	wb2, wh2 := pwm.prepareSendWants(p0, append(cids3, cids[0]), append(cids4, cids2[0]))
 	if !testutil.MatchKeysIgnoreOrder(wb2, cids3) {
 		t.Fatal("Expected 2 want-blocks")
 	}
@@ -154,7 +154,7 @@ func TestPrepareSendWants(t *testing.T) {
 	// Send to p0 as want-blocks: 1 new want-block, 1 old want-have
 	cids5 := testutil.GenerateCids(1)
 	newWantBlockOldWantHave := append(cids5, cids2[0])
-	wb3, wh3 := pwm.PrepareSendWants(p0, newWantBlockOldWantHave, []cid.Cid{})
+	wb3, wh3 := pwm.prepareSendWants(p0, newWantBlockOldWantHave, []cid.Cid{})
 	// If a want was sent as a want-have, it should be ok to now send it as a
 	// want-block
 	if !testutil.MatchKeysIgnoreOrder(wb3, newWantBlockOldWantHave) {
@@ -167,7 +167,7 @@ func TestPrepareSendWants(t *testing.T) {
 	// Send to p0 as want-haves: 1 new want-have, 1 old want-block
 	cids6 := testutil.GenerateCids(1)
 	newWantHaveOldWantBlock := append(cids6, cids[0])
-	wb4, wh4 := pwm.PrepareSendWants(p0, []cid.Cid{}, newWantHaveOldWantBlock)
+	wb4, wh4 := pwm.prepareSendWants(p0, []cid.Cid{}, newWantHaveOldWantBlock)
 	// If a want was previously sent as a want-block, it should not be
 	// possible to now send it as a want-have
 	if !testutil.MatchKeysIgnoreOrder(wh4, cids6) {
@@ -178,7 +178,7 @@ func TestPrepareSendWants(t *testing.T) {
 	}
 
 	// Send 2 want-blocks and 2 want-haves to p1
-	wb5, wh5 := pwm.PrepareSendWants(p1, cids, cids2)
+	wb5, wh5 := pwm.prepareSendWants(p1, cids, cids2)
 	if !testutil.MatchKeysIgnoreOrder(wb5, cids) {
 		t.Fatal("Expected 2 want-blocks")
 	}
@@ -200,24 +200,24 @@ func TestPrepareSendCancels(t *testing.T) {
 	allwb := append(wb1, wb2...)
 	allwh := append(wh1, wh2...)
 
-	pwm.AddPeer(p0)
-	pwm.AddPeer(p1)
+	pwm.addPeer(p0)
+	pwm.addPeer(p1)
 
 	// Send 2 want-blocks and 2 want-haves to p0
-	pwm.PrepareSendWants(p0, wb1, wh1)
+	pwm.prepareSendWants(p0, wb1, wh1)
 	// Send 3 want-blocks and 3 want-haves to p1
 	// (1 overlapping want-block / want-have with p0)
-	pwm.PrepareSendWants(p1, append(wb2, wb1[1]), append(wh2, wh1[1]))
+	pwm.prepareSendWants(p1, append(wb2, wb1[1]), append(wh2, wh1[1]))
 
-	if !testutil.MatchKeysIgnoreOrder(pwm.GetWantBlocks(), allwb) {
+	if !testutil.MatchKeysIgnoreOrder(pwm.getWantBlocks(), allwb) {
 		t.Fatal("Expected 4 cids to be wanted")
 	}
-	if !testutil.MatchKeysIgnoreOrder(pwm.GetWantHaves(), allwh) {
+	if !testutil.MatchKeysIgnoreOrder(pwm.getWantHaves(), allwh) {
 		t.Fatal("Expected 4 cids to be wanted")
 	}
 
 	// Cancel 1 want-block and 1 want-have that were sent to p0
-	res := pwm.PrepareSendCancels([]cid.Cid{wb1[0], wh1[0]})
+	res := pwm.prepareSendCancels([]cid.Cid{wb1[0], wh1[0]})
 	// Should cancel the want-block and want-have
 	if len(res) != 1 {
 		t.Fatal("Expected 1 peer")
@@ -225,16 +225,16 @@ func TestPrepareSendCancels(t *testing.T) {
 	if !testutil.MatchKeysIgnoreOrder(res[p0], []cid.Cid{wb1[0], wh1[0]}) {
 		t.Fatal("Expected 2 cids to be cancelled")
 	}
-	if !testutil.MatchKeysIgnoreOrder(pwm.GetWantBlocks(), append(wb2, wb1[1])) {
+	if !testutil.MatchKeysIgnoreOrder(pwm.getWantBlocks(), append(wb2, wb1[1])) {
 		t.Fatal("Expected 3 want-blocks")
 	}
-	if !testutil.MatchKeysIgnoreOrder(pwm.GetWantHaves(), append(wh2, wh1[1])) {
+	if !testutil.MatchKeysIgnoreOrder(pwm.getWantHaves(), append(wh2, wh1[1])) {
 		t.Fatal("Expected 3 want-haves")
 	}
 
 	// Cancel everything
 	allCids := append(allwb, allwh...)
-	res2 := pwm.PrepareSendCancels(allCids)
+	res2 := pwm.prepareSendCancels(allCids)
 	// Should cancel the remaining want-blocks and want-haves
 	if len(res2) != 2 {
 		t.Fatal("Expected 2 peers", len(res2))
@@ -247,10 +247,10 @@ func TestPrepareSendCancels(t *testing.T) {
 	if !testutil.MatchKeysIgnoreOrder(res2[p1], remainingP2) {
 		t.Fatal("Expected un-cancelled cids to be cancelled")
 	}
-	if len(pwm.GetWantBlocks()) != 0 {
+	if len(pwm.getWantBlocks()) != 0 {
 		t.Fatal("Expected 0 want-blocks")
 	}
-	if len(pwm.GetWantHaves()) != 0 {
+	if len(pwm.getWantHaves()) != 0 {
 		t.Fatal("Expected 0 want-haves")
 	}
 }
@@ -264,10 +264,10 @@ func TestStats(t *testing.T) {
 	cids := testutil.GenerateCids(2)
 	cids2 := testutil.GenerateCids(2)
 
-	pwm.AddPeer(p0)
+	pwm.addPeer(p0)
 
 	// Send 2 want-blocks and 2 want-haves to p0
-	pwm.PrepareSendWants(p0, cids, cids2)
+	pwm.prepareSendWants(p0, cids, cids2)
 
 	if g.count != 2 {
 		t.Fatal("Expected 2 want-blocks")
@@ -275,7 +275,7 @@ func TestStats(t *testing.T) {
 
 	// Send 1 old want-block and 2 new want-blocks to p0
 	cids3 := testutil.GenerateCids(2)
-	pwm.PrepareSendWants(p0, append(cids3, cids[0]), []cid.Cid{})
+	pwm.prepareSendWants(p0, append(cids3, cids[0]), []cid.Cid{})
 
 	if g.count != 4 {
 		t.Fatal("Expected 4 want-blocks")
@@ -284,7 +284,7 @@ func TestStats(t *testing.T) {
 	// Cancel 1 want-block that was sent to p0
 	// and 1 want-block that was not sent
 	cids4 := testutil.GenerateCids(1)
-	pwm.PrepareSendCancels(append(cids4, cids[0]))
+	pwm.prepareSendCancels(append(cids4, cids[0]))
 
 	if g.count != 3 {
 		t.Fatal("Expected 3 want-blocks", g.count)
diff --git a/internal/session/sessionwantsender_test.go b/internal/session/sessionwantsender_test.go
index a791c6c..b679e9c 100644
--- a/internal/session/sessionwantsender_test.go
+++ b/internal/session/sessionwantsender_test.go
@@ -138,7 +138,7 @@ func TestSendWants(t *testing.T) {
 	bpm := bsbpm.New()
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 	onPeersExhausted := func([]cid.Cid) {}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, onPeersExhausted)
 
 	go spm.Run()
 
@@ -176,7 +176,7 @@ func TestSendsWantBlockToOnePeerOnly(t *testing.T) {
 	bpm := bsbpm.New()
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 	onPeersExhausted := func([]cid.Cid) {}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, onPeersExhausted)
 
 	go spm.Run()
 
@@ -234,7 +234,7 @@ func TestReceiveBlock(t *testing.T) {
 	bpm := bsbpm.New()
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 	onPeersExhausted := func([]cid.Cid) {}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, onPeersExhausted)
 
 	go spm.Run()
 
@@ -294,7 +294,7 @@ func TestPeerUnavailable(t *testing.T) {
 	bpm := bsbpm.New()
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 	onPeersExhausted := func([]cid.Cid) {}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, onPeersExhausted)
 
 	go spm.Run()
 
@@ -360,7 +360,7 @@ func TestPeersExhausted(t *testing.T) {
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 
 	ep := exhaustedPeers{}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, ep.onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, ep.onPeersExhausted)
 
 	go spm.Run()
 
@@ -436,7 +436,7 @@ func TestPeersExhaustedLastWaitingPeerUnavailable(t *testing.T) {
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 
 	ep := exhaustedPeers{}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, ep.onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, ep.onPeersExhausted)
 
 	go spm.Run()
 
@@ -484,7 +484,7 @@ func TestPeersExhaustedAllPeersUnavailable(t *testing.T) {
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 
 	ep := exhaustedPeers{}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, ep.onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, ep.onPeersExhausted)
 
 	go spm.Run()
 
@@ -522,7 +522,7 @@ func TestConsecutiveDontHaveLimit(t *testing.T) {
 	bpm := bsbpm.New()
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 	onPeersExhausted := func([]cid.Cid) {}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, onPeersExhausted)
 
 	go spm.Run()
 
@@ -578,7 +578,7 @@ func TestConsecutiveDontHaveLimitInterrupted(t *testing.T) {
 	bpm := bsbpm.New()
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 	onPeersExhausted := func([]cid.Cid) {}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, onPeersExhausted)
 
 	go spm.Run()
 
@@ -633,7 +633,7 @@ func TestConsecutiveDontHaveReinstateAfterRemoval(t *testing.T) {
 	bpm := bsbpm.New()
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 	onPeersExhausted := func([]cid.Cid) {}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, onPeersExhausted)
 
 	go spm.Run()
 
@@ -717,7 +717,7 @@ func TestConsecutiveDontHaveDontRemoveIfHasWantedBlock(t *testing.T) {
 	bpm := bsbpm.New()
 	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
 	onPeersExhausted := func([]cid.Cid) {}
-	spm := newSessionWantSender(context.Background(), sid, pm, fpm, bpm, onSend, onPeersExhausted)
+	spm := newSessionWantSender(sid, pm, fpm, bpm, onSend, onPeersExhausted)
 
 	go spm.Run()