Repository: ipfs/go-bitswap
Pull request #315 (2020-03-24 05:32:42)

chore: make pwm internals private
This makes it easier to tell where module boundaries are.


Patch:
diff --git a/internal/peermanager/peermanager.go b/internal/peermanager/peermanager.go
index 726d4be..5af9887 100644
--- a/internal/peermanager/peermanager.go
+++ b/internal/peermanager/peermanager.go
@@ -94,11 +94,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 +126,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 +137,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 +151,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 +163,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 +175,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 +183,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 +191,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..b4b8748 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,7 +81,7 @@ 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)
 
@@ -124,7 +124,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 +158,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 +174,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 +190,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/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/decision/engine_test.go b/internal/decision/engine_test.go
index 7dac950..bdfa936 100644
--- a/internal/decision/engine_test.go
+++ b/internal/decision/engine_test.go
@@ -935,7 +935,7 @@ func TestSendDontHave(t *testing.T) {
 
 	// Nothing in blockstore, should get DONT_HAVE for entries that wanted it
 	var next envChan
-	next, env := getNextEnvelope(e, next, 5*time.Millisecond)
+	next, env := getNextEnvelope(e, next, 10*time.Millisecond)
 	if env == nil {
 		t.Fatal("expected envelope")
 	}
@@ -965,7 +965,7 @@ func TestSendDontHave(t *testing.T) {
 	e.ReceiveFrom(otherPeer, blks, []cid.Cid{})
 
 	// Envelope should contain 2 HAVEs / 2 blocks
-	_, env = getNextEnvelope(e, next, 5*time.Millisecond)
+	_, env = getNextEnvelope(e, next, 10*time.Millisecond)
 	if env == nil {
 		t.Fatal("expected envelope")
 	}
diff --git a/internal/messagequeue/donthavetimeoutmgr_test.go b/internal/messagequeue/donthavetimeoutmgr_test.go
index 5c0de88..03ceb48 100644
--- a/internal/messagequeue/donthavetimeoutmgr_test.go
+++ b/internal/messagequeue/donthavetimeoutmgr_test.go
@@ -71,7 +71,7 @@ func (tr *timeoutRecorder) clear() {
 func TestDontHaveTimeoutMgrTimeout(t *testing.T) {
 	firstks := testutil.GenerateCids(2)
 	secondks := append(firstks, testutil.GenerateCids(3)...)
-	latency := time.Millisecond * 10
+	latency := time.Millisecond * 20
 	latMultiplier := 2
 	expProcessTime := 5 * time.Millisecond
 	expectedTimeout := expProcessTime + latency*time.Duration(latMultiplier)
@@ -87,7 +87,7 @@ func TestDontHaveTimeoutMgrTimeout(t *testing.T) {
 	dhtm.AddPending(firstks)
 
 	// Wait for less than the expected timeout
-	time.Sleep(expectedTimeout - 5*time.Millisecond)
+	time.Sleep(expectedTimeout - 10*time.Millisecond)
 
 	// At this stage no keys should have timed out
 	if tr.timedOutCount() > 0 {
@@ -98,7 +98,7 @@ func TestDontHaveTimeoutMgrTimeout(t *testing.T) {
 	dhtm.AddPending(secondks)
 
 	// Wait until after the expected timeout
-	time.Sleep(10 * time.Millisecond)
+	time.Sleep(20 * time.Millisecond)
 
 	// At this stage first set of keys should have timed out
 	if tr.timedOutCount() != len(firstks) {
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 1a35c0e..b679e9c 100644
--- a/internal/session/sessionwantsender_test.go
+++ b/internal/session/sessionwantsender_test.go
@@ -82,7 +82,7 @@ func (pm *mockPeerManager) SendWants(ctx context.Context, p peer.ID, wantBlocks
 }
 
 func (pm *mockPeerManager) waitNextWants() map[peer.ID]*sentWants {
-	time.Sleep(5 * time.Millisecond)
+	time.Sleep(10 * time.Millisecond)
 
 	pm.lk.Lock()
 	defer pm.lk.Unlock()
@@ -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()