[management] move network map logic into new design (#4774)

management/internals/controllers/network_map/update_channel/updatechannel.go

@@ -0,0 +1,180 @@
+package update_channel
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/management/server/telemetry"
+)
+
+const channelBufferSize = 100
+
+type PeersUpdateManager struct {
+	// peerChannels is an update channel indexed by Peer.ID
+	peerChannels map[string]chan *network_map.UpdateMessage
+	// channelsMux keeps the mutex to access peerChannels
+	channelsMux *sync.RWMutex
+	// metrics provides method to collect application metrics
+	metrics telemetry.AppMetrics
+}
+
+var _ network_map.PeersUpdateManager = (*PeersUpdateManager)(nil)
+
+// NewPeersUpdateManager returns a new instance of PeersUpdateManager
+func NewPeersUpdateManager(metrics telemetry.AppMetrics) *PeersUpdateManager {
+	return &PeersUpdateManager{
+		peerChannels: make(map[string]chan *network_map.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 *network_map.UpdateMessage) {
+	start := time.Now()
+	var found, dropped bool
+
+	p.channelsMux.RLock()
+
+	defer func() {
+		p.channelsMux.RUnlock()
+		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 *network_map.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 *network_map.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)
+		return
+	}
+
+	log.WithContext(ctx).Debugf("closing updates channel: peer %s has no channel", 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.RLock()
+
+	m := make(map[string]struct{})
+
+	defer func() {
+		p.channelsMux.RUnlock()
+		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.RLock()
+
+	defer func() {
+		p.channelsMux.RUnlock()
+		if p.metrics != nil {
+			p.metrics.UpdateChannelMetrics().CountHasChannelDuration(time.Since(start))
+		}
+	}()
+
+	_, ok := p.peerChannels[peerID]
+
+	return ok
+}
+
+func (p *PeersUpdateManager) CountStreams() int {
+	p.channelsMux.RLock()
+	defer p.channelsMux.RUnlock()
+	return len(p.peerChannels)
+}

management/internals/controllers/network_map/update_channel/updatechannel_test.go

@@ -0,0 +1,80 @@
+package update_channel
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// var peersUpdater *PeersUpdateManager
+
+func TestCreateChannel(t *testing.T) {
+	peer := "test-create"
+	peersUpdater := NewPeersUpdateManager(nil)
+	defer peersUpdater.CloseChannel(context.Background(), peer)
+
+	_ = peersUpdater.CreateChannel(context.Background(), peer)
+	if _, ok := peersUpdater.peerChannels[peer]; !ok {
+		t.Error("Error creating the channel")
+	}
+}
+
+func TestSendUpdate(t *testing.T) {
+	peer := "test-sendupdate"
+	peersUpdater := NewPeersUpdateManager(nil)
+	update1 := &network_map.UpdateMessage{Update: &proto.SyncResponse{
+		NetworkMap: &proto.NetworkMap{
+			Serial: 0,
+		},
+	}}
+	_ = peersUpdater.CreateChannel(context.Background(), peer)
+	if _, ok := peersUpdater.peerChannels[peer]; !ok {
+		t.Error("Error creating the channel")
+	}
+	peersUpdater.SendUpdate(context.Background(), peer, update1)
+	select {
+	case <-peersUpdater.peerChannels[peer]:
+	default:
+		t.Error("Update wasn't send")
+	}
+
+	for range [channelBufferSize]int{} {
+		peersUpdater.SendUpdate(context.Background(), peer, update1)
+	}
+
+	update2 := &network_map.UpdateMessage{Update: &proto.SyncResponse{
+		NetworkMap: &proto.NetworkMap{
+			Serial: 10,
+		},
+	}}
+
+	peersUpdater.SendUpdate(context.Background(), peer, update2)
+	timeout := time.After(5 * time.Second)
+	for range [channelBufferSize]int{} {
+		select {
+		case <-timeout:
+			t.Error("timed out reading previously sent updates")
+		case updateReader := <-peersUpdater.peerChannels[peer]:
+			if updateReader.Update.NetworkMap.Serial == update2.Update.NetworkMap.Serial {
+				t.Error("got the update that shouldn't have been sent")
+			}
+		}
+	}
+
+}
+
+func TestCloseChannel(t *testing.T) {
+	peer := "test-close"
+	peersUpdater := NewPeersUpdateManager(nil)
+	_ = peersUpdater.CreateChannel(context.Background(), peer)
+	if _, ok := peersUpdater.peerChannels[peer]; !ok {
+		t.Error("Error creating the channel")
+	}
+	peersUpdater.CloseChannel(context.Background(), peer)
+	if _, ok := peersUpdater.peerChannels[peer]; ok {
+		t.Error("Error closing the channel")
+	}
+}