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