github.com/koko1123/flow-go-1@v0.29.6/network/p2p/connection/connManager.go (about)

     1  package connection
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/libp2p/go-libp2p/core/connmgr"
     7  	"github.com/libp2p/go-libp2p/core/network"
     8  	"github.com/libp2p/go-libp2p/core/peer"
     9  	"github.com/multiformats/go-multiaddr"
    10  	"github.com/rs/zerolog"
    11  
    12  	"github.com/koko1123/flow-go-1/module"
    13  )
    14  
    15  // ConnManager provides an implementation of Libp2p's ConnManager interface (https://godoc.org/github.com/libp2p/go-libp2p-core/connmgr#ConnManager)
    16  // It is called back by libp2p when certain events occur such as opening/closing a stream, opening/closing connection etc.
    17  // This implementation updates networking metrics when a peer connection is added or removed
    18  type ConnManager struct {
    19  	connmgr.NullConnMgr                  // a null conn mgr provided by libp2p to allow implementing only the functions needed
    20  	n                   network.Notifiee // the notifiee callback provided by libp2p
    21  	log                 zerolog.Logger   // logger to log connection, stream and other statistics about libp2p
    22  	metrics             module.LibP2PConnectionMetrics
    23  
    24  	plk       sync.RWMutex
    25  	protected map[peer.ID]map[string]struct{}
    26  }
    27  
    28  func NewConnManager(log zerolog.Logger, metrics module.LibP2PConnectionMetrics) *ConnManager {
    29  	cn := &ConnManager{
    30  		log:         log.With().Str("component", "conn_manager").Logger(),
    31  		NullConnMgr: connmgr.NullConnMgr{},
    32  		metrics:     metrics,
    33  		protected:   make(map[peer.ID]map[string]struct{}),
    34  	}
    35  	n := &network.NotifyBundle{ListenCloseF: cn.ListenCloseNotifee,
    36  		ListenF:       cn.ListenNotifee,
    37  		ConnectedF:    cn.Connected,
    38  		DisconnectedF: cn.Disconnected}
    39  	cn.n = n
    40  
    41  	return cn
    42  }
    43  
    44  func (cm *ConnManager) Notifee() network.Notifiee {
    45  	return cm.n
    46  }
    47  
    48  // ListenNotifee is called by libp2p when network starts listening on an addr
    49  func (cm *ConnManager) ListenNotifee(n network.Network, m multiaddr.Multiaddr) {
    50  	cm.log.Debug().Str("multiaddress", m.String()).Msg("listen started")
    51  }
    52  
    53  // called by libp2p when network stops listening on an addr
    54  // * This is never called back by libp2p currently and may be a bug on their side
    55  func (cm *ConnManager) ListenCloseNotifee(n network.Network, m multiaddr.Multiaddr) {
    56  	// just log the multiaddress  on which we listen
    57  	cm.log.Debug().Str("multiaddress", m.String()).Msg("listen stopped ")
    58  }
    59  
    60  // Connected is called by libp2p when a connection opened
    61  func (cm *ConnManager) Connected(n network.Network, con network.Conn) {
    62  	cm.logConnectionUpdate(n, con, "connection established")
    63  	cm.updateConnectionMetric(n)
    64  }
    65  
    66  // Disconnected is called by libp2p when a connection closed
    67  func (cm *ConnManager) Disconnected(n network.Network, con network.Conn) {
    68  	cm.logConnectionUpdate(n, con, "connection removed")
    69  	cm.updateConnectionMetric(n)
    70  }
    71  
    72  func (cm *ConnManager) updateConnectionMetric(n network.Network) {
    73  	var totalInbound uint = 0
    74  	var totalOutbound uint = 0
    75  
    76  	for _, conn := range n.Conns() {
    77  		switch conn.Stat().Direction {
    78  		case network.DirInbound:
    79  			totalInbound++
    80  		case network.DirOutbound:
    81  			totalOutbound++
    82  		}
    83  	}
    84  
    85  	cm.metrics.InboundConnections(totalInbound)
    86  	cm.metrics.OutboundConnections(totalOutbound)
    87  }
    88  
    89  func (cm *ConnManager) logConnectionUpdate(n network.Network, con network.Conn, logMsg string) {
    90  	cm.log.Debug().
    91  		Str("remote_peer", con.RemotePeer().String()).
    92  		Str("remote_addrs", con.RemoteMultiaddr().String()).
    93  		Str("local_peer", con.LocalPeer().String()).
    94  		Str("local_addrs", con.LocalMultiaddr().String()).
    95  		Str("direction", con.Stat().Direction.String()).
    96  		Int("total_connections", len(n.Conns())).
    97  		Msg(logMsg)
    98  }
    99  
   100  func (cm *ConnManager) Protect(id peer.ID, tag string) {
   101  	cm.plk.Lock()
   102  	defer cm.plk.Unlock()
   103  
   104  	tags, ok := cm.protected[id]
   105  	if !ok {
   106  		tags = make(map[string]struct{}, 2)
   107  		cm.protected[id] = tags
   108  	}
   109  	tags[tag] = struct{}{}
   110  }
   111  
   112  func (cm *ConnManager) Unprotect(id peer.ID, tag string) (protected bool) {
   113  	cm.plk.Lock()
   114  	defer cm.plk.Unlock()
   115  
   116  	tags, ok := cm.protected[id]
   117  	if !ok {
   118  		return false
   119  	}
   120  	if delete(tags, tag); len(tags) == 0 {
   121  		delete(cm.protected, id)
   122  		return false
   123  	}
   124  	return true
   125  }
   126  
   127  func (cm *ConnManager) IsProtected(id peer.ID, tag string) (protected bool) {
   128  	cm.plk.RLock()
   129  	defer cm.plk.RUnlock()
   130  
   131  	tags, ok := cm.protected[id]
   132  	if !ok {
   133  		return false
   134  	}
   135  
   136  	if tag == "" {
   137  		return true
   138  	}
   139  
   140  	_, protected = tags[tag]
   141  	return protected
   142  }