Repository: ipfs/go-bitswap
Pull request #272 (2020-02-28 00:03:40)

Ensure broadcast when remaining peer becomes unavailable
Fixes https://github.com/ipfs/go-bitswap/issues/259


Patch:
diff --git a/internal/messagequeue/messagequeue.go b/internal/messagequeue/messagequeue.go
index e60d52c..be07400 100644
--- a/internal/messagequeue/messagequeue.go
+++ b/internal/messagequeue/messagequeue.go
@@ -582,7 +582,7 @@ func (mq *MessageQueue) attemptSendAndRecovery(message bsmsg.BitSwapMessage) boo
 		return true
 	case <-time.After(mq.sendErrorBackoff):
 		// wait 100ms in case disconnect notifications are still propagating
-		log.Warning("SendMsg errored but neither 'done' nor context.Done() were set")
+		log.Warn("SendMsg errored but neither 'done' nor context.Done() were set")
 	}
 
 	err = mq.initializeSender()
diff --git a/internal/session/session.go b/internal/session/session.go
index b20db30..c41a65d 100644
--- a/internal/session/session.go
+++ b/internal/session/session.go
@@ -227,9 +227,18 @@ func (s *Session) onWantsSent(p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.C
 }
 
 func (s *Session) onPeersExhausted(ks []cid.Cid) {
+	// We don't want to block the sessionWantSender if the incoming channel
+	// is full. So if we can't immediately send on the incoming channel spin
+	// it off into a go-routine.
 	select {
 	case s.incoming <- op{op: opBroadcast, keys: ks}:
-	case <-s.ctx.Done():
+	default:
+		go func() {
+			select {
+			case s.incoming <- op{op: opBroadcast, keys: ks}:
+			case <-s.ctx.Done():
+			}
+		}()
 	}
 }
 
@@ -287,12 +296,12 @@ func (s *Session) run(ctx context.Context) {
 			case opCancel:
 				s.sw.CancelPending(oper.keys)
 			case opBroadcast:
-				s.handleIdleTick(ctx)
+				s.broadcastWantHaves(ctx, oper.keys)
 			default:
 				panic("unhandled operation")
 			}
 		case <-s.idleTick.C:
-			s.handleIdleTick(ctx)
+			s.broadcastWantHaves(ctx, nil)
 		case <-s.periodicSearchTimer.C:
 			s.handlePeriodicSearch(ctx)
 		case baseTickDelay := <-s.tickDelayReqs:
@@ -304,24 +313,35 @@ func (s *Session) run(ctx context.Context) {
 	}
 }
 
-func (s *Session) handleIdleTick(ctx context.Context) {
-	live := s.sw.PrepareBroadcast()
+// Called when the session hasn't received any blocks for some time, or when
+// all peers in the session have sent DONT_HAVE for a particular set of CIDs.
+// Send want-haves to all connected peers, and search for new peers with the CID.
+func (s *Session) broadcastWantHaves(ctx context.Context, wants []cid.Cid) {
+	// If this broadcast is because of an idle timeout (we haven't received
+	// any blocks for a while) then broadcast all pending wants
+	if wants == nil {
+		wants = s.sw.PrepareBroadcast()
+	}
+
 	// log.Warnf("\n\n\n\n\nSes%d: broadcast %d keys\n\n\n\n\n", s.id, len(live))
 	// log.Infof("Ses%d: broadcast %d keys\n", s.id, len(live))
-	log.Warnf("Ses%d: broadcast %d keys", s.id, len(live))
 
 	// Broadcast a want-have for the live wants to everyone we're connected to
-	s.sprm.RecordPeerRequests(nil, live)
-	s.wm.BroadcastWantHaves(ctx, s.id, live)
+	s.sprm.RecordPeerRequests(nil, wants)
+	s.wm.BroadcastWantHaves(ctx, s.id, wants)
 
 	// do not find providers on consecutive ticks
 	// -- just rely on periodic search widening
-	if len(live) > 0 && (s.consecutiveTicks == 0) {
-		s.sprm.FindMorePeers(ctx, live[0])
+	if len(wants) > 0 && (s.consecutiveTicks == 0) {
+		// Search for providers who have the first want in the list.
+		// Typically if the provider has the first block they will have
+		// the rest of the blocks also.
+		log.Warnf("Ses%d: FindMorePeers with want 0 of %d wants", s.id, len(wants))
+		s.sprm.FindMorePeers(ctx, wants[0])
 	}
 	s.resetIdleTick()
 
-	// If we have live wants
+	// If we have live wants record a consecutive tick
 	if s.sw.HasLiveWants() {
 		s.consecutiveTicks++
 	}
diff --git a/internal/session/sessionwantsender.go b/internal/session/sessionwantsender.go
index 4448f8d..38c6235 100644
--- a/internal/session/sessionwantsender.go
+++ b/internal/session/sessionwantsender.go
@@ -236,10 +236,14 @@ func (spm *sessionWantSender) onChange(changes []change) {
 	}
 
 	// Update peer availability
-	newlyAvailable := spm.processAvailability(availability)
+	newlyAvailable, newlyUnavailable := spm.processAvailability(availability)
 
 	// Update wants
-	spm.processUpdates(updates)
+	dontHaves := spm.processUpdates(updates)
+
+	// Check if there are any wants for which all peers have indicated they
+	// don't have the want
+	spm.checkForExhaustedWants(dontHaves, newlyUnavailable)
 
 	// If there are some connected peers, send any pending wants
 	if spm.peerAvlMgr.haveAvailablePeers() {
@@ -251,8 +255,12 @@ func (spm *sessionWantSender) onChange(changes []change) {
 
 // processAvailability updates the want queue with any changes in
 // peer availability
-func (spm *sessionWantSender) processAvailability(availability map[peer.ID]bool) []peer.ID {
+// It returns the peers that have become
+// - newly available
+// - newly unavailable
+func (spm *sessionWantSender) processAvailability(availability map[peer.ID]bool) (avail []peer.ID, unavail []peer.ID) {
 	var newlyAvailable []peer.ID
+	var newlyUnavailable []peer.ID
 	for p, isNowAvailable := range availability {
 		// Make sure this is a peer that the session is actually interested in
 		if wasAvailable, ok := spm.peerAvlMgr.isAvailable(p); ok {
@@ -264,6 +272,8 @@ func (spm *sessionWantSender) processAvailability(availability map[peer.ID]bool)
 				spm.updateWantsPeerAvailability(p, isNowAvailable)
 				if isNowAvailable {
 					newlyAvailable = append(newlyAvailable, p)
+				} else {
+					newlyUnavailable = append(newlyUnavailable, p)
 				}
 				// Reset the count of consecutive DONT_HAVEs received from the
 				// peer
@@ -272,7 +282,7 @@ func (spm *sessionWantSender) processAvailability(availability map[peer.ID]bool)
 		}
 	}
 
-	return newlyAvailable
+	return newlyAvailable, newlyUnavailable
 }
 
 // isAvailable indicates whether the peer is available and whether
@@ -299,8 +309,9 @@ func (spm *sessionWantSender) trackWant(c cid.Cid) {
 	}
 }
 
-// processUpdates processes incoming blocks and HAVE / DONT_HAVEs
-func (spm *sessionWantSender) processUpdates(updates []update) {
+// processUpdates processes incoming blocks and HAVE / DONT_HAVEs.
+// It returns all DONT_HAVEs.
+func (spm *sessionWantSender) processUpdates(updates []update) []cid.Cid {
 	prunePeers := make(map[peer.ID]struct{})
 	dontHaves := cid.NewSet()
 	for _, upd := range updates {
@@ -354,20 +365,62 @@ func (spm *sessionWantSender) processUpdates(updates []update) {
 		}
 	}
 
+	// If any peers have sent us too many consecutive DONT_HAVEs, remove them
+	// from the session
+	if len(prunePeers) > 0 {
+		go func() {
+			for p := range prunePeers {
+				spm.SignalAvailability(p, false)
+			}
+		}()
+	}
+
+	return dontHaves.Keys()
+}
+
+// checkForExhaustedWants checks if there are any wants for which all peers
+// have sent a DONT_HAVE. We call these "exhausted" wants.
+func (spm *sessionWantSender) checkForExhaustedWants(dontHaves []cid.Cid, newlyUnavailable []peer.ID) {
+	// If there are no new DONT_HAVEs, and no peers became unavailable, then
+	// we don't need to check for exhausted wants
+	if len(dontHaves) == 0 && len(newlyUnavailable) == 0 {
+		return
+	}
+
+	// We need to check each want for which we just received a DONT_HAVE
+	wants := dontHaves
+
+	// If a peer just became unavailable, then we need to check all wants
+	// (because it may be the last peer who hadn't sent a DONT_HAVE for a CID)
+	if len(newlyUnavailable) > 0 {
+		// Collect all pending wants
+		wants = make([]cid.Cid, len(spm.wants))
+		for c := range spm.wants {
+			wants = append(wants, c)
+		}
+
+		// If the last available peer in the session has become unavailable
+		// then we need to broadcast all pending wants
+		if len(spm.peerAvlMgr.availablePeers()) == 0 {
+			spm.processExhaustedWants(wants)
+			return
+		}
+	}
+
 	// If all available peers for a cid sent a DONT_HAVE, signal to the session
 	// that we've exhausted available peers
-	if dontHaves.Len() > 0 {
-		exhausted := spm.bpm.AllPeersDoNotHaveBlock(spm.peerAvlMgr.availablePeers(), dontHaves.Keys())
-		newlyExhausted := spm.newlyExhausted(exhausted)
-		if len(newlyExhausted) > 0 {
-			spm.onPeersExhausted(newlyExhausted)
-		}
+	if len(wants) > 0 {
+		exhausted := spm.bpm.AllPeersDoNotHaveBlock(spm.peerAvlMgr.availablePeers(), wants)
+		spm.processExhaustedWants(exhausted)
 	}
+}
 
-	// If any peers have sent us too many consecutive DONT_HAVEs, remove them
-	// from the session
-	for p := range prunePeers {
-		spm.SignalAvailability(p, false)
+// processExhaustedWants filters the list so that only those wants that haven't
+// already been marked as exhausted are passed to onPeersExhausted()
+func (spm *sessionWantSender) processExhaustedWants(exhausted []cid.Cid) {
+	newlyExhausted := spm.newlyExhausted(exhausted)
+	if len(newlyExhausted) > 0 {
+		spm.onPeersExhausted(newlyExhausted)
 	}
 }
 
diff --git a/workers.go b/workers.go
index fe24305..04dc275 100644
--- a/workers.go
+++ b/workers.go
@@ -147,7 +147,7 @@ func (bs *Bitswap) provideWorker(px process.Process) {
 		defer cancel()
 
 		if err := bs.network.Provide(ctx, k); err != nil {
-			log.Warning(err)
+			log.Warn(err)
 		}
 	}
 


Test patch:
diff --git a/bitswap_with_sessions_test.go b/bitswap_with_sessions_test.go
index 28d3a32..3b5b68e 100644
--- a/bitswap_with_sessions_test.go
+++ b/bitswap_with_sessions_test.go
@@ -12,6 +12,7 @@ import (
 	blocks "github.com/ipfs/go-block-format"
 	cid "github.com/ipfs/go-cid"
 	blocksutil "github.com/ipfs/go-ipfs-blocksutil"
+	delay "github.com/ipfs/go-ipfs-delay"
 	tu "github.com/libp2p/go-libp2p-testing/etc"
 )
 
@@ -216,7 +217,10 @@ func TestFetchAfterDisconnect(t *testing.T) {
 	defer cancel()
 
 	vnet := getVirtualNetwork()
-	ig := testinstance.NewTestInstanceGenerator(vnet, nil, []bitswap.Option{bitswap.ProviderSearchDelay(10 * time.Millisecond)})
+	ig := testinstance.NewTestInstanceGenerator(vnet, nil, []bitswap.Option{
+		bitswap.ProviderSearchDelay(10 * time.Millisecond),
+		bitswap.RebroadcastDelay(delay.Fixed(15 * time.Millisecond)),
+	})
 	defer ig.Close()
 	bgen := blocksutil.NewBlockGenerator()
 
@@ -264,6 +268,8 @@ func TestFetchAfterDisconnect(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	time.Sleep(20 * time.Millisecond)
+
 	// Provide remaining blocks
 	lastBlks := blks[5:]
 	for _, block := range lastBlks {
@@ -276,8 +282,11 @@ func TestFetchAfterDisconnect(t *testing.T) {
 
 	// Should get last 5 blocks
 	for i := 0; i < 5; i++ {
-		b := <-ch
-		got = append(got, b)
+		select {
+		case b := <-ch:
+			got = append(got, b)
+		case <-ctx.Done():
+		}
 	}
 
 	if err := assertBlockLists(got, blks); err != nil {
diff --git a/internal/session/session_test.go b/internal/session/session_test.go
index 21e196f..b3ae26b 100644
--- a/internal/session/session_test.go
+++ b/internal/session/session_test.go
@@ -229,6 +229,49 @@ func TestSessionFindMorePeers(t *testing.T) {
 	}
 }
 
+func TestSessionOnPeersExhausted(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+	fwm := newFakeWantManager()
+	fpm := newFakeSessionPeerManager()
+	sim := bssim.New()
+	bpm := bsbpm.New()
+	notif := notifications.New()
+	defer notif.Shutdown()
+	id := testutil.GenerateSessionID()
+	session := New(ctx, id, fwm, fpm, sim, newFakePeerManager(), bpm, notif, time.Second, delay.Fixed(time.Minute), "")
+	blockGenerator := blocksutil.NewBlockGenerator()
+	blks := blockGenerator.Blocks(broadcastLiveWantsLimit + 5)
+	var cids []cid.Cid
+	for _, block := range blks {
+		cids = append(cids, block.Cid())
+	}
+	_, err := session.GetBlocks(ctx, cids)
+
+	if err != nil {
+		t.Fatal("error getting blocks")
+	}
+
+	// Wait for initial want request
+	receivedWantReq := <-fwm.wantReqs
+
+	// Should have sent out broadcast request for wants
+	if len(receivedWantReq.cids) != broadcastLiveWantsLimit {
+		t.Fatal("did not enqueue correct initial number of wants")
+	}
+
+	// Signal that all peers have send DONT_HAVE for two of the wants
+	session.onPeersExhausted(cids[len(cids)-2:])
+
+	// Wait for want request
+	receivedWantReq = <-fwm.wantReqs
+
+	// Should have sent out broadcast request for wants
+	if len(receivedWantReq.cids) != 2 {
+		t.Fatal("did not enqueue correct initial number of wants")
+	}
+}
+
 func TestSessionFailingToGetFirstBlock(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 	defer cancel()
diff --git a/internal/session/sessionwantsender_test.go b/internal/session/sessionwantsender_test.go
index 75c224d..ecea497 100644
--- a/internal/session/sessionwantsender_test.go
+++ b/internal/session/sessionwantsender_test.go
@@ -347,6 +347,103 @@ func TestPeersExhausted(t *testing.T) {
 	}
 }
 
+// Tests that when
+// - all the peers except one have sent a DONT_HAVE for a CID
+// - the remaining peer becomes unavailable
+// onPeersExhausted should be sent for that CID
+func TestPeersExhaustedLastWaitingPeerUnavailable(t *testing.T) {
+	cids := testutil.GenerateCids(2)
+	peers := testutil.GeneratePeers(2)
+	peerA := peers[0]
+	peerB := peers[1]
+	sid := uint64(1)
+	pm := newMockPeerManager()
+	bpm := bsbpm.New()
+	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
+
+	var exhausted []cid.Cid
+	onPeersExhausted := func(ks []cid.Cid) {
+		exhausted = append(exhausted, ks...)
+	}
+	spm := newSessionWantSender(context.Background(), sid, pm, bpm, onSend, onPeersExhausted)
+
+	go spm.Run()
+
+	// add cid0, cid1
+	spm.Add(cids)
+
+	// peerA: HAVE cid0
+	bpm.ReceiveFrom(peerA, []cid.Cid{cids[0]}, []cid.Cid{})
+	// Note: this also registers peer A as being available
+	spm.Update(peerA, []cid.Cid{}, []cid.Cid{cids[0]}, []cid.Cid{}, true)
+	// peerB: HAVE cid0
+	bpm.ReceiveFrom(peerB, []cid.Cid{cids[0]}, []cid.Cid{})
+	// Note: this also registers peer B as being available
+	spm.Update(peerB, []cid.Cid{}, []cid.Cid{cids[0]}, []cid.Cid{}, true)
+
+	// peerA: DONT_HAVE cid1
+	bpm.ReceiveFrom(peerA, []cid.Cid{}, []cid.Cid{cids[1]})
+	spm.Update(peerA, []cid.Cid{}, []cid.Cid{}, []cid.Cid{cids[0]}, false)
+
+	time.Sleep(5 * time.Millisecond)
+
+	// peerB: becomes unavailable
+	spm.SignalAvailability(peerB, false)
+
+	time.Sleep(5 * time.Millisecond)
+
+	// All remaining peers (peer A) have sent us a DONT_HAVE for cid1,
+	// so expect that onPeersExhausted() will be called with cid1
+	if !testutil.MatchKeysIgnoreOrder(exhausted, []cid.Cid{cids[1]}) {
+		t.Fatal("Wrong keys")
+	}
+}
+
+// Tests that when all the peers are removed from the session
+// onPeersExhausted should be called with all outstanding CIDs
+func TestPeersExhaustedAllPeersUnavailable(t *testing.T) {
+	cids := testutil.GenerateCids(3)
+	peers := testutil.GeneratePeers(2)
+	peerA := peers[0]
+	peerB := peers[1]
+	sid := uint64(1)
+	pm := newMockPeerManager()
+	bpm := bsbpm.New()
+	onSend := func(peer.ID, []cid.Cid, []cid.Cid) {}
+
+	var exhausted []cid.Cid
+	onPeersExhausted := func(ks []cid.Cid) {
+		exhausted = append(exhausted, ks...)
+	}
+	spm := newSessionWantSender(context.Background(), sid, pm, bpm, onSend, onPeersExhausted)
+
+	go spm.Run()
+
+	// add cid0, cid1, cid2
+	spm.Add(cids)
+
+	// peerA: receive block for cid0 (and register peer A with sessionWantSender)
+	spm.Update(peerA, []cid.Cid{cids[0]}, []cid.Cid{}, []cid.Cid{}, true)
+	// peerB: HAVE cid1
+	bpm.ReceiveFrom(peerB, []cid.Cid{cids[0]}, []cid.Cid{})
+	// Note: this also registers peer B as being available
+	spm.Update(peerB, []cid.Cid{}, []cid.Cid{cids[0]}, []cid.Cid{}, true)
+
+	time.Sleep(5 * time.Millisecond)
+
+	// peerA and peerB: become unavailable
+	spm.SignalAvailability(peerA, false)
+	spm.SignalAvailability(peerB, false)
+
+	time.Sleep(5 * time.Millisecond)
+
+	// Expect that onPeersExhausted() will be called with all cids for blocks
+	// that have not been received
+	if !testutil.MatchKeysIgnoreOrder(exhausted, []cid.Cid{cids[1], cids[2]}) {
+		t.Fatal("Wrong keys")
+	}
+}
+
 func TestConsecutiveDontHaveLimit(t *testing.T) {
 	cids := testutil.GenerateCids(peerDontHaveLimit + 10)
 	p := testutil.GeneratePeers(1)[0]