fix: tighten stored gvk review feedback

Signed-off-by: Alex Jun <aljun@nvidia.com>

Author: alexscjundev

fault-remediation/pkg/annotation/annotation.go

@@ -96,6 +96,10 @@ func (m *NodeAnnotationManager) GetRemediationState(
 // UpdateRemediationState updates the node annotation with new remediation state
 func (m *NodeAnnotationManager) UpdateRemediationState(ctx context.Context, nodeName string,
 	group string, crName string, actionName string, resourceRef MaintenanceResourceReference) error {
+	if err := validateMaintenanceResourceReference(resourceRef); err != nil {
+		return fmt.Errorf("invalid maintenance resource reference for node %s group %s: %w", nodeName, group, err)
+	}
+
 	err := retry.RetryOnConflict(conflictBackoff, func() error {
 		// Get current state
 		state, node, err := m.GetRemediationState(ctx, nodeName)
@@ -111,7 +115,7 @@ func (m *NodeAnnotationManager) UpdateRemediationState(ctx context.Context, node
 			ActionName:    actionName,
 			Namespace:     resourceRef.Namespace,
 			Version:       resourceRef.Version,
-			ApiGroup:      resourceRef.ApiGroup,
+			APIGroup:      resourceRef.APIGroup,
 			Kind:          resourceRef.Kind,
 		}
 
@@ -147,11 +151,12 @@ func (m *NodeAnnotationManager) UpdateRemediationState(ctx context.Context, node
 }
 
 // EnsureRemediationStateGVK backfills concrete resource identity for legacy
-// annotation entries using the resourceRefs map, if and only if the GVK is not already set.
+// annotation entries using resourceRefsByAction, if and only if the GVK is not already set.
+// The map must be keyed by ActionName.
 func (m *NodeAnnotationManager) EnsureRemediationStateGVK(
 	ctx context.Context,
 	nodeName string,
-	resourceRefs map[string]MaintenanceResourceReference,
+	resourceRefsByAction map[string]MaintenanceResourceReference,
 ) error {
 	err := retry.RetryOnConflict(conflictBackoff, func() error {
 		state, node, err := m.GetRemediationState(ctx, nodeName)
@@ -163,7 +168,7 @@ func (m *NodeAnnotationManager) EnsureRemediationStateGVK(
 		changed := false
 
 		for group, groupState := range state.EquivalenceGroups {
-			resourceRef, exists := resourceRefs[groupState.ActionName]
+			resourceRef, exists := resourceRefsByAction[groupState.ActionName]
 			if !exists {
 				continue
 			}
@@ -207,14 +212,22 @@ func (m *NodeAnnotationManager) EnsureRemediationStateGVK(
 	return nil
 }
 
+func validateMaintenanceResourceReference(resourceRef MaintenanceResourceReference) error {
+	if resourceRef.APIGroup == "" || resourceRef.Version == "" || resourceRef.Kind == "" {
+		return fmt.Errorf("apiGroup, version, and kind must be non-empty")
+	}
+
+	return nil
+}
+
 func backfillResourceReference(
 	groupState *EquivalenceGroupState,
 	resourceRef MaintenanceResourceReference,
 ) bool {
 	changed := false
 
-	if groupState.ApiGroup == "" && resourceRef.ApiGroup != "" {
-		groupState.ApiGroup = resourceRef.ApiGroup
+	if groupState.APIGroup == "" && resourceRef.APIGroup != "" {
+		groupState.APIGroup = resourceRef.APIGroup
 		changed = true
 	}
 

fault-remediation/pkg/annotation/annotation_interface.go

@@ -37,10 +37,12 @@ type NodeAnnotationManagerInterface interface {
 		actionName string,
 		resourceRef MaintenanceResourceReference,
 	) error
+	// EnsureRemediationStateGVK backfills legacy annotation entries using
+	// resource references keyed by ActionName.
 	EnsureRemediationStateGVK(
 		ctx context.Context,
 		nodeName string,
-		resourceRefs map[string]MaintenanceResourceReference,
+		resourceRefsByAction map[string]MaintenanceResourceReference,
 	) error
 	ClearRemediationState(ctx context.Context, nodeName string) error
 	RemoveGroupsFromState(ctx context.Context, nodeName string, groups []string) error
@@ -56,7 +58,7 @@ type RemediationStateAnnotation struct {
 type MaintenanceResourceReference struct {
 	Namespace string
 	Version   string
-	ApiGroup  string
+	APIGroup  string
 	Kind      string
 }
 
@@ -72,6 +74,6 @@ type EquivalenceGroupState struct {
 	// Concrete resource identity for the CR that was created.
 	Namespace string `json:"namespace,omitempty"`
 	Version   string `json:"version,omitempty"`
-	ApiGroup  string `json:"apiGroup,omitempty"`
+	APIGroup  string `json:"apiGroup,omitempty"`
 	Kind      string `json:"kind,omitempty"`
 }

fault-remediation/pkg/annotation/annotation_test.go

@@ -29,6 +29,14 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 )
 
+func testMaintenanceResourceReference() MaintenanceResourceReference {
+	return MaintenanceResourceReference{
+		Version:  "v1alpha1",
+		APIGroup: "janitor.dgxc.nvidia.com",
+		Kind:     "RebootNode",
+	}
+}
+
 func TestGetRemediationState(t *testing.T) {
 	ctx := context.Background()
 	nodeName := "test-node"
@@ -144,7 +152,7 @@ func TestUpdateRemediationState(t *testing.T) {
 	resourceRef := MaintenanceResourceReference{
 		Namespace: "dgxc-janitor",
 		Version:   "v1alpha1",
-		ApiGroup:  "janitor.dgxc.nvidia.com",
+		APIGroup:  "janitor.dgxc.nvidia.com",
 		Kind:      "RebootNode",
 	}
 	node := &corev1.Node{
@@ -168,10 +176,63 @@ func TestUpdateRemediationState(t *testing.T) {
 	assert.Equal(t, actionName, state.EquivalenceGroups[group].ActionName)
 	assert.Equal(t, resourceRef.Namespace, state.EquivalenceGroups[group].Namespace)
 	assert.Equal(t, resourceRef.Version, state.EquivalenceGroups[group].Version)
-	assert.Equal(t, resourceRef.ApiGroup, state.EquivalenceGroups[group].ApiGroup)
+	assert.Equal(t, resourceRef.APIGroup, state.EquivalenceGroups[group].APIGroup)
 	assert.Equal(t, resourceRef.Kind, state.EquivalenceGroups[group].Kind)
 }
 
+func TestUpdateRemediationStateRejectsIncompleteReference(t *testing.T) {
+	nodeName := "node"
+	node := &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:        nodeName,
+			Annotations: map[string]string{},
+		},
+	}
+	client := fake.NewClientBuilder().WithObjects(node).Build()
+	annotationManager := NodeAnnotationManager{client: client}
+
+	err := annotationManager.UpdateRemediationState(
+		context.TODO(),
+		nodeName,
+		"test",
+		"reboot",
+		"reboot-action",
+		MaintenanceResourceReference{Version: "v1alpha1", Kind: "RebootNode"},
+	)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "apiGroup, version, and kind must be non-empty")
+}
+
+func TestUpdateRemediationStateAllowsClusterScopedReference(t *testing.T) {
+	group := "test"
+	nodeName := "node"
+	resourceRef := MaintenanceResourceReference{
+		Version:  "v1alpha1",
+		APIGroup: "janitor.dgxc.nvidia.com",
+		Kind:     "RebootNode",
+	}
+	node := &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:        nodeName,
+			Annotations: map[string]string{},
+		},
+	}
+	client := fake.NewClientBuilder().WithObjects(node).Build()
+	annotationManager := NodeAnnotationManager{client: client}
+
+	err := annotationManager.UpdateRemediationState(context.TODO(), nodeName, group, "reboot", "reboot-action", resourceRef)
+	require.NoError(t, err)
+
+	updatedNode := &corev1.Node{}
+	require.NoError(t, client.Get(context.TODO(), types.NamespacedName{Name: nodeName}, updatedNode))
+	assert.Contains(t, updatedNode.Annotations[AnnotationKey], `"apiGroup":"janitor.dgxc.nvidia.com"`)
+
+	state, _, err := annotationManager.GetRemediationState(context.TODO(), nodeName)
+	require.NoError(t, err)
+	assert.Empty(t, state.EquivalenceGroups[group].Namespace)
+	assert.Equal(t, resourceRef.APIGroup, state.EquivalenceGroups[group].APIGroup)
+}
+
 func TestEnsureRemediationStateGVKBackfillsLegacyState(t *testing.T) {
 	nodeName := "node"
 	createdAt := time.Now().Add(-time.Hour).UTC()
@@ -195,7 +256,7 @@ func TestEnsureRemediationStateGVKBackfillsLegacyState(t *testing.T) {
 	resourceRef := MaintenanceResourceReference{
 		Namespace: "dgxc-janitor",
 		Version:   "v1alpha1",
-		ApiGroup:  "janitor.dgxc.nvidia.com",
+		APIGroup:  "janitor.dgxc.nvidia.com",
 		Kind:      "RebootNode",
 	}
 	client := fake.NewClientBuilder().WithObjects(node).Build()
@@ -215,7 +276,7 @@ func TestEnsureRemediationStateGVKBackfillsLegacyState(t *testing.T) {
 	assert.Equal(t, createdAt.Unix(), groupState.CreatedAt.Unix())
 	assert.Equal(t, resourceRef.Namespace, groupState.Namespace)
 	assert.Equal(t, resourceRef.Version, groupState.Version)
-	assert.Equal(t, resourceRef.ApiGroup, groupState.ApiGroup)
+	assert.Equal(t, resourceRef.APIGroup, groupState.APIGroup)
 	assert.Equal(t, resourceRef.Kind, groupState.Kind)
 }
 
@@ -248,7 +309,7 @@ func TestEnsureRemediationStateGVKDoesNotOverwriteExistingReference(t *testing.T
 		"RESTART_BM": {
 			Namespace: "new-namespace",
 			Version:   "v2",
-			ApiGroup:  "new.example.com",
+			APIGroup:  "new.example.com",
 			Kind:      "NewKind",
 		},
 	})
@@ -259,7 +320,7 @@ func TestEnsureRemediationStateGVKDoesNotOverwriteExistingReference(t *testing.T
 
 	groupState := state.EquivalenceGroups["restart"]
 	assert.Equal(t, "old-namespace", groupState.Namespace)
-	assert.Equal(t, "old.example.com", groupState.ApiGroup)
+	assert.Equal(t, "old.example.com", groupState.APIGroup)
 	assert.Equal(t, "v1", groupState.Version)
 	assert.Equal(t, "OldKind", groupState.Kind)
 }
@@ -394,10 +455,10 @@ func TestConcurrentUpdateAndRemoveGroupsFromState(t *testing.T) {
 	for i := 0; i < iterations; i++ {
 		// Reset state each iteration
 		err := annotationManager.UpdateRemediationState(context.TODO(), nodeName, "existing-group-1", "old-cr-1",
-			"RESTART_BM", MaintenanceResourceReference{})
+			"RESTART_BM", testMaintenanceResourceReference())
 		require.NoError(t, err)
 		err = annotationManager.UpdateRemediationState(context.TODO(), nodeName, "existing-group-2", "old-cr-2",
-			"COMPONENT_RESET", MaintenanceResourceReference{})
+			"COMPONENT_RESET", testMaintenanceResourceReference())
 		require.NoError(t, err)
 
 		var wg sync.WaitGroup
@@ -411,7 +472,7 @@ func TestConcurrentUpdateAndRemoveGroupsFromState(t *testing.T) {
 		go func() {
 			defer wg.Done()
 			_ = annotationManager.UpdateRemediationState(context.TODO(), nodeName, "new-group", "new-cr",
-				"COMPONENT_RESET", MaintenanceResourceReference{})
+				"COMPONENT_RESET", testMaintenanceResourceReference())
 		}()
 
 		wg.Wait()

fault-remediation/pkg/crstatus/checker.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"log/slog"
 
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -49,6 +50,7 @@ func NewCRStatusChecker(
 	}
 }
 
+// GetCRStateForReference fetches the stored maintenance CR reference and returns its current remediation state.
 func (c *CRStatusChecker) GetCRStateForReference(
 	ctx context.Context,
 	crName string,
@@ -61,7 +63,7 @@ func (c *CRStatusChecker) GetCRStateForReference(
 	}
 
 	gvk := schema.GroupVersionKind{
-		Group:   resourceRef.ApiGroup,
+		Group:   resourceRef.APIGroup,
 		Version: resourceRef.Version,
 		Kind:    resourceRef.Kind,
 	}
@@ -76,8 +78,14 @@ func (c *CRStatusChecker) GetCRStateForReference(
 	key := client.ObjectKey{Name: crName, Namespace: resourceRef.Namespace}
 
 	if err := c.client.Get(ctx, key, obj); err != nil {
-		slog.WarnContext(ctx, "Failed to get CR, allowing create", "crName", crName, "gvk", gvk.String(), "error", err)
-		return CRStateNotFound
+		if apierrors.IsNotFound(err) {
+			slog.InfoContext(ctx, "Stored CR was not found", "crName", crName, "gvk", gvk.String())
+			return CRStateNotFound
+		}
+
+		slog.ErrorContext(ctx, "Failed to get CR, keeping create blocked",
+			"crName", crName, "gvk", gvk.String(), "error", err)
+		return CRStateInProgress
 	}
 
 	return c.checkConditionType(obj, completeConditionType)

fault-remediation/pkg/crstatus/crstatus_test.go

@@ -16,11 +16,14 @@ package crstatus
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+	"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
 
 	"github.com/nvidia/nvsentinel/fault-remediation/pkg/annotation"
 )
@@ -142,11 +145,35 @@ func TestGetCRStateForReferenceUsesStoredReference(t *testing.T) {
 	checker := NewCRStatusChecker(fakeClient, false)
 
 	state := checker.GetCRStateForReference(context.Background(), "stored-cr", annotation.MaintenanceResourceReference{
-		ApiGroup:  "stored.example.com",
+		APIGroup:  "stored.example.com",
 		Version:   "v9",
 		Kind:      "StoredMaintenance",
 		Namespace: "stored-namespace",
 	}, "NodeReady")
 
 	assert.Equal(t, CRStateSucceeded, state)
 }
+
+func TestGetCRStateForReferenceBlocksOnGetError(t *testing.T) {
+	fakeClient := fake.NewClientBuilder().WithInterceptorFuncs(interceptor.Funcs{
+		Get: func(
+			_ context.Context,
+			_ client.WithWatch,
+			_ client.ObjectKey,
+			_ client.Object,
+			_ ...client.GetOption,
+		) error {
+			return fmt.Errorf("api server unavailable")
+		},
+	}).Build()
+	checker := NewCRStatusChecker(fakeClient, false)
+
+	state := checker.GetCRStateForReference(context.Background(), "stored-cr", annotation.MaintenanceResourceReference{
+		APIGroup:  "stored.example.com",
+		Version:   "v9",
+		Kind:      "StoredMaintenance",
+		Namespace: "stored-namespace",
+	}, "NodeReady")
+
+	assert.Equal(t, CRStateInProgress, state)
+}

fault-remediation/pkg/reconciler/reconciler.go

@@ -527,28 +527,28 @@ func (r *FaultRemediationReconciler) ensureRemediationStateGVK(ctx context.Conte
 		return nil
 	}
 
-	resourceRefs := maintenanceResourceReferences(remediationConfig.RemediationActions)
-	if len(resourceRefs) == 0 {
+	resourceRefsByAction := maintenanceResourceReferences(remediationConfig.RemediationActions)
+	if len(resourceRefsByAction) == 0 {
 		return nil
 	}
 
-	return r.annotationManager.EnsureRemediationStateGVK(ctx, nodeName, resourceRefs)
+	return r.annotationManager.EnsureRemediationStateGVK(ctx, nodeName, resourceRefsByAction)
 }
 
 func maintenanceResourceReferences(
 	remediationActions map[string]config.MaintenanceResource,
 ) map[string]annotation.MaintenanceResourceReference {
-	resourceRefs := make(map[string]annotation.MaintenanceResourceReference, len(remediationActions))
+	resourceRefsByAction := make(map[string]annotation.MaintenanceResourceReference, len(remediationActions))
 	for actionName, resource := range remediationActions {
-		resourceRefs[actionName] = annotation.MaintenanceResourceReference{
+		resourceRefsByAction[actionName] = annotation.MaintenanceResourceReference{
 			Namespace: resource.Namespace,
 			Version:   resource.Version,
-			ApiGroup:  resource.ApiGroup,
+			APIGroup:  resource.ApiGroup,
 			Kind:      resource.Kind,
 		}
 	}
 
-	return resourceRefs
+	return resourceRefsByAction
 }
 
 // trySkipEvent returns (result, err, true) when the event should be skipped; otherwise (zero, nil, false).
@@ -1044,10 +1044,10 @@ func resourceReferenceFromState(
 	resourceRef := annotation.MaintenanceResourceReference{
 		Namespace: groupState.Namespace,
 		Version:   groupState.Version,
-		ApiGroup:  groupState.ApiGroup,
+		APIGroup:  groupState.APIGroup,
 		Kind:      groupState.Kind,
 	}
-	if resourceRef.ApiGroup == "" || resourceRef.Version == "" || resourceRef.Kind == "" {
+	if resourceRef.APIGroup == "" || resourceRef.Version == "" || resourceRef.Kind == "" {
 		return annotation.MaintenanceResourceReference{}, false
 	}
 
@@ -1070,6 +1070,9 @@ func (r *FaultRemediationReconciler) completeConditionTypeForStoredState(
 	if !exists {
 		return "", false
 	}
+	if strings.TrimSpace(resource.CompleteConditionType) == "" {
+		return "", false
+	}
 
 	return resource.CompleteConditionType, true
 }