p2p/discutil: new package for discovery iterator utils

p2p/discutil/iter.go

@@ -0,0 +1,124 @@
+// Copyright 2019 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+// Package discutil provides node discovery utilities.
+package discutil
+
+import (
+	"context"
+
+	"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.
+type Iterator interface {
+	NextNode(ctx context.Context) *enode.Node
+}
+
+// ReadNodes reads at most n nodes from the given iterator. The returned slice 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.
+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)
+		if node == nil {
+			continue
+		}
+		prevNode, ok := seen[node.ID()]
+		if ok && prevNode.Seq() > node.Seq() {
+			continue
+		}
+		seen[node.ID()] = node
+	}
+	result := make([]*enode.Node, 0, len(seen))
+	for _, node := range seen {
+		result = append(result, 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
+	})
+}
+
+func nullIterator(context.Context) *enode.Node {
+	return nil
+}
+
+// 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
+		}
+	})
+}
+
+// IterFunc is a function that satisfies the NodeIterator interface.
+type IterFunc func(ctx context.Context) *enode.Node
+
+// 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
+	last    int
+}
+
+// NewMixer creates a Mixer with the given initial sources.
+func NewMixer(sources ...Iterator) *Mixer {
+	return &Mixer{sources: sources}
+}
+
+// AddSource adds a source of nodes.
+func (m *Mixer) AddSource(source Iterator) {
+	m.sources = append(m.sources, source)
+}
+
+// NextNode returns a node from a random source.
+func (m *Mixer) NextNode(ctx context.Context) *enode.Node {
+	if len(m.sources) == 0 {
+		return nil
+	}
+	source := m.nextSource()
+	return source.NextNode(ctx)
+}
+
+func (m *Mixer) nextSource() Iterator {
+	s := m.sources[m.last]
+	m.last = (m.last + 1) % len(m.sources)
+	return s
+}

p2p/discutil/iter_test.go

@@ -0,0 +1,154 @@
+// Copyright 2019 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package discutil
+
+import (
+	"context"
+	"encoding/binary"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/enr"
+)
+
+func TestReadNodes(t *testing.T) {
+	iter := new(genSource)
+	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}
+	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.
+func TestReadNodesCycle(t *testing.T) {
+	iter := &callCountIter{
+		child: CycleNodes([]*enode.Node{
+			testNode(0, 0),
+			testNode(1, 0),
+			testNode(2, 0),
+		}),
+	}
+	nodes := ReadNodes(context.Background(), iter, 10)
+	checkNodes(t, nodes, 3)
+	if iter.count != 100 {
+		t.Fatalf("%d calls to NextNode, want %d", iter.count, 100)
+	}
+}
+
+func checkNodes(t *testing.T, nodes []*enode.Node, wantLen int) {
+	if len(nodes) != wantLen {
+		t.Errorf("slice has %d nodes, want %d", len(nodes), wantLen)
+		return
+	}
+	seen := make(map[enode.ID]bool)
+	for i, e := range nodes {
+		if e == nil {
+			t.Errorf("nil node at index %d", i)
+			return
+		}
+		if seen[e.ID()] {
+			t.Errorf("slice has duplicate node %v", e.ID())
+			return
+		}
+		seen[e.ID()] = true
+	}
+}
+
+type callCountIter struct {
+	child Iterator
+	count int
+}
+
+func (it *callCountIter) NextNode(ctx context.Context) *enode.Node {
+	it.count++
+	return it.child.NextNode(ctx)
+}
+
+// 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.
+	for _, count := range d {
+		if count != len(nodes)/len(sources) {
+			t.Fatalf("ID distribution is unfair: %v", d)
+		}
+	}
+}
+
+// genSource creates fake nodes with numbered IDs based on 'index' and 'gen'
+type genSource struct {
+	index, gen uint32
+}
+
+func (s *genSource) NextNode(ctx context.Context) *enode.Node {
+	n := testNode(uint64(s.index)<<32|uint64(s.gen), 0)
+	s.gen++
+	return n
+}
+
+func testNode(id, seq uint64) *enode.Node {
+	var nodeID enode.ID
+	binary.BigEndian.PutUint64(nodeID[:], id)
+	r := new(enr.Record)
+	r.SetSeq(seq)
+	return enode.SignNull(r, nodeID)
+}
+
+// blockedIter delays NextNodes until the unblock channel receives a value.
+type blockedIter struct {
+	child   Iterator
+	unblock chan struct{}
+}
+
+func (s *blockedIter) NextNode(ctx context.Context) *enode.Node {
+	select {
+	case <-s.unblock:
+		return s.child.NextNode(ctx)
+	case <-ctx.Done():
+		return nil
+	}
+}