[client, management] Support per-peer lazy connection state and default proxy peers to lazy

client/internal/conn_mgr.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"maps"
 	"os"
 	"strconv"
 	"sync"
@@ -14,6 +15,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/peer"
 	"github.com/netbirdio/netbird/client/internal/peerstore"
 	"github.com/netbirdio/netbird/route"
+	mgmProto "github.com/netbirdio/netbird/shared/management/proto"
 )
 
 // lazyForce is the resolved local decision for lazy connections, layered above the
@@ -40,8 +42,14 @@ type ConnMgr struct {
 	iface            lazyconn.WGIface
 	force            lazyForce
 	rosenpassEnabled bool
+	// remoteLazyEnabled caches the account-wide lazy feature flag from management.
+	// It is the default for peers that do not carry a per-peer lazy hint.
+	remoteLazyEnabled bool
 
 	lazyConnMgr *manager.Manager
+	// appliedExcludeList is the exclude set last handed to the lazy manager, kept so an
+	// unchanged set on the next sync skips the O(n) reconciliation.
+	appliedExcludeList map[string]bool
 
 	wg            sync.WaitGroup
 	lazyCtx       context.Context
@@ -59,69 +67,56 @@ func NewConnMgr(engineConfig *EngineConfig, statusRecorder *peer.Status, peerSto
 	return e
 }
 
-// Start initializes the connection manager. It starts the lazy connection manager when a
-// local override forces it on; with no local override it waits for the management feature flag.
+// Start initializes the connection manager. The lazy connection manager always runs so that
+// per-peer lazy defaults (e.g. proxy peers) work even when the account flag is off; the
+// account flag and the local override decide the default lazy state per peer (see
+// PeerLazyDefault). Rosenpass is the only condition that disables it.
 func (e *ConnMgr) Start(ctx context.Context) {
 	if e.lazyConnMgr != nil {
 		log.Errorf("lazy connection manager is already started")
 		return
 	}
 
-	switch e.force {
-	case lazyForceOff:
-		log.Infof("lazy connection manager is disabled by local override (%s or MDM policy)", lazyconn.EnvLazyConn)
-		e.statusRecorder.UpdateLazyConnection(false)
-		return
-	case lazyForceNone:
-		log.Infof("lazy connection manager is managed by the management feature flag")
-		e.statusRecorder.UpdateLazyConnection(false)
-		return
-	}
-
 	if e.rosenpassEnabled {
-		log.Warnf("rosenpass connection manager is enabled, lazy connection manager will not be started")
+		log.Warnf("rosenpass is enabled, lazy connection manager will not be started")
 		e.statusRecorder.UpdateLazyConnection(false)
 		return
 	}
 
 	e.initLazyManager(ctx)
-	e.statusRecorder.UpdateLazyConnection(true)
+	e.statusRecorder.UpdateLazyConnection(e.PeerLazyDefault(mgmProto.LazyState_LazyStateDefault))
 }
 
-// UpdatedRemoteFeatureFlag is called when the remote feature flag is updated.
-// If enabled, it initializes the lazy connection manager and start it. Do not need to call Start() again.
-// If disabled, then it closes the lazy connection manager and open the connections to all peers.
-func (e *ConnMgr) UpdatedRemoteFeatureFlag(ctx context.Context, enabled bool) error {
-	// a local override (NB_LAZY_CONN or local config) takes precedence over management
-	if e.force != lazyForceNone {
-		return nil
+// UpdatedRemoteFeatureFlag caches the account-wide lazy feature flag. The manager itself is
+// not started or stopped here; the per-sync exclude-list reconciliation moves normal peers
+// between the lazy and always-active sets when the flag flips.
+func (e *ConnMgr) UpdatedRemoteFeatureFlag(_ context.Context, enabled bool) error {
+	e.remoteLazyEnabled = enabled
+	if e.isStartedWithLazyMgr() {
+		e.statusRecorder.UpdateLazyConnection(e.PeerLazyDefault(mgmProto.LazyState_LazyStateDefault))
 	}
+	return nil
+}
 
-	if enabled {
-		// if the lazy connection manager is already started, do not start it again
-		if e.lazyConnMgr != nil {
-			return nil
-		}
-
-		if e.rosenpassEnabled {
-			log.Infof("rosenpass connection manager is enabled, lazy connection manager will not be started")
-			e.statusRecorder.UpdateLazyConnection(false)
-			return nil
-		}
+// PeerLazyDefault reports whether a peer should be lazy. The local override
+// (NB_LAZY_CONN/MDM) wins over everything; without a local override the
+// management per-peer state applies (LazyStateLazy/Eager force the decision),
+// and LazyStateDefault follows the account-wide flag.
+func (e *ConnMgr) PeerLazyDefault(state mgmProto.LazyState) bool {
+	switch e.force {
+	case lazyForceOn:
+		return true
+	case lazyForceOff:
+		return false
+	}
 
-		log.Infof("lazy connection manager is enabled by the management feature flag")
-		e.initLazyManager(ctx)
-		e.statusRecorder.UpdateLazyConnection(true)
-		return e.addPeersToLazyConnManager()
-	} else {
-		if e.lazyConnMgr == nil {
-			e.statusRecorder.UpdateLazyConnection(false)
-			return nil
-		}
-		log.Infof("lazy connection manager is disabled by management feature flag")
-		e.closeManager(ctx)
-		e.statusRecorder.UpdateLazyConnection(false)
-		return nil
+	switch state {
+	case mgmProto.LazyState_LazyStateLazy:
+		return true
+	case mgmProto.LazyState_LazyStateEager:
+		return false
+	default:
+		return e.remoteLazyEnabled
 	}
 }
 
@@ -141,6 +136,13 @@ func (e *ConnMgr) SetExcludeList(ctx context.Context, peerIDs map[string]bool) {
 		return
 	}
 
+	// The exclude set is recomputed every sync but rarely changes; skip the O(n)
+	// store lookups and reconciliation when it matches what was already applied.
+	if maps.Equal(peerIDs, e.appliedExcludeList) {
+		return
+	}
+	e.appliedExcludeList = maps.Clone(peerIDs)
+
 	excludedPeers := make([]lazyconn.PeerConfig, 0, len(peerIDs))
 
 	for peerID := range peerIDs {
@@ -176,12 +178,16 @@ func (e *ConnMgr) SetExcludeList(ctx context.Context, peerIDs map[string]bool) {
 	}
 }
 
-func (e *ConnMgr) AddPeerConn(ctx context.Context, peerKey string, conn *peer.Conn) (exists bool) {
+// AddPeerConn registers a peer connection. permanent requests an always-active connection
+// (the peer belongs to the exclude set: a forwarder, or a peer that is not lazy by policy).
+// Non-permanent peers are handed to the lazy manager. The subsequent SetExcludeList call
+// reconciles membership for existing peers across flag flips.
+func (e *ConnMgr) AddPeerConn(ctx context.Context, peerKey string, conn *peer.Conn, permanent bool) (exists bool) {
 	if success := e.peerStore.AddPeerConn(peerKey, conn); !success {
 		return true
 	}
 
-	if !e.isStartedWithLazyMgr() {
+	if !e.isStartedWithLazyMgr() || permanent {
 		if err := conn.Open(ctx); err != nil {
 			conn.Log.Errorf("failed to open connection: %v", err)
 		}
@@ -269,13 +275,15 @@ func (e *ConnMgr) Close() {
 	e.lazyCtxCancel()
 	e.wg.Wait()
 	e.lazyConnMgr = nil
+	e.appliedExcludeList = nil
 }
 
 func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
 	cfg := manager.Config{
 		InactivityThreshold: inactivityThresholdEnv(),
 	}
 	e.lazyConnMgr = manager.NewManager(cfg, engineCtx, e.peerStore, e.iface)
+	e.appliedExcludeList = nil
 
 	e.lazyCtx, e.lazyCtxCancel = context.WithCancel(engineCtx)
 
@@ -286,43 +294,6 @@ func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
 	}()
 }
 
-func (e *ConnMgr) addPeersToLazyConnManager() error {
-	peers := e.peerStore.PeersPubKey()
-	lazyPeerCfgs := make([]lazyconn.PeerConfig, 0, len(peers))
-	for _, peerID := range peers {
-		var peerConn *peer.Conn
-		var exists bool
-		if peerConn, exists = e.peerStore.PeerConn(peerID); !exists {
-			log.Warnf("failed to find peer conn for peerID: %s", peerID)
-			continue
-		}
-
-		lazyPeerCfg := lazyconn.PeerConfig{
-			PublicKey:  peerID,
-			AllowedIPs: peerConn.WgConfig().AllowedIps,
-			PeerConnID: peerConn.ConnID(),
-			Log:        peerConn.Log,
-		}
-		lazyPeerCfgs = append(lazyPeerCfgs, lazyPeerCfg)
-	}
-
-	return e.lazyConnMgr.AddActivePeers(lazyPeerCfgs)
-}
-
-func (e *ConnMgr) closeManager(ctx context.Context) {
-	if e.lazyConnMgr == nil {
-		return
-	}
-
-	e.lazyCtxCancel()
-	e.wg.Wait()
-	e.lazyConnMgr = nil
-
-	for _, peerID := range e.peerStore.PeersPubKey() {
-		e.peerStore.PeerConnOpen(ctx, peerID)
-	}
-}
-
 func (e *ConnMgr) isStartedWithLazyMgr() bool {
 	return e.lazyConnMgr != nil && e.lazyCtxCancel != nil
 }

client/internal/conn_mgr_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/netbirdio/netbird/client/internal/lazyconn"
+	mgmProto "github.com/netbirdio/netbird/shared/management/proto"
 )
 
 func TestResolveLazyForce(t *testing.T) {
@@ -38,3 +39,90 @@ func TestResolveLazyForce(t *testing.T) {
 		})
 	}
 }
+
+func TestPeerLazyDefault(t *testing.T) {
+	tests := []struct {
+		name          string
+		force         lazyForce
+		remoteEnabled bool
+		state         mgmProto.LazyState
+		want          bool
+	}{
+		{name: "force on wins over eager state", force: lazyForceOn, state: mgmProto.LazyState_LazyStateEager, want: true},
+		{name: "force off wins over lazy state", force: lazyForceOff, remoteEnabled: true, state: mgmProto.LazyState_LazyStateLazy, want: false},
+		{name: "none, default, account off -> active", force: lazyForceNone, state: mgmProto.LazyState_LazyStateDefault, want: false},
+		{name: "none, default, account on -> lazy", force: lazyForceNone, remoteEnabled: true, state: mgmProto.LazyState_LazyStateDefault, want: true},
+		{name: "none, lazy state, account off -> lazy", force: lazyForceNone, state: mgmProto.LazyState_LazyStateLazy, want: true},
+		{name: "none, eager state, account on -> active", force: lazyForceNone, remoteEnabled: true, state: mgmProto.LazyState_LazyStateEager, want: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			e := &ConnMgr{force: tt.force, remoteLazyEnabled: tt.remoteEnabled}
+			if got := e.PeerLazyDefault(tt.state); got != tt.want {
+				t.Fatalf("PeerLazyDefault(%v) = %v, want %v", tt.state, got, tt.want)
+			}
+		})
+	}
+}
+
+// TestToExcludedLazyPeers covers the per-peer lazy classification (proxy vs
+// normal, across the force/account-flag matrix). Forwarder-target exclusion is
+// covered by TestToExcludedLazyPeers_ForwardTarget.
+func TestToExcludedLazyPeers(t *testing.T) {
+	const (
+		normalKey = "normal"
+		lazyKey   = "lazy-state"
+		eagerKey  = "eager-state"
+	)
+
+	peers := []*mgmProto.RemotePeerConfig{
+		{WgPubKey: normalKey, AllowedIps: []string{"100.64.0.1/32"}},
+		{WgPubKey: lazyKey, AllowedIps: []string{"100.64.0.2/32"}, LazyState: mgmProto.LazyState_LazyStateLazy},
+		{WgPubKey: eagerKey, AllowedIps: []string{"100.64.0.3/32"}, LazyState: mgmProto.LazyState_LazyStateEager},
+	}
+
+	tests := []struct {
+		name          string
+		force         lazyForce
+		remoteEnabled bool
+		want          map[string]bool
+	}{
+		{
+			name:  "account off: lazy-state peer lazy, normal + eager active",
+			force: lazyForceNone, remoteEnabled: false,
+			want: map[string]bool{normalKey: true, eagerKey: true},
+		},
+		{
+			name:  "account on: only eager-state peer active",
+			force: lazyForceNone, remoteEnabled: true,
+			want: map[string]bool{eagerKey: true},
+		},
+		{
+			name:  "force off: everything active",
+			force: lazyForceOff, remoteEnabled: true,
+			want: map[string]bool{normalKey: true, lazyKey: true, eagerKey: true},
+		},
+		{
+			name:  "force on: nothing active",
+			force: lazyForceOn, remoteEnabled: false,
+			want: map[string]bool{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			e := &Engine{connMgr: &ConnMgr{force: tt.force, remoteLazyEnabled: tt.remoteEnabled}}
+			got := e.toExcludedLazyPeers(nil, peers)
+
+			if len(got) != len(tt.want) {
+				t.Fatalf("toExcludedLazyPeers() = %v, want %v", got, tt.want)
+			}
+			for k := range tt.want {
+				if !got[k] {
+					t.Fatalf("expected peer %s excluded, got %v", k, got)
+				}
+			}
+		})
+	}
+}