[management] prewarm a posture check cache on network map generation (#7093)

management/internals/controllers/network_map/controller/controller.go

@@ -176,6 +176,7 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
 	semaphore := make(chan struct{}, 10)
 
 	c.injectAllProxyPolicies(ctx, account)
+	account.PrecomputePostureValidation(ctx)
 	dnsCache := &cache.DNSConfigCache{}
 	dnsDomain := c.GetDNSDomain(account.Settings)
 	peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain)
@@ -357,6 +358,7 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s
 	// network map that omitted the synth DNS zone, and the agent kept
 	// resolving against the stale or absent record.
 	c.injectAllProxyPolicies(ctx, account)
+	account.PrecomputePostureValidation(ctx)
 	dnsCache := &cache.DNSConfigCache{}
 	dnsDomain := c.GetDNSDomain(account.Settings)
 	peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain)

management/server/account_test.go

@@ -1757,6 +1757,7 @@ func TestAccount_Copy(t *testing.T) {
 				AccountID: "account1",
 			},
 		},
+		PostureValidation: map[string]map[string]bool{"1": {"1": true}},
 	}
 	err := hasNilField(account)
 	if err != nil {

management/server/types/account.go

@@ -91,6 +91,8 @@ type Account struct {
 	Onboarding       AccountOnboarding                `gorm:"foreignKey:AccountID;references:id;constraint:OnDelete:CASCADE"`
 
 	ReverseProxyFreeDomainNonce string
+
+	PostureValidation map[string]map[string]bool `gorm:"-"`
 }
 
 // this class is used by gorm only
@@ -874,6 +876,7 @@ func (a *Account) Copy() *Account {
 		Services:               services,
 		Onboarding:             a.Onboarding,
 		Domains:                domains,
+		PostureValidation:      a.PostureValidation,
 	}
 }
 

management/server/types/account_components.go

@@ -10,6 +10,8 @@ import (
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/management/internals/modules/zones"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/route"
 )
@@ -506,8 +508,8 @@ func (a *Account) getPeersGroupsPoliciesRoutes(
 func (a *Account) getPeersFromGroups(ctx context.Context, groups []string, peerID string, sourcePostureChecksIDs []string,
 	validatedPeersMap map[string]struct{}, postureFailedPeers *map[string]map[string]struct{}) ([]string, bool) {
 	peerInGroups := false
-	filteredPeerIDs := make([]string, 0, len(groups))
-	seenPeerIds := make(map[string]struct{}, len(groups))
+	var filteredPeerIDs []string
+	var seenPeerIds map[string]struct{}
 
 	for _, gid := range groups {
 		group := a.GetGroup(gid)
@@ -547,6 +549,17 @@ func (a *Account) getPeersFromGroups(ctx context.Context, groups []string, peerI
 			return filteredPeerIDs, peerInGroups
 		}
 
+		if seenPeerIds == nil {
+			totalGroupPeers := 0
+			for _, g := range groups {
+				if grp := a.GetGroup(g); grp != nil {
+					totalGroupPeers += len(grp.Peers)
+				}
+			}
+			filteredPeerIDs = make([]string, 0, totalGroupPeers)
+			seenPeerIds = make(map[string]struct{}, totalGroupPeers)
+		}
+
 		for _, pid := range group.Peers {
 			if _, seen := seenPeerIds[pid]; seen {
 				continue
@@ -589,21 +602,109 @@ func (a *Account) validatePostureChecksOnPeerGetFailed(ctx context.Context, sour
 	}
 
 	for _, postureChecksID := range sourcePostureChecksID {
+		if valid, cached := a.cachedPostureCheckResult(postureChecksID, peerID); cached {
+			if !valid {
+				return false, postureChecksID
+			}
+			continue
+		}
+
 		postureChecks := a.GetPostureChecks(postureChecksID)
 		if postureChecks == nil {
 			continue
 		}
 
-		for _, check := range postureChecks.GetChecks() {
-			isValid, _ := check.Check(ctx, *peer)
-			if !isValid {
-				return false, postureChecksID
-			}
+		if !peerPassesPostureChecks(ctx, postureChecks.GetChecks(), peer) {
+			return false, postureChecksID
 		}
 	}
 	return true, ""
 }
 
+// PrecomputePostureValidation evaluates every posture check referenced by an enabled
+// policy once against the peers of that policy's source groups and stores the results,
+// so the per-peer network map calculations that follow look them up instead of
+// re-evaluating checks for every peer pair. It must be called before the account is
+// shared across goroutines; lookups not covered by the precomputed results fall back
+// to direct evaluation.
+func (a *Account) PrecomputePostureValidation(ctx context.Context) {
+	if len(a.PostureChecks) == 0 {
+		a.PostureValidation = nil
+		return
+	}
+
+	checkPeerIDs := make(map[string]map[string]struct{})
+	for _, policy := range a.Policies {
+		if !policy.Enabled || len(policy.SourcePostureChecks) == 0 {
+			continue
+		}
+
+		peerIDs := a.getUniquePeerIDsFromGroupsIDs(ctx, policy.SourceGroups())
+		for _, rule := range policy.Rules {
+			if rule.SourceResource.Type == ResourceTypePeer && rule.SourceResource.ID != "" {
+				peerIDs = append(peerIDs, rule.SourceResource.ID)
+			}
+		}
+
+		for _, postureChecksID := range policy.SourcePostureChecks {
+			set := checkPeerIDs[postureChecksID]
+			if set == nil {
+				set = make(map[string]struct{}, len(peerIDs))
+				checkPeerIDs[postureChecksID] = set
+			}
+			for _, pid := range peerIDs {
+				set[pid] = struct{}{}
+			}
+		}
+	}
+
+	results := make(map[string]map[string]bool, len(checkPeerIDs))
+	for postureChecksID, peerIDs := range checkPeerIDs {
+		results[postureChecksID] = a.evaluatePostureChecksForPeers(ctx, postureChecksID, peerIDs)
+	}
+	a.PostureValidation = results
+}
+
+func (a *Account) evaluatePostureChecksForPeers(ctx context.Context, postureChecksID string, peerIDs map[string]struct{}) map[string]bool {
+	postureChecks := a.GetPostureChecks(postureChecksID)
+	if postureChecks == nil {
+		return nil
+	}
+
+	checks := postureChecks.GetChecks()
+	results := make(map[string]bool, len(peerIDs))
+	for peerID := range peerIDs {
+		peer, ok := a.Peers[peerID]
+		if !ok || peer == nil {
+			continue
+		}
+		results[peerID] = peerPassesPostureChecks(ctx, checks, peer)
+	}
+	return results
+}
+
+func (a *Account) cachedPostureCheckResult(postureChecksID, peerID string) (bool, bool) {
+	results, ok := a.PostureValidation[postureChecksID]
+	if !ok {
+		return false, false
+	}
+	if results == nil {
+		return true, true
+	}
+	valid, found := results[peerID]
+	return valid, found
+}
+
+func peerPassesPostureChecks(ctx context.Context, checks []posture.Check, peer *nbpeer.Peer) bool {
+	for _, check := range checks {
+		isValid, _ := check.Check(ctx, *peer)
+		if !isValid {
+			return false
+		}
+	}
+	return true
+}
+
 func (a *Account) getPostureValidPeersSaveFailed(inputPeers []string, postureChecksIDs []string, validatedPeersMap map[string]struct{}, postureFailedPeers *map[string]map[string]struct{}) []string {
 	var dest []string
 	for _, peerID := range inputPeers {

management/server/types/account_posture_validation_test.go

@@ -0,0 +1,72 @@
+package types_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	nbdns "github.com/netbirdio/netbird/dns"
+	"github.com/netbirdio/netbird/management/server/posture"
+)
+
+func TestPrecomputePostureValidation_MatchesDirectEvaluation(t *testing.T) {
+	account, validatedPeers := scalableTestAccount(60, 5)
+
+	account.PostureChecks = append(account.PostureChecks, &posture.Checks{
+		ID: "posture-check-strict", Name: "Strict version",
+		Checks: posture.ChecksDefinition{
+			NBVersionCheck: &posture.NBVersionCheck{MinVersion: "0.50.0"},
+		},
+	})
+	account.Policies[0].SourcePostureChecks = []string{"posture-check-ver", "posture-check-unknown"}
+	account.Policies[1].SourcePostureChecks = []string{"posture-check-strict"}
+	account.Policies[2].SourcePostureChecks = []string{"posture-check-ver"}
+	account.Policies[2].Enabled = false
+
+	ctx := context.Background()
+	resourcePolicies := account.GetResourcePoliciesMap()
+	routers := account.GetResourceRoutersMap()
+
+	type result struct {
+		peers              map[string]struct{}
+		postureFailedPeers map[string]map[string]struct{}
+	}
+	snapshot := func() map[string]result {
+		results := make(map[string]result, len(account.Peers))
+		for peerID := range account.Peers {
+			components := account.GetPeerNetworkMapComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil)
+			require.NotNil(t, components)
+			peerSet := make(map[string]struct{}, len(components.Peers))
+			for id := range components.Peers {
+				peerSet[id] = struct{}{}
+			}
+			results[peerID] = result{peers: peerSet, postureFailedPeers: components.PostureFailedPeers}
+		}
+		return results
+	}
+
+	direct := snapshot()
+	account.PrecomputePostureValidation(ctx)
+	memoized := snapshot()
+
+	require.Equal(t, len(direct), len(memoized))
+	for peerID, want := range direct {
+		got := memoized[peerID]
+		assert.Equal(t, want.peers, got.peers, "visible peers changed for %s", peerID)
+		assert.Equal(t, want.postureFailedPeers, got.postureFailedPeers, "posture failed peers changed for %s", peerID)
+	}
+}
+
+func TestPrecomputePostureValidation_NoPostureChecks(t *testing.T) {
+	account, validatedPeers := scalableTestAccount(10, 2)
+	account.PostureChecks = nil
+
+	ctx := context.Background()
+	account.PrecomputePostureValidation(ctx)
+
+	components := account.GetPeerNetworkMapComponents(ctx, "peer-0", nbdns.CustomZone{}, nil, validatedPeers, account.GetResourcePoliciesMap(), account.GetResourceRoutersMap(), nil)
+	require.NotNil(t, components)
+	assert.NotEmpty(t, components.Peers)
+}

management/server/types/networkmap_benchmark_test.go

@@ -86,6 +86,43 @@ func BenchmarkNetworkMapGeneration_AllPeers(b *testing.B) {
 			b.ReportAllocs()
 			b.ResetTimer()
 			for range b.N {
+				account.PrecomputePostureValidation(ctx)
+				for _, peerID := range peerIDs {
+					_ = account.GetPeerNetworkMapFromComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil, groupIDToUserIDs)
+				}
+			}
+		})
+	}
+}
+
+// BenchmarkNetworkMapGeneration_AllPeersPostureChecks benchmarks the UpdateAccountPeers
+// hot path with a posture check attached to the account-wide policy, so posture
+// validation runs for every source peer of every target peer's map.
+func BenchmarkNetworkMapGeneration_AllPeersPostureChecks(b *testing.B) {
+	skipCIBenchmark(b)
+	scales := []benchmarkScale{
+		{"500peers_20groups", 500, 20},
+		{"1000peers_50groups", 1000, 50},
+	}
+
+	for _, scale := range scales {
+		account, validatedPeers := scalableTestAccount(scale.peers, scale.groups)
+		account.Policies[0].SourcePostureChecks = []string{"posture-check-ver"}
+		ctx := context.Background()
+
+		peerIDs := make([]string, 0, len(account.Peers))
+		for peerID := range account.Peers {
+			peerIDs = append(peerIDs, peerID)
+		}
+
+		b.Run("components/"+scale.name, func(b *testing.B) {
+			resourcePolicies := account.GetResourcePoliciesMap()
+			routers := account.GetResourceRoutersMap()
+			groupIDToUserIDs := account.GetActiveGroupUsers()
+			b.ReportAllocs()
+			b.ResetTimer()
+			for range b.N {
+				account.PrecomputePostureValidation(ctx)
 				for _, peerID := range peerIDs {
 					_ = account.GetPeerNetworkMapFromComponents(ctx, peerID, nbdns.CustomZone{}, nil, validatedPeers, resourcePolicies, routers, nil, groupIDToUserIDs)
 				}