diff --git a/p2p/discutil/iter.go b/p2p/discutil/iter.go index e879cacc50..e014cd8786 100644 --- a/p2p/discutil/iter.go +++ b/p2p/discutil/iter.go @@ -25,10 +25,11 @@ import ( "github.com/ethereum/go-ethereum/p2p/enode" ) -// Iterator represents a sequence of nodes. The NextNode method returns the next node in -// the sequence. It may return nil if no next node could be found before the context was -// canceled. The isLive return value reports whether the iterator is still open. Once -// closed, iterators keep returning (nil, false). +// Iterator represents a sequence of nodes. +// +// The NextNode method returns the next node in the sequence. It may return nil when no +// node could be found before the context was canceled. The isLive return value reports +// whether the iterator is still open. Once closed, iterators should keep returning (nil, false). // // Implementations are not required to be safe for concurrent use. It is therefore unsafe // to call NextNode from multiple goroutines at the same time. @@ -110,11 +111,10 @@ type mixSource struct { // NewFairMix creates a mixer. // -// The timeout specifies how long the mixer will wait for the 'fair' choice before giving -// up and taking a node from any other source. A good way to set the timeout is deciding -// how long you'd want to wait for a node on average. -// -// Timeout zero is special and makes the mixer completely fair. +// The timeout specifies how long the mixer will wait for the next fairly-chosen source +// before giving up and taking a node from any other source. A good way to set the timeout +// is deciding how long you'd want to wait for a node on average. Passing a negative +// timeout disables the mixer completely fair. func NewFairMix(timeout time.Duration) *FairMix { ctx, cancel := context.WithCancel(context.Background()) m := &FairMix{ @@ -158,25 +158,24 @@ func (m *FairMix) Close() { // NextNode returns a node from a random source. func (m *FairMix) NextNode(ctx context.Context) (*enode.Node, bool) { var timeout <-chan time.Time - if m.timeout > 0 { + if m.timeout >= 0 { timer := time.NewTimer(m.timeout) timeout = timer.C defer timer.Stop() } for { - // Select a source. source := m.pickSource() if source == nil { return m.nextFromAny(ctx) } select { case n, ok := <-source.next: - if ok { - return n, true + if !ok { + // This source has ended. + m.deleteSource(source) + continue } - // This source has ended. Remove it from the list and try again - // with another source. - m.deleteSource(source) + return n, m.isLive() case <-timeout: return m.nextFromAny(ctx) case <-ctx.Done(): @@ -227,13 +226,13 @@ func (m *FairMix) deleteSource(s *mixSource) { } } -// runSource runs a single source in a loop. +// runSource reads a single source in a loop. func (m *FairMix) runSource(s *mixSource) { defer m.wg.Done() + defer close(s.next) for { n, isLive := s.it.NextNode(m.ctx) if !isLive { - close(s.next) return } select { diff --git a/p2p/discutil/iter_test.go b/p2p/discutil/iter_test.go index b52c499210..bc3708af3b 100644 --- a/p2p/discutil/iter_test.go +++ b/p2p/discutil/iter_test.go @@ -19,6 +19,7 @@ package discutil import ( "context" "encoding/binary" + "runtime" "testing" "time" @@ -77,7 +78,6 @@ func checkNodes(t *testing.T, nodes []*enode.Node, wantLen int) { } } - // This test checks fairness of FairMix in the happy case where all sources return nodes // within the context's deadline. func TestFairMix(t *testing.T) { @@ -93,10 +93,10 @@ func testMixerFairness(t *testing.T) { mix.AddSource(&genIter{index: 3}) defer mix.Close() - nodes := ReadNodes(context.Background(), mix, 500) - if len(nodes) != 500 { - t.Fatal("wrong count from ReadNodes:", len(nodes), "want:", 500) - } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + nodes := ReadNodes(ctx, mix, 500) + checkNodes(t, nodes, 500) // Verify that the nodes slice contains an approximately equal number of nodes // from each source. @@ -119,9 +119,8 @@ func TestFairMixNextFromAll(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() nodes := ReadNodes(ctx, mix, 500) - if len(nodes) != 500 { - t.Fatal("wrong count from ReadNodes:", len(nodes), "want:", 500) - } + checkNodes(t, nodes, 500) + d := idPrefixDistribution(nodes) if len(d) > 1 || d[1] != len(nodes) { t.Fatalf("wrong ID distribution: %v", d) @@ -155,7 +154,7 @@ func TestFairMixRemoveSource(t *testing.T) { close(source.unblock) // first NextNode call will return (nil, false) mix.AddSource(source) - ctx, cancel := context.WithTimeout(context.Background(), 100 * time.Millisecond) + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() n, isLive := mix.NextNode(ctx) if n != nil { @@ -169,6 +168,38 @@ func TestFairMixRemoveSource(t *testing.T) { } } +func TestFairMixClose(t *testing.T) { + for i := 0; i < 20 && !t.Failed(); i++ { + testMixerClose(t) + } +} + +func testMixerClose(t *testing.T) { + mix := NewFairMix(-1) + mix.AddSource(cycleNodes{}) + mix.AddSource(cycleNodes{}) + + done := make(chan struct{}) + go func() { + defer close(done) + if _, isLive := mix.NextNode(context.Background()); isLive { + t.Error("NextNode returned isLive == true") + } + }() + // This call is supposed to make it more likely that NextNode is + // actually executing by the time we call Close. + runtime.Gosched() + + mix.Close() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("NextNode didn't unblock on Close") + } + + mix.Close() // shouldn't crash +} + func idPrefixDistribution(nodes []*enode.Node) map[uint32]int { d := make(map[uint32]int) for _, node := range nodes { @@ -225,8 +256,9 @@ func (s *blockedIter) NextNode(ctx context.Context) (*enode.Node, bool) { // cycleNodes is a never-ending interator that cycles through the given slice. type cycleNodes []*enode.Node -func (s cycleNodes) NextNode(context.Context) (*enode.Node, bool) { +func (s cycleNodes) NextNode(ctx context.Context) (*enode.Node, bool) { if len(s) == 0 { + <-ctx.Done() return nil, true } n := s[0]