chore: minor changes

Signed-off-by: Ajay Mishra <[email protected]>

Author: XRFXLP

health-monitors/kubernetes-object-monitor/pkg/annotations/manager.go

@@ -0,0 +1,174 @@
+// Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package annotations
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"maps"
+
+	v1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/client-go/util/retry"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+const (
+	AnnotationKey = "k8s-object-monitor.nvsentinel.nvidia.com/policy-matches"
+)
+
+type PolicyMatchState map[string]string
+
+type Manager struct {
+	client client.Client
+}
+
+func NewManager(c client.Client) *Manager {
+	return &Manager{client: c}
+}
+
+func (m *Manager) AddMatch(ctx context.Context, nodeName, stateKey, targetNode string) error {
+	slog.Debug("Adding match state to annotation", "node", nodeName, "stateKey", stateKey, "targetNode", targetNode)
+
+	return m.updateAnnotation(ctx, nodeName, func(state PolicyMatchState) (PolicyMatchState, bool) {
+		if _, exists := state[stateKey]; exists {
+			return state, false
+		}
+
+		state[stateKey] = targetNode
+
+		return state, true
+	})
+}
+
+func (m *Manager) RemoveMatch(ctx context.Context, nodeName, stateKey string) error {
+	slog.Debug("Removing match state from annotation", "node", nodeName, "stateKey", stateKey)
+
+	err := m.updateAnnotation(ctx, nodeName, func(state PolicyMatchState) (PolicyMatchState, bool) {
+		if _, exists := state[stateKey]; !exists {
+			return state, false
+		}
+
+		delete(state, stateKey)
+
+		return state, true
+	})
+	if apierrors.IsNotFound(err) {
+		return nil
+	}
+
+	return err
+}
+
+func (m *Manager) GetMatches(ctx context.Context, nodeName string) (PolicyMatchState, error) {
+	node := &v1.Node{}
+	if err := m.client.Get(ctx, types.NamespacedName{Name: nodeName}, node); err != nil {
+		if apierrors.IsNotFound(err) {
+			return make(PolicyMatchState), nil
+		}
+
+		return make(PolicyMatchState), fmt.Errorf("failed to get node: %w", err)
+	}
+
+	if node.Annotations == nil {
+		return make(PolicyMatchState), nil
+	}
+
+	annotationStr, exists := node.Annotations[AnnotationKey]
+	if !exists || annotationStr == "" {
+		return make(PolicyMatchState), nil
+	}
+
+	var state PolicyMatchState
+	if err := json.Unmarshal([]byte(annotationStr), &state); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal annotation: %w", err)
+	}
+
+	return state, nil
+}
+
+func (m *Manager) LoadAllMatches(ctx context.Context) (map[string]string, error) {
+	nodeList := &v1.NodeList{}
+	if err := m.client.List(ctx, nodeList); err != nil {
+		return nil, fmt.Errorf("failed to list nodes: %w", err)
+	}
+
+	allMatches := make(map[string]string)
+
+	for _, node := range nodeList.Items {
+		matches, err := m.GetMatches(ctx, node.Name)
+		if err != nil {
+			slog.Warn("Failed to load matches from node", "node", node.Name, "error", err)
+			continue
+		}
+
+		maps.Copy(allMatches, matches)
+	}
+
+	return allMatches, nil
+}
+
+func (m *Manager) updateAnnotation(ctx context.Context,
+	nodeName string,
+	updateFn func(PolicyMatchState) (PolicyMatchState, bool),
+) error {
+	slog.Debug("Updating annotation", "node", nodeName)
+
+	return retry.RetryOnConflict(retry.DefaultBackoff, func() error {
+		node := &v1.Node{}
+		if err := m.client.Get(ctx, types.NamespacedName{Name: nodeName}, node); err != nil {
+			if apierrors.IsNotFound(err) {
+				slog.Warn("Node not found, assuming deleted", "node", nodeName)
+				return nil
+			}
+
+			return fmt.Errorf("failed to get node: %w", err)
+		}
+
+		if node.Annotations == nil {
+			node.Annotations = make(map[string]string)
+		}
+
+		state := make(PolicyMatchState)
+		if existingAnnotation := node.Annotations[AnnotationKey]; existingAnnotation != "" {
+			if err := json.Unmarshal([]byte(existingAnnotation), &state); err != nil {
+				slog.Warn("Failed to parse existing annotation, starting fresh", "node", nodeName, "error", err)
+			}
+		}
+
+		updatedState, changed := updateFn(state)
+		if !changed {
+			return nil
+		}
+
+		if len(updatedState) == 0 {
+			delete(node.Annotations, AnnotationKey)
+		} else {
+			annotationBytes, err := json.Marshal(updatedState)
+			if err != nil {
+				return fmt.Errorf("failed to marshal state: %w", err)
+			}
+
+			node.Annotations[AnnotationKey] = string(annotationBytes)
+		}
+
+		return m.client.Update(ctx, node, &client.UpdateOptions{
+			FieldManager: "kubernetes-object-monitor",
+		})
+	})
+}

health-monitors/kubernetes-object-monitor/pkg/controller/reconciler.go

@@ -17,8 +17,10 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"maps"
 	"sync"
 
+	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/annotations"
 	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/config"
 	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/metrics"
 	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/policy"
@@ -36,6 +38,7 @@ type ResourceReconciler struct {
 	client.Client
 	evaluator     *policy.Evaluator
 	publisher     HealthEventPublisher
+	annotationMgr *annotations.Manager
 	policies      []config.Policy
 	gvk           schema.GroupVersionKind
 	matchStates   map[string]string
@@ -46,19 +49,37 @@ func NewResourceReconciler(
 	c client.Client,
 	evaluator *policy.Evaluator,
 	pub HealthEventPublisher,
+	annotationMgr *annotations.Manager,
 	policies []config.Policy,
 	gvk schema.GroupVersionKind,
 ) *ResourceReconciler {
 	return &ResourceReconciler{
-		Client:      c,
-		evaluator:   evaluator,
-		publisher:   pub,
-		policies:    policies,
-		gvk:         gvk,
-		matchStates: make(map[string]string),
+		Client:        c,
+		evaluator:     evaluator,
+		publisher:     pub,
+		annotationMgr: annotationMgr,
+		policies:      policies,
+		gvk:           gvk,
+		matchStates:   make(map[string]string),
 	}
 }
 
+func (r *ResourceReconciler) LoadState(ctx context.Context) error {
+	allMatches, err := r.annotationMgr.LoadAllMatches(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to load match state: %w", err)
+	}
+
+	r.matchStatesMu.Lock()
+	defer r.matchStatesMu.Unlock()
+
+	maps.Copy(r.matchStates, allMatches)
+
+	slog.Info("Loaded policy match state from annotations", "gvk", r.gvk.String(), "matches", len(allMatches))
+
+	return nil
+}
+
 func (r *ResourceReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
 	obj := &unstructured.Unstructured{}
 	obj.SetGroupVersionKind(r.gvk)
@@ -93,6 +114,8 @@ func (r *ResourceReconciler) handleGetError(ctx context.Context, err error, req
 }
 
 func (r *ResourceReconciler) cleanupDeletedResource(ctx context.Context, req ctrl.Request) {
+	slog.Info("Cleaning up deleted resource", "resource", req.NamespacedName)
+
 	for _, p := range r.policies {
 		if !p.Enabled {
 			continue
@@ -119,9 +142,19 @@ func (r *ResourceReconciler) cleanupDeletedResource(ctx context.Context, req ctr
 				metrics.HealthEventsPublishErrors.WithLabelValues(p.Name, "grpc_error").Inc()
 			}
 
+			slog.Debug("Removing match state for deleted resource",
+				"resource", req.NamespacedName,
+				"node", nodeName,
+				"policy", p.Name,
+				"stateKey", stateKey)
+
 			r.matchStatesMu.Lock()
 			delete(r.matchStates, stateKey)
 			r.matchStatesMu.Unlock()
+
+			if err := r.annotationMgr.RemoveMatch(ctx, nodeName, stateKey); err != nil {
+				slog.Error("Failed to remove match state from annotation", "node", nodeName, "stateKey", stateKey, "error", err)
+			}
 		}
 	}
 }
@@ -179,6 +212,10 @@ func (r *ResourceReconciler) handleUnhealthyTransition(
 	r.matchStates[stateKey] = nodeName
 	r.matchStatesMu.Unlock()
 
+	if err := r.annotationMgr.AddMatch(ctx, nodeName, stateKey, nodeName); err != nil {
+		slog.Error("Failed to persist match state to annotation", "node", nodeName, "stateKey", stateKey, "error", err)
+	}
+
 	metrics.PolicyMatches.WithLabelValues(p.Name, nodeName, r.gvk.Kind).Inc()
 
 	return nil
@@ -199,6 +236,10 @@ func (r *ResourceReconciler) handleHealthyTransition(
 	delete(r.matchStates, stateKey)
 	r.matchStatesMu.Unlock()
 
+	if err := r.annotationMgr.RemoveMatch(ctx, nodeName, stateKey); err != nil {
+		slog.Error("Failed to remove match state from annotation", "node", nodeName, "stateKey", stateKey, "error", err)
+	}
+
 	return nil
 }
 

health-monitors/kubernetes-object-monitor/pkg/controller/reconciler_test.go

@@ -20,6 +20,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/annotations"
 	celenv "github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/cel"
 	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/config"
 	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/controller"
@@ -362,22 +363,23 @@ func TestReconciler_ColdStart(t *testing.T) {
 			postRestartAction: func(t *testing.T, s *testSetup, nodeName string, _ *v1.Node) {
 				updateNodeStatus(t, s, nodeName, v1.ConditionFalse)
 			},
-			expectEvent:   true,
-			expectHealthy: false,
+			expectEvent: false,
 		},
 		{
 			name: "healthy resource",
 			postRestartAction: func(t *testing.T, s *testSetup, nodeName string, _ *v1.Node) {
 				updateNodeStatus(t, s, nodeName, v1.ConditionTrue)
 			},
-			expectEvent: false,
+			expectEvent:   true,
+			expectHealthy: true,
 		},
 		{
 			name: "deleted resource",
 			postRestartAction: func(t *testing.T, s *testSetup, _ string, node *v1.Node) {
 				require.NoError(t, s.k8sClient.Delete(s.ctx, node))
 			},
-			expectEvent: false,
+			expectEvent:   true,
+			expectHealthy: true,
 		},
 	}
 
@@ -492,10 +494,13 @@ func setupTestWithPolicies(t *testing.T, policies []config.Policy) *testSetup {
 		Kind:    "Node",
 	}
 
+	annotationMgr := annotations.NewManager(k8sClient)
+
 	reconciler := controller.NewResourceReconciler(
 		k8sClient,
 		evaluator,
 		mockPub,
+		annotationMgr,
 		policies,
 		gvk,
 	)
@@ -544,10 +549,13 @@ func setupTestWithCRD(t *testing.T, policies []config.Policy, crd *apiextensions
 		Kind:    crd.Spec.Names.Kind,
 	}
 
+	annotationMgr := annotations.NewManager(k8sClient)
+
 	reconciler := controller.NewResourceReconciler(
 		k8sClient,
 		evaluator,
 		mockPub,
+		annotationMgr,
 		policies,
 		gvk,
 	)
@@ -649,14 +657,21 @@ func restartReconciler(t *testing.T, setup *testSetup) *testSetup {
 
 	policies := []config.Policy{defaultNodeNotReadyPolicy()}
 
+	annotationMgr := annotations.NewManager(setup.k8sClient)
+
 	reconciler := controller.NewResourceReconciler(
 		setup.k8sClient,
 		setup.evaluator,
 		mockPub,
+		annotationMgr,
 		policies,
 		gvk,
 	)
 
+	if err := reconciler.LoadState(setup.ctx); err != nil {
+		t.Fatalf("Failed to load state after restart: %v", err)
+	}
+
 	return &testSetup{
 		ctx:        setup.ctx,
 		k8sClient:  setup.k8sClient,

health-monitors/kubernetes-object-monitor/pkg/initializer/initializer.go

@@ -24,6 +24,7 @@ import (
 	"time"
 
 	pb "github.com/nvidia/nvsentinel/data-models/pkg/protos"
+	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/annotations"
 	celenv "github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/cel"
 	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/config"
 	"github.com/nvidia/nvsentinel/health-monitors/kubernetes-object-monitor/pkg/controller"
@@ -150,10 +151,15 @@ func registerControllers(
 	policies []config.Policy,
 	maxConcurrentReconciles int,
 ) error {
+	annotationMgr := annotations.NewManager(mgr.GetClient())
 	gvkPolicies := groupPoliciesByGVK(policies)
 
 	for gvk, policies := range gvkPolicies {
-		reconciler := controller.NewResourceReconciler(mgr.GetClient(), evaluator, pub, policies, gvk)
+		reconciler := controller.NewResourceReconciler(mgr.GetClient(), evaluator, pub, annotationMgr, policies, gvk)
+
+		if err := reconciler.LoadState(context.Background()); err != nil {
+			slog.Warn("Failed to load state for controller, starting fresh", "gvk", gvk.String(), "error", err)
+		}
 
 		if err := ctrl.NewControllerManagedBy(mgr).
 			For(newUnstructuredForGVK(gvk)).