p2p/simulations: add node properties support and utility functions (#20060)

p2p/simulations/adapters/types.go

@@ -101,6 +101,11 @@ type NodeConfig struct {
 	// services registered by calling the RegisterService function)
 	Services []string
 
+	// Properties are the names of the properties this node should hold
+	// within running services (e.g. "bootnode", "lightnode" or any custom values)
+	// These values need to be checked and acted upon by node Services
+	Properties []string
+
 	// Enode
 	node *enode.Node
 
@@ -120,6 +125,7 @@ type nodeConfigJSON struct {
 	PrivateKey      string   `json:"private_key"`
 	Name            string   `json:"name"`
 	Services        []string `json:"services"`
+	Properties      []string `json:"properties"`
 	EnableMsgEvents bool     `json:"enable_msg_events"`
 	Port            uint16   `json:"port"`
 }
@@ -131,6 +137,7 @@ func (n *NodeConfig) MarshalJSON() ([]byte, error) {
 		ID:              n.ID.String(),
 		Name:            n.Name,
 		Services:        n.Services,
+		Properties:      n.Properties,
 		Port:            n.Port,
 		EnableMsgEvents: n.EnableMsgEvents,
 	}
@@ -168,6 +175,7 @@ func (n *NodeConfig) UnmarshalJSON(data []byte) error {
 
 	n.Name = confJSON.Name
 	n.Services = confJSON.Services
+	n.Properties = confJSON.Properties
 	n.Port = confJSON.Port
 	n.EnableMsgEvents = confJSON.EnableMsgEvents
 

p2p/simulations/network.go

@@ -56,6 +56,9 @@ type Network struct {
 	Nodes   []*Node `json:"nodes"`
 	nodeMap map[enode.ID]int
 
+	// Maps a node property string to node indexes of all nodes that hold this property
+	propertyMap map[string][]int
+
 	Conns   []*Conn `json:"conns"`
 	connMap map[string]int
 
@@ -71,6 +74,7 @@ func NewNetwork(nodeAdapter adapters.NodeAdapter, conf *NetworkConfig) *Network
 		NetworkConfig: *conf,
 		nodeAdapter:   nodeAdapter,
 		nodeMap:       make(map[enode.ID]int),
+		propertyMap:   make(map[string][]int),
 		connMap:       make(map[string]int),
 		quitc:         make(chan struct{}),
 	}
@@ -120,9 +124,16 @@ func (net *Network) NewNodeWithConfig(conf *adapters.NodeConfig) (*Node, error)
 		Config: conf,
 	}
 	log.Trace("Node created", "id", conf.ID)
-	net.nodeMap[conf.ID] = len(net.Nodes)
+
+	nodeIndex := len(net.Nodes)
+	net.nodeMap[conf.ID] = nodeIndex
 	net.Nodes = append(net.Nodes, node)
 
+	// Register any node properties with the network-level propertyMap
+	for _, property := range conf.Properties {
+		net.propertyMap[property] = append(net.propertyMap[property], nodeIndex)
+	}
+
 	// emit a "control" event
 	net.events.Send(ControlEvent(node))
 
@@ -410,7 +421,7 @@ func (net *Network) getNode(id enode.ID) *Node {
 	return net.Nodes[i]
 }
 
-// GetNode gets the node with the given name, returning nil if the node does
+// GetNodeByName gets the node with the given name, returning nil if the node does
 // not exist
 func (net *Network) GetNodeByName(name string) *Node {
 	net.lock.RLock()
@@ -427,19 +438,104 @@ func (net *Network) getNodeByName(name string) *Node {
 	return nil
 }
 
-// GetNodes returns the existing nodes
-func (net *Network) GetNodes() (nodes []*Node) {
+// GetNodeIDs returns the IDs of all existing nodes
+// Nodes can optionally be excluded by specifying their enode.ID.
+func (net *Network) GetNodeIDs(excludeIDs ...enode.ID) []enode.ID {
 	net.lock.RLock()
 	defer net.lock.RUnlock()
 
-	return net.getNodes()
+	return net.getNodeIDs(excludeIDs)
+}
+
+func (net *Network) getNodeIDs(excludeIDs []enode.ID) []enode.ID {
+	// Get all curent nodeIDs
+	nodeIDs := make([]enode.ID, 0, len(net.nodeMap))
+	for id := range net.nodeMap {
+		nodeIDs = append(nodeIDs, id)
+	}
+
+	if len(excludeIDs) > 0 {
+		// Return the difference of nodeIDs and excludeIDs
+		return filterIDs(nodeIDs, excludeIDs)
+	} else {
+		return nodeIDs
+	}
+}
+
+// GetNodes returns the existing nodes.
+// Nodes can optionally be excluded by specifying their enode.ID.
+func (net *Network) GetNodes(excludeIDs ...enode.ID) []*Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getNodes(excludeIDs)
+}
+
+func (net *Network) getNodes(excludeIDs []enode.ID) []*Node {
+	if len(excludeIDs) > 0 {
+		nodeIDs := net.getNodeIDs(excludeIDs)
+		return net.getNodesByID(nodeIDs)
+	} else {
+		return net.Nodes
+	}
+}
+
+// GetNodesByID returns existing nodes with the given enode.IDs.
+// If a node doesn't exist with a given enode.ID, it is ignored.
+func (net *Network) GetNodesByID(nodeIDs []enode.ID) []*Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getNodesByID(nodeIDs)
+}
+
+func (net *Network) getNodesByID(nodeIDs []enode.ID) []*Node {
+	nodes := make([]*Node, 0, len(nodeIDs))
+	for _, id := range nodeIDs {
+		node := net.getNode(id)
+		if node != nil {
+			nodes = append(nodes, node)
+		}
 	}
 
-func (net *Network) getNodes() (nodes []*Node) {
-	nodes = append(nodes, net.Nodes...)
 	return nodes
 }
 
+// GetNodesByProperty returns existing nodes that have the given property string registered in their NodeConfig
+func (net *Network) GetNodesByProperty(property string) []*Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getNodesByProperty(property)
+}
+
+func (net *Network) getNodesByProperty(property string) []*Node {
+	nodes := make([]*Node, 0, len(net.propertyMap[property]))
+	for _, nodeIndex := range net.propertyMap[property] {
+		nodes = append(nodes, net.Nodes[nodeIndex])
+	}
+
+	return nodes
+}
+
+// GetNodeIDsByProperty returns existing node's enode IDs that have the given property string registered in the NodeConfig
+func (net *Network) GetNodeIDsByProperty(property string) []enode.ID {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getNodeIDsByProperty(property)
+}
+
+func (net *Network) getNodeIDsByProperty(property string) []enode.ID {
+	nodeIDs := make([]enode.ID, 0, len(net.propertyMap[property]))
+	for _, nodeIndex := range net.propertyMap[property] {
+		node := net.Nodes[nodeIndex]
+		nodeIDs = append(nodeIDs, node.ID())
+	}
+
+	return nodeIDs
+}
+
 // GetRandomUpNode returns a random node on the network, which is running.
 func (net *Network) GetRandomUpNode(excludeIDs ...enode.ID) *Node {
 	net.lock.RLock()
@@ -469,7 +565,7 @@ func (net *Network) GetRandomDownNode(excludeIDs ...enode.ID) *Node {
 }
 
 func (net *Network) getDownNodeIDs() (ids []enode.ID) {
-	for _, node := range net.getNodes() {
+	for _, node := range net.Nodes {
 		if !node.Up() {
 			ids = append(ids, node.ID())
 		}
@@ -477,6 +573,13 @@ func (net *Network) getDownNodeIDs() (ids []enode.ID) {
 	return ids
 }
 
+// GetRandomNode returns a random node on the network, regardless of whether it is running or not
+func (net *Network) GetRandomNode(excludeIDs ...enode.ID) *Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+	return net.getRandomNode(net.getNodeIDs(nil), excludeIDs) // no need to exclude twice
+}
+
 func (net *Network) getRandomNode(ids []enode.ID, excludeIDs []enode.ID) *Node {
 	filtered := filterIDs(ids, excludeIDs)
 
@@ -616,6 +719,7 @@ func (net *Network) Reset() {
 	//re-initialize the maps
 	net.connMap = make(map[string]int)
 	net.nodeMap = make(map[enode.ID]int)
+	net.propertyMap = make(map[string][]int)
 
 	net.Nodes = nil
 	net.Conns = nil
@@ -634,12 +738,14 @@ type Node struct {
 	upMu sync.RWMutex
 }
 
+// Up returns whether the node is currently up (online)
 func (n *Node) Up() bool {
 	n.upMu.RLock()
 	defer n.upMu.RUnlock()
 	return n.up
 }
 
+// SetUp sets the up (online) status of the nodes with the given value
 func (n *Node) SetUp(up bool) {
 	n.upMu.Lock()
 	defer n.upMu.Unlock()

p2p/simulations/network_test.go

@@ -17,6 +17,7 @@
 package simulations
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -393,6 +394,275 @@ func TestNetworkSimulation(t *testing.T) {
 	}
 }
 
+func createTestNodes(count int, network *Network) (nodes []*Node, err error) {
+	for i := 0; i < count; i++ {
+		nodeConf := adapters.RandomNodeConfig()
+		node, err := network.NewNodeWithConfig(nodeConf)
+		if err != nil {
+			return nil, err
+		}
+		if err := network.Start(node.ID()); err != nil {
+			return nil, err
+		}
+
+		nodes = append(nodes, node)
+	}
+
+	return nodes, nil
+}
+
+func createTestNodesWithProperty(property string, count int, network *Network) (propertyNodes []*Node, err error) {
+	for i := 0; i < count; i++ {
+		nodeConf := adapters.RandomNodeConfig()
+		nodeConf.Properties = append(nodeConf.Properties, property)
+
+		node, err := network.NewNodeWithConfig(nodeConf)
+		if err != nil {
+			return nil, err
+		}
+		if err := network.Start(node.ID()); err != nil {
+			return nil, err
+		}
+
+		propertyNodes = append(propertyNodes, node)
+	}
+
+	return propertyNodes, nil
+}
+
+// TestGetNodeIDs creates a set of nodes and attempts to retrieve their IDs,.
+// It then tests again whilst excluding a node ID from being returned.
+// If a node ID is not returned, or more node IDs than expected are returned, the test fails.
+func TestGetNodeIDs(t *testing.T) {
+	adapter := adapters.NewSimAdapter(adapters.Services{
+		"test": newTestService,
+	})
+	network := NewNetwork(adapter, &NetworkConfig{
+		DefaultService: "test",
+	})
+	defer network.Shutdown()
+
+	numNodes := 5
+	nodes, err := createTestNodes(numNodes, network)
+	if err != nil {
+		t.Fatalf("Could not creat test nodes %v", err)
+	}
+
+	gotNodeIDs := network.GetNodeIDs()
+	if len(gotNodeIDs) != numNodes {
+		t.Fatalf("Expected %d nodes, got %d", numNodes, len(gotNodeIDs))
+	}
+
+	for _, node1 := range nodes {
+		match := false
+		for _, node2ID := range gotNodeIDs {
+			if bytes.Equal(node1.ID().Bytes(), node2ID.Bytes()) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("A created node was not returned by GetNodes(), ID: %s", node1.ID().String())
+		}
+	}
+
+	excludeNodeID := nodes[3].ID()
+	gotNodeIDsExcl := network.GetNodeIDs(excludeNodeID)
+	if len(gotNodeIDsExcl) != numNodes-1 {
+		t.Fatalf("Expected one less node ID to be returned")
+	}
+	for _, nodeID := range gotNodeIDsExcl {
+		if bytes.Equal(excludeNodeID.Bytes(), nodeID.Bytes()) {
+			t.Fatalf("GetNodeIDs returned the node ID we excluded, ID: %s", nodeID.String())
+		}
+	}
+}
+
+// TestGetNodes creates a set of nodes and attempts to retrieve them again.
+// It then tests again whilst excluding a node from being returned.
+// If a node is not returned, or more nodes than expected are returned, the test fails.
+func TestGetNodes(t *testing.T) {
+	adapter := adapters.NewSimAdapter(adapters.Services{
+		"test": newTestService,
+	})
+	network := NewNetwork(adapter, &NetworkConfig{
+		DefaultService: "test",
+	})
+	defer network.Shutdown()
+
+	numNodes := 5
+	nodes, err := createTestNodes(numNodes, network)
+	if err != nil {
+		t.Fatalf("Could not creat test nodes %v", err)
+	}
+
+	gotNodes := network.GetNodes()
+	if len(gotNodes) != numNodes {
+		t.Fatalf("Expected %d nodes, got %d", numNodes, len(gotNodes))
+	}
+
+	for _, node1 := range nodes {
+		match := false
+		for _, node2 := range gotNodes {
+			if bytes.Equal(node1.ID().Bytes(), node2.ID().Bytes()) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("A created node was not returned by GetNodes(), ID: %s", node1.ID().String())
+		}
+	}
+
+	excludeNodeID := nodes[3].ID()
+	gotNodesExcl := network.GetNodes(excludeNodeID)
+	if len(gotNodesExcl) != numNodes-1 {
+		t.Fatalf("Expected one less node to be returned")
+	}
+	for _, node := range gotNodesExcl {
+		if bytes.Equal(excludeNodeID.Bytes(), node.ID().Bytes()) {
+			t.Fatalf("GetNodes returned the node we excluded, ID: %s", node.ID().String())
+		}
+	}
+}
+
+// TestGetNodesByID creates a set of nodes and attempts to retrieve a subset of them by ID
+// If a node is not returned, or more nodes than expected are returned, the test fails.
+func TestGetNodesByID(t *testing.T) {
+	adapter := adapters.NewSimAdapter(adapters.Services{
+		"test": newTestService,
+	})
+	network := NewNetwork(adapter, &NetworkConfig{
+		DefaultService: "test",
+	})
+	defer network.Shutdown()
+
+	numNodes := 5
+	nodes, err := createTestNodes(numNodes, network)
+	if err != nil {
+		t.Fatalf("Could not create test nodes: %v", err)
+	}
+
+	numSubsetNodes := 2
+	subsetNodes := nodes[0:numSubsetNodes]
+	var subsetNodeIDs []enode.ID
+	for _, node := range subsetNodes {
+		subsetNodeIDs = append(subsetNodeIDs, node.ID())
+	}
+
+	gotNodesByID := network.GetNodesByID(subsetNodeIDs)
+	if len(gotNodesByID) != numSubsetNodes {
+		t.Fatalf("Expected %d nodes, got %d", numSubsetNodes, len(gotNodesByID))
+	}
+
+	for _, node1 := range subsetNodes {
+		match := false
+		for _, node2 := range gotNodesByID {
+			if bytes.Equal(node1.ID().Bytes(), node2.ID().Bytes()) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("A created node was not returned by GetNodesByID(), ID: %s", node1.ID().String())
+		}
+	}
+}
+
+// TestGetNodesByProperty creates a subset of nodes with a property assigned.
+// GetNodesByProperty is then checked for correctness by comparing the nodes returned to those initially created.
+// If a node with a property is not found, or more nodes than expected are returned, the test fails.
+func TestGetNodesByProperty(t *testing.T) {
+	adapter := adapters.NewSimAdapter(adapters.Services{
+		"test": newTestService,
+	})
+	network := NewNetwork(adapter, &NetworkConfig{
+		DefaultService: "test",
+	})
+	defer network.Shutdown()
+
+	numNodes := 3
+	_, err := createTestNodes(numNodes, network)
+	if err != nil {
+		t.Fatalf("Failed to create nodes: %v", err)
+	}
+
+	numPropertyNodes := 3
+	propertyTest := "test"
+	propertyNodes, err := createTestNodesWithProperty(propertyTest, numPropertyNodes, network)
+	if err != nil {
+		t.Fatalf("Failed to create nodes with property: %v", err)
+	}
+
+	gotNodesByProperty := network.GetNodesByProperty(propertyTest)
+	if len(gotNodesByProperty) != numPropertyNodes {
+		t.Fatalf("Expected %d nodes with a property, got %d", numPropertyNodes, len(gotNodesByProperty))
+	}
+
+	for _, node1 := range propertyNodes {
+		match := false
+		for _, node2 := range gotNodesByProperty {
+			if bytes.Equal(node1.ID().Bytes(), node2.ID().Bytes()) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("A created node with property was not returned by GetNodesByProperty(), ID: %s", node1.ID().String())
+		}
+	}
+}
+
+// TestGetNodeIDsByProperty creates a subset of nodes with a property assigned.
+// GetNodeIDsByProperty is then checked for correctness by comparing the node IDs returned to those initially created.
+// If a node ID with a property is not found, or more nodes IDs than expected are returned, the test fails.
+func TestGetNodeIDsByProperty(t *testing.T) {
+	adapter := adapters.NewSimAdapter(adapters.Services{
+		"test": newTestService,
+	})
+	network := NewNetwork(adapter, &NetworkConfig{
+		DefaultService: "test",
+	})
+	defer network.Shutdown()
+
+	numNodes := 3
+	_, err := createTestNodes(numNodes, network)
+	if err != nil {
+		t.Fatalf("Failed to create nodes: %v", err)
+	}
+
+	numPropertyNodes := 3
+	propertyTest := "test"
+	propertyNodes, err := createTestNodesWithProperty(propertyTest, numPropertyNodes, network)
+	if err != nil {
+		t.Fatalf("Failed to created nodes with property: %v", err)
+	}
+
+	gotNodeIDsByProperty := network.GetNodeIDsByProperty(propertyTest)
+	if len(gotNodeIDsByProperty) != numPropertyNodes {
+		t.Fatalf("Expected %d nodes with a property, got %d", numPropertyNodes, len(gotNodeIDsByProperty))
+	}
+
+	for _, node1 := range propertyNodes {
+		match := false
+		id1 := node1.ID()
+		for _, id2 := range gotNodeIDsByProperty {
+			if bytes.Equal(id1.Bytes(), id2.Bytes()) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("Not all nodes IDs were returned by GetNodeIDsByProperty(), ID: %s", id1.String())
+		}
+	}
+}
+
 func triggerChecks(ctx context.Context, ids []enode.ID, trigger chan enode.ID, interval time.Duration) {
 	tick := time.NewTicker(interval)
 	defer tick.Stop()