p2p/simulations: Add TestNetworkSimulation

Signed-off-by: Lewis Marshall <lewis@lmars.net>

p2p/simulations/http_test.go

@@ -20,6 +20,9 @@ import (
 type testService struct {
 	id discover.NodeID
 
+	// peerCount is incremented once a peer handshake has been performed
+	peerCount int64
+
 	// state stores []byte used to test creating and loading snapshots
 	state atomic.Value
 }
@@ -43,7 +46,10 @@ func (t *testService) APIs() []rpc.API {
 	return []rpc.API{{
 		Namespace: "test",
 		Version:   "1.0",
-		Service:   &TestAPI{state: &t.state},
+		Service: &TestAPI{
+			state:     &t.state,
+			peerCount: &t.peerCount,
+		},
 	}}
 }
 
@@ -56,6 +62,21 @@ func (t *testService) Stop() error {
 }
 
 func (t *testService) Run(_ *p2p.Peer, rw p2p.MsgReadWriter) error {
+	// perform a handshake using an empty struct
+	errc := make(chan error, 2)
+	go func() { errc <- p2p.Send(rw, 0, struct{}{}) }()
+	go func() { errc <- p2p.ExpectMsg(rw, 0, struct{}{}) }()
+	for i := 0; i < 2; i++ {
+		if err := <-errc; err != nil {
+			return err
+		}
+	}
+
+	// track the peer
+	atomic.AddInt64(&t.peerCount, 1)
+	defer atomic.AddInt64(&t.peerCount, -1)
+
+	// block until the peer is dropped
 	for {
 		_, err := rw.ReadMsg()
 		if err != nil {
@@ -68,12 +89,20 @@ func (t *testService) Snapshot() ([]byte, error) {
 	return t.state.Load().([]byte), nil
 }
 
-// TestAPI provides a simple API to get and increment a counter and to
-// subscribe to increment events
+// TestAPI provides a simple API to:
+// * get the peer count
+// * get and set an arbitrary state byte slice
+// * get and increment a counter
+// * subscribe to counter increment events
 type TestAPI struct {
-	state   *atomic.Value
-	counter int64
-	feed    event.Feed
+	state     *atomic.Value
+	peerCount *int64
+	counter   int64
+	feed      event.Feed
+}
+
+func (t *TestAPI) PeerCount() int64 {
+	return atomic.LoadInt64(t.peerCount)
 }
 
 func (t *TestAPI) Get() int64 {

p2p/simulations/network_test.go

@@ -0,0 +1,143 @@
+package simulations
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/ethereum/go-ethereum/p2p/discover"
+	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
+)
+
+// TestNetworkSimulation creates a multi-node simulation network with each node
+// connected in a ring topology, checks that all nodes successfully handshake
+// with each other and that a snapshot fully represents the desired topology
+func TestNetworkSimulation(t *testing.T) {
+	// create simulation network with 20 testService nodes
+	adapter := adapters.NewSimAdapter(adapters.Services{
+		"test": newTestService,
+	})
+	network := NewNetwork(adapter, &NetworkConfig{
+		DefaultService: "test",
+	})
+	defer network.Shutdown()
+	nodeCount := 20
+	ids := make([]discover.NodeID, nodeCount)
+	for i := 0; i < nodeCount; i++ {
+		node, err := network.NewNode()
+		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()
+	}
+
+	// perform a check which connects the nodes in a ring (so each node is
+	// connected to exactly two peers) and then checks that all nodes
+	// performed two handshakes by checking their peerCount
+	action := func(_ context.Context) error {
+		for i, id := range ids {
+			peerID := ids[(i+1)%len(ids)]
+			if err := network.Connect(id, peerID); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+	check := func(ctx context.Context, id discover.NodeID) (bool, error) {
+		// check we haven't run out of time
+		select {
+		case <-ctx.Done():
+			return false, ctx.Err()
+		default:
+		}
+
+		// get the node
+		node := network.GetNode(id)
+		if node == nil {
+			return false, fmt.Errorf("unknown node: %s", id)
+		}
+
+		// check it has exactly two peers
+		client, err := node.Client()
+		if err != nil {
+			return false, err
+		}
+		var peerCount int64
+		if err := client.CallContext(ctx, &peerCount, "test_peerCount"); err != nil {
+			return false, err
+		}
+		switch {
+		case peerCount < 2:
+			return false, nil
+		case peerCount == 2:
+			return true, nil
+		default:
+			return false, fmt.Errorf("unexpected peerCount: %d", peerCount)
+		}
+	}
+
+	timeout := 30 * time.Second
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	// trigger a check every 100ms
+	trigger := make(chan discover.NodeID)
+	go triggerChecks(ctx, ids, trigger, 100*time.Millisecond)
+
+	result := NewSimulation(network).Run(ctx, &Step{
+		Action:  action,
+		Trigger: trigger,
+		Expect: &Expectation{
+			Nodes: ids,
+			Check: check,
+		},
+	})
+	if result.Error != nil {
+		t.Fatalf("simulation failed: %s", result.Error)
+	}
+
+	// take a network snapshot and check it contains the correct topology
+	snap, err := network.Snapshot()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(snap.Nodes) != nodeCount {
+		t.Fatalf("expected snapshot to contain %d nodes, got %d", nodeCount, len(snap.Nodes))
+	}
+	if len(snap.Conns) != nodeCount {
+		t.Fatalf("expected snapshot to contain %d connections, got %d", nodeCount, len(snap.Conns))
+	}
+	for i, id := range ids {
+		conn := snap.Conns[i]
+		if conn.One != id {
+			t.Fatalf("expected conn[%d].One to be %s, got %s", id, conn.One)
+		}
+		peerID := ids[(i+1)%len(ids)]
+		if conn.Other != peerID {
+			t.Fatalf("expected conn[%d].Other to be %s, got %s", peerID, conn.Other)
+		}
+	}
+}
+
+func triggerChecks(ctx context.Context, ids []discover.NodeID, trigger chan discover.NodeID, interval time.Duration) {
+	tick := time.NewTicker(interval)
+	defer tick.Stop()
+	for {
+		select {
+		case <-tick.C:
+			for _, id := range ids {
+				select {
+				case trigger <- id:
+				case <-ctx.Done():
+					return
+				}
+			}
+		case <-ctx.Done():
+			return
+		}
+	}
+}

p2p/simulations/simulation.go

@@ -39,9 +39,18 @@ func (s *Simulation) Run(ctx context.Context, step *Step) (result *StepResult) {
 	}
 
 	// wait for all node expectations to either pass, error or timeout
-	for len(result.Passes) < len(step.Expect.Nodes) {
+	nodes := make(map[discover.NodeID]struct{}, len(step.Expect.Nodes))
+	for _, id := range step.Expect.Nodes {
+		nodes[id] = struct{}{}
+	}
+	for len(result.Passes) < len(nodes) {
 		select {
 		case id := <-step.Trigger:
+			// skip if we aren't checking the node
+			if _, ok := nodes[id]; !ok {
+				continue
+			}
+
 			// skip if the node has already passed
 			if _, ok := result.Passes[id]; ok {
 				continue