p2p/simulation: move connection methods from swarm/network/simulation (#18323)

4 files changed, 571 insertions, 12 deletions

diff --git a/p2p/simulations/connect.go b/p2p/simulations/connect.go
new file mode 100644
index 000000000..606dda056
--- /dev/null
+++ b/p2p/simulations/connect.go
@@ -0,0 +1,176 @@
+// Copyright 2018 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package simulations
+
+import (
+	"errors"
+	"strings"
+
+	"github.com/ethereum/go-ethereum/p2p/enode"
+)
+
+var (
+	ErrNodeNotFound = errors.New("node not found")
+	ErrNoPivotNode  = errors.New("no pivot node set")
+)
+
+// ConnectToPivotNode connects the node with provided NodeID
+// to the pivot node, already set by Network.SetPivotNode method.
+// It is useful when constructing a star network topology
+// when Network adds and removes nodes dynamically.
+func (net *Network) ConnectToPivotNode(id enode.ID) (err error) {
+	pivot := net.GetPivotNode()
+	if pivot == nil {
+		return ErrNoPivotNode
+	}
+	return net.connect(pivot.ID(), id)
+}
+
+// ConnectToLastNode connects the node with provided NodeID
+// to the last node that is up, and avoiding connection to self.
+// It is useful when constructing a chain network topology
+// when Network adds and removes nodes dynamically.
+func (net *Network) ConnectToLastNode(id enode.ID) (err error) {
+	ids := net.getUpNodeIDs()
+	l := len(ids)
+	if l < 2 {
+		return nil
+	}
+	last := ids[l-1]
+	if last == id {
+		last = ids[l-2]
+	}
+	return net.connect(last, id)
+}
+
+// ConnectToRandomNode connects the node with provided NodeID
+// to a random node that is up.
+func (net *Network) ConnectToRandomNode(id enode.ID) (err error) {
+	selected := net.GetRandomUpNode(id)
+	if selected == nil {
+		return ErrNodeNotFound
+	}
+	return net.connect(selected.ID(), id)
+}
+
+// ConnectNodesFull connects all nodes one to another.
+// It provides a complete connectivity in the network
+// which should be rarely needed.
+func (net *Network) ConnectNodesFull(ids []enode.ID) (err error) {
+	if ids == nil {
+		ids = net.getUpNodeIDs()
+	}
+	for i, lid := range ids {
+		for _, rid := range ids[i+1:] {
+			if err = net.connect(lid, rid); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+// ConnectNodesChain connects all nodes in a chain topology.
+// If ids argument is nil, all nodes that are up will be connected.
+func (net *Network) ConnectNodesChain(ids []enode.ID) (err error) {
+	if ids == nil {
+		ids = net.getUpNodeIDs()
+	}
+	l := len(ids)
+	for i := 0; i < l-1; i++ {
+		if err := net.connect(ids[i], ids[i+1]); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// ConnectNodesRing connects all nodes in a ring topology.
+// If ids argument is nil, all nodes that are up will be connected.
+func (net *Network) ConnectNodesRing(ids []enode.ID) (err error) {
+	if ids == nil {
+		ids = net.getUpNodeIDs()
+	}
+	l := len(ids)
+	if l < 2 {
+		return nil
+	}
+	if err := net.ConnectNodesChain(ids); err != nil {
+		return err
+	}
+	return net.connect(ids[l-1], ids[0])
+}
+
+// ConnectNodesStar connects all nodes in a star topology
+// with the center at provided NodeID.
+// If ids argument is nil, all nodes that are up will be connected.
+func (net *Network) ConnectNodesStar(pivot enode.ID, ids []enode.ID) (err error) {
+	if ids == nil {
+		ids = net.getUpNodeIDs()
+	}
+	for _, id := range ids {
+		if pivot == id {
+			continue
+		}
+		if err := net.connect(pivot, id); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// ConnectNodesStarPivot connects all nodes in a star topology
+// with the center at already set pivot node.
+// If ids argument is nil, all nodes that are up will be connected.
+func (net *Network) ConnectNodesStarPivot(ids []enode.ID) (err error) {
+	pivot := net.GetPivotNode()
+	if pivot == nil {
+		return ErrNoPivotNode
+	}
+	return net.ConnectNodesStar(pivot.ID(), ids)
+}
+
+// connect connects two nodes but ignores already connected error.
+func (net *Network) connect(oneID, otherID enode.ID) error {
+	return ignoreAlreadyConnectedErr(net.Connect(oneID, otherID))
+}
+
+func ignoreAlreadyConnectedErr(err error) error {
+	if err == nil || strings.Contains(err.Error(), "already connected") {
+		return nil
+	}
+	return err
+}
+
+// SetPivotNode sets the NodeID of the network's pivot node.
+// Pivot node is just a specific node that should be treated
+// differently then other nodes in test. SetPivotNode and
+// GetPivotNode are just a convenient functions to set and
+// retrieve it.
+func (net *Network) SetPivotNode(id enode.ID) {
+	net.lock.Lock()
+	defer net.lock.Unlock()
+	net.pivotNodeID = id
+}
+
+// GetPivotNode returns NodeID of the pivot node set by
+// Network.SetPivotNode method.
+func (net *Network) GetPivotNode() (node *Node) {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+	return net.getNode(net.pivotNodeID)
+}
diff --git a/p2p/simulations/connect_test.go b/p2p/simulations/connect_test.go
new file mode 100644
index 000000000..36f9442a8
--- /dev/null
+++ b/p2p/simulations/connect_test.go
@@ -0,0 +1,190 @@
+// Copyright 2018 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package simulations
+
+import (
+	"testing"
+
+	"github.com/ethereum/go-ethereum/node"
+	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
+)
+
+func newTestNetwork(t *testing.T, nodeCount int) (*Network, []enode.ID) {
+	adapter := adapters.NewSimAdapter(adapters.Services{
+		"noopwoop": func(ctx *adapters.ServiceContext) (node.Service, error) {
+			return NewNoopService(nil), nil
+		},
+	})
+
+	// create network
+	network := NewNetwork(adapter, &NetworkConfig{
+		DefaultService: "noopwoop",
+	})
+
+	// create and start nodes
+	ids := make([]enode.ID, nodeCount)
+	for i := range ids {
+		conf := adapters.RandomNodeConfig()
+		node, err := network.NewNodeWithConfig(conf)
+		if err != nil {
+			t.Fatalf("error creating node: %s", err)
+		}
+		if err := network.Start(node.ID()); err != nil {
+			t.Fatalf("error starting node: %s", err)
+		}
+		ids[i] = node.ID()
+	}
+
+	if len(network.Conns) > 0 {
+		t.Fatal("no connections should exist after just adding nodes")
+	}
+
+	return network, ids
+}
+
+func TestConnectToPivotNode(t *testing.T) {
+	net, ids := newTestNetwork(t, 2)
+	defer net.Shutdown()
+
+	pivot := ids[0]
+	net.SetPivotNode(pivot)
+
+	other := ids[1]
+	err := net.ConnectToPivotNode(other)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if net.GetConn(pivot, other) == nil {
+		t.Error("pivot and the other node are not connected")
+	}
+}
+
+func TestConnectToLastNode(t *testing.T) {
+	net, ids := newTestNetwork(t, 10)
+	defer net.Shutdown()
+
+	first := ids[0]
+	if err := net.ConnectToLastNode(first); err != nil {
+		t.Fatal(err)
+	}
+
+	last := ids[len(ids)-1]
+	for i, id := range ids {
+		if id == first || id == last {
+			continue
+		}
+
+		if net.GetConn(first, id) != nil {
+			t.Errorf("connection must not exist with node(ind: %v, id: %v)", i, id)
+		}
+	}
+
+	if net.GetConn(first, last) == nil {
+		t.Error("first and last node must be connected")
+	}
+}
+
+func TestConnectToRandomNode(t *testing.T) {
+	net, ids := newTestNetwork(t, 10)
+	defer net.Shutdown()
+
+	err := net.ConnectToRandomNode(ids[0])
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var cc int
+	for i, a := range ids {
+		for _, b := range ids[i:] {
+			if net.GetConn(a, b) != nil {
+				cc++
+			}
+		}
+	}
+
+	if cc != 1 {
+		t.Errorf("expected one connection, got %v", cc)
+	}
+}
+
+func TestConnectNodesFull(t *testing.T) {
+	net, ids := newTestNetwork(t, 12)
+	defer net.Shutdown()
+
+	err := net.ConnectNodesFull(ids)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	VerifyFull(t, net, ids)
+}
+
+func TestConnectNodesChain(t *testing.T) {
+	net, ids := newTestNetwork(t, 10)
+	defer net.Shutdown()
+
+	err := net.ConnectNodesChain(ids)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	VerifyChain(t, net, ids)
+}
+
+func TestConnectNodesRing(t *testing.T) {
+	net, ids := newTestNetwork(t, 10)
+	defer net.Shutdown()
+
+	err := net.ConnectNodesRing(ids)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	VerifyRing(t, net, ids)
+}
+
+func TestConnectNodesStar(t *testing.T) {
+	net, ids := newTestNetwork(t, 10)
+	defer net.Shutdown()
+
+	pivotIndex := 2
+
+	err := net.ConnectNodesStar(ids[pivotIndex], ids)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	VerifyStar(t, net, ids, pivotIndex)
+}
+
+func TestConnectNodesStarPivot(t *testing.T) {
+	net, ids := newTestNetwork(t, 10)
+	defer net.Shutdown()
+
+	pivotIndex := 4
+
+	net.SetPivotNode(ids[pivotIndex])
+
+	err := net.ConnectNodesStarPivot(ids)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	VerifyStar(t, net, ids, pivotIndex)
+}
diff --git a/p2p/simulations/network.go b/p2p/simulations/network.go
index ab9f582c5..a6fac2c2a 100644
--- a/p2p/simulations/network.go
+++ b/p2p/simulations/network.go
@@ -22,6 +22,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math/rand"
 	"sync"
 	"time"
 
@@ -57,6 +58,8 @@ type Network struct {
 	Conns   []*Conn `json:"conns"`
 	connMap map[string]int
 
+	pivotNodeID enode.ID
+
 	nodeAdapter adapters.NodeAdapter
 	events      event.Feed
 	lock        sync.RWMutex
@@ -370,23 +373,32 @@ func (net *Network) DidReceive(sender, receiver enode.ID, proto string, code uin
 // GetNode gets the node with the given ID, returning nil if the node does not
 // exist
 func (net *Network) GetNode(id enode.ID) *Node {
-	net.lock.Lock()
-	defer net.lock.Unlock()
+	net.lock.RLock()
+	defer net.lock.RUnlock()
 	return net.getNode(id)
 }
 
 // GetNode gets the node with the given name, returning nil if the node does
 // not exist
 func (net *Network) GetNodeByName(name string) *Node {
-	net.lock.Lock()
-	defer net.lock.Unlock()
+	net.lock.RLock()
+	defer net.lock.RUnlock()
 	return net.getNodeByName(name)
 }
 
+func (net *Network) getNodeByName(name string) *Node {
+	for _, node := range net.Nodes {
+		if node.Config.Name == name {
+			return node
+		}
+	}
+	return nil
+}
+
 // GetNodes returns the existing nodes
 func (net *Network) GetNodes() (nodes []*Node) {
-	net.lock.Lock()
-	defer net.lock.Unlock()
+	net.lock.RLock()
+	defer net.lock.RUnlock()
 
 	nodes = append(nodes, net.Nodes...)
 	return nodes
@@ -400,20 +412,67 @@ func (net *Network) getNode(id enode.ID) *Node {
 	return net.Nodes[i]
 }
 
-func (net *Network) getNodeByName(name string) *Node {
+// GetRandomUpNode returns a random node on the network, which is running.
+func (net *Network) GetRandomUpNode(excludeIDs ...enode.ID) *Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+	return net.getRandomNode(net.getUpNodeIDs(), excludeIDs)
+}
+
+func (net *Network) getUpNodeIDs() (ids []enode.ID) {
 	for _, node := range net.Nodes {
-		if node.Config.Name == name {
-			return node
+		if node.Up {
+			ids = append(ids, node.ID())
 		}
 	}
-	return nil
+	return ids
+}
+
+// GetRandomDownNode returns a random node on the network, which is stopped.
+func (net *Network) GetRandomDownNode(excludeIDs ...enode.ID) *Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+	return net.getRandomNode(net.getDownNodeIDs(), excludeIDs)
+}
+
+func (net *Network) getDownNodeIDs() (ids []enode.ID) {
+	for _, node := range net.GetNodes() {
+		if !node.Up {
+			ids = append(ids, node.ID())
+		}
+	}
+	return ids
+}
+
+func (net *Network) getRandomNode(ids []enode.ID, excludeIDs []enode.ID) *Node {
+	filtered := filterIDs(ids, excludeIDs)
+
+	l := len(filtered)
+	if l == 0 {
+		return nil
+	}
+	return net.GetNode(filtered[rand.Intn(l)])
+}
+
+func filterIDs(ids []enode.ID, excludeIDs []enode.ID) []enode.ID {
+	exclude := make(map[enode.ID]bool)
+	for _, id := range excludeIDs {
+		exclude[id] = true
+	}
+	var filtered []enode.ID
+	for _, id := range ids {
+		if _, found := exclude[id]; !found {
+			filtered = append(filtered, id)
+		}
+	}
+	return filtered
 }
 
 // GetConn returns the connection which exists between "one" and "other"
 // regardless of which node initiated the connection
 func (net *Network) GetConn(oneID, otherID enode.ID) *Conn {
-	net.lock.Lock()
-	defer net.lock.Unlock()
+	net.lock.RLock()
+	defer net.lock.RUnlock()
 	return net.getConn(oneID, otherID)
 }
 
diff --git a/p2p/simulations/test.go b/p2p/simulations/test.go
new file mode 100644
index 000000000..beeb414e4
--- /dev/null
+++ b/p2p/simulations/test.go
@@ -0,0 +1,134 @@
+package simulations
+
+import (
+	"testing"
+
+	"github.com/ethereum/go-ethereum/p2p"
+	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/enr"
+	"github.com/ethereum/go-ethereum/rpc"
+)
+
+// NoopService is the service that does not do anything
+// but implements node.Service interface.
+type NoopService struct {
+	c map[enode.ID]chan struct{}
+}
+
+func NewNoopService(ackC map[enode.ID]chan struct{}) *NoopService {
+	return &NoopService{
+		c: ackC,
+	}
+}
+
+func (t *NoopService) Protocols() []p2p.Protocol {
+	return []p2p.Protocol{
+		{
+			Name:    "noop",
+			Version: 666,
+			Length:  0,
+			Run: func(peer *p2p.Peer, rw p2p.MsgReadWriter) error {
+				if t.c != nil {
+					t.c[peer.ID()] = make(chan struct{})
+					close(t.c[peer.ID()])
+				}
+				rw.ReadMsg()
+				return nil
+			},
+			NodeInfo: func() interface{} {
+				return struct{}{}
+			},
+			PeerInfo: func(id enode.ID) interface{} {
+				return struct{}{}
+			},
+			Attributes: []enr.Entry{},
+		},
+	}
+}
+
+func (t *NoopService) APIs() []rpc.API {
+	return []rpc.API{}
+}
+
+func (t *NoopService) Start(server *p2p.Server) error {
+	return nil
+}
+
+func (t *NoopService) Stop() error {
+	return nil
+}
+
+func VerifyRing(t *testing.T, net *Network, ids []enode.ID) {
+	t.Helper()
+	n := len(ids)
+	for i := 0; i < n; i++ {
+		for j := i + 1; j < n; j++ {
+			c := net.GetConn(ids[i], ids[j])
+			if i == j-1 || (i == 0 && j == n-1) {
+				if c == nil {
+					t.Errorf("nodes %v and %v are not connected, but they should be", i, j)
+				}
+			} else {
+				if c != nil {
+					t.Errorf("nodes %v and %v are connected, but they should not be", i, j)
+				}
+			}
+		}
+	}
+}
+
+func VerifyChain(t *testing.T, net *Network, ids []enode.ID) {
+	t.Helper()
+	n := len(ids)
+	for i := 0; i < n; i++ {
+		for j := i + 1; j < n; j++ {
+			c := net.GetConn(ids[i], ids[j])
+			if i == j-1 {
+				if c == nil {
+					t.Errorf("nodes %v and %v are not connected, but they should be", i, j)
+				}
+			} else {
+				if c != nil {
+					t.Errorf("nodes %v and %v are connected, but they should not be", i, j)
+				}
+			}
+		}
+	}
+}
+
+func VerifyFull(t *testing.T, net *Network, ids []enode.ID) {
+	t.Helper()
+	n := len(ids)
+	var connections int
+	for i, lid := range ids {
+		for _, rid := range ids[i+1:] {
+			if net.GetConn(lid, rid) != nil {
+				connections++
+			}
+		}
+	}
+
+	want := n * (n - 1) / 2
+	if connections != want {
+		t.Errorf("wrong number of connections, got: %v, want: %v", connections, want)
+	}
+}
+
+func VerifyStar(t *testing.T, net *Network, ids []enode.ID, centerIndex int) {
+	t.Helper()
+	n := len(ids)
+	for i := 0; i < n; i++ {
+		for j := i + 1; j < n; j++ {
+			c := net.GetConn(ids[i], ids[j])
+			if i == centerIndex || j == centerIndex {
+				if c == nil {
+					t.Errorf("nodes %v and %v are not connected, but they should be", i, j)
+				}
+			} else {
+				if c != nil {
+					t.Errorf("nodes %v and %v are connected, but they should not be", i, j)
+				}
+			}
+		}
+	}
+}