From 66c32f8c0a21301d7976abd87139fd8724b14820 Mon Sep 17 00:00:00 2001
From: Hunter Gregory <hunterlgregory@gmail.com>
Date: Mon, 23 May 2022 14:43:20 -0700
Subject: [PATCH] fix race

---
 npm/pkg/dataplane/dataplane.go         | 26 +++++++++++++++++---------
 npm/pkg/dataplane/dataplane_windows.go |  4 +++-
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/npm/pkg/dataplane/dataplane.go b/npm/pkg/dataplane/dataplane.go
index 20f7208f98..2172aa9972 100644
--- a/npm/pkg/dataplane/dataplane.go
+++ b/npm/pkg/dataplane/dataplane.go
@@ -33,7 +33,10 @@ type DataPlane struct {
 	endpointCache  map[string]*NPMEndpoint
 	ioShim         *common.IOShim
 	updatePodCache map[string]*updateNPMPod
-	stopChannel    <-chan struct{}
+	// pendingPolicies includes the policy keys of policies which may
+	// be referenced by ipsets but have not been applied to the kernel yet
+	pendingPolicies map[string]struct{}
+	stopChannel     <-chan struct{}
 }
 
 // TODO this struct could be made unexported
@@ -50,14 +53,15 @@ type NPMEndpoint struct {
 func NewDataPlane(nodeName string, ioShim *common.IOShim, cfg *Config, stopChannel <-chan struct{}) (*DataPlane, error) {
 	metrics.InitializeAll()
 	dp := &DataPlane{
-		Config:         cfg,
-		policyMgr:      policies.NewPolicyManager(ioShim, cfg.PolicyManagerCfg),
-		ipsetMgr:       ipsets.NewIPSetManager(cfg.IPSetManagerCfg, ioShim),
-		endpointCache:  make(map[string]*NPMEndpoint),
-		nodeName:       nodeName,
-		ioShim:         ioShim,
-		updatePodCache: make(map[string]*updateNPMPod),
-		stopChannel:    stopChannel,
+		Config:          cfg,
+		policyMgr:       policies.NewPolicyManager(ioShim, cfg.PolicyManagerCfg),
+		ipsetMgr:        ipsets.NewIPSetManager(cfg.IPSetManagerCfg, ioShim),
+		endpointCache:   make(map[string]*NPMEndpoint),
+		nodeName:        nodeName,
+		ioShim:          ioShim,
+		updatePodCache:  make(map[string]*updateNPMPod),
+		pendingPolicies: make(map[string]struct{}),
+		stopChannel:     stopChannel,
 	}
 
 	err := dp.BootupDataplane()
@@ -202,6 +206,9 @@ func (dp *DataPlane) ApplyDataPlane() error {
 // AddPolicy takes in a translated NPMNetworkPolicy object and applies on dataplane
 func (dp *DataPlane) AddPolicy(policy *policies.NPMNetworkPolicy) error {
 	klog.Infof("[DataPlane] Add Policy called for %s", policy.PolicyKey)
+
+	dp.pendingPolicies[policy.PolicyKey] = struct{}{}
+
 	// Create and add references for Selector IPSets first
 	err := dp.createIPSetsAndReferences(policy.AllPodSelectorIPSets(), policy.PolicyKey, ipsets.SelectorType)
 	if err != nil {
@@ -229,6 +236,7 @@ func (dp *DataPlane) AddPolicy(policy *policies.NPMNetworkPolicy) error {
 	if err != nil {
 		return fmt.Errorf("[DataPlane] error while adding policy: %w", err)
 	}
+	delete(dp.pendingPolicies, policy.PolicyKey)
 	return nil
 }
 
diff --git a/npm/pkg/dataplane/dataplane_windows.go b/npm/pkg/dataplane/dataplane_windows.go
index f8742cd3df..ac6446fbac 100644
--- a/npm/pkg/dataplane/dataplane_windows.go
+++ b/npm/pkg/dataplane/dataplane_windows.go
@@ -162,7 +162,9 @@ func (dp *DataPlane) updatePod(pod *updateNPMPod) error {
 		}
 
 		for policyKey := range selectorReference {
-			toAddPolicies[policyKey] = struct{}{}
+			if _, ok := dp.pendingPolicies[policyKey]; !ok {
+				toAddPolicies[policyKey] = struct{}{}
+			}
 		}
 	}