p2p/discutil: remove context from iterator interface

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
+// The Next method returns the next node in the sequence. The isLive return value reports
-// node could be found before the context was canceled. 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
+// Close may be called concurrently with Next and Node, and interrupts Next if it is blocked.
-// therefore unsafe to call NextNode from multiple goroutines at the same time.
-//
-// Close may be called concurrently with NextNode, and interrupts NextNode.
 type Iterator interface {
-	NextNode(ctx context.Context) (n *enode.Node, isLive bool)
+	Next() bool        // moves to next node
-	Close()
+	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++ {
+	for i := 0; i < n && it.Next(); i++ {
-		node, isLive := it.NextNode(ctx)
-		if !isLive {
-			break
-		}
-		if node == nil {
-			continue
-		}
 		// 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) {
+func (f *filterIter) Next() bool {
-	n, isLive := f.it.NextNode(ctx)
+	for f.Iterator.Next() {
-	if n != nil && !f.check(n) {
+		if f.check(f.Node()) {
-		n = nil
+			return true
+		}
 	}
-	return n, isLive
+	return false
-}
-
-func (f *filterIter) Close() {
-	f.it.Close()
 }
 
 // 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
+	wg      sync.WaitGroup
-	cancelCtx func()
+	fromAny chan *enode.Node
-	wg        sync.WaitGroup
+	timeout time.Duration
-	fromAny   chan *enode.Node
+	cur     *enode.Node
-	timeout   time.Duration
 
 	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,
+		fromAny: make(chan *enode.Node),
-		cancelCtx: cancel,
+		closed:  make(chan struct{}),
-		fromAny:   make(chan *enode.Node),
+		timeout: timeout,
-		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
+// Close shuts down the mixer and all current sources.
-// associated with 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() {
+	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 {
+			if ok {
-				// This source has ended.
+				m.cur = n
-				m.deleteSource(source)
+				return true
-				continue
 			}
-			return n, m.isLive()
+			// This source has ended.
+			m.deleteSource(source)
 		case <-timeout:
-			return m.nextFromAny(ctx)
+			return m.nextFromAny()
-		case <-ctx.Done():
-			return nil, m.isLive()
 		}
 	}
 }
 
+// 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) {
+func (m *FairMix) nextFromAny() bool {
-	select {
+	n, ok := <-m.fromAny
-	case n, ok := <-m.fromAny:
+	if ok {
-		return n, ok
+		m.cur = n
-	case <-ctx.Done():
-		return nil, m.isLive()
 	}
-}
+	return ok
-
-func (m *FairMix) isLive() bool {
-	return m.ctx.Err() == nil
 }
 
 // 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 {
+	for s.it.Next() {
-		n, isLive := s.it.NextNode(m.ctx)
+		n := s.it.Node()
-		if !isLive {
-			return
-		}
 		select {
 		case s.next <- n:
 		case m.fromAny <- n:
-		case <-m.ctx.Done():
+		case <-closed:
 			return
 		}
 	}

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(new(genIter), 10)
-	nodes := ReadNodes(context.Background(), iter, 10)
 	checkNodes(t, nodes, 10)
 }
 
-// This test verifies that ReadNodes checks for context cancelation.
+// This test checks that ReadNodes terminates when reading N nodes from an iterator
-func TestReadNodesCancel(t *testing.T) {
+// which returns less than N nodes in an endless cycle.
-	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.
 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)
+	nodes := ReadNodes(mix, 500)
-	defer cancel()
-	nodes := ReadNodes(ctx, 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)
+	nodes := ReadNodes(mix, 500)
-	defer cancel()
-	nodes := ReadNodes(ctx, 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())
+		mix.Next()
-		ch <- n
+		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{})}
+	source := cycleNodes()
-	close(source.unblock) // first NextNode call will return (nil, false)
+	source.Close()
 	mix.AddSource(source)
 
-	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	if mix.Next() {
-	defer cancel()
+		t.Fatal("Next should've returned false")
-	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))
@@ -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 {
+		if mix.Next() {
-			t.Error("NextNode returned isLive == true")
+			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) {
+func (s *genIter) Next() bool {
-	n := testNode(uint64(s.index)<<32|uint64(s.gen), 0)
+	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.
+// cycleNodes is an interator that cycles through the given slice.
-type blockedIter struct {
+func cycleNodes(nodes ...*enode.Node) Iterator {
-	child   Iterator
+	return &cycleIter{nodes: nodes}
-	unblock chan struct{}
 }
 
-func (s *blockedIter) NextNode(ctx context.Context) (*enode.Node, bool) {
+type cycleIter struct {
-	select {
+	cur   *enode.Node
-	case _, ok := <-s.unblock:
+	mu    sync.Mutex
-		if !ok {
+	index int
-			return nil, false
+	nodes []*enode.Node
-		}
+}
-		return s.child.NextNode(ctx)
+
-	case <-ctx.Done():
+func (s *cycleIter) Next() bool {
-		return nil, true
+	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") }
+func (s *cycleIter) Node() *enode.Node {
-
+	return s.nodes[s.index]
-// 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 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)
+	return it.Iterator.Next()
-}
-
-func (it *callCountIter) Close() {
-	it.child.Close()
 }