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