diff --git a/p2p/enode/iter.go b/p2p/enode/iter.go index 812e1a0e6d..4b7e28929e 100644 --- a/p2p/enode/iter.go +++ b/p2p/enode/iter.go @@ -38,7 +38,7 @@ type SourceIterator interface { } // WithSource attaches a 'source name' to an iterator. -func WithSource(name string, it Iterator) SourceIterator { +func WithSourceName(name string, it Iterator) SourceIterator { return sourceIter{it, name} } @@ -46,7 +46,7 @@ func ensureSourceIter(it Iterator) SourceIterator { if si, ok := it.(SourceIterator); ok { return si } - return WithSource("", it) + return WithSourceName("", it) } type sourceIter struct { diff --git a/p2p/enode/iter_test.go b/p2p/enode/iter_test.go index b736ed450a..577f9c2825 100644 --- a/p2p/enode/iter_test.go +++ b/p2p/enode/iter_test.go @@ -19,6 +19,7 @@ package enode import ( "encoding/binary" "runtime" + "slices" "sync/atomic" "testing" "time" @@ -183,6 +184,53 @@ func TestFairMixRemoveSource(t *testing.T) { } } +// This checks that FairMix correctly returns the name of the source that produced the node. +func TestFairMixSourceName(t *testing.T) { + nodes := make([]*Node, 6) + for i := range nodes { + nodes[i] = testNode(uint64(i), uint64(i)) + } + mix := NewFairMix(-1) + mix.AddSource(WithSourceName("s1", IterNodes(nodes[0:2]))) + mix.AddSource(WithSourceName("s2", IterNodes(nodes[2:4]))) + mix.AddSource(WithSourceName("s3", IterNodes(nodes[4:6]))) + + var names []string + for range nodes { + mix.Next() + names = append(names, mix.NodeSource()) + } + want := []string{"s2", "s3", "s1", "s2", "s3", "s1"} + if !slices.Equal(names, want) { + t.Fatalf("wrong names: %v", names) + } +} + +// This checks that FairMix returns the name of the source that produced the node, +// even when FairMix instances are nested. +func TestFairMixNestedSourceName(t *testing.T) { + nodes := make([]*Node, 6) + for i := range nodes { + nodes[i] = testNode(uint64(i), uint64(i)) + } + mix := NewFairMix(-1) + mix.AddSource(WithSourceName("s1", IterNodes(nodes[0:2]))) + submix := NewFairMix(-1) + submix.AddSource(WithSourceName("s2", IterNodes(nodes[2:4]))) + submix.AddSource(WithSourceName("s3", IterNodes(nodes[4:6]))) + mix.AddSource(submix) + + var names []string + for range nodes { + mix.Next() + names = append(names, mix.NodeSource()) + } + want := []string{"s3", "s1", "s2", "s1", "s3", "s2"} + if !slices.Equal(names, want) { + t.Fatalf("wrong names: %v", names) + } +} + type blockingIter chan struct{} func (it blockingIter) Next() bool {