perf: [WIN-NPM] add all cached NetworkPolicies to a Pod at once (#1893)

* cherry-picking stuff from apply in background POC

* add all policies poc

* add debug prints

* fix deadlock

* fix other GetPolicy deadlock

* update whitespace in yamls

* properly merge

* properly merge 2

* add ACLs in batches

* cleanup errors

* lint and log

* persist state as we add

* refactor into function so we can do UTs on batching

* fix lint

* batch struct

* successful policies

* reduce batch limit to 30

Author: huntergregory

npm/azure-npm.yaml

@@ -148,18 +148,19 @@ metadata:
 data:
   azure-npm.json: |
     {
-        "ResyncPeriodInMinutes": 15,
-        "ListeningPort":         10091,
-        "ListeningAddress":      "0.0.0.0",
-        "ApplyMaxBatches": 100,
+        "ResyncPeriodInMinutes":       15,
+        "ListeningPort":               10091,
+        "ListeningAddress":            "0.0.0.0",
+        "ApplyMaxBatches":             100,
         "ApplyIntervalInMilliseconds": 500,
+        "MaxBatchedACLsPerPod":        30,
         "Toggles": {
             "EnablePrometheusMetrics": true,
             "EnablePprof":             true,
             "EnableHTTPDebugAPI":      true,
             "EnableV2NPM":             true,
             "PlaceAzureChainFirst":    true,
             "ApplyIPSetsOnNeed":       false,
-            "ApplyInBackground": true
+            "ApplyInBackground":       true
         }
     }

npm/cmd/start.go

@@ -123,6 +123,8 @@ func start(config npmconfig.Config, flags npmconfig.Flags) error {
 	stopChannel := wait.NeverStop
 	if config.Toggles.EnableV2NPM {
 		// update the dataplane config
+		npmV2DataplaneCfg.MaxBatchedACLsPerPod = config.MaxBatchedACLsPerPod
+
 		npmV2DataplaneCfg.ApplyInBackground = config.Toggles.ApplyInBackground
 		if config.ApplyMaxBatches > 0 {
 			npmV2DataplaneCfg.ApplyMaxBatches = config.ApplyMaxBatches

npm/config/config.go

@@ -3,12 +3,13 @@ package npmconfig
 import "github.com/Azure/azure-container-networking/npm/util"
 
 const (
-	defaultResyncPeriod    = 15
-	defaultApplyMaxBatches = 100
-	defaultApplyInterval   = 500
-	defaultListeningPort   = 10091
-	defaultGrpcPort        = 10092
-	defaultGrpcServicePort = 9002
+	defaultResyncPeriod         = 15
+	defaultApplyMaxBatches      = 100
+	defaultApplyInterval        = 500
+	defaultMaxBatchedACLsPerPod = 30
+	defaultListeningPort        = 10091
+	defaultGrpcPort             = 10092
+	defaultGrpcServicePort      = 9002
 	// ConfigEnvPath is what's used by viper to load config path
 	ConfigEnvPath = "NPM_CONFIG"
 
@@ -32,6 +33,7 @@ var DefaultConfig = Config{
 	WindowsNetworkName:          util.AzureNetworkName,
 	ApplyMaxBatches:             defaultApplyMaxBatches,
 	ApplyIntervalInMilliseconds: defaultApplyInterval,
+	MaxBatchedACLsPerPod:        defaultMaxBatchedACLsPerPod,
 
 	Toggles: Toggles{
 		EnablePrometheusMetrics: true,
@@ -62,9 +64,13 @@ type Config struct {
 	// It can also be the empty string, which results in the default value of 'azure'.
 	WindowsNetworkName string `json:"WindowsNetworkName,omitempty"`
 	// Apply options for Windows only. Relevant when ApplyInBackground is true.
-	ApplyMaxBatches             int     `json:"ApplyDataPlaneMaxBatches,omitempty"`
-	ApplyIntervalInMilliseconds int     `json:"ApplyDataPlaneMaxWaitInMilliseconds,omitempty"`
-	Toggles                     Toggles `json:"Toggles,omitempty"`
+	ApplyMaxBatches             int `json:"ApplyDataPlaneMaxBatches,omitempty"`
+	ApplyIntervalInMilliseconds int `json:"ApplyDataPlaneMaxWaitInMilliseconds,omitempty"`
+	// MaxBatchedACLsPerPod is the maximum number of ACLs that can be added to a Pod at once in Windows.
+	// The zero value is valid.
+	// A NetworkPolicy's ACLs are always in the same batch, and there will be at least one NetworkPolicy per batch.
+	MaxBatchedACLsPerPod int     `json:"MaxBatchedACLsPerPod,omitempty"`
+	Toggles              Toggles `json:"Toggles,omitempty"`
 }
 
 type Toggles struct {

npm/examples/windows/azure-npm.yaml

@@ -140,18 +140,19 @@ metadata:
 data:
   azure-npm.json: |
     {
-        "ResyncPeriodInMinutes": 15,
-        "ListeningPort":         10091,
-        "ListeningAddress":      "0.0.0.0",
-        "ApplyMaxBatches": 100,
+        "ResyncPeriodInMinutes":       15,
+        "ListeningPort":               10091,
+        "ListeningAddress":            "0.0.0.0",
+        "ApplyMaxBatches":             100,
         "ApplyIntervalInMilliseconds": 500,
+        "MaxBatchedACLsPerPod":        30,
         "Toggles": {
             "EnablePrometheusMetrics": true,
             "EnablePprof":             true,
             "EnableHTTPDebugAPI":      true,
             "EnableV2NPM":             true,
             "PlaceAzureChainFirst":    true,
             "ApplyIPSetsOnNeed":       false,
-            "ApplyInBackground": true
+            "ApplyInBackground":       true
         }
     }

npm/pkg/dataplane/dataplane_windows.go

@@ -217,6 +217,9 @@ func (dp *DataPlane) updatePod(pod *updateNPMPod) error {
 	}
 
 	// for every ipset we're adding to the endpoint, consider adding to the endpoint every policy that the set touches
+	// add policy if:
+	// 1. it's not already there
+	// 2. the pod IP is part of every set that the policy requires (every set in the pod selector)
 	toAddPolicies := make(map[string]struct{})
 	for _, setName := range pod.IPSetsToAdd {
 		/*
@@ -239,48 +242,42 @@ func (dp *DataPlane) updatePod(pod *updateNPMPod) error {
 		}
 
 		for policyKey := range selectorReference {
-			if dp.policyMgr.PolicyExists(policyKey) {
-				toAddPolicies[policyKey] = struct{}{}
-			} else {
+			if _, ok := endpoint.netPolReference[policyKey]; ok {
+				continue
+			}
+
+			policy, ok := dp.policyMgr.GetPolicy(policyKey)
+			if !ok {
 				klog.Infof("[DataPlane] while updating pod, policy is referenced but does not exist. pod: [%s], policy: [%s], set [%s]", pod.PodKey, policyKey, setName)
+				continue
 			}
-		}
-	}
 
-	// for all of these policies, add the policy to the endpoint if:
-	// 1. it's not already there
-	// 2. the pod IP is part of every set that the policy requires (every set in the pod selector)
-	for policyKey := range toAddPolicies {
-		if _, ok := endpoint.netPolReference[policyKey]; ok {
-			continue
-		}
+			selectorIPSets := dp.getSelectorIPSets(policy)
+			ok, err := dp.ipsetMgr.DoesIPSatisfySelectorIPSets(pod.PodIP, pod.PodKey, selectorIPSets)
+			if err != nil {
+				return fmt.Errorf("[DataPlane] error getting IPs satisfying selector ipsets: %w", err)
+			}
+			if !ok {
+				continue
+			}
 
-		// TODO Also check if the endpoint reference in policy for this Ip is right
-		policy, ok := dp.policyMgr.GetPolicy(policyKey)
-		if !ok {
-			return fmt.Errorf("policy with name %s does not exist", policyKey)
-		}
-
-		selectorIPSets := dp.getSelectorIPSets(policy)
-		ok, err := dp.ipsetMgr.DoesIPSatisfySelectorIPSets(pod.PodIP, pod.PodKey, selectorIPSets)
-		if err != nil {
-			return err
-		}
-		if !ok {
-			continue
+			toAddPolicies[policyKey] = struct{}{}
 		}
+	}
 
-		// Apply the network policy
-		endpointList := map[string]string{
-			endpoint.ip: endpoint.id,
-		}
-		err = dp.policyMgr.AddPolicy(policy, endpointList)
-		if err != nil {
-			return err
-		}
+	if len(toAddPolicies) == 0 {
+		return nil
+	}
 
+	successfulPolicies, err := dp.policyMgr.AddAllPolicies(toAddPolicies, endpoint.id, endpoint.ip)
+	for policyKey := range successfulPolicies {
 		endpoint.netPolReference[policyKey] = struct{}{}
 	}
+	if err != nil {
+		return fmt.Errorf("failed to add all policies while updating pod. endpoint: %+v. policies: %+v. err: %w", endpoint, toAddPolicies, err)
+	}
+
+	klog.Infof("[DataPlane] updatedPod complete. podKey: %s. endpoint: %+v", pod.PodKey, endpoint)
 
 	return nil
 }

npm/pkg/dataplane/policies/policymanager.go

@@ -35,6 +35,10 @@ type PolicyManagerCfg struct {
 	PolicyMode PolicyManagerMode
 	// PlaceAzureChainFirst only affects Linux
 	PlaceAzureChainFirst bool
+	// MaxBatchedACLsPerPod is the maximum number of ACLs that can be added to a Pod at once in Windows.
+	// The zero value is valid.
+	// A NetworkPolicy's ACLs are always in the same batch, and there will be at least one NetworkPolicy per batch.
+	MaxBatchedACLsPerPod int
 }
 
 type PolicyMap struct {

npm/pkg/dataplane/policies/policymanager_windows.go

@@ -87,6 +87,11 @@ var baseACLsForCalicoCNI = []*NPMACLPolSettings{
 	},
 }
 
+type aclBatch struct {
+	rules    []*NPMACLPolSettings
+	policies []string
+}
+
 type staleChains struct{} // unused in Windows
 
 type shouldResetAllACLs bool
@@ -125,6 +130,114 @@ func (pMgr *PolicyManager) reconcile() {
 	// not implemented
 }
 
+// AddAllPolicies is used in Windows to add all NetworkPolicies to an endpoint.
+// Will make a series of sequential HNS ADD calls based on MaxBatchedACLsPerPod.
+// A NetworkPolicy's ACLs are always in the same batch, and there will be at least one NetworkPolicy per batch.
+func (pMgr *PolicyManager) AddAllPolicies(policyKeys map[string]struct{}, epToModifyID, epToModifyIP string) (map[string]struct{}, error) {
+	pMgr.policyMap.Lock()
+	defer pMgr.policyMap.Unlock()
+
+	klog.Infof("[PolicyManagerWindows] adding all policies. epID: %s. epIP: %s. policyKeys: %+v", epToModifyID, epToModifyIP, policyKeys)
+
+	batches, err := pMgr.batchPolicies(policyKeys, epToModifyID, epToModifyIP)
+	if err != nil {
+		return nil, fmt.Errorf("error while batching policies for endpoint. err: %w", err)
+	}
+
+	successfulPolicies := make(map[string]struct{})
+
+	for i, batch := range batches {
+		klog.Infof("[PolicyManagerWindows] processing batch %d out of %d for adding all policies to endpoint. endpoint ID: %s. policyBatch: %+v", i+1, len(batches), epToModifyID, batch.policies)
+
+		epPolicyRequest, err := getEPPolicyReqFromACLSettings(batch.rules)
+		if err != nil {
+			return successfulPolicies, fmt.Errorf("error while applying all policies for batch %d out of %d. ruleBatch: %+v. err: %w", i+1, len(batches), batch, err)
+		}
+
+		klog.Infof("[PolicyManager] applying all rules to endpoint for batch %d out of %d. endpoint ID: %s", i+1, len(batches), epToModifyID)
+		err = pMgr.applyPoliciesToEndpointID(epToModifyID, epPolicyRequest)
+		if err != nil {
+			return successfulPolicies, fmt.Errorf("failed to add all policies on endpoint for batch %d out of %d. ruleBatch: %+v. err: %w", i+1, len(batches), batch, err)
+		}
+
+		klog.Infof("[PolicyManager] finished applying all rules to endpoint for batch %d out of %d. endpoint ID: %s, policyBatch: %+v", i+1, len(batches), epToModifyID, batch.policies)
+		for _, policyKey := range batch.policies {
+			policy, ok := pMgr.policyMap.cache[policyKey]
+			if ok {
+				policy.PodEndpoints[epToModifyIP] = epToModifyID
+				successfulPolicies[policyKey] = struct{}{}
+			} else {
+				klog.Errorf("[PolicyManagerWindows] unexpected error: policy not found after adding all policies for batch %d out of %d. policyKey: %s. epID: %s",
+					i+1, len(batches), policyKey, epToModifyID)
+				metrics.SendErrorLogAndMetric(util.IptmID, "[PolicyManagerWindows] unexpected error: policy not found after adding all policies for batch %d out of %d. policyKey: %s. epID: %s",
+					i+1, len(batches), policyKey, epToModifyID)
+			}
+		}
+	}
+
+	return successfulPolicies, nil
+}
+
+// batchPolicies returns a list of batches
+func (pMgr *PolicyManager) batchPolicies(policyKeys map[string]struct{}, epToModifyID, epToModifyIP string) ([]*aclBatch, error) {
+	batches := make([]*aclBatch, 0)
+	for policyKey := range policyKeys {
+		policy, ok := pMgr.policyMap.cache[policyKey]
+		if !ok {
+			klog.Infof("[PolicyManagerWindows] policy not found while adding all policies. policyKey: %s. epID: %s", policyKey, epToModifyID)
+			delete(policyKeys, policyKey)
+			continue
+		}
+
+		// 1. remove stale endpoints from policy.PodEndpoints and skip adding to endpoints that already have the policy
+		if policy.PodEndpoints == nil {
+			policy.PodEndpoints = make(map[string]string)
+		}
+
+		epID, ok := policy.PodEndpoints[epToModifyIP]
+		if ok {
+			if epID == epToModifyID {
+				klog.Infof("[PolicyManagerWindows] while adding all policies, will not add policy %s to endpoint since it already exists there. endpoint IP: %s, endpoint ID: %s",
+					policy.PolicyKey, epToModifyIP, epToModifyID)
+				delete(policyKeys, policyKey)
+				continue
+			}
+
+			// If the expected ID is not same as epID, there is a chance that old pod got deleted
+			// and same IP is used by new pod with new endpoint.
+			// so we should delete the non-existent endpoint from policy reference
+			klog.Infof("[PolicyManagerWindows] while adding all policies, removing deleted endpoint from policy's current endpoints. policy: %s, endpoint IP: %s, new ID: %s, previous ID: %s",
+				policy.PolicyKey, epToModifyIP, epToModifyID, epID)
+			delete(policy.PodEndpoints, epToModifyIP)
+		}
+
+		// 2. add this policy's rules to a batch
+		policyRules, err := pMgr.getSettingsFromACL(policy)
+		if err != nil {
+			return batches, fmt.Errorf("error while getting settings while applying all policies. err: %w", err)
+		}
+
+		if len(batches) > 0 {
+			batch := batches[len(batches)-1]
+			if len(batch.rules)+len(policyRules) <= pMgr.MaxBatchedACLsPerPod {
+				batch.rules = append(batch.rules, policyRules...)
+				batch.policies = append(batch.policies, policy.PolicyKey)
+				continue
+			}
+		}
+
+		// create a new batch
+		// either this is the first NetPol we've seen, or adding this NetPol's rules to the previous batch would exceed the max rules per batch
+		batch := &aclBatch{
+			rules:    policyRules,
+			policies: []string{policy.PolicyKey},
+		}
+		batches = append(batches, batch)
+	}
+
+	return batches, nil
+}
+
 // AddBaseACLsForCalicoCNI attempts to add base ACLs for Calico CNI.
 func (pMgr *PolicyManager) AddBaseACLsForCalicoCNI(epID string) {
 	epPolicyRequest, err := getEPPolicyReqFromACLSettings(baseACLsForCalicoCNI)