From 0aeb1e29b5131e836d44b9a3f589644264c15f61 Mon Sep 17 00:00:00 2001
From: Ivan Kolodiazhnyi <ikolodiazhny@nvidia.com>
Date: Mon, 3 Jul 2023 18:49:57 +0300
Subject: [PATCH] Implement DrainAnnotationPredicate to watch only for a needed
 annotation

---
 controllers/drain_controller.go |  5 ++--
 controllers/helper.go           | 47 +++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/controllers/drain_controller.go b/controllers/drain_controller.go
index 0147db36bf..f43d8ac237 100644
--- a/controllers/drain_controller.go
+++ b/controllers/drain_controller.go
@@ -16,7 +16,6 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/event"
 	"sigs.k8s.io/controller-runtime/pkg/handler"
 	"sigs.k8s.io/controller-runtime/pkg/log"
-	"sigs.k8s.io/controller-runtime/pkg/predicate"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 	"sigs.k8s.io/controller-runtime/pkg/source"
 
@@ -92,7 +91,7 @@ func (dr *DrainReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctr
 			}
 			drainingNodes++
 		} else {
-			reqLogger.Info("Too many nodes to be draining at the moment. Skipping node %s", node.Name)
+			reqLogger.Info("Too many nodes to be draining at the moment. Skipping node %s", "node", node.Name)
 			return reconcile.Result{}, nil
 		}
 	}
@@ -121,7 +120,7 @@ func (dr *DrainReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	}
 
 	// Watch for spec and annotation changes
-	nodePredicates := builder.WithPredicates(predicate.AnnotationChangedPredicate{})
+	nodePredicates := builder.WithPredicates(DrainAnnotationPredicate{})
 
 	return ctrl.NewControllerManagedBy(mgr).
 		For(&sriovnetworkv1.SriovOperatorConfig{}).
diff --git a/controllers/helper.go b/controllers/helper.go
index a6d93c2ac8..56da29c58e 100644
--- a/controllers/helper.go
+++ b/controllers/helper.go
@@ -18,10 +18,15 @@ package controllers
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"os"
 	"strings"
 
+	"sigs.k8s.io/controller-runtime/pkg/event"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+	"sigs.k8s.io/controller-runtime/pkg/predicate"
+
 	constants "github.com/k8snetworkplumbingwg/sriov-network-operator/pkg/consts"
 )
 
@@ -32,6 +37,48 @@ var webhooks = map[string](string){
 
 var namespace = os.Getenv("NAMESPACE")
 
+type DrainAnnotationPredicate struct {
+	predicate.Funcs
+}
+
+func (DrainAnnotationPredicate) Create(e event.CreateEvent) bool {
+	logger := log.FromContext(context.TODO())
+	if e.Object == nil {
+		logger.Error(nil, "Create event has no object for create", "event", e)
+		return false
+	}
+
+	if _, hasAnno := e.Object.GetAnnotations()[constants.NodeDrainAnnotation]; hasAnno {
+		logger.Error(nil, "Create event: node has no drain annotation", "event", e)
+		return true
+	}
+	return false
+}
+
+func (DrainAnnotationPredicate) Update(e event.UpdateEvent) bool {
+	logger := log.FromContext(context.TODO())
+	if e.ObjectOld == nil {
+		logger.Error(nil, "Update event has no old object to update", "event", e)
+		return false
+	}
+	if e.ObjectNew == nil {
+		logger.Error(nil, "Update event has no new object for update", "event", e)
+		return false
+	}
+
+	oldAnno, hasOldAnno := e.ObjectOld.GetAnnotations()[constants.NodeDrainAnnotation]
+	newAnno, hasNewAnno := e.ObjectNew.GetAnnotations()[constants.NodeDrainAnnotation]
+
+	if !hasOldAnno || !hasNewAnno {
+		logger.Error(nil, "Update event: can not compare annotations", "old", hasOldAnno)
+		logger.Error(nil, "Update event: can not compare annotations", "new", hasNewAnno)
+		logger.Error(nil, "Update event: can not compare annotations", "event", e)
+		return false
+	}
+
+	return oldAnno == newAnno
+}
+
 func GetImagePullSecrets() []string {
 	imagePullSecrets := os.Getenv("IMAGE_PULL_SECRETS")
 	if imagePullSecrets != "" {