mirror of
https://github.com/netbirdio/netbird.git
synced 2026-04-16 15:26:40 +00:00
120 lines
3.0 KiB
Go
120 lines
3.0 KiB
Go
package peer
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
|
|
"errors"
|
|
|
|
log "github.com/sirupsen/logrus"
|
|
|
|
"github.com/netbirdio/netbird/shared/signal/proto"
|
|
"github.com/netbirdio/netbird/signal/metrics"
|
|
)
|
|
|
|
var (
|
|
ErrPeerAlreadyRegistered = errors.New("peer already registered")
|
|
)
|
|
|
|
// Peer representation of a connected Peer
|
|
type Peer struct {
|
|
// a unique id of the Peer (e.g. sha256 fingerprint of the Wireguard public key)
|
|
Id string
|
|
|
|
StreamID int64
|
|
|
|
// a gRpc connection stream to the Peer
|
|
Stream proto.SignalExchange_ConnectStreamServer
|
|
|
|
// registration time
|
|
RegisteredAt time.Time
|
|
|
|
Cancel context.CancelFunc
|
|
}
|
|
|
|
// NewPeer creates a new instance of a connected Peer
|
|
func NewPeer(id string, stream proto.SignalExchange_ConnectStreamServer, cancel context.CancelFunc) *Peer {
|
|
return &Peer{
|
|
Id: id,
|
|
Stream: stream,
|
|
StreamID: time.Now().UnixNano(),
|
|
RegisteredAt: time.Now(),
|
|
Cancel: cancel,
|
|
}
|
|
}
|
|
|
|
// Registry that holds all currently connected Peers
|
|
type Registry struct {
|
|
// Peer.key -> Peer
|
|
Peers sync.Map
|
|
// regMutex ensures that registration and de-registrations are safe
|
|
regMutex sync.Mutex
|
|
metrics *metrics.AppMetrics
|
|
}
|
|
|
|
// NewRegistry creates a new connected Peer registry
|
|
func NewRegistry(metrics *metrics.AppMetrics) *Registry {
|
|
return &Registry{
|
|
regMutex: sync.Mutex{},
|
|
metrics: metrics,
|
|
}
|
|
}
|
|
|
|
// Get gets a peer from the registry
|
|
func (registry *Registry) Get(peerId string) (*Peer, bool) {
|
|
if load, ok := registry.Peers.Load(peerId); ok {
|
|
return load.(*Peer), ok
|
|
}
|
|
return nil, false
|
|
|
|
}
|
|
|
|
func (registry *Registry) IsPeerRegistered(peerId string) bool {
|
|
if _, ok := registry.Peers.Load(peerId); ok {
|
|
return ok
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Register registers peer in the registry
|
|
func (registry *Registry) Register(peer *Peer) error {
|
|
start := time.Now()
|
|
|
|
// can be that peer already exists, but it is fine (e.g. reconnect)
|
|
p, loaded := registry.Peers.LoadOrStore(peer.Id, peer)
|
|
if loaded {
|
|
pp := p.(*Peer)
|
|
if peer.StreamID > pp.StreamID {
|
|
log.Tracef("peer [%s] is already registered [new streamID %d, previous StreamID %d]. Will override stream.",
|
|
peer.Id, peer.StreamID, pp.StreamID)
|
|
if swapped := registry.Peers.CompareAndSwap(peer.Id, pp, peer); !swapped {
|
|
return registry.Register(peer)
|
|
}
|
|
pp.Cancel()
|
|
log.Debugf("peer re-registered [%s]", peer.Id)
|
|
return nil
|
|
}
|
|
return ErrPeerAlreadyRegistered
|
|
}
|
|
|
|
log.Debugf("peer registered [%s]", peer.Id)
|
|
registry.metrics.ActivePeers.Add(context.Background(), 1)
|
|
|
|
// record time as milliseconds
|
|
registry.metrics.RegistrationDelay.Record(context.Background(), float64(time.Since(start).Nanoseconds())/1e6)
|
|
|
|
registry.metrics.Registrations.Add(context.Background(), 1)
|
|
|
|
return nil
|
|
}
|
|
|
|
// Deregister Peer from the Registry (usually once it disconnects)
|
|
func (registry *Registry) Deregister(peer *Peer) {
|
|
if deleted := registry.Peers.CompareAndDelete(peer.Id, peer); deleted {
|
|
registry.metrics.ActivePeers.Add(context.Background(), -1)
|
|
log.Debugf("peer deregistered [%s]", peer.Id)
|
|
registry.metrics.Deregistrations.Add(context.Background(), 1)
|
|
}
|
|
}
|