1package peer 2 3import ( 4 "sync" 5) 6 7// PeerSet is a threadsafe set of peers. 8type Set struct { 9 lk sync.RWMutex 10 ps map[ID]struct{} 11 12 size int 13} 14 15func NewSet() *Set { 16 ps := new(Set) 17 ps.ps = make(map[ID]struct{}) 18 ps.size = -1 19 return ps 20} 21 22func NewLimitedSet(size int) *Set { 23 ps := new(Set) 24 ps.ps = make(map[ID]struct{}) 25 ps.size = size 26 return ps 27} 28 29func (ps *Set) Add(p ID) { 30 ps.lk.Lock() 31 ps.ps[p] = struct{}{} 32 ps.lk.Unlock() 33} 34 35func (ps *Set) Contains(p ID) bool { 36 ps.lk.RLock() 37 _, ok := ps.ps[p] 38 ps.lk.RUnlock() 39 return ok 40} 41 42func (ps *Set) 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 52func (ps *Set) TryAdd(p 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} 62 63func (ps *Set) Peers() []ID { 64 ps.lk.Lock() 65 out := make([]ID, 0, len(ps.ps)) 66 for p, _ := range ps.ps { 67 out = append(out, p) 68 } 69 ps.lk.Unlock() 70 return out 71} 72