add p2p Service Peers method, fix libp2p Disconnect and improve tests

pkg/p2p/libp2p/libp2p.go

@@ -176,7 +176,7 @@ func New(ctx context.Context, o Options) (*Service, error) {
 			return
 		}
 
-		s.peers.add(peerID, i.Address)
+		s.peers.add(stream.Conn(), i.Address)
 		s.metrics.HandledStreamCount.Inc()
 		s.logger.Infof("peer %s connected", i.Address)
 	})
@@ -280,7 +280,7 @@ func (s *Service) Connect(ctx context.Context, addr ma.Multiaddr) (overlay swarm
 		return swarm.Address{}, err
 	}
 
-	s.peers.add(info.ID, i.Address)
+	s.peers.add(stream.Conn(), i.Address)
 	s.metrics.CreatedConnectionCount.Inc()
 	s.logger.Infof("peer %s connected", i.Address)
 	return i.Address, nil
@@ -302,6 +302,10 @@ func (s *Service) disconnect(peerID libp2ppeer.ID) error {
 	return nil
 }
 
+func (s *Service) Peers() []p2p.Peer {
+	return s.peers.peers()
+}
+
 func (s *Service) NewStream(ctx context.Context, overlay swarm.Address, protocolName, protocolVersion, streamName string) (p2p.Stream, error) {
 	peerID, found := s.peers.peerID(overlay)
 	if !found {

pkg/p2p/libp2p/peer.go

@@ -5,26 +5,31 @@
 package libp2p
 
 import (
+	"bytes"
+	"sort"
 	"sync"
 
+	"github.com/ethersphere/bee/pkg/p2p"
 	"github.com/ethersphere/bee/pkg/swarm"
 	"github.com/libp2p/go-libp2p-core/network"
 	libp2ppeer "github.com/libp2p/go-libp2p-core/peer"
 )
 
 type peerRegistry struct {
-	peers    map[string]libp2ppeer.ID        // map overlay address to underlay peer id
-	overlays map[libp2ppeer.ID]swarm.Address // map underlay peer id to overlay address
-	mu       sync.RWMutex
+	underlays   map[string]libp2ppeer.ID                    // map overlay address to underlay peer id
+	overlays    map[libp2ppeer.ID]swarm.Address             // map underlay peer id to overlay address
+	connections map[libp2ppeer.ID]map[network.Conn]struct{} // list of connections for safe removal on Disconnect notification
+	mu          sync.RWMutex
 
 	network.Notifiee // peerRegistry can be the receiver for network.Notify
 }
 
 func newPeerRegistry() *peerRegistry {
 	return &peerRegistry{
-		peers:    make(map[string]libp2ppeer.ID),
-		overlays: make(map[libp2ppeer.ID]swarm.Address),
-		Notifiee: new(network.NoopNotifiee),
+		underlays:   make(map[string]libp2ppeer.ID),
+		overlays:    make(map[libp2ppeer.ID]swarm.Address),
+		connections: make(map[libp2ppeer.ID]map[network.Conn]struct{}),
+		Notifiee:    new(network.NoopNotifiee),
 	}
 }
 
@@ -33,16 +38,61 @@ func (r *peerRegistry) Exists(overlay swarm.Address) (found bool) {
 	return found
 }
 
-func (r *peerRegistry) add(peerID libp2ppeer.ID, overlay swarm.Address) {
+// Disconnect removes the peer from registry in disconnect.
+// peerRegistry has to be set by network.Network.Notify().
+func (r *peerRegistry) Disconnected(_ network.Network, c network.Conn) {
+	peerID := c.RemotePeer()
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	// remove only the related connection,
+	// not eventually newly created one for the same peer
+	if _, ok := r.connections[peerID][c]; !ok {
+		return
+	}
+
+	overlay := r.overlays[peerID]
+	delete(r.overlays, peerID)
+	delete(r.underlays, encodeunderlaysKey(overlay))
+
+	delete(r.connections[peerID], c)
+	if len(r.connections[peerID]) == 0 {
+		delete(r.connections, peerID)
+	}
+}
+
+func (r *peerRegistry) peers() []p2p.Peer {
+	r.mu.Lock()
+	peers := make([]p2p.Peer, 0, len(r.overlays))
+	for _, a := range r.overlays {
+		peers = append(peers, p2p.Peer{
+			Address: a,
+		})
+	}
+	r.mu.Unlock()
+	sort.Slice(peers, func(i, j int) bool {
+		return bytes.Compare(peers[i].Address.Bytes(), peers[j].Address.Bytes()) == -1
+	})
+	return peers
+}
+
+func (r *peerRegistry) add(c network.Conn, overlay swarm.Address) {
+	peerID := c.RemotePeer()
+
 	r.mu.Lock()
-	r.peers[encodePeersKey(overlay)] = peerID
+	r.underlays[encodeunderlaysKey(overlay)] = peerID
 	r.overlays[peerID] = overlay
+	if _, ok := r.connections[peerID]; !ok {
+		r.connections[peerID] = make(map[network.Conn]struct{})
+	}
+	r.connections[peerID][c] = struct{}{}
 	r.mu.Unlock()
 }
 
 func (r *peerRegistry) peerID(overlay swarm.Address) (peerID libp2ppeer.ID, found bool) {
 	r.mu.RLock()
-	peerID, found = r.peers[encodePeersKey(overlay)]
+	peerID, found = r.underlays[encodeunderlaysKey(overlay)]
 	r.mu.RUnlock()
 	return peerID, found
 }
@@ -58,16 +108,11 @@ func (r *peerRegistry) remove(peerID libp2ppeer.ID) {
 	r.mu.Lock()
 	overlay := r.overlays[peerID]
 	delete(r.overlays, peerID)
-	delete(r.peers, encodePeersKey(overlay))
+	delete(r.underlays, encodeunderlaysKey(overlay))
+	delete(r.connections, peerID)
 	r.mu.Unlock()
 }
 
-func encodePeersKey(overlay swarm.Address) string {
+func encodeunderlaysKey(overlay swarm.Address) string {
 	return string(overlay.Bytes())
 }
-
-// Disconnect removes the peer from registry in disconnect.
-// peerRegistry has to be set by network.Network.Notify().
-func (r *peerRegistry) Disconnected(_ network.Network, c network.Conn) {
-	r.remove(c.RemotePeer())
-}

pkg/p2p/mock/mock.go

@@ -17,6 +17,7 @@ type Service struct {
 	addProtocolFunc func(p2p.ProtocolSpec) error
 	connectFunc     func(ctx context.Context, addr ma.Multiaddr) (overlay swarm.Address, err error)
 	disconnectFunc  func(overlay swarm.Address) error
+	peersFunc       func() []p2p.Peer
 }
 
 func WithAddProtocolFunc(f func(p2p.ProtocolSpec) error) Option {
@@ -37,6 +38,12 @@ func WithDisconnectFunc(f func(overlay swarm.Address) error) Option {
 	})
 }
 
+func WithPeersFunc(f func() []p2p.Peer) Option {
+	return optionFunc(func(s *Service) {
+		s.peersFunc = f
+	})
+}
+
 func New(opts ...Option) *Service {
 	s := new(Service)
 	for _, o := range opts {
@@ -47,25 +54,32 @@ func New(opts ...Option) *Service {
 
 func (s *Service) AddProtocol(spec p2p.ProtocolSpec) error {
 	if s.addProtocolFunc == nil {
-		return errors.New("AddProtocol function not configured")
+		return errors.New("function AddProtocol not configured")
 	}
 	return s.addProtocolFunc(spec)
 }
 
 func (s *Service) Connect(ctx context.Context, addr ma.Multiaddr) (overlay swarm.Address, err error) {
 	if s.connectFunc == nil {
-		return swarm.Address{}, errors.New("Connect function not configured")
+		return swarm.Address{}, errors.New("function Connect not configured")
 	}
 	return s.connectFunc(ctx, addr)
 }
 
 func (s *Service) Disconnect(overlay swarm.Address) error {
 	if s.disconnectFunc == nil {
-		return errors.New("Disconnect function not configured")
+		return errors.New("function Disconnect not configured")
 	}
 	return s.disconnectFunc(overlay)
 }
 
+func (s *Service) Peers() []p2p.Peer {
+	if s.peersFunc == nil {
+		return nil
+	}
+	return s.peersFunc()
+}
+
 type Option interface {
 	apply(*Service)
 }

pkg/p2p/p2p.go

@@ -16,6 +16,7 @@ type Service interface {
 	AddProtocol(ProtocolSpec) error
 	Connect(ctx context.Context, addr ma.Multiaddr) (overlay swarm.Address, err error)
 	Disconnect(overlay swarm.Address) error
+	Peers() []Peer
 }
 
 type Streamer interface {