forks/go-ethereum
1
0
Fork 1
mirror of https://github.com/ethereum/go-ethereum.git synced 2026-08-20 10:52:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

p2p/discutil: remove context from iterator interface

Browse source
This commit is contained in:
Felix Lange 2019-07-09 13:49:42 +02:00
parent abc7dd8644
commit 5ee78b76b6
2 changed files with 129 additions and 157 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

130
p2p/discutil/iter.go
View file

@ -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
}
}

156
p2p/discutil/iter_test.go
View file

@ -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()
}