github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/p2p/node/internal/protocolPeerCache.go (about)

     1  package internal
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/libp2p/go-libp2p/core/event"
     8  	"github.com/libp2p/go-libp2p/core/host"
     9  	libp2pnet "github.com/libp2p/go-libp2p/core/network"
    10  	"github.com/libp2p/go-libp2p/core/peer"
    11  	"github.com/libp2p/go-libp2p/core/protocol"
    12  	"github.com/rs/zerolog"
    13  
    14  	p2plogging "github.com/onflow/flow-go/network/p2p/logging"
    15  )
    16  
    17  // ProtocolPeerCache store a mapping from protocol ID to peers who support that protocol
    18  type ProtocolPeerCache struct {
    19  	protocolPeers map[protocol.ID]map[peer.ID]struct{}
    20  	sync.RWMutex
    21  }
    22  
    23  func NewProtocolPeerCache(logger zerolog.Logger, h host.Host) (*ProtocolPeerCache, error) {
    24  	sub, err := h.EventBus().
    25  		Subscribe([]interface{}{new(event.EvtPeerIdentificationCompleted), new(event.EvtPeerProtocolsUpdated)})
    26  	if err != nil {
    27  		return nil, fmt.Errorf("could not subscribe to peer protocol update events: %w", err)
    28  	}
    29  	p := &ProtocolPeerCache{protocolPeers: make(map[protocol.ID]map[peer.ID]struct{})}
    30  	h.Network().Notify(&libp2pnet.NotifyBundle{
    31  		DisconnectedF: func(n libp2pnet.Network, c libp2pnet.Conn) {
    32  			peer := c.RemotePeer()
    33  			if len(n.ConnsToPeer(peer)) == 0 {
    34  				p.RemovePeer(peer)
    35  			}
    36  		},
    37  	})
    38  	go p.consumeSubscription(logger, h, sub)
    39  
    40  	return p, nil
    41  }
    42  
    43  func (p *ProtocolPeerCache) RemovePeer(peerID peer.ID) {
    44  	p.Lock()
    45  	defer p.Unlock()
    46  	for pid, peers := range p.protocolPeers {
    47  		delete(peers, peerID)
    48  		if len(peers) == 0 {
    49  			delete(p.protocolPeers, pid)
    50  		}
    51  	}
    52  }
    53  
    54  func (p *ProtocolPeerCache) AddProtocols(peerID peer.ID, protocols []protocol.ID) {
    55  	p.Lock()
    56  	defer p.Unlock()
    57  	for _, pid := range protocols {
    58  		peers, ok := p.protocolPeers[pid]
    59  		if !ok {
    60  			peers = make(map[peer.ID]struct{})
    61  			p.protocolPeers[pid] = peers
    62  		}
    63  		peers[peerID] = struct{}{}
    64  	}
    65  }
    66  
    67  func (p *ProtocolPeerCache) RemoveProtocols(peerID peer.ID, protocols []protocol.ID) {
    68  	p.Lock()
    69  	defer p.Unlock()
    70  	for _, pid := range protocols {
    71  		peers := p.protocolPeers[pid]
    72  		delete(peers, peerID)
    73  		if len(peers) == 0 {
    74  			delete(p.protocolPeers, pid)
    75  		}
    76  	}
    77  }
    78  
    79  func (p *ProtocolPeerCache) GetPeers(pid protocol.ID) map[peer.ID]struct{} {
    80  	p.RLock()
    81  	defer p.RUnlock()
    82  
    83  	// it is not safe to return a reference to the map, so we make a copy
    84  	peersCopy := make(map[peer.ID]struct{}, len(p.protocolPeers[pid]))
    85  	for peerID := range p.protocolPeers[pid] {
    86  		peersCopy[peerID] = struct{}{}
    87  	}
    88  	return peersCopy
    89  }
    90  
    91  func (p *ProtocolPeerCache) consumeSubscription(logger zerolog.Logger, h host.Host, sub event.Subscription) {
    92  	defer sub.Close()
    93  	logger.Debug().Msg("starting peer protocol event subscription loop")
    94  	for e := range sub.Out() {
    95  		logger.Debug().Interface("event", e).Msg("received new peer protocol event")
    96  		switch evt := e.(type) {
    97  		case event.EvtPeerIdentificationCompleted:
    98  			protocols, err := h.Peerstore().GetProtocols(evt.Peer)
    99  			if err != nil {
   100  				logger.Err(err).Str("peer_id", p2plogging.PeerId(evt.Peer)).Msg("failed to get protocols for peer")
   101  				continue
   102  			}
   103  			p.AddProtocols(evt.Peer, protocols)
   104  		case event.EvtPeerProtocolsUpdated:
   105  			p.AddProtocols(evt.Peer, evt.Added)
   106  			p.RemoveProtocols(evt.Peer, evt.Removed)
   107  		}
   108  	}
   109  	logger.Debug().Msg("exiting peer protocol event subscription loop")
   110  }