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

Author: chadsr

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.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)
+		}
+	}
+
+	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.getNodes()
+	return net.getNodesByProperty(property)
 }
 
-func (net *Network) getNodes() (nodes []*Node) {
-	nodes = append(nodes, net.Nodes...)
+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,14 +565,21 @@ 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())
 		}
 	}
 	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()