diff --git a/client/internal/conn_mgr.go b/client/internal/conn_mgr.go index c44986928..7f620209e 100644 --- a/client/internal/conn_mgr.go +++ b/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() } diff --git a/client/internal/lazyconn/manager/manager.go b/client/internal/lazyconn/manager/manager.go index cb9662bc9..6283e7092 100644 --- a/client/internal/lazyconn/manager/manager.go +++ b/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 + } + } +}