package server import ( "context" "sync" "time" log "github.com/sirupsen/logrus" "github.com/netbirdio/netbird/management/proto" "github.com/netbirdio/netbird/management/server/telemetry" ) const channelBufferSize = 100 type UpdateMessage struct { Update *proto.SyncResponse NetworkMap *NetworkMap } type PeersUpdateManager struct { // peerChannels is an update channel indexed by Peer.ID peerChannels map[string]chan *UpdateMessage // channelsMux keeps the mutex to access peerChannels channelsMux *sync.RWMutex // metrics provides method to collect application metrics metrics telemetry.AppMetrics } // NewPeersUpdateManager returns a new instance of PeersUpdateManager func NewPeersUpdateManager(metrics telemetry.AppMetrics) *PeersUpdateManager { return &PeersUpdateManager{ peerChannels: make(map[string]chan *UpdateMessage), channelsMux: &sync.RWMutex{}, metrics: metrics, } } // SendUpdate sends update message to the peer's channel func (p *PeersUpdateManager) SendUpdate(ctx context.Context, peerID string, update *UpdateMessage) { start := time.Now() var found, dropped bool p.channelsMux.Lock() defer func() { p.channelsMux.Unlock() if p.metrics != nil { p.metrics.UpdateChannelMetrics().CountSendUpdateDuration(time.Since(start), found, dropped) } }() if channel, ok := p.peerChannels[peerID]; ok { found = true select { case channel <- update: log.WithContext(ctx).Debugf("update was sent to channel for peer %s", peerID) default: dropped = true log.WithContext(ctx).Warnf("channel for peer %s is %d full or closed", peerID, len(channel)) } } else { log.WithContext(ctx).Debugf("peer %s has no channel", peerID) } } // CreateChannel creates a go channel for a given peer used to deliver updates relevant to the peer. func (p *PeersUpdateManager) CreateChannel(ctx context.Context, peerID string) chan *UpdateMessage { start := time.Now() closed := false p.channelsMux.Lock() defer func() { p.channelsMux.Unlock() if p.metrics != nil { p.metrics.UpdateChannelMetrics().CountCreateChannelDuration(time.Since(start), closed) } }() if channel, ok := p.peerChannels[peerID]; ok { closed = true delete(p.peerChannels, peerID) close(channel) } // mbragin: todo shouldn't it be more? or configurable? channel := make(chan *UpdateMessage, channelBufferSize) p.peerChannels[peerID] = channel log.WithContext(ctx).Debugf("opened updates channel for a peer %s", peerID) return channel } func (p *PeersUpdateManager) closeChannel(ctx context.Context, peerID string) { if channel, ok := p.peerChannels[peerID]; ok { delete(p.peerChannels, peerID) close(channel) } log.WithContext(ctx).Debugf("closed updates channel of a peer %s", peerID) } // CloseChannels closes updates channel for each given peer func (p *PeersUpdateManager) CloseChannels(ctx context.Context, peerIDs []string) { start := time.Now() p.channelsMux.Lock() defer func() { p.channelsMux.Unlock() if p.metrics != nil { p.metrics.UpdateChannelMetrics().CountCloseChannelsDuration(time.Since(start), len(peerIDs)) } }() for _, id := range peerIDs { p.closeChannel(ctx, id) } } // CloseChannel closes updates channel of a given peer func (p *PeersUpdateManager) CloseChannel(ctx context.Context, peerID string) { start := time.Now() p.channelsMux.Lock() defer func() { p.channelsMux.Unlock() if p.metrics != nil { p.metrics.UpdateChannelMetrics().CountCloseChannelDuration(time.Since(start)) } }() p.closeChannel(ctx, peerID) } // GetAllConnectedPeers returns a copy of the connected peers map func (p *PeersUpdateManager) GetAllConnectedPeers() map[string]struct{} { start := time.Now() p.channelsMux.Lock() m := make(map[string]struct{}) defer func() { p.channelsMux.Unlock() if p.metrics != nil { p.metrics.UpdateChannelMetrics().CountGetAllConnectedPeersDuration(time.Since(start), len(m)) } }() for ID := range p.peerChannels { m[ID] = struct{}{} } return m } // HasChannel returns true if peers has channel in update manager, otherwise false func (p *PeersUpdateManager) HasChannel(peerID string) bool { start := time.Now() p.channelsMux.Lock() defer func() { p.channelsMux.Unlock() if p.metrics != nil { p.metrics.UpdateChannelMetrics().CountHasChannelDuration(time.Since(start)) } }() _, ok := p.peerChannels[peerID] return ok }