6 log "github.com/sirupsen/logrus"
7 "gopkg.in/fatih/set.v0"
9 "github.com/bytom/errors"
10 "github.com/bytom/p2p"
11 "github.com/bytom/protocol/bc"
12 "github.com/bytom/protocol/bc/types"
16 errClosed = errors.New("peer set is closed")
17 errAlreadyRegistered = errors.New("peer is already registered")
18 errNotRegistered = errors.New("peer is not registered")
21 const defaultVersion = 1
25 version int // Protocol version negotiated
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
35 func newPeer(height uint64, hash *bc.Hash, Peer *p2p.Peer) *peer {
37 version: defaultVersion,
42 knownBlocks: set.New(),
46 func (p *peer) GetStatus() (height uint64, hash *bc.Hash) {
49 return p.height, p.hash
52 func (p *peer) SetStatus(height uint64, hash *bc.Hash) {
60 func (p *peer) requestBlockByHash(hash *bc.Hash) error {
61 msg := &BlockRequestMessage{RawHash: hash.Byte32()}
62 p.Peer.TrySend(BlockchainChannel, struct{ BlockchainMessage }{msg})
66 func (p *peer) requestBlockByHeight(height uint64) error {
67 msg := &BlockRequestMessage{Height: height}
68 p.Peer.TrySend(BlockchainChannel, struct{ BlockchainMessage }{msg})
72 func (p *peer) SendTransactions(txs []*types.Tx) error {
73 for _, tx := range txs {
74 msg, err := NewTransactionNotifyMessage(tx)
76 return errors.New("Failed construction tx msg")
79 p.knownTxs.Add(hash.String())
80 p.Peer.TrySend(BlockchainChannel, struct{ BlockchainMessage }{msg})
85 func (p *peer) getPeer() *p2p.Peer {
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) {
98 // If we reached the memory allowance, drop a previously known transaction hash
99 for p.knownTxs.Size() >= maxKnownTxs {
102 p.knownTxs.Add(hash.String())
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) {
111 // If we reached the memory allowance, drop a previously known block hash
112 for p.knownBlocks.Size() >= maxKnownBlocks {
115 p.knownBlocks.Add(hash.String())
118 type peerSet struct {
119 peers map[string]*peer
124 // newPeerSet creates a new peer set to track the active participants.
125 func newPeerSet() *peerSet {
127 peers: make(map[string]*peer),
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 {
135 defer ps.lock.Unlock()
140 if _, ok := ps.peers[p.id]; ok {
141 return errAlreadyRegistered
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 {
151 defer ps.lock.Unlock()
153 if _, ok := ps.peers[id]; !ok {
154 return errNotRegistered
160 // Peer retrieves the registered peer with the given id.
161 func (ps *peerSet) Peer(id string) *peer {
163 defer ps.lock.RUnlock()
168 // Len returns if the current number of peers in the set.
169 func (ps *peerSet) Len() int {
171 defer ps.lock.RUnlock()
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) {
180 defer ps.lock.RUnlock()
182 if peer, ok := ps.peers[peerID]; ok {
183 peer.MarkTransaction(hash)
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) {
191 defer ps.lock.RUnlock()
193 if peer, ok := ps.peers[peerID]; ok {
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 {
202 defer ps.lock.RUnlock()
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)
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 {
217 defer ps.lock.RUnlock()
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)
228 // BestPeer retrieves the known peer with the currently highest total difficulty.
229 func (ps *peerSet) BestPeer() (*p2p.Peer, uint64) {
231 defer ps.lock.RUnlock()
233 var bestPeer *p2p.Peer
234 var bestHeight uint64
236 for _, p := range ps.peers {
237 if bestPeer == nil || p.height > bestHeight {
238 bestPeer, bestHeight = p.Peer, p.height
242 return bestPeer, bestHeight
245 // Close disconnects all peers.
246 // No new peers can be registered after Close has returned.
247 func (ps *peerSet) Close() {
249 defer ps.lock.Unlock()
251 for _, p := range ps.peers {
257 func (ps *peerSet) AddPeer(peer *p2p.Peer) {
259 defer ps.lock.Unlock()
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")
267 log.WithField("ID", peer.Key).Warning("Add existing peer to blockKeeper")
270 func (ps *peerSet) RemovePeer(peerID string) {
272 defer ps.lock.Unlock()
274 delete(ps.peers, peerID)
275 log.WithField("ID", peerID).Info("Delete peer from peerset")
278 func (ps *peerSet) SetPeerStatus(peerID string, height uint64, hash *bc.Hash) {
280 defer ps.lock.Unlock()
282 if peer, ok := ps.peers[peerID]; ok {
283 peer.SetStatus(height, hash)
287 func (ps *peerSet) requestBlockByHash(peerID string, hash *bc.Hash) error {
289 defer ps.lock.Unlock()
291 peer, ok := ps.peers[peerID]
293 return errors.New("Can't find peer. ")
295 return peer.requestBlockByHash(hash)
298 func (ps *peerSet) requestBlockByHeight(peerID string, height uint64) error {
300 defer ps.lock.Unlock()
302 peer, ok := ps.peers[peerID]
304 return errors.New("Can't find peer. ")
306 return peer.requestBlockByHeight(height)
309 func (ps *peerSet) BroadcastMinedBlock(block *types.Block) error {
310 msg, err := NewMinedBlockMessage(block)
312 return errors.New("Failed construction block msg")
315 peers := ps.PeersWithoutBlock(&hash)
316 for _, peer := range peers {
317 ps.MarkBlock(peer.Key, &hash)
318 peer.Send(BlockchainChannel, struct{ BlockchainMessage }{msg})
323 func (ps *peerSet) BroadcastTx(tx *types.Tx) error {
324 msg, err := NewTransactionNotifyMessage(tx)
326 return errors.New("Failed construction tx msg")
328 peers := ps.PeersWithoutTx(&tx.ID)
329 for _, peer := range peers {
330 ps.peers[peer.Key].MarkTransaction(&tx.ID)
331 peer.Send(BlockchainChannel, struct{ BlockchainMessage }{msg})