p2p/discutil: fix hang in NextNode when FairMix is closed

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 {

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]