p2p/discutil: new Iterator interface and improve mixer

p2p/discutil/iter.go

@@ -19,25 +19,33 @@ package discutil
 
 import (
 	"context"
+	"sync"
+	"time"
 
 	"github.com/ethereum/go-ethereum/p2p/enode"
 )
 
-// Iterator represents an infinite 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. Implementations are not required to be safe for concurrent use.
+// 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).
+//
+// Implementations are not required to be safe for concurrent use. It is therefore unsafe
+// to call NextNode from multiple goroutines at the same time.
 type Iterator interface {
-	NextNode(ctx context.Context) *enode.Node
+	NextNode(ctx context.Context) (n *enode.Node, isLive bool)
 }
 
-// ReadNodes reads at most n nodes from the given iterator. The returned slice contains no
+// ReadNodes reads at most n nodes from the given iterator. The return value contains no
 // duplicates and no nil values. To prevent looping indefinitely for small repeating node
-// sequences, e.g. when reading from a CycleNodes iterator with a slice length < n, this
-// function calls NextNode at most n times.
+// sequences, this function calls NextNode at most n times.
 func ReadNodes(ctx context.Context, it Iterator, n int) []*enode.Node {
 	seen := make(map[enode.ID]*enode.Node, n)
 	for i := 0; i < n && ctx.Err() == nil; i++ {
-		node := it.NextNode(ctx)
+		node, isLive := it.NextNode(ctx)
+		if !isLive {
+			break
+		}
 		if node == nil {
 			continue
 		}
@@ -54,71 +62,185 @@ func ReadNodes(ctx context.Context, it Iterator, n int) []*enode.Node {
 	return result
 }
 
-// CycleNodes returns a never-ending interator that cycles through the given slice.
-func CycleNodes(nodes []*enode.Node) Iterator {
-	if len(nodes) == 0 {
-		return IterFunc(nullIterator)
-	}
-	index := 0
-	return IterFunc(func(context.Context) *enode.Node {
-		n := nodes[index]
-		index = (index + 1) % len(nodes)
-		return n
-	})
+// Filter wraps an iterator such that NextNode only returns nodes for which
+// the 'check' function returns true.
+func Filter(it Iterator, check func(*enode.Node) bool) Iterator {
+	return &filterIter{it, check}
 }
 
-func nullIterator(context.Context) *enode.Node {
-	return nil
+type filterIter struct {
+	it    Iterator
+	check func(*enode.Node) bool
 }
 
-// IterChan returns a NodeIterator wrapping the given channel.
-func IterChan(ch <-chan *enode.Node) Iterator {
-	return IterFunc(func(ctx context.Context) *enode.Node {
-		select {
-		case n := <-ch:
-			return n
-		case <-ctx.Done():
-			return nil
+func (f *filterIter) NextNode(ctx context.Context) (*enode.Node, bool) {
+	n, isLive := f.it.NextNode(ctx)
+	if n != nil && !f.check(n) {
+		n = nil
 	}
-	})
+	return n, isLive
 }
 
-// IterFunc is a function that satisfies the NodeIterator interface.
-type IterFunc func(ctx context.Context) *enode.Node
+// FairMix aggregates multiple node iterators. The mixer itself is an iterator which ends
+// only when Close is called. Source iterators added via AddSource are removed from the mix
+// when they end.
+//
+// The distribution of nodes returned by NextNode is approximately fair, i.e. FairMix
+// attempts to draw from all sources equally often. However, if a certain source is slow
+// and doesn't return a node within the configured timeout, a node from any other source
+// will be returned.
+//
+// It's safe to call AddSource and Close concurrently with NextNode.
+type FairMix struct {
+	ctx       context.Context
+	cancelCtx func()
+	wg        sync.WaitGroup
+	fromAny   chan *enode.Node
+	timeout   time.Duration
 
-// NextNode calls the function.
-func (fn IterFunc) NextNode(ctx context.Context) *enode.Node {
-	return fn(ctx)
-}
-
-// Mixer aggregates multiple node iterators. The distribution of nodes drawn from the mixer
-// is fair, i.e. all iterators are drawn from equally often.
-type Mixer struct {
-	sources []Iterator
+	mu      sync.Mutex
+	sources []*mixSource
 	last    int
 }
 
-// NewMixer creates a Mixer with the given initial sources.
-func NewMixer(sources ...Iterator) *Mixer {
-	return &Mixer{sources: sources}
+type mixSource struct {
+	it   Iterator
+	next chan *enode.Node
+}
+
+// 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.
+func NewFairMix(timeout time.Duration) *FairMix {
+	ctx, cancel := context.WithCancel(context.Background())
+	m := &FairMix{
+		ctx:       ctx,
+		cancelCtx: cancel,
+		fromAny:   make(chan *enode.Node),
+		timeout:   timeout,
+	}
+	return m
 }
 
 // AddSource adds a source of nodes.
-func (m *Mixer) AddSource(source Iterator) {
+func (m *FairMix) AddSource(it Iterator) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if !m.isLive() {
+		return
+	}
+	m.wg.Add(1)
+	source := &mixSource{it, make(chan *enode.Node)}
 	m.sources = append(m.sources, source)
+	go m.runSource(source)
+}
+
+// Close shuts down the mixer. Calling this is required to release resources
+// associated with the mixer.
+func (m *FairMix) Close() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if !m.isLive() {
+		return
+	}
+	m.cancelCtx()
+	m.wg.Wait()
+	m.sources = nil
+	close(m.fromAny)
 }
 
 // NextNode returns a node from a random source.
-func (m *Mixer) NextNode(ctx context.Context) *enode.Node {
+func (m *FairMix) NextNode(ctx context.Context) (*enode.Node, bool) {
+	var timeout <-chan time.Time
+	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
+			}
+			// This source has ended. Remove it from the list and try again
+			// with another source.
+			m.deleteSource(source)
+		case <-timeout:
+			return m.nextFromAny(ctx)
+		case <-ctx.Done():
+			return nil, m.isLive()
+		}
+	}
+}
+
+// nextFromAny is used when there are no sources or when the 'fair' choice
+// doesn't turn up a node quickly enough.
+func (m *FairMix) nextFromAny(ctx context.Context) (*enode.Node, bool) {
+	select {
+	case n, ok := <-m.fromAny:
+		return n, ok
+	case <-ctx.Done():
+		return nil, m.isLive()
+	}
+}
+
+func (m *FairMix) isLive() bool {
+	return m.ctx.Err() == nil
+}
+
+// pickSource chooses the next source to read from, cycling through them in order.
+func (m *FairMix) pickSource() *mixSource {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	if len(m.sources) == 0 {
 		return nil
 	}
-	source := m.nextSource()
-	return source.NextNode(ctx)
+	m.last = (m.last + 1) % len(m.sources)
+	return m.sources[m.last]
 }
 
-func (m *Mixer) nextSource() Iterator {
-	s := m.sources[m.last]
-	m.last = (m.last + 1) % len(m.sources)
-	return s
+// deleteSource deletes a source.
+func (m *FairMix) deleteSource(s *mixSource) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	for i := range m.sources {
+		if m.sources[i] == s {
+			copy(m.sources[i:], m.sources[i+1:])
+			m.sources[len(m.sources)-1] = nil
+			m.sources = m.sources[:len(m.sources)-1]
+			break
+		}
+	}
+}
+
+// runSource runs a single source in a loop.
+func (m *FairMix) runSource(s *mixSource) {
+	defer m.wg.Done()
+	for {
+		n, isLive := s.it.NextNode(m.ctx)
+		if !isLive {
+			close(s.next)
+			return
+		}
+		select {
+		case s.next <- n:
+		case m.fromAny <- n:
+		case <-m.ctx.Done():
+			return
+		}
+	}
 }

p2p/discutil/iter_test.go

@@ -20,39 +20,40 @@ import (
 	"context"
 	"encoding/binary"
 	"testing"
+	"time"
 
 	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/enr"
 )
 
 func TestReadNodes(t *testing.T) {
-	iter := new(genSource)
+	iter := new(genIter)
 	nodes := ReadNodes(context.Background(), iter, 10)
 	checkNodes(t, nodes, 10)
 }
 
 // This test verifies that ReadNodes checks for context cancelation.
 func TestReadNodesCancel(t *testing.T) {
-	iter := &blockedIter{new(genSource), nil}
+	iter := &blockedIter{new(genIter), nil}
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 	nodes := ReadNodes(ctx, iter, 10)
 	checkNodes(t, nodes, 0)
 }
 
-// This test checks that ReadNodes terminates when reading N nodes from an iterator
-// which returns less than N nodes in an endless cycle.
+// // This test checks that ReadNodes terminates when reading N nodes from an iterator
+// // which returns less than N nodes in an endless cycle.
 func TestReadNodesCycle(t *testing.T) {
 	iter := &callCountIter{
-		child: CycleNodes([]*enode.Node{
+		child: cycleNodes{
 			testNode(0, 0),
 			testNode(1, 0),
 			testNode(2, 0),
-		}),
+		},
 	}
 	nodes := ReadNodes(context.Background(), iter, 10)
 	checkNodes(t, nodes, 3)
-	if iter.count != 100 {
+	if iter.count != 10 {
 		t.Fatalf("%d calls to NextNode, want %d", iter.count, 100)
 	}
 }
@@ -76,58 +77,123 @@ func checkNodes(t *testing.T, nodes []*enode.Node, wantLen int) {
 	}
 }
 
-type callCountIter struct {
-	child Iterator
-	count 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) {
+	for i := 0; i < 500; i++ {
+		testMixerFairness(t)
+	}
 }
 
-func (it *callCountIter) NextNode(ctx context.Context) *enode.Node {
-	it.count++
-	return it.child.NextNode(ctx)
+func testMixerFairness(t *testing.T) {
+	mix := NewFairMix(1 * time.Second)
+	mix.AddSource(&genIter{index: 1})
+	mix.AddSource(&genIter{index: 2})
+	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)
 	}
 
-// This test ensures Mixer doesn't crash for NextNode with no sources.
-func TestMixerEmpty(t *testing.T) {
-	mix := NewMixer()
-	_ = mix.NextNode(context.Background())
-}
-
-// This test checks fairness of Mixer for the simple case of three non-overlapping sources
-// which return nodes immediately.
-func TestMixerFairSimple(t *testing.T) {
-	sources := []Iterator{&genSource{index: 1}, &genSource{index: 2}, &genSource{index: 3}}
-	mix := NewMixer(sources...)
-	nodes := ReadNodes(context.Background(), mix, 198)
-	if len(nodes) != 198 {
-		t.Fatal("wrong count from ReadNodes:", len(nodes), "want:", 198)
-	}
-
-	// Compute distribution.
-	d := make(map[uint32]int)
-	for i, node := range nodes {
-		if node == nil {
-			t.Fatalf("node %d is nil", i)
-		}
-		id := node.ID()
-		d[binary.BigEndian.Uint32(id[:4])]++
-	}
-	// Verify that the nodes slice contains an equal number of nodes from each source.
+	// Verify that the nodes slice contains an approximately equal number of nodes
+	// from each source.
+	d := idPrefixDistribution(nodes)
 	for _, count := range d {
-		if count != len(nodes)/len(sources) {
+		if approxEqual(count, len(nodes)/3, 30) {
 			t.Fatalf("ID distribution is unfair: %v", d)
 		}
 	}
 }
 
-// genSource creates fake nodes with numbered IDs based on 'index' and 'gen'
-type genSource struct {
+// This test checks that FairMix falls back to an alternative source when
+// the 'fair' choice doesn't return a node within the context's deadline.
+func TestFairMixNextFromAll(t *testing.T) {
+	mix := NewFairMix(1 * time.Millisecond)
+	mix.AddSource(&genIter{index: 1})
+	mix.AddSource(&blockedIter{child: &genIter{index: 2}})
+	defer mix.Close()
+
+	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)
+	}
+	d := idPrefixDistribution(nodes)
+	if len(d) > 1 || d[1] != len(nodes) {
+		t.Fatalf("wrong ID distribution: %v", d)
+	}
+}
+
+// This test ensures FairMix works for NextNode with no sources.
+func TestFairMixEmpty(t *testing.T) {
+	var (
+		mix   = NewFairMix(1 * time.Second)
+		testN = testNode(1, 1)
+		ch    = make(chan *enode.Node)
+	)
+	defer mix.Close()
+
+	go func() {
+		n, _ := mix.NextNode(context.Background())
+		ch <- n
+	}()
+
+	mix.AddSource(cycleNodes{testN})
+	if n := <-ch; n != testN {
+		t.Errorf("got wrong node: %v", n)
+	}
+}
+
+// This test checks closing a source while NextNode runs.
+func TestFairMixRemoveSource(t *testing.T) {
+	mix := NewFairMix(1 * time.Second)
+	source := &blockedIter{child: &genIter{index: 1}, unblock: make(chan struct{})}
+	close(source.unblock) // first NextNode call will return (nil, false)
+	mix.AddSource(source)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100 * time.Millisecond)
+	defer cancel()
+	n, isLive := mix.NextNode(ctx)
+	if n != nil {
+		t.Fatal("NextNode returned a node but shouldn't")
+	}
+	if !isLive {
+		t.Fatal("NextNode returned isLive == false")
+	}
+	if len(mix.sources) != 0 {
+		t.Fatalf("have %d sources, want zero", len(mix.sources))
+	}
+}
+
+func idPrefixDistribution(nodes []*enode.Node) map[uint32]int {
+	d := make(map[uint32]int)
+	for _, node := range nodes {
+		id := node.ID()
+		d[binary.BigEndian.Uint32(id[:4])]++
+	}
+	return d
+}
+
+func approxEqual(x, y, ε int) bool {
+	if y > x {
+		x, y = y, x
+	}
+	return x-y > ε
+}
+
+// genIter creates fake nodes with numbered IDs based on 'index' and 'gen'
+type genIter struct {
 	index, gen uint32
 }
 
-func (s *genSource) NextNode(ctx context.Context) *enode.Node {
+func (s *genIter) NextNode(ctx context.Context) (*enode.Node, bool) {
 	n := testNode(uint64(s.index)<<32|uint64(s.gen), 0)
 	s.gen++
-	return n
+	return n, true
 }
 
 func testNode(id, seq uint64) *enode.Node {
@@ -144,11 +210,38 @@ type blockedIter struct {
 	unblock chan struct{}
 }
 
-func (s *blockedIter) NextNode(ctx context.Context) *enode.Node {
+func (s *blockedIter) NextNode(ctx context.Context) (*enode.Node, bool) {
 	select {
-	case <-s.unblock:
+	case _, ok := <-s.unblock:
+		if !ok {
+			return nil, false
+		}
 		return s.child.NextNode(ctx)
 	case <-ctx.Done():
-		return nil
+		return nil, true
 	}
 }
+
+// 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) {
+	if len(s) == 0 {
+		return nil, true
+	}
+	n := s[0]
+	copy(s[:], s[1:])
+	s[len(s)-1] = n
+	return n, true
+}
+
+// callCountIter counts calls to NextNode.
+type callCountIter struct {
+	child Iterator
+	count int
+}
+
+func (it *callCountIter) NextNode(ctx context.Context) (*enode.Node, bool) {
+	it.count++
+	return it.child.NextNode(ctx)
+}