netsync/peer.go

   1 package netsync
   2
   3 import (
   4         "sync"
   5
   6         log "github.com/sirupsen/logrus"
   7         "gopkg.in/fatih/set.v0"
   8
   9         "github.com/bytom/errors"
  10         "github.com/bytom/p2p"
  11         "github.com/bytom/protocol/bc"
  12         "github.com/bytom/protocol/bc/types"
  13 )
  14
  15 var (
  16         errClosed            = errors.New("peer set is closed")
  17         errAlreadyRegistered = errors.New("peer is already registered")
  18         errNotRegistered     = errors.New("peer is not registered")
  19 )
  20
  21 const defaultVersion = 1
  22
  23 type peer struct {
  24         mtx     sync.RWMutex
  25         version int // Protocol version negotiated
  26         id      string
  27         height  uint64
  28         hash    *bc.Hash
  29         *p2p.Peer
  30
  31         knownTxs    *set.Set // Set of transaction hashes known to be known by this peer
  32         knownBlocks *set.Set // Set of block hashes known to be known by this peer
  33 }
  34
  35 func newPeer(height uint64, hash *bc.Hash, Peer *p2p.Peer) *peer {
  36         return &peer{
  37                 version:     defaultVersion,
  38                 height:      height,
  39                 hash:        hash,
  40                 Peer:        Peer,
  41                 knownTxs:    set.New(),
  42                 knownBlocks: set.New(),
  43         }
  44 }
  45
  46 func (p *peer) GetStatus() (height uint64, hash *bc.Hash) {
  47         p.mtx.RLock()
  48         defer p.mtx.RUnlock()
  49         return p.height, p.hash
  50 }
  51
  52 func (p *peer) SetStatus(height uint64, hash *bc.Hash) {
  53         p.mtx.Lock()
  54         defer p.mtx.Unlock()
  55
  56         p.height = height
  57         p.hash = hash
  58 }
  59
  60 func (p *peer) requestBlockByHash(hash *bc.Hash) error {
  61         msg := &BlockRequestMessage{RawHash: hash.Byte32()}
  62         p.Peer.TrySend(BlockchainChannel, struct{ BlockchainMessage }{msg})
  63         return nil
  64 }
  65
  66 func (p *peer) requestBlockByHeight(height uint64) error {
  67         msg := &BlockRequestMessage{Height: height}
  68         p.Peer.TrySend(BlockchainChannel, struct{ BlockchainMessage }{msg})
  69         return nil
  70 }
  71
  72 func (p *peer) SendTransactions(txs []*types.Tx) error {
  73         for _, tx := range txs {
  74                 msg, err := NewTransactionNotifyMessage(tx)
  75                 if err != nil {
  76                         return errors.New("Failed construction tx msg")
  77                 }
  78                 hash := &tx.ID
  79                 p.knownTxs.Add(hash.String())
  80                 p.Peer.TrySend(BlockchainChannel, struct{ BlockchainMessage }{msg})
  81         }
  82         return nil
  83 }
  84
  85 func (p *peer) getPeer() *p2p.Peer {
  86         p.mtx.RLock()
  87         defer p.mtx.RUnlock()
  88
  89         return p.Peer
  90 }
  91
  92 // MarkTransaction marks a transaction as known for the peer, ensuring that it
  93 // will never be propagated to this particular peer.
  94 func (p *peer) MarkTransaction(hash *bc.Hash) {
  95         p.mtx.Lock()
  96         defer p.mtx.Unlock()
  97
  98         // If we reached the memory allowance, drop a previously known transaction hash
  99         for p.knownTxs.Size() >= maxKnownTxs {
 100                 p.knownTxs.Pop()
 101         }
 102         p.knownTxs.Add(hash.String())
 103 }
 104
 105 // MarkBlock marks a block as known for the peer, ensuring that the block will
 106 // never be propagated to this particular peer.
 107 func (p *peer) MarkBlock(hash *bc.Hash) {
 108         p.mtx.Lock()
 109         defer p.mtx.Unlock()
 110
 111         // If we reached the memory allowance, drop a previously known block hash
 112         for p.knownBlocks.Size() >= maxKnownBlocks {
 113                 p.knownBlocks.Pop()
 114         }
 115         p.knownBlocks.Add(hash.String())
 116 }
 117
 118 type peerSet struct {
 119         peers  map[string]*peer
 120         lock   sync.RWMutex
 121         closed bool
 122 }
 123
 124 // newPeerSet creates a new peer set to track the active participants.
 125 func newPeerSet() *peerSet {
 126         return &peerSet{
 127                 peers: make(map[string]*peer),
 128         }
 129 }
 130
 131 // Register injects a new peer into the working set, or returns an error if the
 132 // peer is already known.
 133 func (ps *peerSet) Register(p *peer) error {
 134         ps.lock.Lock()
 135         defer ps.lock.Unlock()
 136
 137         if ps.closed {
 138                 return errClosed
 139         }
 140         if _, ok := ps.peers[p.id]; ok {
 141                 return errAlreadyRegistered
 142         }
 143         ps.peers[p.id] = p
 144         return nil
 145 }
 146
 147 // Unregister removes a remote peer from the active set, disabling any further
 148 // actions to/from that particular entity.
 149 func (ps *peerSet) Unregister(id string) error {
 150         ps.lock.Lock()
 151         defer ps.lock.Unlock()
 152
 153         if _, ok := ps.peers[id]; !ok {
 154                 return errNotRegistered
 155         }
 156         delete(ps.peers, id)
 157         return nil
 158 }
 159
 160 // Peer retrieves the registered peer with the given id.
 161 func (ps *peerSet) Peer(id string) *peer {
 162         ps.lock.RLock()
 163         defer ps.lock.RUnlock()
 164
 165         return ps.peers[id]
 166 }
 167
 168 // Len returns if the current number of peers in the set.
 169 func (ps *peerSet) Len() int {
 170         ps.lock.RLock()
 171         defer ps.lock.RUnlock()
 172
 173         return len(ps.peers)
 174 }
 175
 176 // MarkTransaction marks a transaction as known for the peer, ensuring that it
 177 // will never be propagated to this particular peer.
 178 func (ps *peerSet) MarkTransaction(peerID string, hash *bc.Hash) {
 179         ps.lock.RLock()
 180         defer ps.lock.RUnlock()
 181
 182         if peer, ok := ps.peers[peerID]; ok {
 183                 peer.MarkTransaction(hash)
 184         }
 185 }
 186
 187 // MarkBlock marks a block as known for the peer, ensuring that the block will
 188 // never be propagated to this particular peer.
 189 func (ps *peerSet) MarkBlock(peerID string, hash *bc.Hash) {
 190         ps.lock.RLock()
 191         defer ps.lock.RUnlock()
 192
 193         if peer, ok := ps.peers[peerID]; ok {
 194                 peer.MarkBlock(hash)
 195         }
 196 }
 197
 198 // PeersWithoutBlock retrieves a list of peers that do not have a given block in
 199 // their set of known hashes.
 200 func (ps *peerSet) PeersWithoutBlock(hash *bc.Hash) []*peer {
 201         ps.lock.RLock()
 202         defer ps.lock.RUnlock()
 203
 204         list := make([]*peer, 0, len(ps.peers))
 205         for _, p := range ps.peers {
 206                 if !p.knownBlocks.Has(hash.String()) {
 207                         list = append(list, p)
 208                 }
 209         }
 210         return list
 211 }
 212
 213 // PeersWithoutTx retrieves a list of peers that do not have a given transaction
 214 // in their set of known hashes.
 215 func (ps *peerSet) PeersWithoutTx(hash *bc.Hash) []*peer {
 216         ps.lock.RLock()
 217         defer ps.lock.RUnlock()
 218
 219         list := make([]*peer, 0, len(ps.peers))
 220         for _, p := range ps.peers {
 221                 if !p.knownTxs.Has(hash.String()) {
 222                         list = append(list, p)
 223                 }
 224         }
 225         return list
 226 }
 227
 228 // BestPeer retrieves the known peer with the currently highest total difficulty.
 229 func (ps *peerSet) BestPeer() (*p2p.Peer, uint64) {
 230         ps.lock.RLock()
 231         defer ps.lock.RUnlock()
 232
 233         var bestPeer *p2p.Peer
 234         var bestHeight uint64
 235
 236         for _, p := range ps.peers {
 237                 if bestPeer == nil || p.height > bestHeight {
 238                         bestPeer, bestHeight = p.Peer, p.height
 239                 }
 240         }
 241
 242         return bestPeer, bestHeight
 243 }
 244
 245 // Close disconnects all peers.
 246 // No new peers can be registered after Close has returned.
 247 func (ps *peerSet) Close() {
 248         ps.lock.Lock()
 249         defer ps.lock.Unlock()
 250
 251         for _, p := range ps.peers {
 252                 p.CloseConn()
 253         }
 254         ps.closed = true
 255 }
 256
 257 func (ps *peerSet) AddPeer(peer *p2p.Peer) {
 258         ps.lock.Lock()
 259         defer ps.lock.Unlock()
 260
 261         if _, ok := ps.peers[peer.Key]; !ok {
 262                 keeperPeer := newPeer(0, nil, peer)
 263                 ps.peers[peer.Key] = keeperPeer
 264                 log.WithFields(log.Fields{"ID": peer.Key}).Info("Add new peer to blockKeeper")
 265                 return
 266         }
 267         log.WithField("ID", peer.Key).Warning("Add existing peer to blockKeeper")
 268 }
 269
 270 func (ps *peerSet) RemovePeer(peerID string) {
 271         ps.lock.Lock()
 272         defer ps.lock.Unlock()
 273
 274         delete(ps.peers, peerID)
 275         log.WithField("ID", peerID).Info("Delete peer from peerset")
 276 }
 277
 278 func (ps *peerSet) SetPeerStatus(peerID string, height uint64, hash *bc.Hash) {
 279         ps.lock.Lock()
 280         defer ps.lock.Unlock()
 281
 282         if peer, ok := ps.peers[peerID]; ok {
 283                 peer.SetStatus(height, hash)
 284         }
 285 }
 286
 287 func (ps *peerSet) requestBlockByHash(peerID string, hash *bc.Hash) error {
 288         ps.lock.Lock()
 289         defer ps.lock.Unlock()
 290
 291         peer, ok := ps.peers[peerID]
 292         if !ok {
 293                 return errors.New("Can't find peer. ")
 294         }
 295         return peer.requestBlockByHash(hash)
 296 }
 297
 298 func (ps *peerSet) requestBlockByHeight(peerID string, height uint64) error {
 299         ps.lock.Lock()
 300         defer ps.lock.Unlock()
 301
 302         peer, ok := ps.peers[peerID]
 303         if !ok {
 304                 return errors.New("Can't find peer. ")
 305         }
 306         return peer.requestBlockByHeight(height)
 307 }
 308
 309 func (ps *peerSet) BroadcastMinedBlock(block *types.Block) error {
 310         msg, err := NewMinedBlockMessage(block)
 311         if err != nil {
 312                 return errors.New("Failed construction block msg")
 313         }
 314         hash := block.Hash()
 315         peers := ps.PeersWithoutBlock(&hash)
 316         for _, peer := range peers {
 317                 ps.MarkBlock(peer.Key, &hash)
 318                 peer.Send(BlockchainChannel, struct{ BlockchainMessage }{msg})
 319         }
 320         return nil
 321 }
 322
 323 func (ps *peerSet) BroadcastNewStatus(block *types.Block) error {
 324         return ps.BroadcastMinedBlock(block)
 325 }
 326
 327 func (ps *peerSet) BroadcastTx(tx *types.Tx) error {
 328         msg, err := NewTransactionNotifyMessage(tx)
 329         if err != nil {
 330                 return errors.New("Failed construction tx msg")
 331         }
 332         peers := ps.PeersWithoutTx(&tx.ID)
 333         for _, peer := range peers {
 334                 ps.peers[peer.Key].MarkTransaction(&tx.ID)
 335                 peer.Send(BlockchainChannel, struct{ BlockchainMessage }{msg})
 336         }
 337         return nil
 338 }