p2p/simulations: Add simulation network functionality for bootnodes & lightnodes

p2p/simulations/network.go

@@ -56,6 +56,10 @@ type Network struct {
 	Nodes   []*Node `json:"nodes"`
 	nodeMap map[enode.ID]int
 
+	// Node subtypes are also mapped separately, so they can be distinguished quickly
+	bootNodeMap  map[enode.ID]int
+	lightNodeMap map[enode.ID]int
+
 	Conns   []*Conn `json:"conns"`
 	connMap map[string]int
 
@@ -71,6 +75,8 @@ func NewNetwork(nodeAdapter adapters.NodeAdapter, conf *NetworkConfig) *Network
 		NetworkConfig: *conf,
 		nodeAdapter:   nodeAdapter,
 		nodeMap:       make(map[enode.ID]int),
+		bootNodeMap:   make(map[enode.ID]int),
+		lightNodeMap:  make(map[enode.ID]int),
 		connMap:       make(map[string]int),
 		quitc:         make(chan struct{}),
 	}
@@ -120,7 +126,15 @@ 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)
+	if conf.BootNode {
+		net.bootNodeMap[conf.ID] = nodeIndex
+	} else if conf.LightNode {
+		net.lightNodeMap[conf.ID] = nodeIndex
+	}
+
+	net.nodeMap[conf.ID] = nodeIndex
 	net.Nodes = append(net.Nodes, node)
 
 	// emit a "control" event
@@ -427,19 +441,164 @@ func (net *Network) getNodeByName(name string) *Node {
 	return nil
 }
 
-// GetNodes returns the existing nodes
+// GetNodes returns the existing nodes.
-func (net *Network) GetNodes() (nodes []*Node) {
+// 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()
+	return net.getNodes(excludeIDs)
 }
 
-func (net *Network) getNodes() (nodes []*Node) {
+func (net *Network) getNodes(excludeIDs []enode.ID) []*Node {
-	nodes = append(nodes, net.Nodes...)
+	if len(excludeIDs) > 0 {
+		// Get all curent nodeIDs
+		nodeIDs := make([]enode.ID, 0, len(net.nodeMap))
+		for id := range net.nodeMap {
+			nodeIDs = append(nodeIDs, id)
+		}
+
+		// Return the difference of nodeIDs and excludeIDs
+		filteredIDs := filterIDs(nodeIDs, excludeIDs)
+		return net.getNodesByID(filteredIDs)
+	} 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
 }
 
+// GetBootNodes returns all configured bootnodes in the network.
+func (net *Network) GetBootNodes() []*Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getBootNodes()
+}
+
+func (net *Network) getBootNodes() []*Node {
+	bootNodes := make([]*Node, 0, len(net.bootNodeMap))
+	for _, i := range net.bootNodeMap {
+		bootNodes = append(bootNodes, net.Nodes[i])
+	}
+
+	return bootNodes
+}
+
+// GetBootNodeIDs returns a slice of all bootnode enode.ID
+func (net *Network) GetBootNodeIDs() []enode.ID {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getBootNodeIDs()
+}
+
+func (net *Network) getBootNodeIDs() []enode.ID {
+	bootNodeIDs := make([]enode.ID, 0, len(net.bootNodeMap))
+	for id := range net.bootNodeMap {
+		bootNodeIDs = append(bootNodeIDs, id)
+	}
+
+	return bootNodeIDs
+}
+
+// GetLightNodes returns all configured light nodes in the network.
+func (net *Network) GetLightNodes() []*Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getLightNodes()
+}
+
+func (net *Network) getLightNodes() []*Node {
+	lightNodes := make([]*Node, 0, len(net.lightNodeMap))
+	for _, i := range net.lightNodeMap {
+		lightNodes = append(lightNodes, net.Nodes[i])
+	}
+
+	return lightNodes
+}
+
+// GetLightNodeIDs returns a slice of all light node enode.ID
+func (net *Network) GetLightNodeIDs() []enode.ID {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getLightNodeIDs()
+}
+
+func (net *Network) getLightNodeIDs() []enode.ID {
+	lightNodeIDs := make([]enode.ID, 0, len(net.lightNodeMap))
+	for id := range net.lightNodeMap {
+		lightNodeIDs = append(lightNodeIDs, id)
+	}
+
+	return lightNodeIDs
+}
+
+// GetFullNodes returns all configured full nodes in the network.
+// This excludes bootnodes and lightnodes.
+func (net *Network) GetFullNodes() []*Node {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getFullNodes()
+}
+
+// Collect the enode.IDs of all nodes types that are not full nodes and provide them to getNodes for exclusion
+func (net *Network) getFullNodes() []*Node {
+	excludeNodeCount := len(net.lightNodeMap) + len(net.bootNodeMap)
+	excludeIDs := make([]enode.ID, 0, excludeNodeCount)
+	for ID := range net.lightNodeMap {
+		excludeIDs = append(excludeIDs, ID)
+	}
+
+	for ID := range net.bootNodeMap {
+		excludeIDs = append(excludeIDs, ID)
+	}
+
+	return net.getNodes(excludeIDs)
+}
+
+// GetFullNodeIDs returns a slice of all full node enode.ID
+func (net *Network) GetFullNodeIDs() []enode.ID {
+	net.lock.RLock()
+	defer net.lock.RUnlock()
+
+	return net.getFullNodeIDs()
+}
+
+func (net *Network) getFullNodeIDs() []enode.ID {
+	// The number of full nodes is the total number minus all sub mapping counts
+	fullNodeCount := len(net.nodeMap) - len(net.lightNodeMap) - len(net.bootNodeMap)
+	fullNodeIDs := make([]enode.ID, 0, fullNodeCount)
+
+	for _, node := range net.getFullNodes() {
+		fullNodeIDs = append(fullNodeIDs, node.ID())
+	}
+
+	return fullNodeIDs
+}
+
 // GetRandomUpNode returns a random node on the network, which is running.
 func (net *Network) GetRandomUpNode(excludeIDs ...enode.ID) *Node {
 	net.lock.RLock()
@@ -469,7 +628,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())
 		}
@@ -616,6 +775,8 @@ func (net *Network) Reset() {
 	//re-initialize the maps
 	net.connMap = make(map[string]int)
 	net.nodeMap = make(map[enode.ID]int)
+	net.bootNodeMap = make(map[enode.ID]int)
+	net.lightNodeMap = make(map[enode.ID]int)
 
 	net.Nodes = nil
 	net.Conns = nil

p2p/simulations/network_test.go

@@ -17,6 +17,7 @@
 package simulations
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -393,6 +394,191 @@ func TestNetworkSimulation(t *testing.T) {
 	}
 }
 
+// TestMultiNodeRetrieval creates a multi-node simulation network.
+// Full nodes, bootnodes and lightnodes are created.
+// Functions for retrieving specific subgroups of nodes are then tested for correctness.
+func TestMultiNodeRetrieval(t *testing.T) {
+	adapter := adapters.NewSimAdapter(adapters.Services{
+		"test": newTestService,
+	})
+	network := NewNetwork(adapter, &NetworkConfig{
+		DefaultService: "test",
+	})
+	defer network.Shutdown()
+
+	// Create a bootnode
+	bootNodeConf := adapters.RandomNodeConfig()
+	bootNodeConf.BootNode = true
+	bootNode, err := network.NewNodeWithConfig(bootNodeConf)
+	if err != nil {
+		t.Fatalf("error creating bootnode: %s", err)
+	}
+	if err := network.Start(bootNode.ID()); err != nil {
+		t.Fatalf("error starting bootnode: %s", err)
+	}
+
+	// Create 20 light nodes
+	lightNodeCount := 20
+	lightNodes := make(map[enode.ID]*Node, lightNodeCount)
+	for i := 0; i < lightNodeCount; i++ {
+		conf := adapters.RandomNodeConfig()
+		conf.LightNode = true
+		node, err := network.NewNodeWithConfig(conf)
+		if err != nil {
+			t.Fatalf("error creating light node: %s", err)
+		}
+		if err := network.Start(node.ID()); err != nil {
+			t.Fatalf("error starting light node: %s", err)
+		}
+		lightNodes[node.ID()] = node
+	}
+
+	// Create 20 full nodes
+	fullNodeCount := 20
+	fullNodes := make(map[enode.ID]*Node, fullNodeCount)
+	for i := 0; i < fullNodeCount; i++ {
+		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)
+		}
+		fullNodes[node.ID()] = node
+	}
+
+	// Check that network.GetBootNodes returns the boot node we created and only that bootnode
+	bootNodes := network.GetBootNodes()
+	if len(bootNodes) == 0 {
+		t.Fatal("GetBootNodes returned empty when size of one was expected")
+	}
+
+	for _, bn := range bootNodes {
+		if !bytes.Equal(bn.ID().Bytes(), bootNode.ID().Bytes()) {
+			t.Fatalf("Found an unexpected node in GetBootNodes: %s", bn.String())
+		}
+	}
+
+	// Check that the boot node's ID is the only one returned by GetBoodNodeIDs()
+	// If a non-matching ID is found, the test fails
+	bootNodeIDs := network.GetBootNodeIDs()
+	if len(bootNodeIDs) == 0 {
+		t.Fatal("GetBootNodeIDs returned empty when one ID was expected")
+	}
+
+	for _, id := range bootNodeIDs {
+		if !bytes.Equal(id.Bytes(), bootNode.ID().Bytes()) {
+			t.Fatalf("Found an unexpected enode.ID in GetBootNodeIDs: %s", id.String())
+		}
+	}
+
+	// Check that each of lightNodes (the light nodes we just created) are available from the GetLightNodes method.
+	// If a light node isn't found in GetLightNodes, the test fails.
+	for _, ln1 := range lightNodes {
+		match := false
+		lightNode1IDBytes := ln1.ID().Bytes()
+
+		for _, ln2 := range network.GetLightNodes() {
+			lightNode2IDBytes := ln2.ID().Bytes()
+			if bytes.Equal(lightNode1IDBytes, lightNode2IDBytes) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("A created light node was not returned by GetLightNodes(), ID: %s", ln1.ID().String())
+		}
+	}
+
+	// Check that the IDs of each of fullNodes are returned by GetFullNodeIDs()
+	// If a full not isn't found in GetFullNodeIDs(), the test fails
+	lightNodeIDs := network.GetLightNodeIDs()
+	for id1 := range lightNodes {
+		match := false
+		for _, id2 := range lightNodeIDs {
+			if bytes.Equal(id1.Bytes(), id2.Bytes()) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("Not all light nodes were returned by GetLightNodeIDs(), ID: %s", id1.String())
+		}
+	}
+
+	// Check that each of fullNodes (the full nodes we just created) are available from the GetFullNodes method.
+	// If a full node isn't found in GetFullNodes, the test fails.
+	for _, fn1 := range fullNodes {
+		match := false
+		fullNode1IDBytes := fn1.ID().Bytes()
+
+		for _, fn2 := range network.GetFullNodes() {
+			fullNode2IDBytes := fn2.ID().Bytes()
+			if bytes.Equal(fullNode1IDBytes, fullNode2IDBytes) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("A created full node was not returned by GetFullNodes(), ID: %s", fn1.ID().String())
+		}
+	}
+
+	// Check that the IDs of each of fullNodes are returned by GetFullNodeIDs()
+	// If a full not isn't found in GetFullNodeIDs(), the test fails
+	fullNodeIDs := network.GetFullNodeIDs()
+	for id1 := range fullNodes {
+		match := false
+		for _, id2 := range fullNodeIDs {
+			if bytes.Equal(id1.Bytes(), id2.Bytes()) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("Not all full nodes were returned by GetFullNodeIDs(), ID: %s", id1.String())
+		}
+	}
+
+	// Get all nodes, excluding the bootnode by passing the bootnode ID.
+	// Checks that the bootnode is excluded as expected and fails the test if not.
+	nodesExclBootNode := network.GetNodes(bootNode.ID())
+	for _, node := range nodesExclBootNode {
+		if bytes.Equal(node.ID().Bytes(), bootNode.ID().Bytes()) {
+			t.Fatalf("Bootnode still found in GetNodes when it has been explicitly excluded.")
+		}
+	}
+
+	// Get all node IDs and call GetNodesByID using them.
+	// The test then confirms that the nodes returned from GetNodes() match those returned from GetNodesByID(allIDs)
+	var nodeIDs []enode.ID
+	for _, node := range network.GetNodes() {
+		nodeIDs = append(nodeIDs, node.ID())
+	}
+
+	nodesByID := network.GetNodesByID(nodeIDs)
+	for _, node1 := range network.GetNodes() {
+		match := false
+		node1IDBytes := node1.ID().Bytes()
+		for _, node2 := range nodesByID {
+			node2IDBytes := node2.ID().Bytes()
+			if bytes.Equal(node1IDBytes, node2IDBytes) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			t.Fatalf("A node was found in GetNodes() that was not returned by GetNodesByID() for all node IDs")
+		}
+	}
+}
+
 func triggerChecks(ctx context.Context, ids []enode.ID, trigger chan enode.ID, interval time.Duration) {
 	tick := time.NewTicker(interval)
 	defer tick.Stop()