diff --git a/p2p/discutil/iter.go b/p2p/discutil/iter.go index 8081a91a6e..cebd2236fd 100644 --- a/p2p/discutil/iter.go +++ b/p2p/discutil/iter.go @@ -18,7 +18,6 @@ package discutil import ( - "context" "sync" "time" @@ -27,33 +26,24 @@ import ( // 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 +// The Next method returns the next node in the sequence. The isLive return value reports // whether the iterator is still open. Once closed, iterators should keep returning (nil, false). // -// Implementations of NextNode are not required to be safe for concurrent use. It is -// therefore unsafe to call NextNode from multiple goroutines at the same time. -// -// Close may be called concurrently with NextNode, and interrupts NextNode. +// Close may be called concurrently with Next and Node, and interrupts Next if it is blocked. type Iterator interface { - NextNode(ctx context.Context) (n *enode.Node, isLive bool) - Close() + Next() bool // moves to next node + Node() *enode.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(ctx context.Context, it Iterator, n int) []*enode.Node { +func ReadNodes(it Iterator, n int) []*enode.Node { seen := make(map[enode.ID]*enode.Node, n) - for i := 0; i < n && ctx.Err() == nil; i++ { - node, isLive := it.NextNode(ctx) - if !isLive { - break - } - if node == nil { - continue - } + 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 @@ -74,20 +64,17 @@ func Filter(it Iterator, check func(*enode.Node) bool) Iterator { } type filterIter struct { - it Iterator + Iterator check func(*enode.Node) bool } -func (f *filterIter) NextNode(ctx context.Context) (*enode.Node, bool) { - n, isLive := f.it.NextNode(ctx) - if n != nil && !f.check(n) { - n = nil +func (f *filterIter) Next() bool { + for f.Iterator.Next() { + if f.check(f.Node()) { + return true + } } - return n, isLive -} - -func (f *filterIter) Close() { - f.it.Close() + return false } // FairMix aggregates multiple node iterators. The mixer itself is an iterator which ends @@ -101,13 +88,13 @@ func (f *filterIter) Close() { // // 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 + wg sync.WaitGroup + fromAny chan *enode.Node + timeout time.Duration + cur *enode.Node mu sync.Mutex + closed chan struct{} sources []*mixSource last int } @@ -124,12 +111,10 @@ type mixSource struct { // 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{ - ctx: ctx, - cancelCtx: cancel, - fromAny: make(chan *enode.Node), - timeout: timeout, + fromAny: make(chan *enode.Node), + closed: make(chan struct{}), + timeout: timeout, } return m } @@ -139,32 +124,38 @@ func (m *FairMix) AddSource(it Iterator) { m.mu.Lock() defer m.mu.Unlock() - if !m.isLive() { + if m.closed == nil { return } m.wg.Add(1) source := &mixSource{it, make(chan *enode.Node)} m.sources = append(m.sources, source) - go m.runSource(source) + go m.runSource(m.closed, source) } -// Close shuts down the mixer. Calling this is required to release resources -// associated with the mixer. +// 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.isLive() { + if m.closed == nil { return } - m.cancelCtx() + for _, s := range m.sources { + s.it.Close() + } + close(m.closed) m.wg.Wait() - m.sources = nil close(m.fromAny) + m.sources = nil + m.closed = nil } // NextNode returns a node from a random source. -func (m *FairMix) NextNode(ctx context.Context) (*enode.Node, bool) { +func (m *FairMix) Next() bool { + m.cur = nil + var timeout <-chan time.Time if m.timeout >= 0 { timer := time.NewTimer(m.timeout) @@ -174,37 +165,35 @@ func (m *FairMix) NextNode(ctx context.Context) (*enode.Node, bool) { for { source := m.pickSource() if source == nil { - return m.nextFromAny(ctx) + return m.nextFromAny() } select { case n, ok := <-source.next: - if !ok { - // This source has ended. - m.deleteSource(source) - continue + if ok { + m.cur = n + return true } - return n, m.isLive() + // This source has ended. + m.deleteSource(source) case <-timeout: - return m.nextFromAny(ctx) - case <-ctx.Done(): - return nil, m.isLive() + return m.nextFromAny() } } } +// Node returns the current node. +func (m *FairMix) Node() *enode.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(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) nextFromAny() bool { + n, ok := <-m.fromAny + if ok { + m.cur = n } -} - -func (m *FairMix) isLive() bool { - return m.ctx.Err() == nil + return ok } // pickSource chooses the next source to read from, cycling through them in order. @@ -235,18 +224,15 @@ func (m *FairMix) deleteSource(s *mixSource) { } // runSource reads a single source in a loop. -func (m *FairMix) runSource(s *mixSource) { +func (m *FairMix) runSource(closed chan struct{}, s *mixSource) { defer m.wg.Done() defer close(s.next) - for { - n, isLive := s.it.NextNode(m.ctx) - if !isLive { - return - } + for s.it.Next() { + n := s.it.Node() select { case s.next <- n: case m.fromAny <- n: - case <-m.ctx.Done(): + case <-closed: return } } diff --git a/p2p/discutil/iter_test.go b/p2p/discutil/iter_test.go index 42f6f213d4..4d65d67a22 100644 --- a/p2p/discutil/iter_test.go +++ b/p2p/discutil/iter_test.go @@ -17,9 +17,10 @@ package discutil import ( - "context" "encoding/binary" "runtime" + "sync" + "sync/atomic" "testing" "time" @@ -28,34 +29,24 @@ import ( ) func TestReadNodes(t *testing.T) { - iter := new(genIter) - nodes := ReadNodes(context.Background(), iter, 10) + nodes := ReadNodes(new(genIter), 10) checkNodes(t, nodes, 10) } -// This test verifies that ReadNodes checks for context cancelation. -func TestReadNodesCancel(t *testing.T) { - 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{ + Iterator: cycleNodes( testNode(0, 0), testNode(1, 0), testNode(2, 0), - }, + ), } - nodes := ReadNodes(context.Background(), iter, 10) + nodes := ReadNodes(iter, 10) checkNodes(t, nodes, 3) if iter.count != 10 { - t.Fatalf("%d calls to NextNode, want %d", iter.count, 100) + t.Fatalf("%d calls to Next, want %d", iter.count, 100) } } @@ -93,9 +84,7 @@ func testMixerFairness(t *testing.T) { mix.AddSource(&genIter{index: 3}) defer mix.Close() - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - nodes := ReadNodes(ctx, mix, 500) + nodes := ReadNodes(mix, 500) checkNodes(t, nodes, 500) // Verify that the nodes slice contains an approximately equal number of nodes @@ -109,16 +98,14 @@ func testMixerFairness(t *testing.T) { } // 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. +// 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(&blockedIter{child: &genIter{index: 2}}) + mix.AddSource(cycleNodes()) defer mix.Close() - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - nodes := ReadNodes(ctx, mix, 500) + nodes := ReadNodes(mix, 500) checkNodes(t, nodes, 500) d := idPrefixDistribution(nodes) @@ -127,7 +114,7 @@ func TestFairMixNextFromAll(t *testing.T) { } } -// This test ensures FairMix works for NextNode with no sources. +// This test ensures FairMix works for Next with no sources. func TestFairMixEmpty(t *testing.T) { var ( mix = NewFairMix(1 * time.Second) @@ -137,31 +124,25 @@ func TestFairMixEmpty(t *testing.T) { defer mix.Close() go func() { - n, _ := mix.NextNode(context.Background()) - ch <- n + mix.Next() + ch <- mix.Node() }() - mix.AddSource(cycleNodes{testN}) + 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. +// This test checks closing a source while Next 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) + source := cycleNodes() + source.Close() 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 mix.Next() { + t.Fatal("Next should've returned false") } if len(mix.sources) != 0 { t.Fatalf("have %d sources, want zero", len(mix.sources)) @@ -176,14 +157,14 @@ func TestFairMixClose(t *testing.T) { func testMixerClose(t *testing.T) { mix := NewFairMix(-1) - mix.AddSource(cycleNodes{}) - mix.AddSource(cycleNodes{}) + 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") + if mix.Next() { + t.Error("Next returned true") } }() // This call is supposed to make it more likely that NextNode is @@ -194,7 +175,7 @@ func testMixerClose(t *testing.T) { select { case <-done: case <-time.After(3 * time.Second): - t.Fatal("NextNode didn't unblock on Close") + t.Fatal("Next didn't unblock on Close") } mix.Close() // shouldn't crash @@ -218,16 +199,28 @@ func approxEqual(x, y, ε int) bool { // genIter creates fake nodes with numbered IDs based on 'index' and 'gen' type genIter struct { + node *enode.Node index, gen uint32 } -func (s *genIter) NextNode(ctx context.Context) (*enode.Node, bool) { - n := testNode(uint64(s.index)<<32|uint64(s.gen), 0) +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 n, true + return true } -func (s *genIter) Close() { panic("called") } +func (s *genIter) Node() *enode.Node { + return s.node +} + +func (s *genIter) Close() { + s.index = ^uint32(0) +} func testNode(id, seq uint64) *enode.Node { var nodeID enode.ID @@ -237,53 +230,46 @@ func testNode(id, seq uint64) *enode.Node { return enode.SignNull(r, nodeID) } -// blockedIter delays NextNodes until the unblock channel receives a value. -type blockedIter struct { - child Iterator - unblock chan struct{} +// cycleNodes is an interator that cycles through the given slice. +func cycleNodes(nodes ...*enode.Node) Iterator { + return &cycleIter{nodes: nodes} } -func (s *blockedIter) NextNode(ctx context.Context) (*enode.Node, bool) { - select { - case _, ok := <-s.unblock: - if !ok { - return nil, false - } - return s.child.NextNode(ctx) - case <-ctx.Done(): - return nil, true +type cycleIter struct { + cur *enode.Node + mu sync.Mutex + index int + nodes []*enode.Node +} + +func (s *cycleIter) Next() bool { + s.mu.Lock() + defer s.mu.Unlock() + if len(s.nodes) == 0 { + return false } + s.cur = s.nodes[s.index] + s.index = (s.index + 1) % len(s.nodes) + return true } -func (s *blockedIter) Close() { panic("called") } - -// cycleNodes is a never-ending interator that cycles through the given slice. -type cycleNodes []*enode.Node - -func (s cycleNodes) NextNode(ctx context.Context) (*enode.Node, bool) { - if len(s) == 0 { - <-ctx.Done() - return nil, true - } - n := s[0] - copy(s[:], s[1:]) - s[len(s)-1] = n - return n, true +func (s *cycleIter) Node() *enode.Node { + return s.nodes[s.index] } -func (s cycleNodes) Close() { panic("called") } +func (s *cycleIter) Close() { + s.mu.Lock() + s.nodes = nil + s.mu.Unlock() +} // callCountIter counts calls to NextNode. type callCountIter struct { - child Iterator + Iterator count int } -func (it *callCountIter) NextNode(ctx context.Context) (*enode.Node, bool) { +func (it *callCountIter) Next() bool { it.count++ - return it.child.NextNode(ctx) -} - -func (it *callCountIter) Close() { - it.child.Close() + return it.Iterator.Next() }