p2p/simulations, swarm/network: Custom services in snapshot (#17991)

* p2p/simulations: Add custom services to simnodes + remove sim down conn objs

* p2p/simulation, swarm/network: Add selective services to discovery sim

* p2p/simulations, swarm/network: Remove useless comments

* p2p/simulations, swarm/network: Clean up mess from rebase

* p2p/simulation: Add sleep to prevent connect flakiness in http test

* p2p/simulations: added concurrent goroutines to prevent sleeps on simulation connect/disconnect

* p2p/simulations, swarm/network/simulations: address pr comments

* reinstated dummy service

* fixed http snapshot test

3 files changed, 118 insertions, 12 deletions

diff --git a/p2p/simulations/http_test.go b/p2p/simulations/http_test.go
index 7bdbad1b5..d9513caaa 100644
--- a/p2p/simulations/http_test.go
+++ b/p2p/simulations/http_test.go
@@ -18,6 +18,7 @@ package simulations
 
 import (
 	"context"
+	"flag"
 	"fmt"
 	"math/rand"
 	"net/http/httptest"
@@ -28,13 +29,26 @@ import (
 	"time"
 
 	"github.com/ethereum/go-ethereum/event"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/node"
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
 	"github.com/ethereum/go-ethereum/rpc"
+	colorable "github.com/mattn/go-colorable"
 )
 
+var (
+	loglevel = flag.Int("loglevel", 2, "verbosity of logs")
+)
+
+func init() {
+	flag.Parse()
+
+	log.PrintOrigins(true)
+	log.Root().SetHandler(log.LvlFilterHandler(log.Lvl(*loglevel), log.StreamHandler(colorable.NewColorableStderr(), log.TerminalFormat(true))))
+}
+
 // testService implements the node.Service interface and provides protocols
 // and APIs which are useful for testing nodes in a simulation network
 type testService struct {
@@ -584,9 +598,26 @@ func TestHTTPNodeRPC(t *testing.T) {
 // TestHTTPSnapshot tests creating and loading network snapshots
 func TestHTTPSnapshot(t *testing.T) {
 	// start the server
-	_, s := testHTTPServer(t)
+	network, s := testHTTPServer(t)
 	defer s.Close()
 
+	var eventsDone = make(chan struct{})
+	count := 1
+	eventsDoneChan := make(chan *Event)
+	eventSub := network.Events().Subscribe(eventsDoneChan)
+	go func() {
+		defer eventSub.Unsubscribe()
+		for event := range eventsDoneChan {
+			if event.Type == EventTypeConn && !event.Control {
+				count--
+				if count == 0 {
+					eventsDone <- struct{}{}
+					return
+				}
+			}
+		}
+	}()
+
 	// create a two-node network
 	client := NewClient(s.URL)
 	nodeCount := 2
@@ -620,7 +651,7 @@ func TestHTTPSnapshot(t *testing.T) {
 		}
 		states[i] = state
 	}
-
+	<-eventsDone
 	// create a snapshot
 	snap, err := client.CreateSnapshot()
 	if err != nil {
@@ -634,9 +665,23 @@ func TestHTTPSnapshot(t *testing.T) {
 	}
 
 	// create another network
-	_, s = testHTTPServer(t)
+	network2, s := testHTTPServer(t)
 	defer s.Close()
 	client = NewClient(s.URL)
+	count = 1
+	eventSub = network2.Events().Subscribe(eventsDoneChan)
+	go func() {
+		defer eventSub.Unsubscribe()
+		for event := range eventsDoneChan {
+			if event.Type == EventTypeConn && !event.Control {
+				count--
+				if count == 0 {
+					eventsDone <- struct{}{}
+					return
+				}
+			}
+		}
+	}()
 
 	// subscribe to events so we can check them later
 	events := make(chan *Event, 100)
@@ -651,6 +696,7 @@ func TestHTTPSnapshot(t *testing.T) {
 	if err := client.LoadSnapshot(snap); err != nil {
 		t.Fatalf("error loading snapshot: %s", err)
 	}
+	<-eventsDone
 
 	// check the nodes and connection exists
 	net, err := client.GetNetwork()
@@ -676,6 +722,9 @@ func TestHTTPSnapshot(t *testing.T) {
 	if conn.Other.String() != nodes[1].ID {
 		t.Fatalf("expected connection to have other=%q, got other=%q", nodes[1].ID, conn.Other)
 	}
+	if !conn.Up {
+		t.Fatal("should be up")
+	}
 
 	// check the node states were restored
 	for i, node := range nodes {
diff --git a/p2p/simulations/network.go b/p2p/simulations/network.go
index 200015ff3..92ccfde81 100644
--- a/p2p/simulations/network.go
+++ b/p2p/simulations/network.go
@@ -644,11 +644,18 @@ type NodeSnapshot struct {
 
 // Snapshot creates a network snapshot
 func (net *Network) Snapshot() (*Snapshot, error) {
+	return net.snapshot(nil, nil)
+}
+
+func (net *Network) SnapshotWithServices(addServices []string, removeServices []string) (*Snapshot, error) {
+	return net.snapshot(addServices, removeServices)
+}
+
+func (net *Network) snapshot(addServices []string, removeServices []string) (*Snapshot, error) {
 	net.lock.Lock()
 	defer net.lock.Unlock()
 	snap := &Snapshot{
 		Nodes: make([]NodeSnapshot, len(net.Nodes)),
-		Conns: make([]Conn, len(net.Conns)),
 	}
 	for i, node := range net.Nodes {
 		snap.Nodes[i] = NodeSnapshot{Node: *node}
@@ -660,9 +667,40 @@ func (net *Network) Snapshot() (*Snapshot, error) {
 			return nil, err
 		}
 		snap.Nodes[i].Snapshots = snapshots
+		for _, addSvc := range addServices {
+			haveSvc := false
+			for _, svc := range snap.Nodes[i].Node.Config.Services {
+				if svc == addSvc {
+					haveSvc = true
+					break
+				}
+			}
+			if !haveSvc {
+				snap.Nodes[i].Node.Config.Services = append(snap.Nodes[i].Node.Config.Services, addSvc)
+			}
+		}
+		if len(removeServices) > 0 {
+			var cleanedServices []string
+			for _, svc := range snap.Nodes[i].Node.Config.Services {
+				haveSvc := false
+				for _, rmSvc := range removeServices {
+					if rmSvc == svc {
+						haveSvc = true
+						break
+					}
+				}
+				if !haveSvc {
+					cleanedServices = append(cleanedServices, svc)
+				}
+
+			}
+			snap.Nodes[i].Node.Config.Services = cleanedServices
+		}
 	}
-	for i, conn := range net.Conns {
-		snap.Conns[i] = *conn
+	for _, conn := range net.Conns {
+		if conn.Up {
+			snap.Conns = append(snap.Conns, *conn)
+		}
 	}
 	return snap, nil
 }
diff --git a/swarm/network/simulations/discovery/discovery_test.go b/swarm/network/simulations/discovery/discovery_test.go
index 3c3affe58..cd5456b73 100644
--- a/swarm/network/simulations/discovery/discovery_test.go
+++ b/swarm/network/simulations/discovery/discovery_test.go
@@ -85,11 +85,12 @@ func getDbStore(nodeID string) (*state.DBStore, error) {
 }
 
 var (
-	nodeCount    = flag.Int("nodes", 10, "number of nodes to create (default 10)")
-	initCount    = flag.Int("conns", 1, "number of originally connected peers	 (default 1)")
-	snapshotFile = flag.String("snapshot", "", "create snapshot")
-	loglevel     = flag.Int("loglevel", 3, "verbosity of logs")
-	rawlog       = flag.Bool("rawlog", false, "remove terminal formatting from logs")
+	nodeCount       = flag.Int("nodes", 10, "number of nodes to create (default 10)")
+	initCount       = flag.Int("conns", 1, "number of originally connected peers	 (default 1)")
+	snapshotFile    = flag.String("snapshot", "", "path to create snapshot file in")
+	loglevel        = flag.Int("loglevel", 3, "verbosity of logs")
+	rawlog          = flag.Bool("rawlog", false, "remove terminal formatting from logs")
+	serviceOverride = flag.String("services", "", "remove or add services to the node snapshot; prefix with \"+\" to add, \"-\" to remove; example: +pss,-discovery")
 )
 
 func init() {
@@ -306,7 +307,25 @@ func discoverySimulation(nodes, conns int, adapter adapters.NodeAdapter) (*simul
 	}
 
 	if *snapshotFile != "" {
-		snap, err := net.Snapshot()
+		var err error
+		var snap *simulations.Snapshot
+		if len(*serviceOverride) > 0 {
+			var addServices []string
+			var removeServices []string
+			for _, osvc := range strings.Split(*serviceOverride, ",") {
+				if strings.Index(osvc, "+") == 0 {
+					addServices = append(addServices, osvc[1:])
+				} else if strings.Index(osvc, "-") == 0 {
+					removeServices = append(removeServices, osvc[1:])
+				} else {
+					panic("stick to the rules, you know what they are")
+				}
+			}
+			snap, err = net.SnapshotWithServices(addServices, removeServices)
+		} else {
+			snap, err = net.Snapshot()
+		}
+
 		if err != nil {
 			return nil, errors.New("no shapshot dude")
 		}