diff --git a/network/p2p/p2p.go b/network/p2p/p2p.go index 8d5f3e4f38..4b4671c2c0 100644 --- a/network/p2p/p2p.go +++ b/network/p2p/p2p.go @@ -62,6 +62,8 @@ type Service interface { ID() peer.ID // return peer.ID for self IDSigner() *PeerIDChallengeSigner AddrInfo() peer.AddrInfo // return addrInfo for self + NetworkNotify(network.Notifiee) + NetworkStopNotify(network.Notifiee) DialPeersUntilTargetCount(targetConnCount int) ClosePeer(peer.ID) error @@ -227,6 +229,7 @@ func (s *serviceImpl) Start() error { // Close shuts down the P2P service func (s *serviceImpl) Close() error { + s.host.Network().StopNotify(s.streams) return s.host.Close() } @@ -424,3 +427,13 @@ func addressFilter(addrs []multiaddr.Multiaddr) []multiaddr.Multiaddr { func (s *serviceImpl) GetHTTPClient(addrInfo *peer.AddrInfo, connTimeStore limitcaller.ConnectionTimeStore, queueingTimeout time.Duration) (*http.Client, error) { return makeHTTPClientWithRateLimit(addrInfo, s, connTimeStore, queueingTimeout) } + +// NetworkNotify registers a notifiee with the host's network +func (s *serviceImpl) NetworkNotify(n network.Notifiee) { + s.host.Network().Notify(n) +} + +// NetworkStopNotify unregisters a notifiee with the host's network +func (s *serviceImpl) NetworkStopNotify(n network.Notifiee) { + s.host.Network().StopNotify(n) +} diff --git a/network/p2pNetwork_test.go b/network/p2pNetwork_test.go index d78b9a7573..556f99c238 100644 --- a/network/p2pNetwork_test.go +++ b/network/p2pNetwork_test.go @@ -375,6 +375,12 @@ func (s *mockService) GetHTTPClient(addrInfo *peer.AddrInfo, connTimeStore limit return nil, nil } +func (s *mockService) NetworkNotify(notifiee network.Notifiee) { +} + +func (s *mockService) NetworkStopNotify(notifiee network.Notifiee) { +} + func makeMockService(id peer.ID, addrs []ma.Multiaddr) *mockService { return &mockService{ id: id, @@ -1357,8 +1363,23 @@ func TestP2PEnableGossipService_BothDisable(t *testing.T) { relayCfg := cfg relayCfg.NetAddress = "127.0.0.1:0" + var netAConnected atomic.Bool + var netBConnected atomic.Bool + notifiee1 := &network.NotifyBundle{ + ConnectedF: func(n network.Network, c network.Conn) { + netAConnected.Store(true) + }, + } + notifiee2 := &network.NotifyBundle{ + ConnectedF: func(n network.Network, c network.Conn) { + netBConnected.Store(true) + }, + } + netA, err := NewP2PNetwork(log.With("net", "netA"), relayCfg, "", nil, genesisID, config.Devtestnet, &nopeNodeInfo{}, nil) require.NoError(t, err) + netA.service.NetworkNotify(notifiee1) + defer netA.service.NetworkStopNotify(notifiee1) netA.Start() defer netA.Stop() @@ -1374,11 +1395,13 @@ func TestP2PEnableGossipService_BothDisable(t *testing.T) { netB, err := NewP2PNetwork(log.With("net", "netB"), nodeCfg, "", phoneBookAddresses, genesisID, config.Devtestnet, &nopeNodeInfo{}, nil) require.NoError(t, err) + netB.service.NetworkNotify(notifiee2) + defer netB.service.NetworkStopNotify(notifiee2) netB.Start() defer netB.Stop() require.Eventually(t, func() bool { - return len(netA.service.Conns()) > 0 && len(netB.service.Conns()) > 0 + return netAConnected.Load() && netBConnected.Load() }, 1*time.Second, 50*time.Millisecond) require.False(t, netA.hasPeers())