@@ -20,6 +20,7 @@ import (
20
20
"bytes"
21
21
"context"
22
22
"encoding/json"
23
+ "errors"
23
24
"fmt"
24
25
"sync"
25
26
"time"
@@ -705,8 +706,11 @@ func (net *Network) snapshot(addServices []string, removeServices []string) (*Sn
705
706
return snap , nil
706
707
}
707
708
709
+ var snapshotLoadTimeout = 120 * time .Second
710
+
708
711
// Load loads a network snapshot
709
712
func (net * Network ) Load (snap * Snapshot ) error {
713
+ // Start nodes.
710
714
for _ , n := range snap .Nodes {
711
715
if _ , err := net .NewNodeWithConfig (n .Node .Config ); err != nil {
712
716
return err
@@ -718,6 +722,69 @@ func (net *Network) Load(snap *Snapshot) error {
718
722
return err
719
723
}
720
724
}
725
+
726
+ // Prepare connection events counter.
727
+ allConnected := make (chan struct {}) // closed when all connections are established
728
+ done := make (chan struct {}) // ensures that the event loop goroutine is terminated
729
+ defer close (done )
730
+
731
+ // Subscribe to event channel.
732
+ // It needs to be done outside of the event loop goroutine (created below)
733
+ // to ensure that the event channel is blocking before connect calls are made.
734
+ events := make (chan * Event )
735
+ sub := net .Events ().Subscribe (events )
736
+ defer sub .Unsubscribe ()
737
+
738
+ go func () {
739
+ // Expected number of connections.
740
+ total := len (snap .Conns )
741
+ // Set of all established connections from the snapshot, not other connections.
742
+ // Key array element 0 is the connection One field value, and element 1 connection Other field.
743
+ connections := make (map [[2 ]enode.ID ]struct {}, total )
744
+
745
+ for {
746
+ select {
747
+ case e := <- events :
748
+ // Ignore control events as they do not represent
749
+ // connect or disconnect (Up) state change.
750
+ if e .Control {
751
+ continue
752
+ }
753
+ // Detect only connection events.
754
+ if e .Type != EventTypeConn {
755
+ continue
756
+ }
757
+ connection := [2 ]enode.ID {e .Conn .One , e .Conn .Other }
758
+ // Nodes are still not connected or have been disconnected.
759
+ if ! e .Conn .Up {
760
+ // Delete the connection from the set of established connections.
761
+ // This will prevent false positive in case disconnections happen.
762
+ delete (connections , connection )
763
+ log .Warn ("load snapshot: unexpected disconnection" , "one" , e .Conn .One , "other" , e .Conn .Other )
764
+ continue
765
+ }
766
+ // Check that the connection is from the snapshot.
767
+ for _ , conn := range snap .Conns {
768
+ if conn .One == e .Conn .One && conn .Other == e .Conn .Other {
769
+ // Add the connection to the set of established connections.
770
+ connections [connection ] = struct {}{}
771
+ if len (connections ) == total {
772
+ // Signal that all nodes are connected.
773
+ close (allConnected )
774
+ return
775
+ }
776
+
777
+ break
778
+ }
779
+ }
780
+ case <- done :
781
+ // Load function returned, terminate this goroutine.
782
+ return
783
+ }
784
+ }
785
+ }()
786
+
787
+ // Start connecting.
721
788
for _ , conn := range snap .Conns {
722
789
723
790
if ! net .GetNode (conn .One ).Up || ! net .GetNode (conn .Other ).Up {
@@ -729,6 +796,14 @@ func (net *Network) Load(snap *Snapshot) error {
729
796
return err
730
797
}
731
798
}
799
+
800
+ select {
801
+ // Wait until all connections from the snapshot are established.
802
+ case <- allConnected :
803
+ // Make sure that we do not wait forever.
804
+ case <- time .After (snapshotLoadTimeout ):
805
+ return errors .New ("snapshot connections not established" )
806
+ }
732
807
return nil
733
808
}
734
809
0 commit comments