Prevent allocate/free ENIs when node is marked noSchedule (#1927)

* Prevent allocate/free ENIs when node is marked noSchedule * Update UTs
aws · Jul 11, 2022 · b01d356 · b01d356
1 parent f979630
commit b01d356
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 10 deletions.
diff --git a/pkg/ipamd/ipamd.go b/pkg/ipamd/ipamd.go
@@ -691,7 +691,7 @@ func (c *IPAMContext) decreaseDatastorePool(interval time.Duration) {
 
 // tryFreeENI always tries to free one ENI
 func (c *IPAMContext) tryFreeENI() {
-	if c.isTerminating() {
+	if c.isTerminating() || c.isNodeNonSchedulable() {
 		log.Debug("AWS CNI is terminating, not detaching any ENIs")
 		return
 	}
@@ -780,7 +780,7 @@ func (c *IPAMContext) increaseDatastorePool(ctx context.Context) {
 		}
 	}
 
-	if c.isTerminating() {
+	if c.isTerminating() || c.isNodeNonSchedulable() {
 		log.Debug("AWS CNI is terminating, will not try to attach any new IPs or ENIs right now")
 		return
 	}
@@ -1806,6 +1806,34 @@ func (c *IPAMContext) isTerminating() bool {
 	return atomic.LoadInt32(&c.terminating) > 0
 }
 
+func (c *IPAMContext) isNodeNonSchedulable() bool {
+	ctx := context.TODO()
+
+	request := types.NamespacedName{
+		Name: c.myNodeName,
+	}
+
+	node := &corev1.Node{}
+	// Find my node
+	err := c.cachedK8SClient.Get(ctx, request, node)
+	if err != nil {
+		log.Errorf("Failed to get node while determining schedulability: %v", err)
+		return false
+	}
+	log.Debugf("Node found %q - no of taints - %d", node.Name, len(node.Spec.Taints))
+	taintToMatch := &corev1.Taint{
+		Key:    "node.kubernetes.io/unschedulable",
+		Effect: corev1.TaintEffectNoSchedule,
+	}
+	for _, taint := range node.Spec.Taints {
+		if taint.MatchTaint(taintToMatch) {
+			return true
+		}
+	}
+
+	return false
+}
+
 // GetConfigForDebug returns the active values of the configuration env vars (for debugging purposes).
 func GetConfigForDebug() map[string]interface{} {
 	return map[string]interface{}{

diff --git a/pkg/ipamd/ipamd_test.go b/pkg/ipamd/ipamd_test.go
@@ -718,14 +718,16 @@ func TestTryAddIPToENI(t *testing.T) {
 
 	warmIPTarget := 3
 	mockContext := &IPAMContext{
-		awsClient:     m.awsutils,
-		maxIPsPerENI:  14,
-		maxENI:        4,
-		warmENITarget: 1,
-		warmIPTarget:  warmIPTarget,
-		networkClient: m.network,
-		primaryIP:     make(map[string]string),
-		terminating:   int32(0),
+		rawK8SClient:    m.rawK8SClient,
+		cachedK8SClient: m.cachedK8SClient,
+		awsClient:       m.awsutils,
+		maxIPsPerENI:    14,
+		maxENI:          4,
+		warmENITarget:   1,
+		warmIPTarget:    warmIPTarget,
+		networkClient:   m.network,
+		primaryIP:       make(map[string]string),
+		terminating:     int32(0),
 	}
 
 	mockContext.dataStore = testDatastore()
@@ -766,6 +768,17 @@ func TestTryAddIPToENI(t *testing.T) {
 	m.awsutils.EXPECT().GetPrimaryENI().Return(primaryENIid)
 	m.network.EXPECT().SetupENINetwork(gomock.Any(), secMAC, secDevice, secSubnet)
 
+	mockContext.myNodeName = myNodeName
+
+	//Create a Fake Node
+	fakeNode := v1.Node{
+		TypeMeta:   metav1.TypeMeta{Kind: "Node"},
+		ObjectMeta: metav1.ObjectMeta{Name: myNodeName},
+		Spec:       v1.NodeSpec{},
+		Status:     v1.NodeStatus{},
+	}
+	_ = m.cachedK8SClient.Create(ctx, &fakeNode)
+
 	mockContext.increaseDatastorePool(ctx)
 }