swarm: snapshot load improvement

p2p/simulations/network.go

@@ -20,7 +20,9 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"strings"
 	"sync"
 	"time"
 
@@ -707,6 +709,7 @@ func (net *Network) snapshot(addServices []string, removeServices []string) (*Sn
 
 // Load loads a network snapshot
 func (net *Network) Load(snap *Snapshot) error {
+	// Start nodes.
 	for _, n := range snap.Nodes {
 		if _, err := net.NewNodeWithConfig(n.Node.Config); err != nil {
 			return err
@@ -718,17 +721,102 @@ func (net *Network) Load(snap *Snapshot) error {
 			return err
 		}
 	}
+
+	// Prepare connection events counter.
+	allConnected := make(chan struct{})     // closed when all connections are established
+	eventLoopStarted := make(chan struct{}) // ensures that event loop is started before it is closed
+	done := make(chan struct{})             // ensures that the event loop goroutine is terminated
+	defer close(done)
+
+	go func() {
+		// Subscribe to event channel.
+		events := make(chan *Event)
+		sub := net.Events().Subscribe(events)
+		defer sub.Unsubscribe()
+
+		// Expected number of connections.
+		total := len(snap.Conns)
+		// Set of all established connections from the snapshot, not other connections.
+		// Key array element 0 is the connection One field value, and element 1 connection Other field.
+		connections := make(map[[2]enode.ID]struct{}, total)
+
+		// once is a closed channel that is read in the event loop below
+		// only once.
+		// It ensures that eventLoopStarted is closed which signals that
+		// it is safe to call connect method on the network without the
+		// possibility to miss a few first connection events.
+		once := make(chan struct{})
+		// Close once channel so that it can be read from in the event loop.
+		close(once)
+
+		for {
+			select {
+			case e := <-events:
+				// Detect only connection events.
+				if e.Type != EventTypeConn {
+					continue
+				}
+				connection := [2]enode.ID{e.Conn.One, e.Conn.Other}
+				// Nodes are still not connected or have been disconnected.
+				if !e.Conn.Up {
+					// Delete the connection from the set of established connections.
+					// This will prevent false positive in case disconnections happen.
+					delete(connections, connection)
+					continue
+				}
+				// Check that the connection is from the snapshot.
+				for _, conn := range snap.Conns {
+					if conn.One == e.Conn.One && conn.Other == e.Conn.Other {
+						// Add the connection to the set of established connections.
+						connections[connection] = struct{}{}
+						if len(connections) == total {
+							// Signal that all nodes are connected.
+							close(allConnected)
+							return
+						}
+
+						break
+					}
+				}
+			case <-once:
+				// Set once to nil as nil channel never blocks forever.
+				// This ensures that this for loop never gets into this part
+				// of the code again.
+				once = nil
+				// Proceed with connecting the nodes, as we are ready to
+				// detect events.
+				close(eventLoopStarted)
+			case <-done:
+				// Load function returned, terminate this goroutine.
+				return
+			}
+		}
+	}()
+
+	// Do not proceed until the goroutine with the event loop actually is ready
+	// to receive events.
+	<-eventLoopStarted
+
+	// Start connecting.
 	for _, conn := range snap.Conns {
 
 		if !net.GetNode(conn.One).Up || !net.GetNode(conn.Other).Up {
 			//in this case, at least one of the nodes of a connection is not up,
 			//so it would result in the snapshot `Load` to fail
 			continue
 		}
-		if err := net.Connect(conn.One, conn.Other); err != nil {
+		if err := net.Connect(conn.One, conn.Other); err != nil && !strings.Contains(err.Error(), "already connected") {
 			return err
 		}
 	}
+
+	select {
+	// Wait until all connections from the snapshot are established.
+	case <-allConnected:
+	// Make sure that we do not wait forever.
+	case <-time.After(120 * time.Second):
+		return errors.New("snapshot connections not established")
+	}
 	return nil
 }
 

swarm/network/hive.go

@@ -165,8 +165,8 @@ func (h *Hive) Run(p *BzzPeer) error {
 			// otherwise just send depth to new peer
 			dp.NotifyDepth(depth)
 		}
+		NotifyPeer(p.BzzAddr, h.Kademlia)
 	}
-	NotifyPeer(p.BzzAddr, h.Kademlia)
 	defer h.Off(dp)
 	return dp.Run(dp.HandleMsg)
 }

swarm/network/simulation/kademlia_test.go

@@ -33,7 +33,6 @@ func TestWaitTillHealthy(t *testing.T) {
 		"bzz": func(ctx *adapters.ServiceContext, b *sync.Map) (node.Service, func(), error) {
 			addr := network.NewAddr(ctx.Config.Node())
 			hp := network.NewHiveParams()
-			hp.Discovery = false
 			config := &network.BzzConfig{
 				OverlayAddr:  addr.Over(),
 				UnderlayAddr: addr.Under(),