feat(controller): re-apply runtime-required taint on reboot when autoTaintNewNodes=true (#273)

When REAPPLY_ON_REBOOT=true, runtimeRequired=true, and autoTaintNewNodes=true
are all set, a rebooted node must be re-tainted so workloads cannot schedule on
it until Skyhook finishes re-applying. Previously, node.Reset() cleared the
nodeState/cordon/status annotations but left the autoTaint annotation intact,
so HasSkyhookAnnotations() remained true and HandleAutoTaint never re-tainted
the node.

The fix appends the runtime-required taint to the node inside the existing
ReapplyOnReboot block, after node.Reset() and before the StrategicMerge patch,
so the taint re-application and state reset land in one atomic apiserver write.
Taint removal on completion is unchanged (HandleRuntimeRequiredTaint).

Closes #180

Signed-off-by: Brian Lockwood <lockwobr@gmail.com>

Author: lockwobr

docs/runtime_required.md

@@ -37,6 +37,8 @@ When enabled, the operator automatically applies the runtime-required taint to n
 
 A node is considered "new" if it has no Skyhook annotations. This works for both initial cluster setup (day 0) and nodes joining an existing cluster (day 2+). Nodes that have already been processed by Skyhook (and had their taint removed after completion) will not be re-tainted because they retain their Skyhook annotations.
 
+**Exception: reboot with `REAPPLY_ON_REBOOT=true`.** When the operator is configured with `REAPPLY_ON_REBOOT=true` and a Skyhook has both `runtimeRequired: true` and `autoTaintNewNodes: true`, a node whose boot ID changes is treated as new for taint purposes. The runtime-required taint is re-applied alongside the state reset in the same atomic operation, ensuring no workloads can schedule on the rebooted node before Skyhook finishes re-applying. The taint is removed again by the normal completion path once all runtime-required Skyhooks finish on that node.
+
 ## What runtimeRequired: true will NOT do
 
 1. Without `autoTaintNewNodes: true`, it will NOT add the taint to any nodes targeted by a SCR with `runtimeRequired: true`

operator/internal/controller/skyhook_controller.go

@@ -680,6 +680,19 @@ func (r *SkyhookReconciler) TrackReboots(ctx context.Context, clusterState *clus
 					r.recorder.Eventf(node.GetNode(), nil, EventTypeNormal, EventsReasonNodeReboot, "ResetNodeState", "detected reboot, resetting node for [%s] to be reapplied", node.GetSkyhook().Name)
 					node.Reset()
 
+					// Re-apply the runtime-required taint so workloads cannot schedule on the
+					// rebooted node until Skyhook finishes re-applying. The original auto-taint
+					// annotation survives Reset() and remains the record that this taint is
+					// operator-managed; no annotation update is needed.
+					if skyhook.GetSkyhook().Spec.RuntimeRequired && skyhook.GetSkyhook().Spec.AutoTaintNewNodes {
+						taintToAdd := r.opts.GetRuntimeRequiredTaint()
+						newNode, updated, _ := taints.AddOrUpdateTaint(node.GetNode(), &taintToAdd)
+						if updated {
+							node.GetNode().Spec.Taints = newNode.Spec.Taints
+							log.FromContext(ctx).Info("re-applying runtime-required taint after reboot", "node", node.GetNode().Name, "taint", taintToAdd.Key)
+						}
+					}
+
 					// Persist the reset before recording the new boot id. We Patch rather than
 					// Update because a busy node's resourceVersion churns constantly under other
 					// controllers, and a full Update would lose that optimistic-concurrency race; a

operator/internal/controller/skyhook_controller_test.go

@@ -3795,3 +3795,149 @@ var _ = Describe("ProcessInterrupt skipped-package promotion", func() {
 			"a skipped package must keep waiting while the preempting interrupt is still pending")
 	})
 })
+
+var _ = Describe("TrackReboots auto-taint on reboot", func() {
+	const (
+		defaultRuntimeRequiredTaint = "skyhook.nvidia.com=runtime-required:NoSchedule"
+		oldBootID                   = "boot-A"
+		newBootID                   = "boot-B"
+	)
+
+	It("should re-apply the runtime-required taint when all three conditions are met", func() {
+		const (
+			skyhookName = "auto-taint-reboot-sh"
+			nodeName    = "auto-taint-reboot-node"
+		)
+		nodeLabel := map[string]string{"auto-taint-reboot-test": "yes"}
+		autoTaintAnnotationKey := fmt.Sprintf("%s/autoTaint_skyhook.nvidia.com", v1alpha1.METADATA_PREFIX)
+
+		// Node arrives without the taint (it was removed after previous completion)
+		// but retains the autoTaint annotation from the original auto-taint.
+		node := &corev1.Node{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   nodeName,
+				Labels: nodeLabel,
+				Annotations: map[string]string{
+					autoTaintAnnotationKey: "true",
+				},
+			},
+		}
+		Expect(k8sClient.Create(ctx, node)).To(Succeed())
+		DeferCleanup(func() { _ = k8sClient.Delete(ctx, node) })
+
+		node.Status.NodeInfo.BootID = newBootID
+		Expect(k8sClient.Status().Update(ctx, node)).To(Succeed())
+		snapshotNode := node.DeepCopy()
+
+		pkgRef := v1alpha1.PackageRef{Name: "pkg1", Version: "1.0.0"}
+		skyhook := v1alpha1.Skyhook{
+			ObjectMeta: metav1.ObjectMeta{Name: skyhookName},
+			Spec: v1alpha1.SkyhookSpec{
+				NodeSelector:      metav1.LabelSelector{MatchLabels: nodeLabel},
+				RuntimeRequired:   true,
+				AutoTaintNewNodes: true,
+				Packages: v1alpha1.Packages{
+					pkgRef.Name: {PackageRef: pkgRef, Image: "ghcr.io/org/pkg1"},
+				},
+			},
+			Status: v1alpha1.SkyhookStatus{
+				NodeBootIds: map[string]string{nodeName: oldBootID},
+			},
+		}
+
+		clusterState, err := BuildState(
+			&v1alpha1.SkyhookList{Items: []v1alpha1.Skyhook{skyhook}},
+			&corev1.NodeList{Items: []corev1.Node{*snapshotNode}},
+			&v1alpha1.DeploymentPolicyList{},
+		)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(clusterState.skyhooks[0].GetNodes()).To(HaveLen(1))
+
+		r := &SkyhookReconciler{
+			Client:   k8sClient,
+			recorder: operator.recorder,
+			opts: SkyhookOperatorOptions{
+				ReapplyOnReboot:      true,
+				RuntimeRequiredTaint: defaultRuntimeRequiredTaint,
+			},
+		}
+
+		_, err = r.TrackReboots(ctx, clusterState)
+		// Status().Update fails because the Skyhook is in-memory only; the node patch is what matters.
+		if err != nil {
+			Expect(err.Error()).ToNot(ContainSubstring("node after reboot"),
+				"the node-state reset/taint write must not fail")
+		}
+
+		live := &corev1.Node{}
+		Expect(k8sClient.Get(ctx, types.NamespacedName{Name: nodeName}, live)).To(Succeed())
+		Expect(live.Spec.Taints).To(ContainElement(
+			Equal(corev1.Taint{Key: "skyhook.nvidia.com", Value: "runtime-required", Effect: corev1.TaintEffectNoSchedule}),
+		), "runtime-required taint must be re-applied after reboot when all three conditions are met")
+	})
+
+	It("should NOT re-apply the taint when AutoTaintNewNodes is false", func() {
+		const (
+			skyhookName = "no-auto-taint-reboot-sh"
+			nodeName    = "no-auto-taint-reboot-node"
+		)
+		nodeLabel := map[string]string{"no-auto-taint-reboot-test": "yes"}
+
+		node := &corev1.Node{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:   nodeName,
+				Labels: nodeLabel,
+			},
+		}
+		Expect(k8sClient.Create(ctx, node)).To(Succeed())
+		DeferCleanup(func() { _ = k8sClient.Delete(ctx, node) })
+
+		node.Status.NodeInfo.BootID = newBootID
+		Expect(k8sClient.Status().Update(ctx, node)).To(Succeed())
+		snapshotNode := node.DeepCopy()
+
+		pkgRef := v1alpha1.PackageRef{Name: "pkg1", Version: "1.0.0"}
+		skyhook := v1alpha1.Skyhook{
+			ObjectMeta: metav1.ObjectMeta{Name: skyhookName},
+			Spec: v1alpha1.SkyhookSpec{
+				NodeSelector:      metav1.LabelSelector{MatchLabels: nodeLabel},
+				RuntimeRequired:   true,
+				AutoTaintNewNodes: false, // guard: disabled
+				Packages: v1alpha1.Packages{
+					pkgRef.Name: {PackageRef: pkgRef, Image: "ghcr.io/org/pkg1"},
+				},
+			},
+			Status: v1alpha1.SkyhookStatus{
+				NodeBootIds: map[string]string{nodeName: oldBootID},
+			},
+		}
+
+		clusterState, err := BuildState(
+			&v1alpha1.SkyhookList{Items: []v1alpha1.Skyhook{skyhook}},
+			&corev1.NodeList{Items: []corev1.Node{*snapshotNode}},
+			&v1alpha1.DeploymentPolicyList{},
+		)
+		Expect(err).ToNot(HaveOccurred())
+
+		r := &SkyhookReconciler{
+			Client:   k8sClient,
+			recorder: operator.recorder,
+			opts: SkyhookOperatorOptions{
+				ReapplyOnReboot:      true,
+				RuntimeRequiredTaint: defaultRuntimeRequiredTaint,
+			},
+		}
+
+		_, err = r.TrackReboots(ctx, clusterState)
+		if err != nil {
+			Expect(err.Error()).ToNot(ContainSubstring("node after reboot"),
+				"the node-state reset write must not fail")
+		}
+
+		live := &corev1.Node{}
+		Expect(k8sClient.Get(ctx, types.NamespacedName{Name: nodeName}, live)).To(Succeed())
+		Expect(live.Spec.Taints).ToNot(ContainElement(
+			HaveField("Key", "skyhook.nvidia.com"),
+		), "taint must NOT be re-applied when AutoTaintNewNodes is false")
+	})
+})