From ed4de8a78abee1a9b8a6201712514cf7f3d2113f Mon Sep 17 00:00:00 2001 From: Felix Lange Date: Fri, 5 Jul 2019 20:18:31 +0200 Subject: [PATCH] p2p/discutil: new Iterator interface and improve mixer --- p2p/discutil/iter.go | 234 +++++++++++++++++++++++++++++--------- p2p/discutil/iter_test.go | 185 ++++++++++++++++++++++-------- 2 files changed, 317 insertions(+), 102 deletions(-) diff --git a/p2p/discutil/iter.go b/p2p/discutil/iter.go index c4e2c7523e..e879cacc50 100644 --- a/p2p/discutil/iter.go +++ b/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) +// 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} +} + +type filterIter struct { + it 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 } - index := 0 - return IterFunc(func(context.Context) *enode.Node { - n := nodes[index] - index = (index + 1) % len(nodes) - return n - }) + return n, isLive } -func nullIterator(context.Context) *enode.Node { - return nil -} +// 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 -// 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 + 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 + } + } } diff --git a/p2p/discutil/iter_test.go b/p2p/discutil/iter_test.go index a0ea96e5e3..b52c499210 100644 --- a/p2p/discutil/iter_test.go +++ b/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() -// 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) + nodes := ReadNodes(context.Background(), mix, 500) + if len(nodes) != 500 { + t.Fatal("wrong count from ReadNodes:", len(nodes), "want:", 500) } - // 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) +}