diff --git a/internal/statesync/dispatcher.go b/internal/statesync/dispatcher.go index 48de5d7f2..55fe7be10 100644 --- a/internal/statesync/dispatcher.go +++ b/internal/statesync/dispatcher.go @@ -57,7 +57,7 @@ func (d *dispatcher) LightBlock(ctx context.Context, height int64) (*types.Light // fetch the next peer id in the list and request a light block from that // peer - peer := d.availablePeers.Pop() + peer := d.availablePeers.Pop(ctx) lb, err := d.lightBlock(ctx, height, peer) return lb, peer, err } @@ -272,7 +272,7 @@ func (l *peerlist) Len() int { return len(l.peers) } -func (l *peerlist) Pop() types.NodeID { +func (l *peerlist) Pop(ctx context.Context) types.NodeID { l.mtx.Lock() if len(l.peers) == 0 { // if we don't have any peers in the list we block until a peer is @@ -281,8 +281,13 @@ func (l *peerlist) Pop() types.NodeID { l.waiting = append(l.waiting, wait) // unlock whilst waiting so that the list can be appended to l.mtx.Unlock() - peer := <-wait - return peer + select { + case peer := <-wait: + return peer + + case <-ctx.Done(): + return "" + } } peer := l.peers[0] diff --git a/internal/statesync/dispatcher_test.go b/internal/statesync/dispatcher_test.go index 0e33ab8be..9bec94c88 100644 --- a/internal/statesync/dispatcher_test.go +++ b/internal/statesync/dispatcher_test.go @@ -95,7 +95,7 @@ func TestPeerListBasic(t *testing.T) { half := numPeers / 2 for i := 0; i < half; i++ { - assert.Equal(t, peerSet[i], peerList.Pop()) + assert.Equal(t, peerSet[i], peerList.Pop(ctx)) } assert.Equal(t, half, peerList.Len()) @@ -104,7 +104,7 @@ func TestPeerListBasic(t *testing.T) { peerList.Remove(peerSet[half]) half++ - assert.Equal(t, peerSet[half], peerList.Pop()) + assert.Equal(t, peerSet[half], peerList.Pop(ctx)) } @@ -117,7 +117,7 @@ func TestPeerListConcurrent(t *testing.T) { // peer list hasn't been populated each these go routines should block for i := 0; i < numPeers/2; i++ { go func() { - _ = peerList.Pop() + _ = peerList.Pop(ctx) wg.Done() }() } @@ -132,7 +132,7 @@ func TestPeerListConcurrent(t *testing.T) { // we request the second half of the peer set for i := 0; i < numPeers/2; i++ { go func() { - _ = peerList.Pop() + _ = peerList.Pop(ctx) wg.Done() }() }