Use GroupMemberSet.Merge to reduce CPU usage and memory footprint (#2467

)

When getting the union of multiple AddressGroups or AppliedToGroups,
agent used the GroupMemberSet.Union method, which always creates an
intermediate set and copies the original items when merging another
set. This wasted significant time and memory when the groups had
massive members. This patch adds the GroupMemberSet.Merge method which
avoids the above overhead.

benchmark comparison when merging two groups with 10K members and 100
members:

name                            old time/op    new time/op    delta
RuleCacheUnionAddressGroups-48    4.27ms ±15%    2.15ms ±15%  -49.72%  (p=0.008 n=5+5)

name                            old alloc/op   new alloc/op   delta
RuleCacheUnionAddressGroups-48    1.95MB ± 0%    0.97MB ± 0%  -50.00%  (p=0.008 n=5+5)

name                            old allocs/op  new allocs/op  delta
RuleCacheUnionAddressGroups-48       565 ± 0%       282 ± 0%  -50.16%  (p=0.000 n=4+5)

Signed-off-by: Quan Tian <qtian@vmware.com>

pkg/agent/controller/networkpolicy/cache.go

@@ -787,7 +787,7 @@ func (c *ruleCache) unionAddressGroups(groupNames []string) (v1beta.GroupMemberS
 			klog.V(2).Infof("AddressGroup %v was not found", groupName)
 			return nil, false
 		}
-		set = set.Union(curSet)
+		set.Merge(curSet)
 	}
 	return set, true
 }
@@ -807,7 +807,7 @@ func (c *ruleCache) unionAppliedToGroups(groupNames []string) (v1beta.GroupMembe
 			continue
 		}
 		anyExists = true
-		set = set.Union(curSet)
+		set.Merge(curSet)
 	}
 	return set, anyExists
 }

pkg/agent/controller/networkpolicy/cache_test.go

@@ -15,6 +15,7 @@
 package networkpolicy
 
 import (
+	"fmt"
 	"net"
 	"reflect"
 	"testing"
@@ -1186,3 +1187,28 @@ func TestRuleCacheProcessPodUpdates(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkRuleCacheUnionAddressGroups(b *testing.B) {
+	var addressGroupMembers1, addressGroupMembers2 []*v1beta2.GroupMember
+	// addressGroup1 includes 10K members.
+	for i := 0; i < 100; i++ {
+		for j := 0; j < 100; j++ {
+			addressGroupMembers1 = append(addressGroupMembers1, newAddressGroupMember(fmt.Sprintf("1.1.%d.%d", i, j)))
+		}
+	}
+	addressGroup1 := v1beta2.NewGroupMemberSet(addressGroupMembers1...)
+	// addressGroup2 includes 10 members.
+	for i := 0; i < 10; i++ {
+		addressGroupMembers2 = append(addressGroupMembers2, newAddressGroupMember(fmt.Sprintf("2.2.2.%d", i)))
+	}
+	addressGroup2 := v1beta2.NewGroupMemberSet(addressGroupMembers2...)
+	c, _, _ := newFakeRuleCache()
+	c.addressSetByGroup["addressGroup1"] = addressGroup1
+	c.addressSetByGroup["addressGroup2"] = addressGroup2
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		c.unionAddressGroups([]string{"addressGroup1", "addressGroup2"})
+	}
+}

pkg/apis/controlplane/sets.go

@@ -118,6 +118,21 @@ func (s GroupMemberSet) Union(o GroupMemberSet) GroupMemberSet {
 	return result
 }
 
+// Merge merges the other set into the set.
+// For example:
+// s1 = {a1, a2, a3}
+// s2 = {a1, a2, a4, a5}
+// s1.Merge(s2) = {a1, a2, a3, a4, a5}
+// s1 = {a1, a2, a3, a4, a5}
+//
+// It should be used instead of s1.Union(s2) when constructing a new set is not required.
+func (s GroupMemberSet) Merge(o GroupMemberSet) GroupMemberSet {
+	for key, item := range o {
+		s[key] = item
+	}
+	return s
+}
+
 // IsSuperset returns true if and only if s1 is a superset of s2.
 func (s GroupMemberSet) IsSuperset(o GroupMemberSet) bool {
 	for key := range o {

pkg/apis/controlplane/v1beta1/sets.go

@@ -91,6 +91,21 @@ func (s GroupMemberPodSet) Union(o GroupMemberPodSet) GroupMemberPodSet {
 	return result
 }
 
+// Merge merges the other set into the set.
+// For example:
+// s1 = {a1, a2, a3}
+// s2 = {a1, a2, a4, a5}
+// s1.Merge(s2) = {a1, a2, a3, a4, a5}
+// s1 = {a1, a2, a3, a4, a5}
+//
+// It should be used instead of s1.Union(s2) when constructing a new set is not required.
+func (s GroupMemberSet) Merge(o GroupMemberSet) GroupMemberSet {
+	for key, item := range o {
+		s[key] = item
+	}
+	return s
+}
+
 // IsSuperset returns true if and only if s1 is a superset of s2.
 func (s GroupMemberPodSet) IsSuperset(o GroupMemberPodSet) bool {
 	for key := range o {

pkg/apis/controlplane/v1beta2/sets.go

@@ -118,6 +118,21 @@ func (s GroupMemberSet) Union(o GroupMemberSet) GroupMemberSet {
 	return result
 }
 
+// Merge merges the other set into the set.
+// For example:
+// s1 = {a1, a2, a3}
+// s2 = {a1, a2, a4, a5}
+// s1.Merge(s2) = {a1, a2, a3, a4, a5}
+// s1 = {a1, a2, a3, a4, a5}
+//
+// It should be used instead of s1.Union(s2) when constructing a new set is not required.
+func (s GroupMemberSet) Merge(o GroupMemberSet) GroupMemberSet {
+	for key, item := range o {
+		s[key] = item
+	}
+	return s
+}
+
 // IsSuperset returns true if and only if s1 is a superset of s2.
 func (s GroupMemberSet) IsSuperset(o GroupMemberSet) bool {
 	for key := range o {

pkg/controller/egress/store/egressgroup.go

@@ -75,11 +75,11 @@ func (event *egressGroupEvent) ToWatchEvent(selectors *storage.Selectors, isInit
 		} else {
 			currMembers = controlplane.GroupMemberSet{}
 			for _, members := range event.CurrGroup.GroupMemberByNode {
-				currMembers = currMembers.Union(members)
+				currMembers.Merge(members)
 			}
 			prevMembers = controlplane.GroupMemberSet{}
 			for _, members := range event.PrevGroup.GroupMemberByNode {
-				prevMembers = prevMembers.Union(members)
+				prevMembers.Merge(members)
 			}
 		}
 		for _, member := range currMembers.Difference(prevMembers) {

pkg/controller/networkpolicy/networkpolicy_controller.go

@@ -1129,7 +1129,7 @@ func (n *NetworkPolicyController) getClusterGroupMemberSet(group *antreatypes.Gr
 		childGroup, found, _ := n.internalGroupStore.Get(childName)
 		if found {
 			child := childGroup.(*antreatypes.Group)
-			groupMemberSet = groupMemberSet.Union(n.getMemberSetForGroupType(clusterGroupType, child.Name))
+			groupMemberSet.Merge(n.getMemberSetForGroupType(clusterGroupType, child.Name))
 		}
 	}
 	return groupMemberSet

pkg/controller/networkpolicy/store/appliedtogroup.go

@@ -76,11 +76,11 @@ func (event *appliedToGroupEvent) ToWatchEvent(selectors *storage.Selectors, isI
 		} else {
 			currMembers = controlplane.GroupMemberSet{}
 			for _, members := range event.CurrGroup.GroupMemberByNode {
-				currMembers = currMembers.Union(members)
+				currMembers.Merge(members)
 			}
 			prevMembers = controlplane.GroupMemberSet{}
 			for _, members := range event.PrevGroup.GroupMemberByNode {
-				prevMembers = prevMembers.Union(members)
+				prevMembers.Merge(members)
 			}
 		}
 		for _, member := range currMembers.Difference(prevMembers) {