fix: cover cancellation completion marker for fault-remediation (#1335)

Signed-off-by: Kaivalya Dabhadkar <kdabhadkar@nvidia.com>
Co-authored-by: rohansav <rohansavar@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 3 people

fault-remediation/pkg/reconciler/reconciler.go

@@ -155,8 +155,7 @@ func (r *FaultRemediationReconciler) Reconcile(
 	nodeQuarantined := healthEventWithStatus.HealthEventStatus.NodeQuarantined
 
 	if nodeQuarantined == string(model.UnQuarantined) || nodeQuarantined == string(model.Cancelled) {
-		return r.handleCancellationEvent(ctx, nodeName, model.Status(nodeQuarantined), r.Watcher, event.ResumeToken,
-			r.healthEventStore)
+		return r.handleCancellationEvent(ctx, nodeName, model.Status(nodeQuarantined), r.Watcher, *event, r.healthEventStore)
 	}
 
 	return r.handleRemediationEvent(ctx, &healthEventWithStatus, *event, r.Watcher, r.healthEventStore)
@@ -371,13 +370,15 @@ func (r *FaultRemediationReconciler) performRemediation(ctx context.Context,
 	return crName, nil
 }
 
-// handleCancellationEvent handles node unquarantine and cancellation events by clearing annotations
+// handleCancellationEvent handles node unquarantine and cancellation events by clearing
+// annotations and writing a completion marker so the event is excluded from cold start
+// queries on future restarts.
 func (r *FaultRemediationReconciler) handleCancellationEvent(
 	ctx context.Context,
 	nodeName string,
 	status model.Status,
 	watcherInstance datastore.ChangeStreamWatcher,
-	resumeToken []byte,
+	eventWithToken datastore.EventWithToken,
 	healthEventStore datastore.HealthEventStore,
 ) (ctrl.Result, error) {
 	ctx, span := tracing.StartSpan(ctx, "fault_remediation.cancellation_event")
@@ -418,7 +419,16 @@ func (r *FaultRemediationReconciler) handleCancellationEvent(
 		return ctrl.Result{}, fmt.Errorf("failed to clear remediation state for node: %w", err)
 	}
 
-	if err := safeMarkProcessed(context.Background(), watcherInstance, resumeToken, nodeName); err != nil {
+	if err := r.updateNodeRemediatedStatus(ctx, healthEventStore, eventWithToken, true); err != nil {
+		slog.ErrorContext(ctx, "Failed to write completion marker for cancellation event",
+			"node", nodeName,
+			"error", err)
+		tracing.RecordError(span, err)
+
+		return ctrl.Result{}, fmt.Errorf("failed to write completion marker for cancellation event: %w", err)
+	}
+
+	if err := safeMarkProcessed(ctx, watcherInstance, eventWithToken.ResumeToken, nodeName); err != nil {
 		return ctrl.Result{}, err
 	}
 
@@ -1128,9 +1138,12 @@ func (r *FaultRemediationReconciler) HandleColdStart(ctx context.Context) {
 		query.Or(
 			// Quarantined + drained but not yet remediated
 			unresolvedRemediationReadyEventsCondition(""),
-			// Cancelled/unquarantined events (need cleanup)
-			query.In("healtheventstatus.nodequarantined",
-				[]interface{}{string(model.UnQuarantined), string(model.Cancelled)}),
+			// Cancelled/unquarantined events that haven't been marked complete
+			query.And(
+				query.In("healtheventstatus.nodequarantined",
+					[]interface{}{string(model.UnQuarantined), string(model.Cancelled)}),
+				query.Eq("healtheventstatus.faultremediated", nil),
+			),
 		),
 	)
 

fault-remediation/pkg/reconciler/reconciler_e2e_test.go

@@ -2103,6 +2103,241 @@ func TestHandleColdStart_CancellationFlow(t *testing.T) {
 	_ = testDynamic.Resource(gvr).Delete(ctx, crName, metav1.DeleteOptions{})
 }
 
+// TestCancellationEvent_WritesCompletionMarker verifies that handleCancellationEvent writes
+// faultRemediated=true to the health event store, making cold start self-terminating for
+// cancellation events.
+func TestCancellationEvent_WritesCompletionMarker(t *testing.T) {
+	ctx, cancel := context.WithTimeout(testContext, 30*time.Second)
+	defer cancel()
+
+	nodeName := testutils.GenerateTestNodeName("test-cancel-marker")
+	createTestNode(ctx, nodeName, nil, map[string]string{"test": "label"})
+	defer func() {
+		_ = testClient.CoreV1().Nodes().Delete(ctx, nodeName, metav1.DeleteOptions{})
+	}()
+
+	cleanupNodeAnnotations(ctx, t, nodeName)
+
+	gvr := schema.GroupVersionResource{
+		Group:    "janitor.dgxc.nvidia.com",
+		Version:  "v1alpha1",
+		Resource: "rebootnodes",
+	}
+
+	remediationClient, err := createTestRemediationClient(false, restartRemediationActions)
+	require.NoError(t, err)
+
+	var capturedStatuses []datastore.HealthEventStatus
+	var capturedIDs []string
+	var mu sync.Mutex
+
+	localStore := &MockHealthEventStore{}
+	localStore.UpdateHealthEventStatusFn = func(ctx context.Context, id string, status datastore.HealthEventStatus) error {
+		mu.Lock()
+		capturedIDs = append(capturedIDs, id)
+		capturedStatuses = append(capturedStatuses, status)
+		mu.Unlock()
+		return nil
+	}
+
+	localWatcher := NewMockChangeStreamWatcher()
+	cfg := ReconcilerConfig{
+		RemediationClient: remediationClient,
+		StateManager:      statemanager.NewStateManager(testClient),
+		UpdateMaxRetries:  3,
+		UpdateRetryDelay:  100 * time.Millisecond,
+	}
+	localReconciler := NewFaultRemediationReconciler(nil, localWatcher, localStore, cfg, false)
+
+	// Pre-set the node state to match what fault-quarantine + node-drainer would have done.
+	_, err = cfg.StateManager.UpdateNVSentinelStateNodeLabel(ctx, nodeName,
+		statemanager.QuarantinedLabelValue, false)
+	require.NoError(t, err)
+	_, err = cfg.StateManager.UpdateNVSentinelStateNodeLabel(ctx, nodeName,
+		statemanager.DrainingLabelValue, false)
+	require.NoError(t, err)
+	_, err = cfg.StateManager.UpdateNVSentinelStateNodeLabel(ctx, nodeName,
+		statemanager.DrainSucceededLabelValue, false)
+	require.NoError(t, err)
+
+	// Step 1: Send a quarantine event to establish remediation state
+	t.Log("Step 1: Processing quarantine event to create CR and annotation")
+	eventID := "cancel-marker-event-1"
+	quarantineEvent := createQuarantineEvent(eventID, nodeName, protos.RecommendedAction_RESTART_BM)
+	quarantineEventToken := datastore.EventWithToken{
+		Event:       map[string]interface{}(quarantineEvent),
+		ResumeToken: []byte("cancel-marker-token-1"),
+	}
+	_, err = localReconciler.Reconcile(ctx, &quarantineEventToken)
+	require.NoError(t, err)
+
+	var crName string
+	require.Eventually(t, func() bool {
+		state, _, err := localReconciler.annotationManager.GetRemediationState(ctx, nodeName)
+		if err != nil {
+			return false
+		}
+		if grp, ok := state.EquivalenceGroups["restart"]; ok {
+			crName = grp.MaintenanceCR
+			return crName != ""
+		}
+		return false
+	}, 5*time.Second, 100*time.Millisecond, "CR and annotation should be created")
+
+	mu.Lock()
+	capturedIDs = nil
+	capturedStatuses = nil
+	mu.Unlock()
+
+	// Step 2: Send the cancellation event and capture only the cancellation store updates.
+	t.Log("Step 2: Sending cancellation event and capturing store updates")
+	cancelledEvent := createCancelledEvent(eventID, nodeName, protos.RecommendedAction_RESTART_BM)
+	cancelledEventToken := datastore.EventWithToken{
+		Event:       map[string]interface{}(cancelledEvent),
+		ResumeToken: []byte("cancel-marker-token-2"),
+	}
+	_, err = localReconciler.Reconcile(ctx, &cancelledEventToken)
+	require.NoError(t, err)
+
+	// Step 3: Wait for remediation state to be cleared (proves cancellation was processed)
+	require.Eventually(t, func() bool {
+		node, err := testClient.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
+		if err != nil {
+			return false
+		}
+		_, hasAnnotation := node.Annotations[annotation.AnnotationKey]
+		return !hasAnnotation
+	}, 5*time.Second, 100*time.Millisecond, "Remediation annotation should be cleared")
+
+	// Step 4: Verify that UpdateHealthEventStatus was called with FaultRemediated=true
+	t.Log("Step 4: Verifying completion marker was written")
+	mu.Lock()
+	defer mu.Unlock()
+
+	require.NotEmpty(t, capturedStatuses, "UpdateHealthEventStatus should have been called for cancellation")
+
+	var foundCompletionMarker bool
+	for i, status := range capturedStatuses {
+		if capturedIDs[i] == eventID && status.FaultRemediated != nil && *status.FaultRemediated {
+			foundCompletionMarker = true
+			t.Logf("Completion marker written for event ID: %s (FaultRemediated=true)", capturedIDs[i])
+			assert.NotNil(t, status.LastRemediationTimestamp,
+				"LastRemediationTimestamp should be set when FaultRemediated=true")
+			break
+		}
+	}
+	require.True(t, foundCompletionMarker,
+		"handleCancellationEvent must write FaultRemediated=true for the cancelled event")
+
+	_ = testDynamic.Resource(gvr).Delete(ctx, crName, metav1.DeleteOptions{})
+}
+
+// TestColdStartCancellationQueryRequiresCompletionMarker verifies that the cold start
+// cancellation query includes the same completion gate as the remediation query leg.
+// Without this faultremediated==nil predicate, every historical cancellation can replay
+// on every restart.
+func TestColdStartCancellationQueryRequiresCompletionMarker(t *testing.T) {
+	ctx, cancel := context.WithTimeout(testContext, 30*time.Second)
+	defer cancel()
+
+	var capturedQuery map[string]interface{}
+	coldStartCallCount := 0
+
+	localStore := &MockHealthEventStore{}
+	localStore.FindHealthEventsByQueryBatchedFn = func(_ context.Context, builder datastore.QueryBuilder, _ int, fn func([]datastore.HealthEventWithStatus) error) error {
+		coldStartCallCount++
+		capturedQuery = builder.ToMongo()
+		return fn([]datastore.HealthEventWithStatus{})
+	}
+
+	remediationClient, err := createTestRemediationClient(false, restartRemediationActions)
+	require.NoError(t, err)
+	cfg := ReconcilerConfig{
+		RemediationClient: remediationClient,
+		StateManager:      statemanager.NewStateManager(testClient),
+		UpdateMaxRetries:  3,
+		UpdateRetryDelay:  100 * time.Millisecond,
+	}
+	localReconciler := NewFaultRemediationReconciler(nil, NewMockChangeStreamWatcher(), localStore, cfg, false)
+
+	localReconciler.HandleColdStart(ctx)
+
+	assert.Equal(t, 1, coldStartCallCount,
+		"Cold start should query the store exactly once")
+	require.NotNil(t, capturedQuery, "Cold start query should be captured")
+	require.True(t, coldStartCancellationQueryHasCompletionGate(capturedQuery),
+		"Cold start cancellation query must include faultremediated==nil gate")
+}
+
+func coldStartCancellationQueryHasCompletionGate(q map[string]interface{}) bool {
+	orConditions, ok := q["$or"].([]interface{})
+	if !ok {
+		return false
+	}
+
+	for _, rawCondition := range orConditions {
+		condition, ok := rawCondition.(map[string]interface{})
+		if !ok {
+			continue
+		}
+
+		if isCancellationQueryLeg(condition) && hasFaultRemediatedNilGate(condition) {
+			return true
+		}
+
+		andConditions, ok := condition["$and"].([]interface{})
+		if !ok {
+			continue
+		}
+
+		hasCancellationFilter := false
+		hasCompletionGate := false
+		for _, rawAndCondition := range andConditions {
+			andCondition, ok := rawAndCondition.(map[string]interface{})
+			if !ok {
+				continue
+			}
+			hasCancellationFilter = hasCancellationFilter || isCancellationQueryLeg(andCondition)
+			hasCompletionGate = hasCompletionGate || hasFaultRemediatedNilGate(andCondition)
+		}
+		if hasCancellationFilter && hasCompletionGate {
+			return true
+		}
+	}
+
+	return false
+}
+
+func isCancellationQueryLeg(condition map[string]interface{}) bool {
+	nodeQuarantined, ok := condition["healtheventstatus.nodequarantined"].(map[string]interface{})
+	if !ok {
+		return false
+	}
+
+	values, ok := nodeQuarantined["$in"].([]interface{})
+	if !ok {
+		return false
+	}
+
+	foundUnquarantined := false
+	foundCancelled := false
+	for _, value := range values {
+		switch value {
+		case string(model.UnQuarantined):
+			foundUnquarantined = true
+		case string(model.Cancelled):
+			foundCancelled = true
+		}
+	}
+
+	return foundUnquarantined && foundCancelled
+}
+
+func hasFaultRemediatedNilGate(condition map[string]interface{}) bool {
+	value, ok := condition["healtheventstatus.faultremediated"]
+	return ok && value == nil
+}
+
 // TestCustomAction_E2E tests the full reconciler flow with a custom remediation action.
 // Exercises the complete path: raw MongoDB event → Reconcile → GetEffectiveActionName
 // → template rendering → CR creation.

tests/fault_management_test.go

@@ -547,14 +547,20 @@ func TestManualUncordonWithFaultRemediation(t *testing.T) {
 		WithLabel("suite", "fault-management-manual-intervention-fr")
 
 	var testCtx *helpers.RemediationTestContext
+	var mongoPod string
 	var podNames []string
 	testNamespace := "immediate-test"
+	const cancellationMarkerMessage = "XID 79 fatal error for FR cancellation marker"
 
 	feature.Setup(func(ctx context.Context, t *testing.T, c *envconf.Config) context.Context {
 		var newCtx context.Context
 		newCtx, testCtx = helpers.SetupFaultRemediationTest(ctx, t, c, testNamespace)
 		newCtx = helpers.ApplyNodeDrainerConfig(newCtx, t, c, "data/nd-all-modes.yaml")
 
+		client, err := c.NewClient()
+		require.NoError(t, err)
+		mongoPod, _ = helpers.TryGetMongoDBPrimaryPodName(ctx, t, client)
+
 		return newCtx
 	})
 
@@ -569,7 +575,7 @@ func TestManualUncordonWithFaultRemediation(t *testing.T) {
 		t.Log("Sending fatal health event to trigger quarantine")
 		fatalEvent := helpers.NewHealthEvent(testCtx.NodeName).
 			WithErrorCode("79").
-			WithMessage("XID 79 fatal error").
+			WithMessage(cancellationMarkerMessage).
 			WithRecommendedAction(24)
 		helpers.SendHealthEvent(ctx, t, fatalEvent)
 
@@ -639,6 +645,25 @@ func TestManualUncordonWithFaultRemediation(t *testing.T) {
 		return ctx
 	})
 
+	feature.Assess("verify FR marks cancelled event remediated", func(ctx context.Context, t *testing.T, c *envconf.Config) context.Context {
+		if mongoPod == "" {
+			t.Log("Skipping direct faultRemediated MongoDB assertion for non-MongoDB datastore")
+			return ctx
+		}
+
+		client, err := c.NewClient()
+		require.NoError(t, err)
+
+		t.Log("Waiting for FR to write faultRemediated=true on the Cancelled HealthEvent")
+		require.Eventually(t, func() bool {
+			return helpers.CancelledHealthEventFaultRemediated(ctx, t, client.RESTConfig(), client, mongoPod,
+				testCtx.NodeName, cancellationMarkerMessage)
+		}, helpers.EventuallyWaitTimeout, helpers.WaitInterval,
+			"FR should write faultRemediated=true for the cancelled HealthEvent")
+
+		return ctx
+	})
+
 	feature.Teardown(func(ctx context.Context, t *testing.T, c *envconf.Config) context.Context {
 		client, err := c.NewClient()
 		require.NoError(t, err)