Skip to content

Commit 4016287

Browse files
vicentefbvicentefb
committed
feat: Added Disruption control for Sandbox
feat: added PDB to Sandbox spec updated rbac generated file nit nit refacted pdb into its own controller nit sandbox controller test
1 parent 0ab6451 commit 4016287

File tree

7 files changed

+337
-1
lines changed

7 files changed

+337
-1
lines changed

‎api/v1alpha1/sandbox_types.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,10 @@ func (c ConditionType) String() string { return string(c) }
2727
const (
2828
// SandboxConditionReady indicates readiness for Sandbox
2929
SandboxConditionReady ConditionType = "Ready"
30+
31+
// PDBRequiredAnnotation is the annotation that the sandbox-policy-controller
32+
// looks for to create a PodDisruptionBudget for a Sandbox.
33+
PDBRequiredAnnotation = "policy.agents.x-k8s.io/pdb"
3034
)
3135

3236
type PodMetadata struct {

‎cmd/agent-sandbox-controller/main.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,12 @@ func main() {
6969
setupLog.Error(err, "unable to create controller", "controller", "Sandbox")
7070
os.Exit(1)
7171
}
72+
if err = (&controllers.SandboxPolicyReconciler{
73+
Client: mgr.GetClient(),
74+
}).SetupWithManager(mgr); err != nil {
75+
setupLog.Error(err, "unable to create controller", "controller", "SandboxPolicy")
76+
os.Exit(1)
77+
}
7278
//+kubebuilder:scaffold:builder
7379

7480
if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {

‎codegen.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,5 +17,5 @@
1717
package agentsandbox
1818

1919
// Generate CRDs and RBAC rules
20-
//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
20+
//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
2121
//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/sandboxpolicy_controller.go‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
package controllers
2+
3+
import (
4+
"context"
5+
"fmt"
6+
7+
corev1 "k8s.io/api/core/v1"
8+
policyv1 "k8s.io/api/policy/v1"
9+
k8serrors "k8s.io/apimachinery/pkg/api/errors"
10+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
11+
"k8s.io/apimachinery/pkg/types"
12+
"k8s.io/apimachinery/pkg/util/intstr"
13+
ctrl "sigs.k8s.io/controller-runtime"
14+
"sigs.k8s.io/controller-runtime/pkg/client"
15+
"sigs.k8s.io/controller-runtime/pkg/log"
16+
17+
sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
18+
)
19+
20+
// SandboxPolicyReconciler reconciles a Sandbox object for policy purposes
21+
type SandboxPolicyReconciler struct {
22+
client.Client
23+
}
24+
25+
const (
26+
safeToEvictAnnotation = "cluster-autoscaler.kubernetes.io/safe-to-evict"
27+
)
28+
29+
//+kubebuilder:rbac:groups=agents.x-k8s.io,resources=sandboxes,verbs=get;list;watch
30+
//+kubebuilder:rbac:groups=policy,resources=poddisruptionbudgets,verbs=get;list;watch;create;update;patch;delete
31+
//+kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch;patch;update
32+
33+
func (r *SandboxPolicyReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
34+
log := log.FromContext(ctx)
35+
36+
sandbox := &sandboxv1alpha1.Sandbox{}
37+
if err := r.Get(ctx, req.NamespacedName, sandbox); err != nil {
38+
return ctrl.Result{}, client.IgnoreNotFound(err)
39+
}
40+
41+
// Reconcile the PodDisruptionBudget
42+
if err := r.reconcilePDB(ctx, sandbox); err != nil {
43+
log.Error(err, "Failed to reconcile PDB")
44+
return ctrl.Result{}, err
45+
}
46+
47+
// Reconcile the safe-to-evict annotation on the Pod
48+
if err := r.reconcilePodAnnotation(ctx, sandbox); err != nil {
49+
log.Error(err, "Failed to reconcile Pod annotation")
50+
return ctrl.Result{}, err
51+
}
52+
53+
return ctrl.Result{}, nil
54+
}
55+
56+
// reconcilePDB manages the PDB for a Sandbox based on its annotation.
57+
func (r *SandboxPolicyReconciler) reconcilePDB(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox) error {
58+
log := log.FromContext(ctx)
59+
pdb := &policyv1.PodDisruptionBudget{}
60+
pdbName := types.NamespacedName{Name: sandbox.Name, Namespace: sandbox.Namespace}
61+
nameHash := NameHash(sandbox.Name)
62+
63+
// Check if the annotation requests a PDB
64+
pdbRequested := sandbox.Annotations[sandboxv1alpha1.PDBRequiredAnnotation] == "true"
65+
66+
if !pdbRequested {
67+
// If PDB is not requested, ensure it is deleted.
68+
if err := r.Get(ctx, pdbName, pdb); err != nil {
69+
return client.IgnoreNotFound(err) // PDB doesn't exist, which is correct.
70+
}
71+
log.Info("Deleting PDB as policy annotation is not set", "PDB.Name", pdb.Name)
72+
return r.Delete(ctx, pdb)
73+
}
74+
75+
// If PDB is requested, ensure it exists.
76+
if err := r.Get(ctx, pdbName, pdb); err != nil {
77+
if !k8serrors.IsNotFound(err) {
78+
return fmt.Errorf("PDB Get Failed: %w", err)
79+
}
80+
// PDB does not exist, so create it.
81+
log.Info("Creating a new PodDisruptionBudget", "PDB.Name", sandbox.Name)
82+
minAvailable := intstr.FromInt(1)
83+
newPDB := &policyv1.PodDisruptionBudget{
84+
ObjectMeta: metav1.ObjectMeta{Name: sandbox.Name, Namespace: sandbox.Namespace},
85+
Spec: policyv1.PodDisruptionBudgetSpec{
86+
MinAvailable: &minAvailable,
87+
Selector: &metav1.LabelSelector{MatchLabels: map[string]string{sandboxLabel: nameHash}},
88+
},
89+
}
90+
if err := ctrl.SetControllerReference(sandbox, newPDB, r.Scheme()); err != nil {
91+
return err
92+
}
93+
return r.Create(ctx, newPDB)
94+
}
95+
return nil // PDB exists, which is correct.
96+
}
97+
98+
// reconcilePodAnnotation manages the safe-to-evict annotation on the Pod.
99+
func (r *SandboxPolicyReconciler) reconcilePodAnnotation(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox) error {
100+
pod := &corev1.Pod{}
101+
if err := r.Get(ctx, client.ObjectKeyFromObject(sandbox), pod); err != nil {
102+
return client.IgnoreNotFound(err) // Pod may not have been created yet.
103+
}
104+
105+
pdbRequested := sandbox.Annotations[sandboxv1alpha1.PDBRequiredAnnotation] == "true"
106+
patch := client.MergeFrom(pod.DeepCopy())
107+
108+
if pdbRequested {
109+
if pod.Annotations[safeToEvictAnnotation] == "false" {
110+
return nil // Annotation is already correct.
111+
}
112+
if pod.Annotations == nil {
113+
pod.Annotations = make(map[string]string)
114+
}
115+
pod.Annotations[safeToEvictAnnotation] = "false"
116+
} else {
117+
if _, exists := pod.Annotations[safeToEvictAnnotation]; !exists {
118+
return nil // Annotation is already absent.
119+
}
120+
delete(pod.Annotations, safeToEvictAnnotation)
121+
}
122+
123+
return r.Patch(ctx, pod, patch)
124+
}
125+
126+
// SetupWithManager sets up the controller with the Manager.
127+
func (r *SandboxPolicyReconciler) SetupWithManager(mgr ctrl.Manager) error {
128+
return ctrl.NewControllerManagedBy(mgr).
129+
For(&sandboxv1alpha1.Sandbox{}).
130+
Owns(&policyv1.PodDisruptionBudget{}).
131+
Named("sandboxpolicy").
132+
Complete(r)
133+
}

‎controllers/sandboxpolicy_controller_test.go‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
// Copyright 2025 The Kubernetes Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package controllers
16+
17+
import (
18+
"context"
19+
"testing"
20+
21+
"github.com/stretchr/testify/require"
22+
corev1 "k8s.io/api/core/v1"
23+
policyv1 "k8s.io/api/policy/v1"
24+
k8serrors "k8s.io/apimachinery/pkg/api/errors"
25+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
26+
"k8s.io/apimachinery/pkg/types"
27+
sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
28+
ctrl "sigs.k8s.io/controller-runtime"
29+
"sigs.k8s.io/controller-runtime/pkg/client/fake"
30+
)
31+
32+
// TestSandboxPolicyReconciler_Reconcile covers creation and no-op scenarios.
33+
func TestSandboxPolicyReconciler_Reconcile(t *testing.T) {
34+
sandboxName := "test-sandbox"
35+
sandboxNs := "default"
36+
sandboxKey := types.NamespacedName{Name: sandboxName, Namespace: sandboxNs}
37+
38+
// A pod that would be created by the main sandbox-controller
39+
existingPod := &corev1.Pod{
40+
ObjectMeta: metav1.ObjectMeta{
41+
Name: sandboxName,
42+
Namespace: sandboxNs,
43+
Labels: map[string]string{
44+
// Assume this label is present for the selector to match
45+
sandboxLabel: NameHash(sandboxName),
46+
},
47+
},
48+
}
49+
50+
testCases := []struct {
51+
name string
52+
initialSandbox *sandboxv1alpha1.Sandbox
53+
expectPDB bool
54+
expectAnnotation bool
55+
}{
56+
{
57+
name: "should create PDB and add annotation when annotation is present",
58+
initialSandbox: &sandboxv1alpha1.Sandbox{
59+
ObjectMeta: metav1.ObjectMeta{
60+
Name: sandboxName,
61+
Namespace: sandboxNs,
62+
Annotations: map[string]string{
63+
sandboxv1alpha1.PDBRequiredAnnotation: "true",
64+
},
65+
},
66+
},
67+
expectPDB: true,
68+
expectAnnotation: true,
69+
},
70+
{
71+
name: "should do nothing when annotation is absent",
72+
initialSandbox: &sandboxv1alpha1.Sandbox{
73+
ObjectMeta: metav1.ObjectMeta{
74+
Name: sandboxName,
75+
Namespace: sandboxNs,
76+
},
77+
},
78+
expectPDB: false,
79+
expectAnnotation: false,
80+
},
81+
}
82+
83+
for _, tc := range testCases {
84+
t.Run(tc.name, func(t *testing.T) {
85+
fakeClient := fake.NewClientBuilder().
86+
WithScheme(Scheme).
87+
WithRuntimeObjects(tc.initialSandbox.DeepCopy(), existingPod.DeepCopy()).
88+
Build()
89+
90+
r := &SandboxPolicyReconciler{Client: fakeClient}
91+
req := ctrl.Request{NamespacedName: sandboxKey}
92+
93+
_, err := r.Reconcile(context.Background(), req)
94+
require.NoError(t, err)
95+
96+
// Check PDB state
97+
pdb := &policyv1.PodDisruptionBudget{}
98+
err = r.Get(context.Background(), sandboxKey, pdb)
99+
if tc.expectPDB {
100+
require.NoError(t, err, "Expected PDB to be created")
101+
} else {
102+
require.True(t, k8serrors.IsNotFound(err), "Expected PDB to not exist")
103+
}
104+
105+
// Check Pod annotation state
106+
pod := &corev1.Pod{}
107+
err = r.Get(context.Background(), sandboxKey, pod)
108+
require.NoError(t, err)
109+
110+
_, hasAnnotation := pod.Annotations[safeToEvictAnnotation]
111+
require.Equal(t, tc.expectAnnotation, hasAnnotation)
112+
})
113+
}
114+
}
115+
116+
// TestSandboxPolicyReconciler_Cleanup tests the update/cleanup scenario.
117+
func TestSandboxPolicyReconciler_Cleanup(t *testing.T) {
118+
sandboxName := "cleanup-sandbox"
119+
sandboxNs := "default"
120+
sandboxKey := types.NamespacedName{Name: sandboxName, Namespace: sandboxNs}
121+
122+
initialSandbox := &sandboxv1alpha1.Sandbox{
123+
ObjectMeta: metav1.ObjectMeta{
124+
Name: sandboxName,
125+
Namespace: sandboxNs,
126+
Annotations: map[string]string{
127+
sandboxv1alpha1.PDBRequiredAnnotation: "true",
128+
},
129+
},
130+
}
131+
initialPod := &corev1.Pod{
132+
ObjectMeta: metav1.ObjectMeta{
133+
Name: sandboxName,
134+
Namespace: sandboxNs,
135+
Labels: map[string]string{sandboxLabel: NameHash(sandboxName)},
136+
},
137+
}
138+
139+
fakeClient := fake.NewClientBuilder().
140+
WithScheme(Scheme).
141+
WithRuntimeObjects(initialSandbox, initialPod).
142+
Build()
143+
144+
r := &SandboxPolicyReconciler{Client: fakeClient}
145+
req := ctrl.Request{NamespacedName: sandboxKey}
146+
147+
// First reconcile: Create the PDB and add annotation
148+
_, err := r.Reconcile(context.Background(), req)
149+
require.NoError(t, err)
150+
151+
// Verify PDB was created
152+
pdb := &policyv1.PodDisruptionBudget{}
153+
require.NoError(t, r.Get(context.Background(), sandboxKey, pdb), "PDB should exist after first reconcile")
154+
155+
// Verify Pod has annotation
156+
pod := &corev1.Pod{}
157+
require.NoError(t, r.Get(context.Background(), sandboxKey, pod))
158+
require.Equal(t, "false", pod.Annotations[safeToEvictAnnotation])
159+
160+
// Update the Sandbox to remove the annotation
161+
updatedSandbox := &sandboxv1alpha1.Sandbox{}
162+
require.NoError(t, r.Get(context.Background(), sandboxKey, updatedSandbox))
163+
delete(updatedSandbox.Annotations, sandboxv1alpha1.PDBRequiredAnnotation)
164+
require.NoError(t, r.Update(context.Background(), updatedSandbox))
165+
166+
// Second reconcile: Should clean up the PDB and annotation
167+
_, err = r.Reconcile(context.Background(), req)
168+
require.NoError(t, err)
169+
170+
// Verify PDB was deleted
171+
err = r.Get(context.Background(), sandboxKey, pdb)
172+
require.True(t, k8serrors.IsNotFound(err), "PDB should be deleted after annotation is removed")
173+
174+
// Verify Pod annotation was removed
175+
require.NoError(t, r.Get(context.Background(), sandboxKey, pod))
176+
_, hasAnnotation := pod.Annotations[safeToEvictAnnotation]
177+
require.False(t, hasAnnotation, "Pod annotation should be removed")
178+
}

‎examples/sandbox.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@ kind: Sandbox
33
metadata:
44
name: sandbox-example
55
namespace: sandbox-ns
6+
# To request a PDB
7+
annotations:
8+
policy.agents.x-k8s.io/pdb: "true"
69
spec:
710
podTemplate:
811
metadata:

‎k8s/rbac.generated.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,3 +38,15 @@ rules:
3838
- get
3939
- patch
4040
- update
41+
- apiGroups:
42+
- policy
43+
resources:
44+
- poddisruptionbudgets
45+
verbs:
46+
- create
47+
- delete
48+
- get
49+
- list
50+
- patch
51+
- update
52+
- watch

0 commit comments

Comments
 (0)