Adding self taint toleration for wasp-agent deployment

manifests/openshift/ds.yaml

@@ -56,6 +56,9 @@ spec:
       priorityClassName: system-node-critical
       serviceAccountName: wasp
       terminationGracePeriodSeconds: 5
+      tolerations:
+        - effect: NoSchedule
+          key: waspEvictionTaint
       volumes:
         - hostPath:
             path: /

pkg/taints/taints.go

@@ -0,0 +1,85 @@
+package taints
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"github.com/openshift-virtualization/wasp-agent/pkg/client"
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+)
+
+const (
+	WaspTaint = "waspEvictionTaint"
+)
+
+func GenerateTaint() v1.Taint {
+	return v1.Taint{
+		Key:    WaspTaint,
+		Effect: v1.TaintEffectNoSchedule,
+	}
+}
+
+func GenerateToleration() v1.Toleration {
+	return v1.Toleration{
+		Key:    WaspTaint,
+		Effect: v1.TaintEffectNoSchedule,
+	}
+}
+
+func NodeHasEvictionTaint(node *v1.Node) bool {
+	// Check if the node has the specified taint
+	for _, taint := range node.Spec.Taints {
+		if taint.Key == WaspTaint && taint.Effect == v1.TaintEffectNoSchedule {
+			return true
+		}
+	}
+	return false
+}
+
+func AddWaspEvictionTaint(waspCli client.WaspClient, node *v1.Node) error {
+
+	taints := append(node.Spec.Taints, GenerateTaint())
+
+	taintsPatch, err := json.Marshal(map[string]interface{}{
+		"spec": map[string]interface{}{
+			"taints": taints,
+		},
+	})
+	if err != nil {
+		return fmt.Errorf("failed to marshal taints patch: %v", err)
+	}
+
+	_, err = waspCli.CoreV1().Nodes().Patch(context.TODO(), node.Name, types.StrategicMergePatchType, taintsPatch, metav1.PatchOptions{})
+	if err != nil {
+		return fmt.Errorf("failed to patch node: %v", err)
+	}
+
+	return nil
+}
+
+func RemoveWaspEvictionTaint(waspCli client.WaspClient, node *v1.Node) error {
+	var newTaints []v1.Taint
+	for _, taint := range node.Spec.Taints {
+		if taint.Key != WaspTaint {
+			newTaints = append(newTaints, taint)
+		}
+	}
+
+	taintsPatch, err := json.Marshal(map[string]interface{}{
+		"spec": map[string]interface{}{
+			"taints": newTaints,
+		},
+	})
+	if err != nil {
+		return fmt.Errorf("failed to marshal taints patch: %v", err)
+	}
+
+	_, err = waspCli.CoreV1().Nodes().Patch(context.TODO(), node.Name, types.StrategicMergePatchType, taintsPatch, metav1.PatchOptions{})
+	if err != nil {
+		return fmt.Errorf("failed to patch node: %v", err)
+	}
+
+	return nil
+}

pkg/wasp/eviction-controller/eviction_controller.go

@@ -2,10 +2,10 @@ package eviction_controller
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"github.com/openshift-virtualization/wasp-agent/pkg/client"
 	"github.com/openshift-virtualization/wasp-agent/pkg/log"
+	wasp_taints "github.com/openshift-virtualization/wasp-agent/pkg/taints"
 	pod_evictor "github.com/openshift-virtualization/wasp-agent/pkg/wasp/pod-evictor"
 	pod_filter "github.com/openshift-virtualization/wasp-agent/pkg/wasp/pod-filter"
 	pod_ranker "github.com/openshift-virtualization/wasp-agent/pkg/wasp/pod-ranker"
@@ -14,7 +14,6 @@ import (
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/fields"
-	"k8s.io/apimachinery/pkg/types"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	"k8s.io/apimachinery/pkg/util/wait"
 	v1lister "k8s.io/client-go/listers/core/v1"
@@ -24,7 +23,6 @@ import (
 
 const (
 	timeToWaitForCacheSync = 10 * time.Second
-	WaspTaint              = "waspEvictionTaint"
 )
 
 type EvictionController struct {
@@ -98,16 +96,16 @@ func (ctrl *EvictionController) handleMemorySwapEviction() {
 		log.Log.Infof(err.Error())
 		return
 	}
-	evicting := nodeHasEvictionTaint(node)
+	evicting := wasp_taints.NodeHasEvictionTaint(node)
 
 	switch {
 	case evicting && !shouldEvict:
-		err := removeWaspEvictionTaint(ctrl.waspCli, node)
+		err := wasp_taints.RemoveWaspEvictionTaint(ctrl.waspCli, node)
 		if err != nil {
 			log.Log.Infof(err.Error())
 		}
 	case !evicting && shouldEvict:
-		err := addWaspEvictionTaint(ctrl.waspCli, node)
+		err := wasp_taints.AddWaspEvictionTaint(ctrl.waspCli, node)
 		if err != nil {
 			log.Log.Infof(err.Error())
 		}
@@ -143,66 +141,6 @@ func (ctrl *EvictionController) handleMemorySwapEviction() {
 	ctrl.statsCollector.FlushStats()
 }
 
-func nodeHasEvictionTaint(node *v1.Node) bool {
-	// Check if the node has the specified taint
-	for _, taint := range node.Spec.Taints {
-		if taint.Key == WaspTaint && taint.Effect == v1.TaintEffectNoSchedule {
-			return true
-		}
-	}
-	return false
-}
-
-func addWaspEvictionTaint(waspCli client.WaspClient, node *v1.Node) error {
-	taint := v1.Taint{
-		Key:    WaspTaint,
-		Effect: v1.TaintEffectNoSchedule,
-	}
-
-	taints := append(node.Spec.Taints, taint)
-
-	taintsPatch, err := json.Marshal(map[string]interface{}{
-		"spec": map[string]interface{}{
-			"taints": taints,
-		},
-	})
-	if err != nil {
-		return fmt.Errorf("failed to marshal taints patch: %v", err)
-	}
-
-	_, err = waspCli.CoreV1().Nodes().Patch(context.TODO(), node.Name, types.StrategicMergePatchType, taintsPatch, metav1.PatchOptions{})
-	if err != nil {
-		return fmt.Errorf("failed to patch node: %v", err)
-	}
-
-	return nil
-}
-
-func removeWaspEvictionTaint(waspCli client.WaspClient, node *v1.Node) error {
-	var newTaints []v1.Taint
-	for _, taint := range node.Spec.Taints {
-		if taint.Key != WaspTaint {
-			newTaints = append(newTaints, taint)
-		}
-	}
-
-	taintsPatch, err := json.Marshal(map[string]interface{}{
-		"spec": map[string]interface{}{
-			"taints": newTaints,
-		},
-	})
-	if err != nil {
-		return fmt.Errorf("failed to marshal taints patch: %v", err)
-	}
-
-	_, err = waspCli.CoreV1().Nodes().Patch(context.TODO(), node.Name, types.StrategicMergePatchType, taintsPatch, metav1.PatchOptions{})
-	if err != nil {
-		return fmt.Errorf("failed to patch node: %v", err)
-	}
-
-	return nil
-}
-
 func waitForSyncedStore(timeout <-chan time.Time, informerSynced func() bool) bool {
 	for !informerSynced() {
 		select {

pkg/wasp/resources/operator/operator.go

@@ -2,6 +2,7 @@ package operator
 
 import (
 	"fmt"
+	"github.com/openshift-virtualization/wasp-agent/pkg/taints"
 
 	"github.com/openshift-virtualization/wasp-agent/pkg/monitoring/rules"
 	utils2 "github.com/openshift-virtualization/wasp-agent/pkg/util"
@@ -170,6 +171,9 @@ func createWaspDaemonSet(namespace, swapUtilizationTHresholdFactor, maxAverageSw
 	container.Env = createDaemonSetEnvVar(swapUtilizationTHresholdFactor, maxAverageSwapInPagesPerSecond, maxAverageSwapOutPagesPerSecond, averageWindowSizeSeconds, verbosity)
 
 	labels := resources.WithLabels(map[string]string{"name": "wasp"}, utils2.DaemonSetLabels)
+	tolerations := []corev1.Toleration{}
+	tolerations = append(tolerations, taints.GenerateToleration())
+
 	ds := &appsv1.DaemonSet{
 		TypeMeta: metav1.TypeMeta{
 			Kind:       "DaemonSet",
@@ -220,6 +224,7 @@ func createWaspDaemonSet(namespace, swapUtilizationTHresholdFactor, maxAverageSw
 						},
 					},
 					PriorityClassName: "system-node-critical",
+					Tolerations:       tolerations,
 				},
 			},
 			UpdateStrategy: appsv1.DaemonSetUpdateStrategy{