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

     1  package connection
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	mrand "math/rand"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/rs/zerolog"
    11  
    12  	"github.com/koko1123/flow-go-1/module/component"
    13  	"github.com/koko1123/flow-go-1/module/irrecoverable"
    14  	"github.com/koko1123/flow-go-1/network/p2p"
    15  )
    16  
    17  // DefaultPeerUpdateInterval is default duration for which the peer manager waits in between attempts to update peer connections
    18  var DefaultPeerUpdateInterval = 10 * time.Minute
    19  
    20  var _ p2p.PeerManager = (*PeerManager)(nil)
    21  var _ component.Component = (*PeerManager)(nil)
    22  
    23  // PeerManager adds and removes connections to peers periodically and on request
    24  type PeerManager struct {
    25  	component.Component
    26  
    27  	logger             zerolog.Logger
    28  	peersProvider      p2p.PeersProvider // callback to retrieve list of peers to connect to
    29  	peerRequestQ       chan struct{}     // a channel to queue a peer update request
    30  	connector          p2p.Connector     // connector to connect or disconnect from peers
    31  	peerUpdateInterval time.Duration     // interval the peer manager runs on
    32  
    33  	peersProviderMu sync.RWMutex
    34  }
    35  
    36  // NewPeerManager creates a new peer manager which calls the peersProvider callback to get a list of peers to connect to
    37  // and it uses the connector to actually connect or disconnect from peers.
    38  func NewPeerManager(logger zerolog.Logger, updateInterval time.Duration, connector p2p.Connector) *PeerManager {
    39  	pm := &PeerManager{
    40  		logger:             logger,
    41  		connector:          connector,
    42  		peerRequestQ:       make(chan struct{}, 1),
    43  		peerUpdateInterval: updateInterval,
    44  	}
    45  
    46  	pm.Component = component.NewComponentManagerBuilder().
    47  		AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
    48  			// makes sure that peer update request is invoked once before returning
    49  			pm.RequestPeerUpdate()
    50  
    51  			ready()
    52  			pm.updateLoop(ctx)
    53  		}).
    54  		AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
    55  			ready()
    56  			pm.periodicLoop(ctx)
    57  		}).
    58  		Build()
    59  
    60  	return pm
    61  }
    62  
    63  // updateLoop triggers an update peer request when it has been requested
    64  func (pm *PeerManager) updateLoop(ctx irrecoverable.SignalerContext) {
    65  	for {
    66  		select {
    67  		case <-ctx.Done():
    68  			return
    69  		default:
    70  		}
    71  
    72  		select {
    73  		case <-ctx.Done():
    74  			return
    75  		case <-pm.peerRequestQ:
    76  			pm.updatePeers(ctx)
    77  		}
    78  	}
    79  }
    80  
    81  // periodicLoop periodically triggers an update peer request
    82  func (pm *PeerManager) periodicLoop(ctx irrecoverable.SignalerContext) {
    83  	// add a random delay to initial launch to avoid synchronizing this
    84  	// potentially expensive operation across the network
    85  	delay := time.Duration(mrand.Int63n(pm.peerUpdateInterval.Nanoseconds()))
    86  
    87  	ticker := time.NewTicker(pm.peerUpdateInterval)
    88  	defer ticker.Stop()
    89  
    90  	select {
    91  	case <-ctx.Done():
    92  		return
    93  	case <-time.After(delay):
    94  	}
    95  
    96  	for {
    97  		select {
    98  		case <-ctx.Done():
    99  			return
   100  		case <-ticker.C:
   101  			pm.RequestPeerUpdate()
   102  		}
   103  	}
   104  }
   105  
   106  // RequestPeerUpdate requests an update to the peer connections of this node.
   107  // If a peer update has already been requested (either as a periodic request or an on-demand request) and is outstanding,
   108  // then this call is a no-op.
   109  func (pm *PeerManager) RequestPeerUpdate() {
   110  	select {
   111  	case pm.peerRequestQ <- struct{}{}:
   112  	default:
   113  	}
   114  }
   115  
   116  // updatePeers updates the peers by connecting to all the nodes provided by the peersProvider callback and disconnecting from
   117  // previous nodes that are no longer in the new list of nodes.
   118  func (pm *PeerManager) updatePeers(ctx context.Context) {
   119  	pm.peersProviderMu.RLock()
   120  	defer pm.peersProviderMu.RUnlock()
   121  
   122  	if pm.peersProvider == nil {
   123  		pm.logger.Error().Msg("peers provider not set")
   124  		return
   125  	}
   126  
   127  	// get all the peer ids to connect to
   128  	peers := pm.peersProvider()
   129  
   130  	pm.logger.Trace().
   131  		Str("peers", fmt.Sprintf("%v", peers)).
   132  		Msg("connecting to peers")
   133  
   134  	// ask the connector to connect to all peers in the list
   135  	pm.connector.UpdatePeers(ctx, peers)
   136  }
   137  
   138  // ForceUpdatePeers initiates an update to the peer connections of this node immediately
   139  func (pm *PeerManager) ForceUpdatePeers(ctx context.Context) {
   140  	pm.updatePeers(ctx)
   141  }
   142  
   143  // SetPeersProvider sets the peers provider
   144  // SetPeersProvider may be called at most once
   145  func (pm *PeerManager) SetPeersProvider(peersProvider p2p.PeersProvider) {
   146  	pm.peersProviderMu.Lock()
   147  	defer pm.peersProviderMu.Unlock()
   148  
   149  	if pm.peersProvider != nil {
   150  		pm.logger.Fatal().Msg("peers provider already set")
   151  	}
   152  
   153  	pm.peersProvider = peersProvider
   154  }