github.com/arieschain/arieschain@v0.0.0-20191023063405-37c074544356/p2p/testing/peerpool.go (about)

     1  package testing
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"github.com/quickchainproject/quickchain/log"
     8  	"github.com/quickchainproject/quickchain/p2p/discover"
     9  )
    10  
    11  type TestPeer interface {
    12  	ID() discover.NodeID
    13  	Drop(error)
    14  }
    15  
    16  // TestPeerPool is an example peerPool to demonstrate registration of peer connections
    17  type TestPeerPool struct {
    18  	lock  sync.Mutex
    19  	peers map[discover.NodeID]TestPeer
    20  }
    21  
    22  func NewTestPeerPool() *TestPeerPool {
    23  	return &TestPeerPool{peers: make(map[discover.NodeID]TestPeer)}
    24  }
    25  
    26  func (self *TestPeerPool) Add(p TestPeer) {
    27  	self.lock.Lock()
    28  	defer self.lock.Unlock()
    29  	log.Trace(fmt.Sprintf("pp add peer  %v", p.ID()))
    30  	self.peers[p.ID()] = p
    31  
    32  }
    33  
    34  func (self *TestPeerPool) Remove(p TestPeer) {
    35  	self.lock.Lock()
    36  	defer self.lock.Unlock()
    37  	delete(self.peers, p.ID())
    38  }
    39  
    40  func (self *TestPeerPool) Has(id discover.NodeID) bool {
    41  	self.lock.Lock()
    42  	defer self.lock.Unlock()
    43  	_, ok := self.peers[id]
    44  	return ok
    45  }
    46  
    47  func (self *TestPeerPool) Get(id discover.NodeID) TestPeer {
    48  	self.lock.Lock()
    49  	defer self.lock.Unlock()
    50  	return self.peers[id]
    51  }