diff options
Diffstat (limited to 'swarm/network/simulation/events.go')
| -rw-r--r-- | swarm/network/simulation/events.go | 11 |
1 files changed, 11 insertions, 0 deletions
diff --git a/swarm/network/simulation/events.go b/swarm/network/simulation/events.go index f9cfadb73..980a9a756 100644 --- a/swarm/network/simulation/events.go +++ b/swarm/network/simulation/events.go @@ -18,6 +18,7 @@ package simulation import ( "context" + "sync" "github.com/ethereum/go-ethereum/p2p/discover" @@ -71,24 +72,32 @@ func (f *PeerEventsFilter) MsgCode(c uint64) *PeerEventsFilter { func (s *Simulation) PeerEvents(ctx context.Context, ids []discover.NodeID, filters ...*PeerEventsFilter) <-chan PeerEvent { eventC := make(chan PeerEvent) + // wait group to make sure all subscriptions to admin peerEvents are established + // before this function returns. + var subsWG sync.WaitGroup for _, id := range ids { s.shutdownWG.Add(1) + subsWG.Add(1) go func(id discover.NodeID) { defer s.shutdownWG.Done() client, err := s.Net.GetNode(id).Client() if err != nil { + subsWG.Done() eventC <- PeerEvent{NodeID: id, Error: err} return } events := make(chan *p2p.PeerEvent) sub, err := client.Subscribe(ctx, "admin", events, "peerEvents") if err != nil { + subsWG.Done() eventC <- PeerEvent{NodeID: id, Error: err} return } defer sub.Unsubscribe() + subsWG.Done() + for { select { case <-ctx.Done(): @@ -153,5 +162,7 @@ func (s *Simulation) PeerEvents(ctx context.Context, ids []discover.NodeID, filt }(id) } + // wait all subscriptions + subsWG.Wait() return eventC } |