refactor(sessions): minor cleanup

Encapsulate functions for readability, and move code for understanding

session/session.go

@@ -333,6 +333,7 @@ func (s *Session) cidIsWanted(c cid.Cid) bool {
 func (s *Session) receiveBlock(ctx context.Context, blk blocks.Block) {
 	c := blk.Cid()
 	if s.cidIsWanted(c) {
+		s.srs.RecordUniqueBlock()
 		tval, ok := s.liveWants[c]
 		if ok {
 			s.latTotal += time.Since(tval)
@@ -363,10 +364,6 @@ func (s *Session) updateReceiveCounters(ctx context.Context, blk blkRecv) {
 	ks := blk.blk.Cid()
 	if s.pastWants.Has(ks) {
 		s.srs.RecordDuplicateBlock()
-	} else {
-		if s.cidIsWanted(ks) {
-			s.srs.RecordUniqueBlock()
-		}
 	}
 }
 

session/session_test.go

@@ -174,7 +174,7 @@ func TestSessionFindMorePeers(t *testing.T) {
 	wantReqs := make(chan wantReq, 1)
 	cancelReqs := make(chan wantReq, 1)
 	fwm := &fakeWantManager{wantReqs, cancelReqs}
-	fpm := &fakePeerManager{findMorePeersRequested: make(chan struct{})}
+	fpm := &fakePeerManager{findMorePeersRequested: make(chan struct{}, 1)}
 	frs := &fakeRequestSplitter{}
 	id := testutil.GenerateSessionID()
 	session := New(ctx, id, fwm, fpm, frs)

sessionpeermanager/sessionpeermanager.go

@@ -134,6 +134,25 @@ func (spm *SessionPeerManager) insertOptimizedPeer(p peer.ID) {
 	spm.optimizedPeersArr = append([]peer.ID{p}, spm.optimizedPeersArr...)
 }
 
+func (spm *SessionPeerManager) removeOptimizedPeer(p peer.ID) {
+	for i := 0; i < len(spm.optimizedPeersArr); i++ {
+		if spm.optimizedPeersArr[i] == p {
+			spm.optimizedPeersArr = append(spm.optimizedPeersArr[:i], spm.optimizedPeersArr[i+1:]...)
+			return
+		}
+	}
+}
+
+func (spm *SessionPeerManager) removeUnoptimizedPeer(p peer.ID) {
+	for i := 0; i < len(spm.unoptimizedPeersArr); i++ {
+		if spm.unoptimizedPeersArr[i] == p {
+			spm.unoptimizedPeersArr[i] = spm.unoptimizedPeersArr[len(spm.unoptimizedPeersArr)-1]
+			spm.unoptimizedPeersArr = spm.unoptimizedPeersArr[:len(spm.unoptimizedPeersArr)-1]
+			return
+		}
+	}
+}
+
 type peerFoundMessage struct {
 	p peer.ID
 }
@@ -160,24 +179,10 @@ func (prm *peerResponseMessage) handle(spm *SessionPeerManager) {
 		spm.tagPeer(p)
 	} else {
 		if isOptimized {
-			if spm.optimizedPeersArr[0] == p {
-				return
-			}
-			for i := 0; i < len(spm.optimizedPeersArr); i++ {
-				if spm.optimizedPeersArr[i] == p {
-					spm.optimizedPeersArr = append(spm.optimizedPeersArr[:i], spm.optimizedPeersArr[i+1:]...)
-					break
-				}
-			}
+			spm.removeOptimizedPeer(p)
 		} else {
 			spm.activePeers[p] = true
-			for i := 0; i < len(spm.unoptimizedPeersArr); i++ {
-				if spm.unoptimizedPeersArr[i] == p {
-					spm.unoptimizedPeersArr[i] = spm.unoptimizedPeersArr[len(spm.unoptimizedPeersArr)-1]
-					spm.unoptimizedPeersArr = spm.unoptimizedPeersArr[:len(spm.unoptimizedPeersArr)-1]
-					break
-				}
-			}
+			spm.removeUnoptimizedPeer(p)
 		}
 	}
 	spm.insertOptimizedPeer(p)