mirror of
https://github.com/ipfs/kubo.git
synced 2026-02-28 22:08:01 +08:00
89f5cd4c94
I think it's time to move a lot of the peer-to-peer networking but-not-ipfs-specific things into its own package: p2p. This could in the future be split off into its own library. The first thing to go is the peer.
62 lines
1.1 KiB
Go
62 lines
1.1 KiB
Go
package peerset
|
|
|
|
import (
|
|
peer "github.com/jbenet/go-ipfs/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
|
|
}
|