github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/util/peerset/peerset.go (about)

     1  package peerset
     2  
     3  import (
     4  	peer "github.com/ipfs/go-ipfs/p2p/peer"
     5  	"sync"
     6  )
     7  
     8  // PeerSet is a threadsafe set of peers
     9  type PeerSet struct {
    10  	ps   map[peer.ID]struct{}
    11  	lk   sync.RWMutex
    12  	size int
    13  }
    14  
    15  func New() *PeerSet {
    16  	ps := new(PeerSet)
    17  	ps.ps = make(map[peer.ID]struct{})
    18  	ps.size = -1
    19  	return ps
    20  }
    21  
    22  func NewLimited(size int) *PeerSet {
    23  	ps := new(PeerSet)
    24  	ps.ps = make(map[peer.ID]struct{})
    25  	ps.size = size
    26  	return ps
    27  }
    28  
    29  func (ps *PeerSet) Add(p peer.ID) {
    30  	ps.lk.Lock()
    31  	ps.ps[p] = struct{}{}
    32  	ps.lk.Unlock()
    33  }
    34  
    35  func (ps *PeerSet) Contains(p peer.ID) bool {
    36  	ps.lk.RLock()
    37  	_, ok := ps.ps[p]
    38  	ps.lk.RUnlock()
    39  	return ok
    40  }
    41  
    42  func (ps *PeerSet) Size() int {
    43  	ps.lk.RLock()
    44  	defer ps.lk.RUnlock()
    45  	return len(ps.ps)
    46  }
    47  
    48  // TryAdd Attempts to add the given peer into the set.
    49  // This operation can fail for one of two reasons:
    50  // 1) The given peer is already in the set
    51  // 2) The number of peers in the set is equal to size
    52  func (ps *PeerSet) TryAdd(p peer.ID) bool {
    53  	var success bool
    54  	ps.lk.Lock()
    55  	if _, ok := ps.ps[p]; !ok && (len(ps.ps) < ps.size || ps.size == -1) {
    56  		success = true
    57  		ps.ps[p] = struct{}{}
    58  	}
    59  	ps.lk.Unlock()
    60  	return success
    61  }