feat: Added Disruption control for Sandbox

feat: added PDB to Sandbox spec

updated rbac generated file

nit

nit

Author: vicentefb

api/v1alpha1/sandbox_types.go

@@ -90,6 +90,14 @@ type PersistentVolumeClaimTemplate struct {
 	Spec corev1.PersistentVolumeClaimSpec `json:"spec" protobuf:"bytes,3,opt,name=spec"`
 }
 
+// ResilienceLevel defines the desired level of resilience for a Sandbox.
+type ResilienceLevel string
+
+const (
+	// ResilienceLevelHigh indicates the Sandbox should be protected from voluntary disruptions.
+	ResilienceLevelHigh ResilienceLevel = "High"
+)
+
 // SandboxSpec defines the desired state of Sandbox
 type SandboxSpec struct {
 	// The following markers will use OpenAPI v3 schema to validate the value
@@ -117,6 +125,11 @@ type SandboxSpec struct {
 	// +kubebuilder:validation:Maximum=1
 	// +optional
 	Replicas *int32 `json:"replicas,omitempty"`
+	// Resilience defines the desired level of resilience for the Sandbox Pod.
+	// When set to "High", a PodDisruptionBudget is created to prevent voluntary
+	// disruptions and an annotation is added to prevent cluster-autoscaler evictions.
+	// +optional
+	Resilience ResilienceLevel `json:"resilience,omitempty"`
 }
 
 // SandboxStatus defines the observed state of Sandbox.

codegen.go

@@ -17,5 +17,5 @@
 package agentsandbox
 
 // Generate CRDs and RBAC rules
-//go:generate go tool -modfile=tools.mod sigs.k8s.io/controller-tools/cmd/controller-gen object crd:maxDescLen=0 paths="./api/..." output:crd:dir=k8s/crds output:rbac:dir=k8s rbac:roleName=agent-sandbox-controller,fileName=rbac.generated.yaml
+//go:generate go tool -modfile=tools.mod sigs.k8s.io/controller-tools/cmd/controller-gen rbac:roleName=agent-sandbox-controller,fileName=rbac.generated.yaml crd:maxDescLen=0 paths="./..." output:crd:dir=k8s/crds output:rbac:dir=k8s
 //go:generate go tool -modfile=tools.mod sigs.k8s.io/controller-tools/cmd/controller-gen object crd:maxDescLen=0 paths="./extensions/..." output:crd:dir=k8s/crds output:rbac:dir=k8s rbac:roleName=agent-sandbox-controller,fileName=rbac.generated.yaml

controllers/sandbox_controller.go

@@ -23,11 +23,13 @@ import (
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
+	policyv1 "k8s.io/api/policy/v1"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/intstr"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
 	ctrl "sigs.k8s.io/controller-runtime"
@@ -42,6 +44,8 @@ import (
 
 const (
 	sandboxLabel = "agents.x-k8s.io/sandbox-name-hash"
+	// safeToEvictAnnotation is used to mark pods that should not be evicted by a PDB
+	safeToEvictAnnotation = "cluster-autoscaler.kubernetes.io/safe-to-evict"
 )
 
 var (
@@ -66,6 +70,7 @@ type SandboxReconciler struct {
 //+kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;watch;create;update;patch;delete
 //+kubebuilder:rbac:groups=core,resources=services,verbs=get;list;watch;create;update;patch;delete
 //+kubebuilder:rbac:groups=core,resources=persistentvolumeclaims,verbs=get;list;watch;create;update;patch;delete
+//+kubebuilder:rbac:groups=policy,resources=poddisruptionbudgets,verbs=get;list;watch;create;update;patch;delete
 
 // Reconcile is part of the main kubernetes reconciliation loop which aims to
 // move the current state of the cluster closer to the desired state.
@@ -131,6 +136,10 @@ func (r *SandboxReconciler) reconcileChildResources(ctx context.Context, sandbox
 	err := r.reconcilePVCs(ctx, sandbox)
 	allErrors = errors.Join(allErrors, err)
 
+	// Reconcile PDB
+	err = r.reconcilePDB(ctx, sandbox, nameHash)
+	allErrors = errors.Join(allErrors, err)
+
 	// Reconcile Pod
 	pod, err := r.reconcilePod(ctx, sandbox, nameHash)
 	allErrors = errors.Join(allErrors, err)
@@ -332,6 +341,10 @@ func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1
 		annotations[k] = v
 	}
 
+	if sandbox.Spec.Resilience == sandboxv1alpha1.ResilienceLevelHigh {
+		annotations[safeToEvictAnnotation] = "false"
+	}
+
 	mutatedSpec := sandbox.Spec.PodTemplate.Spec.DeepCopy()
 
 	for _, pvcTemplate := range sandbox.Spec.VolumeClaimTemplates {
@@ -365,6 +378,59 @@ func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1
 	return pod, nil
 }
 
+func (r *SandboxReconciler) reconcilePDB(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox, nameHash string) error {
+	log := log.FromContext(ctx)
+	pdb := &policyv1.PodDisruptionBudget{}
+	pdbName := types.NamespacedName{Name: sandbox.Name, Namespace: sandbox.Namespace}
+
+	// If resilience is not "High", ensure the PDB is deleted.
+	if sandbox.Spec.Resilience != sandboxv1alpha1.ResilienceLevelHigh {
+		if err := r.Get(ctx, pdbName, pdb); err != nil {
+			if k8serrors.IsNotFound(err) {
+				return nil // PDB doesn't exist, which is the desired state.
+			}
+			return err
+		}
+		log.Info("Deleting PDB as resilience level is not High", "PDB.Name", pdb.Name)
+		return r.Delete(ctx, pdb)
+	}
+
+	// If resilience is "High", ensure the PDB exists.
+	if err := r.Get(ctx, pdbName, pdb); err != nil {
+		if !k8serrors.IsNotFound(err) {
+			log.Error(err, "Failed to get PDB")
+			return fmt.Errorf("PDB Get Failed: %w", err)
+		}
+
+		// PDB does not exist, so create it.
+		log.Info("Creating a new PodDisruptionBudget", "PDB.Namespace", sandbox.Namespace, "PDB.Name", sandbox.Name)
+		minAvailable := intstr.FromInt(1) // For a single-pod Sandbox, minAvailable=1 is appropriate
+		newPDB := &policyv1.PodDisruptionBudget{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      sandbox.Name,
+				Namespace: sandbox.Namespace,
+			},
+			Spec: policyv1.PodDisruptionBudgetSpec{
+				MinAvailable: &minAvailable,
+				Selector: &metav1.LabelSelector{
+					MatchLabels: map[string]string{
+						sandboxLabel: nameHash,
+					},
+				},
+			},
+		}
+
+		if err := ctrl.SetControllerReference(sandbox, newPDB, r.Scheme); err != nil {
+			return fmt.Errorf("SetControllerReference for PDB failed: %w", err)
+		}
+
+		return r.Create(ctx, newPDB)
+	}
+
+	log.Info("Found PDB", "PDB.Name", pdb.Name)
+	return nil
+}
+
 func (r *SandboxReconciler) reconcilePVCs(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox) error {
 	log := log.FromContext(ctx)
 	for _, pvcTemplate := range sandbox.Spec.VolumeClaimTemplates {

controllers/sandbox_controller_test.go

@@ -22,11 +22,13 @@ import (
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
+	policyv1 "k8s.io/api/policy/v1"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/utils/ptr"
 	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
 	ctrl "sigs.k8s.io/controller-runtime"
@@ -405,6 +407,88 @@ func TestReconcile(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "sandbox with high resilience creates PDB and adds pod annotation",
+			sandboxSpec: sandboxv1alpha1.SandboxSpec{
+				Resilience: sandboxv1alpha1.ResilienceLevelHigh,
+				PodTemplate: sandboxv1alpha1.PodTemplate{
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{{Name: "test-container"}},
+					},
+				},
+			},
+			// Verify Sandbox status
+			wantStatus: sandboxv1alpha1.SandboxStatus{
+				Service:       sandboxName,
+				ServiceFQDN:   "sandbox-name.sandbox-ns.svc.cluster.local",
+				Replicas:      1,
+				LabelSelector: "agents.x-k8s.io/sandbox-name-hash=ab179450",
+				Conditions: []metav1.Condition{
+					{
+						Type:               "Ready",
+						Status:             "False",
+						ObservedGeneration: 1,
+						Reason:             "DependenciesNotReady",
+						Message:            "Pod exists with phase: ; Service Exists",
+					},
+				},
+			},
+			wantObjs: []client.Object{
+				// Verify Pod has the new annotation
+				&corev1.Pod{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:            sandboxName,
+						Namespace:       sandboxNs,
+						ResourceVersion: "1",
+						Labels: map[string]string{
+							"agents.x-k8s.io/sandbox-name-hash": "ab179450",
+						},
+						Annotations: map[string]string{
+							"cluster-autoscaler.kubernetes.io/safe-to-evict": "false",
+						},
+						OwnerReferences: []metav1.OwnerReference{sandboxControllerRef(sandboxName)},
+					},
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{{Name: "test-container"}},
+					},
+				},
+				// Verify Service
+				&corev1.Service{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:            sandboxName,
+						Namespace:       sandboxNs,
+						ResourceVersion: "1",
+						Labels: map[string]string{
+							"agents.x-k8s.io/sandbox-name-hash": "ab179450",
+						},
+						OwnerReferences: []metav1.OwnerReference{sandboxControllerRef(sandboxName)},
+					},
+					Spec: corev1.ServiceSpec{
+						Selector: map[string]string{
+							"agents.x-k8s.io/sandbox-name-hash": "ab179450",
+						},
+						ClusterIP: "None",
+					},
+				},
+				// Verify the new PDB
+				&policyv1.PodDisruptionBudget{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:            sandboxName,
+						Namespace:       sandboxNs,
+						ResourceVersion: "1",
+						OwnerReferences: []metav1.OwnerReference{sandboxControllerRef(sandboxName)},
+					},
+					Spec: policyv1.PodDisruptionBudgetSpec{
+						MinAvailable: ptr.To(intstr.FromInt(1)),
+						Selector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								"agents.x-k8s.io/sandbox-name-hash": "ab179450",
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, tc := range testCases {
@@ -595,15 +679,68 @@ func TestReconcilePod(t *testing.T) {
 			require.Equal(t, tc.wantPod, pod)
 			// Validate the Pod from the "cluster" (fake client)
 			if tc.wantPod != nil {
+				// If we expect a pod, verify it exists and matches.
 				livePod := &corev1.Pod{}
-				err = r.Get(t.Context(), types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}, livePod)
+				err = r.Get(t.Context(), types.NamespacedName{Name: tc.wantPod.Name, Namespace: tc.wantPod.Namespace}, livePod)
 				require.NoError(t, err)
 				require.Equal(t, tc.wantPod, livePod)
 			} else {
+				// If we don't expect a pod (it was deleted), verify it's gone.
 				livePod := &corev1.Pod{}
-				err = r.Get(t.Context(), types.NamespacedName{Name: sandboxName, Namespace: sandboxNs}, livePod)
+				err = r.Get(t.Context(), types.NamespacedName{Name: tc.sandbox.Name, Namespace: tc.sandbox.Namespace}, livePod)
 				require.True(t, k8serrors.IsNotFound(err))
 			}
 		})
 	}
 }
+
+// This test simulates updating a Sandbox and ensures the controller correctly deletes the now-unneeded PDB
+func TestReconcile_ResilienceCleanup(t *testing.T) {
+	sandboxName := "sandbox-name"
+	sandboxNs := "sandbox-ns"
+
+	// Initial Sandbox with High Resilience
+	sb := &sandboxv1alpha1.Sandbox{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:       sandboxName,
+			Namespace:  sandboxNs,
+			Generation: 1,
+		},
+		Spec: sandboxv1alpha1.SandboxSpec{
+			Resilience: sandboxv1alpha1.ResilienceLevelHigh,
+			PodTemplate: sandboxv1alpha1.PodTemplate{
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{{Name: "test-container"}},
+				},
+			},
+		},
+	}
+
+	r := SandboxReconciler{
+		Client: newFakeClient(sb),
+		Scheme: Scheme,
+	}
+	req := ctrl.Request{NamespacedName: types.NamespacedName{Name: sandboxName, Namespace: sandboxNs}}
+
+	_, err := r.Reconcile(t.Context(), req)
+	require.NoError(t, err)
+
+	// Verify PDB was created
+	pdb := &policyv1.PodDisruptionBudget{}
+	require.NoError(t, r.Get(t.Context(), req.NamespacedName, pdb), "PDB should exist after first reconcile")
+
+	// Update Sandbox to remove resilience
+	liveSandbox := &sandboxv1alpha1.Sandbox{}
+	require.NoError(t, r.Get(t.Context(), req.NamespacedName, liveSandbox))
+	liveSandbox.Spec.Resilience = "" // Remove resilience
+	liveSandbox.Generation = 2
+	require.NoError(t, r.Update(t.Context(), liveSandbox))
+
+	// Re-run reconcile
+	_, err = r.Reconcile(t.Context(), req)
+	require.NoError(t, err)
+
+	// Verify PDB was deleted
+	err = r.Get(t.Context(), req.NamespacedName, pdb)
+	require.True(t, k8serrors.IsNotFound(err), "PDB should be deleted after resilience is removed")
+}

examples/sandbox.yaml

@@ -4,6 +4,7 @@ metadata:
   name: sandbox-example
   namespace: sandbox-ns
 spec:
+  resilience: High
   podTemplate:
     metadata:
       labels:

k8s/crds/agents.x-k8s.io_sandboxes.yaml

@@ -3819,6 +3819,8 @@ spec:
                 maximum: 1
                 minimum: 0
                 type: integer
+              resilience:
+                type: string
               shutdownTime:
                 format: date-time
                 type: string

k8s/rbac.generated.yaml

@@ -38,3 +38,15 @@ rules:
   - get
   - patch
   - update
+- apiGroups:
+  - policy
+  resources:
+  - poddisruptionbudgets
+  verbs:
+  - create
+  - delete
+  - get
+  - list
+  - patch
+  - update
+  - watch