diff --git a/p2p/enode/iter.go b/p2p/enode/iter.go new file mode 100644 index 0000000000..2c76ed9c8f --- /dev/null +++ b/p2p/enode/iter.go @@ -0,0 +1,283 @@ +// 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 . + +package enode + +import ( + "sync" + "time" +) + +// Iterator represents a sequence of nodes. The Next method moves to the next node in the +// sequence. It returns false when the sequence has ended or the iterator is closed. Close +// may be called concurrently with Next and Node, and interrupts Next if it is blocked. +type Iterator interface { + Next() bool // moves to next node + Node() *Node // returns current node + Close() // ends the iterator +} + +// 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, this function calls NextNode at most n times. +func ReadNodes(it Iterator, n int) []*Node { + seen := make(map[ID]*Node, n) + for i := 0; i < n && it.Next(); i++ { + // Remove duplicates, keeping the node with higher seq. + node := it.Node() + prevNode, ok := seen[node.ID()] + if ok && prevNode.Seq() > node.Seq() { + continue + } + seen[node.ID()] = node + } + result := make([]*Node, 0, len(seen)) + for _, node := range seen { + result = append(result, node) + } + return result +} + +// IterNodes makes an iterator which runs through the given nodes once. +func IterNodes(nodes []*Node) Iterator { + return &sliceIter{nodes: nodes, index: -1} +} + +// CycleNodes makes an iterator which cycles through the given nodes indefinitely. +func CycleNodes(nodes []*Node) Iterator { + return &sliceIter{nodes: nodes, index: -1, cycle: true} +} + +type sliceIter struct { + mu sync.Mutex + nodes []*Node + index int + cycle bool +} + +func (it *sliceIter) Next() bool { + it.mu.Lock() + defer it.mu.Unlock() + + if len(it.nodes) == 0 { + return false + } + it.index++ + if it.index == len(it.nodes) { + if it.cycle { + it.index = 0 + } else { + it.nodes = nil + return false + } + } + return true +} + +func (it *sliceIter) Node() *Node { + if len(it.nodes) == 0 { + return nil + } + return it.nodes[it.index] +} + +func (it *sliceIter) Close() { + it.mu.Lock() + defer it.mu.Unlock() + + it.nodes = nil +} + +// Filter wraps an iterator such that NextNode only returns nodes for which +// the 'check' function returns true. +func Filter(it Iterator, check func(*Node) bool) Iterator { + return &filterIter{it, check} +} + +type filterIter struct { + Iterator + check func(*Node) bool +} + +func (f *filterIter) Next() bool { + for f.Iterator.Next() { + if f.check(f.Node()) { + return true + } + } + return false +} + +// 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 { + wg sync.WaitGroup + fromAny chan *Node + timeout time.Duration + cur *Node + + mu sync.Mutex + closed chan struct{} + sources []*mixSource + last int +} + +type mixSource struct { + it Iterator + next chan *Node +} + +// NewFairMix creates a mixer. +// +// 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 makes the mixer completely fair. +func NewFairMix(timeout time.Duration) *FairMix { + m := &FairMix{ + fromAny: make(chan *Node), + closed: make(chan struct{}), + timeout: timeout, + } + return m +} + +// AddSource adds a source of nodes. +func (m *FairMix) AddSource(it Iterator) { + m.mu.Lock() + defer m.mu.Unlock() + + if m.closed == nil { + return + } + m.wg.Add(1) + source := &mixSource{it, make(chan *Node)} + m.sources = append(m.sources, source) + go m.runSource(m.closed, source) +} + +// Close shuts down the mixer and all current sources. +// Calling this is required to release resources associated with the mixer. +func (m *FairMix) Close() { + m.mu.Lock() + defer m.mu.Unlock() + + if m.closed == nil { + return + } + for _, s := range m.sources { + s.it.Close() + } + close(m.closed) + m.wg.Wait() + close(m.fromAny) + m.sources = nil + m.closed = nil +} + +// NextNode returns a node from a random source. +func (m *FairMix) Next() bool { + m.cur = nil + + var timeout <-chan time.Time + if m.timeout >= 0 { + timer := time.NewTimer(m.timeout) + timeout = timer.C + defer timer.Stop() + } + for { + source := m.pickSource() + if source == nil { + return m.nextFromAny() + } + select { + case n, ok := <-source.next: + if ok { + m.cur = n + return true + } + // This source has ended. + m.deleteSource(source) + case <-timeout: + return m.nextFromAny() + } + } +} + +// Node returns the current node. +func (m *FairMix) Node() *Node { + return m.cur +} + +// 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() bool { + n, ok := <-m.fromAny + if ok { + m.cur = n + } + return ok +} + +// 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 + } + m.last = (m.last + 1) % len(m.sources) + return m.sources[m.last] +} + +// 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 reads a single source in a loop. +func (m *FairMix) runSource(closed chan struct{}, s *mixSource) { + defer m.wg.Done() + defer close(s.next) + for s.it.Next() { + n := s.it.Node() + select { + case s.next <- n: + case m.fromAny <- n: + case <-closed: + return + } + } +} diff --git a/p2p/enode/iter_test.go b/p2p/enode/iter_test.go new file mode 100644 index 0000000000..6009661f3c --- /dev/null +++ b/p2p/enode/iter_test.go @@ -0,0 +1,291 @@ +// 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 . + +package enode + +import ( + "encoding/binary" + "runtime" + "sync/atomic" + "testing" + "time" + + "github.com/ethereum/go-ethereum/p2p/enr" +) + +func TestReadNodes(t *testing.T) { + nodes := ReadNodes(new(genIter), 10) + checkNodes(t, nodes, 10) +} + +// 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{ + Iterator: CycleNodes([]*Node{ + testNode(0, 0), + testNode(1, 0), + testNode(2, 0), + }), + } + nodes := ReadNodes(iter, 10) + checkNodes(t, nodes, 3) + if iter.count != 10 { + t.Fatalf("%d calls to Next, want %d", iter.count, 100) + } +} + +func TestFilterNodes(t *testing.T) { + nodes := make([]*Node, 100) + for i := range nodes { + nodes[i] = testNode(uint64(i), uint64(i)) + } + + it := Filter(IterNodes(nodes), func(n *Node) bool { + return n.Seq() >= 50 + }) + for i := 50; i < len(nodes); i++ { + if !it.Next() { + t.Fatal("Next returned false") + } + if it.Node() != nodes[i] { + t.Fatalf("iterator returned wrong node %v\nwant %v", it.Node(), nodes[i]) + } + } + if it.Next() { + t.Fatal("Next returned true after underlying iterator has ended") + } +} + +func checkNodes(t *testing.T, nodes []*Node, wantLen int) { + if len(nodes) != wantLen { + t.Errorf("slice has %d nodes, want %d", len(nodes), wantLen) + return + } + seen := make(map[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 + } +} + +// 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 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(mix, 500) + checkNodes(t, nodes, 500) + + // Verify that the nodes slice contains an approximately equal number of nodes + // from each source. + d := idPrefixDistribution(nodes) + for _, count := range d { + if approxEqual(count, len(nodes)/3, 30) { + t.Fatalf("ID distribution is unfair: %v", d) + } + } +} + +// This test checks that FairMix falls back to an alternative source when +// the 'fair' choice doesn't return a node within the timeout. +func TestFairMixNextFromAll(t *testing.T) { + mix := NewFairMix(1 * time.Millisecond) + mix.AddSource(&genIter{index: 1}) + mix.AddSource(CycleNodes(nil)) + defer mix.Close() + + nodes := ReadNodes(mix, 500) + checkNodes(t, nodes, 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 Next with no sources. +func TestFairMixEmpty(t *testing.T) { + var ( + mix = NewFairMix(1 * time.Second) + testN = testNode(1, 1) + ch = make(chan *Node) + ) + defer mix.Close() + + go func() { + mix.Next() + ch <- mix.Node() + }() + + mix.AddSource(CycleNodes([]*Node{testN})) + if n := <-ch; n != testN { + t.Errorf("got wrong node: %v", n) + } +} + +// This test checks closing a source while Next runs. +func TestFairMixRemoveSource(t *testing.T) { + mix := NewFairMix(1 * time.Second) + source := make(blockingIter) + mix.AddSource(source) + + sig := make(chan *Node) + go func() { + <-sig + mix.Next() + sig <- mix.Node() + }() + + sig <- nil + runtime.Gosched() + source.Close() + + wantNode := testNode(0, 0) + mix.AddSource(CycleNodes([]*Node{wantNode})) + n := <-sig + + if len(mix.sources) != 1 { + t.Fatalf("have %d sources, want one", len(mix.sources)) + } + if n != wantNode { + t.Fatalf("mixer returned wrong node") + } +} + +type blockingIter chan struct{} + +func (it blockingIter) Next() bool { + <-it + return false +} + +func (it blockingIter) Node() *Node { + return nil +} + +func (it blockingIter) Close() { + close(it) +} + +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(nil)) + mix.AddSource(CycleNodes(nil)) + + done := make(chan struct{}) + go func() { + defer close(done) + if mix.Next() { + t.Error("Next returned 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("Next didn't unblock on Close") + } + + mix.Close() // shouldn't crash +} + +func idPrefixDistribution(nodes []*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 { + node *Node + index, gen uint32 +} + +func (s *genIter) Next() bool { + index := atomic.LoadUint32(&s.index) + if index == ^uint32(0) { + s.node = nil + return false + } + s.node = testNode(uint64(index)<<32|uint64(s.gen), 0) + s.gen++ + return true +} + +func (s *genIter) Node() *Node { + return s.node +} + +func (s *genIter) Close() { + s.index = ^uint32(0) +} + +func testNode(id, seq uint64) *Node { + var nodeID ID + binary.BigEndian.PutUint64(nodeID[:], id) + r := new(enr.Record) + r.SetSeq(seq) + return SignNull(r, nodeID) +} + +// callCountIter counts calls to NextNode. +type callCountIter struct { + Iterator + count int +} + +func (it *callCountIter) Next() bool { + it.count++ + return it.Iterator.Next() +}