Handle env var

client/internal/conn_mgr.go

@@ -16,6 +16,12 @@ const (
 	envDisableLazyConn = "NB_LAZY_CONN_DISABLE"
 )
 
+// ConnMgr coordinates both lazy connections (established on-demand) and permanent peer connections.
+//
+// The connection manager is responsible for:
+// - Managing lazy connections via the lazyConnManager
+// - Maintaining a list of excluded peers that should always have permanent connections
+// - Handling connection establishment based on peer signaling
 type ConnMgr struct {
 	lazyConnMgr *lazyConnManager.Manager
 	peerStore   *peerstore.Store
@@ -24,19 +30,25 @@ type ConnMgr struct {
 }
 
 func NewConnMgr(peerStore *peerstore.Store, iface lazyconn.WGIface) *ConnMgr {
+	var lazyConnMgr *lazyConnManager.Manager
+	if os.Getenv(envDisableLazyConn) != "true" {
+		lazyConnMgr = lazyConnManager.NewManager(iface)
+	}
+
 	e := &ConnMgr{
 		peerStore:   peerStore,
-		lazyConnMgr: lazyConnManager.NewManager(iface),
+		lazyConnMgr: lazyConnMgr,
 	}
 	return e
 }
 
 func (e *ConnMgr) Start(ctx context.Context) {
-	if os.Getenv(envDisableLazyConn) == "true" {
+	if e.lazyConnMgr == nil {
 		log.Infof("lazy connection manager is disabled")
 		return
 	}
-	go e.lazyConnMgr.Start()
+
+	e.lazyConnMgr.Start(ctx)
 	go e.receiveLazyConnEvents(ctx)
 }
 
@@ -49,6 +61,11 @@ func (e *ConnMgr) AddPeerConn(peerKey string, conn *peer.Conn) (exists bool) {
 		return true
 	}
 
+	if e.lazyConnMgr == nil {
+		conn.Open()
+		return
+	}
+
 	_, exists = e.excludes[peerKey]
 	if exists {
 		conn.Open()
@@ -72,6 +89,10 @@ func (e *ConnMgr) OnSignalMsg(peerKey string) (*peer.Conn, bool) {
 		return nil, false
 	}
 
+	if e.lazyConnMgr == nil {
+		return conn, true
+	}
+
 	if ok := e.lazyConnMgr.RemovePeer(peerKey); ok {
 		conn.Open()
 	}
@@ -84,10 +105,17 @@ func (e *ConnMgr) RemovePeerConn(peerKey string) {
 		conn.Close()
 	}
 
+	if e.lazyConnMgr == nil {
+		return
+	}
+
 	e.lazyConnMgr.RemovePeer(peerKey)
 }
 
 func (e *ConnMgr) Close() {
+	if e.lazyConnMgr == nil {
+		return
+	}
 	// todo wait for receiveLazyConnEvents to finish
 	e.lazyConnMgr.Close()
 }

client/internal/lazyconn/manager/manager.go

@@ -10,12 +10,17 @@ import (
 	"github.com/netbirdio/netbird/client/internal/lazyconn/listener"
 )
 
+// Manager manages lazy connections
+// This is not a thread safe implementation, do not call exported functions concurrently
 type Manager struct {
 	PeerActivityChan chan string
 
-	listenerMgr  *listener.Manager
-	managedPeers map[string]lazyconn.PeerConfig
-	mu           sync.Mutex
+	listenerMgr    *listener.Manager
+	managedPeers   map[string]lazyconn.PeerConfig
+	managedPeersMu sync.Mutex
+
+	ctxCancel context.CancelFunc
+	wg        sync.WaitGroup
 }
 
 func NewManager(wgIface lazyconn.WGIface) *Manager {
@@ -27,38 +32,28 @@ func NewManager(wgIface lazyconn.WGIface) *Manager {
 	return m
 }
 
-func (m *Manager) Start() {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+// Start to listen for traffic start events
+func (m *Manager) Start(parentCtx context.Context) {
+	ctx, cancel := context.WithCancel(parentCtx)
+	m.ctxCancel = cancel
 
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case peerID := <-m.listenerMgr.TrafficStartChan:
-			m.mu.Lock()
-			_, ok := m.managedPeers[peerID]
-			if !ok {
-				log.Debugf("lazy peer is not managed: %s", peerID)
-				m.mu.Unlock()
-				continue
-			}
+	m.wg.Add(1)
 
-			//m.watcher.AddPeer(peerID)
-			log.Infof("lazy peer is active: %s", peerID)
-			m.notifyPeerAction(ctx, peerID)
-			m.mu.Unlock()
-		}
-	}
+	go func() {
+		defer m.wg.Done()
+		defer cancel()
+		m.receiveLazyConnEvents(ctx)
+	}()
 }
 
 func (m *Manager) AddPeer(peer lazyconn.PeerConfig) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
+	m.managedPeersMu.Lock()
+	defer m.managedPeersMu.Unlock()
 
 	log.Debugf("adding lazy peer: %s", peer.PublicKey)
 
 	if _, ok := m.managedPeers[peer.PublicKey]; ok {
+		log.Warnf("peer already managed: %s", peer.PublicKey)
 		return nil
 	}
 
@@ -71,8 +66,8 @@ func (m *Manager) AddPeer(peer lazyconn.PeerConfig) error {
 }
 
 func (m *Manager) RemovePeer(peerID string) bool {
-	m.mu.Lock()
-	defer m.mu.Unlock()
+	m.managedPeersMu.Lock()
+	defer m.managedPeersMu.Unlock()
 
 	if _, ok := m.managedPeers[peerID]; !ok {
 		return false
@@ -85,12 +80,31 @@ func (m *Manager) RemovePeer(peerID string) bool {
 	return false
 }
 
+// Close the manager and all the listeners
+// block until all routine are done and cleanup the exported Channels
 func (m *Manager) Close() {
-	m.mu.Lock()
-	defer m.mu.Unlock()
+	m.managedPeersMu.Lock()
+	defer m.managedPeersMu.Unlock()
+
+	m.ctxCancel()
 
 	m.listenerMgr.Close()
+	m.wg.Wait()
 	m.managedPeers = make(map[string]lazyconn.PeerConfig)
+
+	// clean up the channel for the future reuse
+	m.drainPeerActivityChan()
+}
+
+func (m *Manager) receiveLazyConnEvents(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case peerID := <-m.listenerMgr.TrafficStartChan:
+			m.notifyPeerAction(ctx, peerID)
+		}
+	}
 }
 
 func (m *Manager) notifyPeerAction(ctx context.Context, peerID string) {
@@ -99,3 +113,13 @@ func (m *Manager) notifyPeerAction(ctx context.Context, peerID string) {
 	case m.PeerActivityChan <- peerID:
 	}
 }
+
+func (m *Manager) drainPeerActivityChan() {
+	for {
+		select {
+		case <-m.PeerActivityChan:
+		default:
+			return
+		}
+	}
+}