github.com/aergoio/aergo@v1.3.1/p2p/rolemanager.go (about)

     1  /*
     2   * @file
     3   * @copyright defined in aergo/LICENSE.txt
     4   */
     5  
     6  package p2p
     7  
     8  import (
     9  	"github.com/aergoio/aergo-lib/log"
    10  	"github.com/aergoio/aergo/p2p/p2pcommon"
    11  	"github.com/aergoio/aergo/p2p/p2putil"
    12  	"github.com/aergoio/aergo/types"
    13  	"strings"
    14  	"sync"
    15  )
    16  
    17  type RaftRoleManager struct {
    18  	p2ps      *P2P
    19  	logger    *log.Logger
    20  	raftBP    map[types.PeerID]bool
    21  	raftMutex sync.Mutex
    22  }
    23  
    24  func (rm *RaftRoleManager) UpdateBP(toAdd []types.PeerID, toRemove []types.PeerID) {
    25  	rm.raftMutex.Lock()
    26  	defer rm.raftMutex.Unlock()
    27  	changes := make([]p2pcommon.AttrModifier,0, len(toAdd)+len(toRemove))
    28  	for _, pid := range toRemove {
    29  		delete(rm.raftBP, pid)
    30  		changes = append(changes, p2pcommon.AttrModifier{pid, p2pcommon.Watcher})
    31  		rm.logger.Debug().Str(p2putil.LogPeerID, p2putil.ShortForm(pid)).Msg("raftBP removed")
    32  	}
    33  	for _, pid := range toAdd {
    34  		rm.raftBP[pid] = true
    35  		changes = append(changes, p2pcommon.AttrModifier{pid, p2pcommon.BlockProducer})
    36  		rm.logger.Debug().Str(p2putil.LogPeerID, p2putil.ShortForm(pid)).Msg("raftBP added")
    37  	}
    38  	rm.p2ps.pm.UpdatePeerRole(changes)
    39  }
    40  
    41  func (rm *RaftRoleManager) SelfRole() p2pcommon.PeerRole {
    42  	return rm.p2ps.selfRole
    43  }
    44  
    45  func (rm *RaftRoleManager) GetRole(pid types.PeerID) p2pcommon.PeerRole {
    46  	rm.raftMutex.Lock()
    47  	defer rm.raftMutex.Unlock()
    48  	if _, found := rm.raftBP[pid]; found {
    49  		// TODO check if leader or follower
    50  		return p2pcommon.BlockProducer
    51  	} else {
    52  		return p2pcommon.Watcher
    53  	}
    54  }
    55  
    56  func (rm *RaftRoleManager) NotifyNewBlockMsg(mo p2pcommon.MsgOrder, peers []p2pcommon.RemotePeer) (skipped, sent int) {
    57  	// TODO filter to only contain bp and trusted node.
    58  	for _, neighbor := range peers {
    59  		if neighbor != nil && neighbor.State() == types.RUNNING &&
    60  			neighbor.Role() == p2pcommon.Watcher {
    61  			sent++
    62  			neighbor.SendMessage(mo)
    63  		} else {
    64  			skipped++
    65  		}
    66  	}
    67  	return
    68  }
    69  
    70  type DefaultRoleManager struct {
    71  	p2ps *P2P
    72  }
    73  
    74  func (rm *DefaultRoleManager) UpdateBP(toAdd []types.PeerID, toRemove []types.PeerID) {
    75  	changes := make([]p2pcommon.AttrModifier,0, len(toAdd)+len(toRemove))
    76  	for _, pid := range toRemove {
    77  		changes = append(changes, p2pcommon.AttrModifier{pid, p2pcommon.Watcher})
    78  	}
    79  	for _, pid := range toAdd {
    80  		changes = append(changes, p2pcommon.AttrModifier{pid, p2pcommon.BlockProducer})
    81  	}
    82  	rm.p2ps.pm.UpdatePeerRole(changes)
    83  }
    84  
    85  func (rm *DefaultRoleManager) SelfRole() p2pcommon.PeerRole {
    86  	return rm.p2ps.selfRole
    87  }
    88  
    89  func (rm *DefaultRoleManager) GetRole(pid types.PeerID) p2pcommon.PeerRole {
    90  	prettyID := pid.Pretty()
    91  	bps := rm.p2ps.consacc.ConsensusInfo().Bps
    92  	for _, bp := range bps {
    93  		if strings.Contains(bp, prettyID) {
    94  			return p2pcommon.BlockProducer
    95  		}
    96  	}
    97  	return p2pcommon.Watcher
    98  }
    99  
   100  func (rm *DefaultRoleManager) NotifyNewBlockMsg(mo p2pcommon.MsgOrder, peers []p2pcommon.RemotePeer) (skipped, sent int) {
   101  	// TODO filter to only contain bp and trusted node.
   102  	for _, neighbor := range peers {
   103  		if neighbor != nil && neighbor.State() == types.RUNNING {
   104  			sent++
   105  			neighbor.SendMessage(mo)
   106  		} else {
   107  			skipped++
   108  		}
   109  	}
   110  	return
   111  }
   112