[management] do in place expand of loaded snapshots (#8177)

management/server/affectedpeers/resolver.go

@@ -267,10 +267,10 @@ func (c Change) isEmpty() bool {
 }
 
 // Expand returns the deduplicated affected peer IDs from the preloaded Snapshot,
-// no store access. Run after the producing tx commits. Logs the full walk at
-// trace level for diagnosing a miscalculation.
+// no store access. Run after the producing tx commits. A nil Snapshot expands to
+// nothing. Logs the full walk at trace level for diagnosing a miscalculation.
 func (snap *Snapshot) Expand(ctx context.Context, accountID string, c Change) []string {
-	if c.isEmpty() {
+	if snap == nil || c.isEmpty() {
 		return nil
 	}
 	r := newResolver(ctx, snap, accountID, c)

management/server/affectedpeers/resolver_test.go

@@ -1,6 +1,7 @@
 package affectedpeers
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -77,6 +78,11 @@ func TestPolicyGroupsAndPeers_NonPeerResource(t *testing.T) {
 	assert.Empty(t, peers, "domain resource type should not produce direct peer IDs")
 }
 
+func TestExpand_NilSnapshot(t *testing.T) {
+	var snap *Snapshot
+	assert.Nil(t, snap.Expand(context.Background(), "acc", Change{ChangedGroupIDs: []string{"g"}}))
+}
+
 func TestChangeIsEmpty(t *testing.T) {
 	assert.True(t, Change{}.isEmpty())
 	assert.False(t, Change{ChangedGroupIDs: []string{"g"}}.isEmpty())

management/server/group.go

@@ -6,12 +6,12 @@ import (
 	"fmt"
 	"slices"
 
-	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
-	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
 
 	nbdns "github.com/netbirdio/netbird/dns"
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/affectedpeers"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
@@ -249,8 +249,7 @@ func (am *DefaultAccountManager) CreateGroups(ctx context.Context, accountID, us
 	}
 
 	var eventsToStore []func()
-	var snaps []*affectedpeers.Snapshot
-	var changes []affectedpeers.Change
+	affected := make(affectedPeerSet)
 
 	var globalErr error
 	createdCount := 0
@@ -281,7 +280,7 @@ func (am *DefaultAccountManager) CreateGroups(ctx context.Context, accountID, us
 			return err
 		})
 		if err != nil {
-			log.WithContext(ctx).Errorf("failed to update group %s: %v", newGroup.ID, err)
+			log.WithContext(ctx).Errorf("failed to create group %s: %v", newGroup.ID, err)
 			if createdCount == 0 {
 				return err
 			}
@@ -290,15 +289,14 @@ func (am *DefaultAccountManager) CreateGroups(ctx context.Context, accountID, us
 			continue
 		}
 		createdCount++
-		snaps = append(snaps, snap)
-		changes = append(changes, change)
+		affected.add(snap.Expand(ctx, accountID, change))
 	}
 
 	for _, storeEvent := range eventsToStore {
 		storeEvent()
 	}
 
-	go am.dispatchAffected(ctx, accountID, snaps, changes)
+	go am.dispatchAffectedPeerIDs(ctx, accountID, affected.ids())
 
 	return globalErr
 }
@@ -317,8 +315,7 @@ func (am *DefaultAccountManager) UpdateGroups(ctx context.Context, accountID, us
 	}
 
 	var eventsToStore []func()
-	var snaps []*affectedpeers.Snapshot
-	var changes []affectedpeers.Change
+	affected := make(affectedPeerSet)
 
 	var globalErr error
 	for _, newGroup := range groups {
@@ -333,15 +330,14 @@ func (am *DefaultAccountManager) UpdateGroups(ctx context.Context, accountID, us
 			continue
 		}
 		eventsToStore = append(eventsToStore, events...)
-		snaps = append(snaps, snap)
-		changes = append(changes, change)
+		affected.add(snap.Expand(ctx, accountID, change))
 	}
 
 	for _, storeEvent := range eventsToStore {
 		storeEvent()
 	}
 
-	go am.dispatchAffected(ctx, accountID, snaps, changes)
+	go am.dispatchAffectedPeerIDs(ctx, accountID, affected.ids())
 
 	return globalErr
 }
@@ -538,7 +534,10 @@ func (am *DefaultAccountManager) DeleteGroups(ctx context.Context, accountID, us
 		am.StoreEvent(ctx, userID, group.ID, accountID, activity.GroupDeleted, group.EventMeta())
 	}
 
-	go am.dispatchAffected(ctx, accountID, []*affectedpeers.Snapshot{snap, ipv6Snap}, []affectedpeers.Change{change, ipv6Change})
+	affected := make(affectedPeerSet)
+	affected.add(snap.Expand(ctx, accountID, change))
+	affected.add(ipv6Snap.Expand(ctx, accountID, ipv6Change))
+	go am.dispatchAffectedPeerIDs(ctx, accountID, affected.ids())
 
 	return allErrors
 }

management/server/group_test.go

@@ -1386,3 +1386,101 @@ func TestDefaultAccountManager_GroupPeersMustBelongToAccount(t *testing.T) {
 		assert.Equal(t, []string{peer1.ID}, stored.Peers, "rejected direct adds must not change membership")
 	})
 }
+
+func TestGroupsBatchAccountPeersUpdate(t *testing.T) {
+	runPeerUpdateTest(t, func(t *testing.T) {
+		manager, updateManager, account, peer1, peer2, peer3 := setupNetworkMapTest(t)
+
+		for _, group := range []*types.Group{
+			{ID: "groupA", Name: "GroupA", Peers: []string{peer1.ID, peer2.ID}},
+			{ID: "groupB", Name: "GroupB", Peers: []string{peer3.ID}},
+		} {
+			require.NoError(t, manager.CreateGroup(context.Background(), account.Id, userID, group))
+		}
+
+		updMsg := updateManager.CreateChannel(context.Background(), peer1.ID)
+		t.Cleanup(func() {
+			updateManager.CloseChannel(context.Background(), peer1.ID)
+		})
+
+		step(t, "creating unlinked groups", func(t *testing.T) {
+			settleAffectedUpdates(updMsg)
+
+			done := make(chan struct{})
+			go func() {
+				peerShouldNotReceiveUpdate(t, updMsg)
+				close(done)
+			}()
+
+			err := manager.CreateGroups(context.Background(), account.Id, userID, []*types.Group{
+				{ID: "groupC", Name: "GroupC", Peers: []string{peer1.ID}},
+				{ID: "groupD", Name: "GroupD"},
+			})
+			require.NoError(t, err)
+
+			select {
+			case <-done:
+			case <-time.After(time.Second):
+				t.Error("timeout waiting for peerShouldNotReceiveUpdate")
+			}
+		})
+
+		_, err := manager.SavePolicy(context.Background(), account.Id, userID, &types.Policy{
+			Enabled: true,
+			Rules: []*types.PolicyRule{
+				{
+					Enabled:       true,
+					Sources:       []string{"groupA"},
+					Destinations:  []string{"groupB"},
+					Bidirectional: true,
+					Action:        types.PolicyTrafficActionAccept,
+				},
+			},
+		}, true)
+		require.NoError(t, err)
+
+		step(t, "updating a batch with a linked group", func(t *testing.T) {
+			settleAffectedUpdates(updMsg)
+
+			done := make(chan struct{})
+			go func() {
+				peerShouldReceiveUpdate(t, updMsg)
+				close(done)
+			}()
+
+			err := manager.UpdateGroups(context.Background(), account.Id, userID, []*types.Group{
+				{ID: "groupD", Name: "GroupD renamed"},
+				{ID: "groupB", Name: "GroupB renamed"},
+			})
+			require.NoError(t, err)
+
+			select {
+			case <-done:
+			case <-time.After(peerUpdateTimeout):
+				t.Error("timeout waiting for peerShouldReceiveUpdate")
+			}
+		})
+
+		step(t, "updating a batch where one item fails", func(t *testing.T) {
+			settleAffectedUpdates(updMsg)
+
+			done := make(chan struct{})
+			go func() {
+				peerShouldReceiveUpdate(t, updMsg)
+				close(done)
+			}()
+
+			err := manager.UpdateGroups(context.Background(), account.Id, userID, []*types.Group{
+				{ID: "missing", Name: "Missing"},
+				{ID: "groupB", Name: "GroupB"},
+			})
+			require.Error(t, err)
+
+			select {
+			case <-done:
+			case <-time.After(peerUpdateTimeout):
+				t.Error("timeout waiting for peerShouldReceiveUpdate")
+			}
+		})
+	})
+}

management/server/peer.go

@@ -1596,30 +1596,20 @@ func (am *DefaultAccountManager) UpdateAccountPeers(ctx context.Context, account
 // ExpandAndUpdateAffected expands a Snapshot (loaded INSIDE the now-committed
 // transaction) into the affected peers and dispatches the network-map refresh.
 // Pure in-memory work plus dispatch, so it runs AFTER commit — the fan-out walk
-// never holds the write lock, over the consistent in-tx snapshot. Exported so the
-// networks sub-package managers (which hold only account.Manager) share it.
+// never holds the write lock, over the consistent in-tx snapshot. The walk runs
+// before the dispatch goroutine starts, so the snapshot is released on return.
+// Exported so the networks sub-package managers (which hold only account.Manager)
+// share it.
 func (am *DefaultAccountManager) ExpandAndUpdateAffected(ctx context.Context, accountID string, snap *affectedpeers.Snapshot, change affectedpeers.Change) {
-	go am.dispatchAffected(ctx, accountID, []*affectedpeers.Snapshot{snap}, []affectedpeers.Change{change})
+	go am.dispatchAffectedPeerIDs(ctx, accountID, snap.Expand(ctx, accountID, change))
 }
 
-// dispatchAffected expands one or more (snapshot, change) pairs — collected across
-// one or several transactions — unions their affected peers, and dispatches a
-// single network-map refresh. Each snapshot must already be loaded inside its
-// transaction; this runs AFTER commit (pure in-memory + dispatch). It is spawned
-// in a goroutine that outlives the request, so it detaches from the request
-// context's cancellation up front.
-func (am *DefaultAccountManager) dispatchAffected(ctx context.Context, accountID string, snaps []*affectedpeers.Snapshot, changes []affectedpeers.Change) {
+// dispatchAffectedPeerIDs dispatches a single network-map refresh for already
+// expanded affected peer IDs. It may run in a goroutine that outlives the
+// request, so it detaches from the request context's cancellation up front.
+func (am *DefaultAccountManager) dispatchAffectedPeerIDs(ctx context.Context, accountID string, affectedPeerIDs []string) {
 	ctx = context.WithoutCancel(ctx)
 
-	var lists [][]string
-	for i, snap := range snaps {
-		if snap == nil {
-			continue
-		}
-		lists = append(lists, snap.Expand(ctx, accountID, changes[i]))
-	}
-
-	affectedPeerIDs := unionStrings(lists...)
 	if len(affectedPeerIDs) == 0 {
 		log.WithContext(ctx).Tracef("no affected peers for account %s", accountID)
 		return
@@ -1629,21 +1619,19 @@ func (am *DefaultAccountManager) dispatchAffected(ctx context.Context, accountID
 	_ = am.networkMapController.UpdateAffectedPeers(ctx, accountID, affectedPeerIDs)
 }
 
-// unionStrings concatenates the given string lists into one deduplicated slice,
-// preserving first-occurrence order.
-func unionStrings(lists ...[]string) []string {
-	seen := make(map[string]struct{})
-	var out []string
-	for _, list := range lists {
-		for _, id := range list {
-			if _, ok := seen[id]; ok {
-				continue
-			}
-			seen[id] = struct{}{}
-			out = append(out, id)
-		}
+// affectedPeerSet collects the affected peer IDs of several snapshot expansions,
+// deduplicating as they are added so it stays bounded by the account's peers
+// however many items a batch has.
+type affectedPeerSet map[string]struct{}
+
+func (s affectedPeerSet) add(peerIDs []string) {
+	for _, id := range peerIDs {
+		s[id] = struct{}{}
 	}
-	return out
+}
+
+func (s affectedPeerSet) ids() []string {
+	return maps.Keys(s)
 }
 
 // affectedPeerIDsFromNetworkMap returns the peer IDs referenced by a peer's

management/server/user.go

@@ -593,8 +593,7 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 		return nil, err
 	}
 
-	var snaps []*affectedpeers.Snapshot
-	var changes []affectedpeers.Change
+	affected := make(affectedPeerSet)
 	var peersToExpire []*nbpeer.Peer
 	var addUserEvents []func()
 	var usersToSave = make([]*types.User, 0, len(updates))
@@ -629,6 +628,8 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 			return nil, status.Errorf(status.InvalidArgument, "provided user update is nil")
 		}
 
+		var itemSnap *affectedpeers.Snapshot
+		var itemChange affectedpeers.Change
 		err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error {
 			change, updatedUser, userPeersToExpire, userEvents, err := am.processUserUpdate(
 				ctx, transaction, groupsMap, accountID, initiatorUserID, initiatorUser, update, addIfNotExists, settings,
@@ -647,8 +648,7 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 				return err
 			}
 
-			snaps = append(snaps, snap)
-			changes = append(changes, change)
+			itemSnap, itemChange = snap, change
 			usersToSave = append(usersToSave, updatedUser)
 			addUserEvents = append(addUserEvents, userEvents...)
 			peersToExpire = append(peersToExpire, userPeersToExpire...)
@@ -662,7 +662,9 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 			}
 			globalErr = errors.Join(globalErr, err)
 			// continue when updating multiple users
+			continue
 		}
+		affected.add(itemSnap.Expand(ctx, accountID, itemChange))
 	}
 
 	var updatedUsersInfo = make([]*types.UserInfo, 0, len(usersToSave))
@@ -693,7 +695,7 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 		if err = am.Store.IncrementNetworkSerial(ctx, accountID); err != nil {
 			return nil, fmt.Errorf("failed to increment network serial: %w", err)
 		}
-		go am.dispatchAffected(ctx, accountID, snaps, changes)
+		go am.dispatchAffectedPeerIDs(ctx, accountID, affected.ids())
 	}
 
 	return updatedUsersInfo, globalErr