Skip to content

Commit 85201c2

Browse files
mattia-eleuterilexfrei
authored andcommitted
feat(scheduler): set scheduling-class label on mutated pods
The lineage-controller-webhook now sets a scheduler.cozystack.io/scheduling-class label on mutated pods (in addition to the existing annotation of the same name). This unlocks SchedulingClass authors writing podAffinity rules whose labelSelector targets pods across different applications that share the same class. A natural pattern becomes: define a class whose name is the affinity-group identifier, attach it to several applications, and have the scheduler co-locate them at the chosen topology key (region, zone, hostname). Behavior preserved when the SchedulingClass CR is missing: pods are left untouched (no scheduler injection, no label) so they do not stay Pending on a non-existent scheduler. Tests cover label/annotation/schedulerName injection, existing label preservation, non-Pod no-op, no-class no-op, and missing-CR no-op. Signed-off-by: mattia-eleuteri <[email protected]>
1 parent f6a615f commit 85201c2

2 files changed

Lines changed: 202 additions & 2 deletions

File tree

‎internal/lineagecontrollerwebhook/webhook.go‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -230,8 +230,14 @@ func hasAllLineageLabels(o *unstructured.Unstructured) bool {
230230
return true
231231
}
232232

233-
// applySchedulingClass injects schedulerName and scheduling-class annotation
234-
// into Pods whose namespace carries the scheduler.cozystack.io/scheduling-class label.
233+
// applySchedulingClass injects schedulerName, scheduling-class annotation, and
234+
// scheduling-class label into Pods whose namespace carries the
235+
// scheduler.cozystack.io/scheduling-class label.
236+
//
237+
// The label (in addition to the annotation) lets SchedulingClass authors write
238+
// podAffinity rules whose labelSelector matches pods across applications that
239+
// share the same class.
240+
//
235241
// If the referenced SchedulingClass CR does not exist (e.g. the scheduler
236242
// package is not installed), the injection is silently skipped so that pods
237243
// are not left Pending.
@@ -283,6 +289,13 @@ func (h *LineageControllerWebhook) applySchedulingClass(ctx context.Context, obj
283289
annotations[schedulerapi.SchedulingClassAnnotation] = schedulingClass
284290
obj.SetAnnotations(annotations)
285291

292+
podLabels := obj.GetLabels()
293+
if podLabels == nil {
294+
podLabels = make(map[string]string)
295+
}
296+
podLabels[schedulerapi.SchedulingClassLabel] = schedulingClass
297+
obj.SetLabels(podLabels)
298+
286299
return nil
287300
}
288301

‎internal/lineagecontrollerwebhook/webhook_test.go‎

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@ func TestHandle_PatchTouchesOnlyLineageAndScheduling(t *testing.T) {
120120
"add /metadata/labels/apps.cozystack.io~1application.name": "demo",
121121
"add /metadata/labels/internal.cozystack.io~1tenantresource": "false",
122122
"add /metadata/annotations/scheduler.cozystack.io~1scheduling-class": "gpu",
123+
"add /metadata/labels/scheduler.cozystack.io~1scheduling-class": "gpu",
123124
"add /spec/schedulerName": schedulerapi.SchedulerName,
124125
}
125126

@@ -273,3 +274,189 @@ func TestHandle_ReadmissionShortCircuit(t *testing.T) {
273274
})
274275
}
275276
}
277+
278+
const testSchedulingClass = "co-region-test"
279+
280+
// newSchedulingClassCR builds an unstructured SchedulingClass CR the fake
281+
// dynamic client can serve.
282+
func newSchedulingClassCR(name string) *unstructured.Unstructured {
283+
u := &unstructured.Unstructured{}
284+
u.SetGroupVersionKind(schema.GroupVersionKind{
285+
Group: schedulerapi.Group,
286+
Version: schedulerapi.Version,
287+
Kind: "SchedulingClass",
288+
})
289+
u.SetName(name)
290+
return u
291+
}
292+
293+
// newWebhookTestEnv builds a LineageControllerWebhook wired with fakes,
294+
// pre-populated with a Namespace carrying the scheduling-class label and a
295+
// SchedulingClass CR by that name.
296+
func newWebhookTestEnv(t *testing.T, namespace string, classOnNamespace string, classExists bool) *LineageControllerWebhook {
297+
t.Helper()
298+
299+
testScheme := runtime.NewScheme()
300+
_ = clientgoscheme.AddToScheme(testScheme)
301+
302+
ns := &corev1.Namespace{
303+
ObjectMeta: metav1.ObjectMeta{
304+
Name: namespace,
305+
},
306+
}
307+
if classOnNamespace != "" {
308+
ns.Labels = map[string]string{
309+
schedulerapi.SchedulingClassLabel: classOnNamespace,
310+
}
311+
}
312+
313+
c := fake.NewClientBuilder().
314+
WithScheme(testScheme).
315+
WithObjects(ns).
316+
Build()
317+
318+
gvr := schema.GroupVersionResource{
319+
Group: schedulerapi.Group,
320+
Version: schedulerapi.Version,
321+
Resource: schedulerapi.Resource,
322+
}
323+
324+
dynObjs := []runtime.Object{}
325+
if classExists {
326+
dynObjs = append(dynObjs, newSchedulingClassCR(classOnNamespace))
327+
}
328+
dyn := dynamicfake.NewSimpleDynamicClientWithCustomListKinds(
329+
testScheme,
330+
map[schema.GroupVersionResource]string{gvr: "SchedulingClassList"},
331+
dynObjs...,
332+
)
333+
334+
return &LineageControllerWebhook{
335+
Client: c,
336+
Scheme: testScheme,
337+
dynClient: dyn,
338+
}
339+
}
340+
341+
// newPod returns a minimal unstructured Pod for mutation tests.
342+
func newPod(namespace, name string) *unstructured.Unstructured {
343+
u := &unstructured.Unstructured{}
344+
u.SetGroupVersionKind(schema.GroupVersionKind{Version: "v1", Kind: "Pod"})
345+
u.SetNamespace(namespace)
346+
u.SetName(name)
347+
return u
348+
}
349+
350+
func TestApplySchedulingClass_SetsLabelAnnotationAndSchedulerName(t *testing.T) {
351+
ctx := context.Background()
352+
h := newWebhookTestEnv(t, "tenant-acme", testSchedulingClass, true)
353+
pod := newPod("tenant-acme", "db-0")
354+
355+
if err := h.applySchedulingClass(ctx, pod, nil, "tenant-acme"); err != nil {
356+
t.Fatalf("applySchedulingClass returned error: %v", err)
357+
}
358+
359+
gotLabel := pod.GetLabels()[schedulerapi.SchedulingClassLabel]
360+
if gotLabel != testSchedulingClass {
361+
t.Errorf("expected pod label %s=%s, got %q", schedulerapi.SchedulingClassLabel, testSchedulingClass, gotLabel)
362+
}
363+
364+
gotAnnotation := pod.GetAnnotations()[schedulerapi.SchedulingClassAnnotation]
365+
if gotAnnotation != testSchedulingClass {
366+
t.Errorf("expected pod annotation %s=%s, got %q", schedulerapi.SchedulingClassAnnotation, testSchedulingClass, gotAnnotation)
367+
}
368+
369+
gotSched, _, _ := unstructured.NestedString(pod.Object, "spec", "schedulerName")
370+
if gotSched != schedulerapi.SchedulerName {
371+
t.Errorf("expected pod spec.schedulerName=%s, got %q", schedulerapi.SchedulerName, gotSched)
372+
}
373+
}
374+
375+
func TestApplySchedulingClass_PreservesExistingLabels(t *testing.T) {
376+
ctx := context.Background()
377+
h := newWebhookTestEnv(t, "tenant-acme", testSchedulingClass, true)
378+
pod := newPod("tenant-acme", "db-0")
379+
pod.SetLabels(map[string]string{
380+
"app.kubernetes.io/name": "postgres",
381+
"role": "primary",
382+
})
383+
384+
if err := h.applySchedulingClass(ctx, pod, nil, "tenant-acme"); err != nil {
385+
t.Fatalf("applySchedulingClass returned error: %v", err)
386+
}
387+
388+
labels := pod.GetLabels()
389+
if labels[schedulerapi.SchedulingClassLabel] != testSchedulingClass {
390+
t.Errorf("scheduling-class label not set")
391+
}
392+
if labels["app.kubernetes.io/name"] != "postgres" {
393+
t.Errorf("existing label app.kubernetes.io/name was clobbered: %q", labels["app.kubernetes.io/name"])
394+
}
395+
if labels["role"] != "primary" {
396+
t.Errorf("existing label role was clobbered: %q", labels["role"])
397+
}
398+
}
399+
400+
func TestApplySchedulingClass_NonPodIsNoOp(t *testing.T) {
401+
ctx := context.Background()
402+
h := newWebhookTestEnv(t, "tenant-acme", testSchedulingClass, true)
403+
404+
svc := &unstructured.Unstructured{}
405+
svc.SetGroupVersionKind(schema.GroupVersionKind{Version: "v1", Kind: "Service"})
406+
svc.SetNamespace("tenant-acme")
407+
svc.SetName("svc-0")
408+
409+
if err := h.applySchedulingClass(ctx, svc, nil, "tenant-acme"); err != nil {
410+
t.Fatalf("applySchedulingClass returned error: %v", err)
411+
}
412+
413+
if _, ok := svc.GetLabels()[schedulerapi.SchedulingClassLabel]; ok {
414+
t.Errorf("non-Pod object should not be mutated")
415+
}
416+
if _, ok := svc.GetAnnotations()[schedulerapi.SchedulingClassAnnotation]; ok {
417+
t.Errorf("non-Pod object should not be mutated")
418+
}
419+
}
420+
421+
func TestApplySchedulingClass_NoClassOnNamespaceIsNoOp(t *testing.T) {
422+
ctx := context.Background()
423+
h := newWebhookTestEnv(t, "tenant-acme", "", false)
424+
pod := newPod("tenant-acme", "db-0")
425+
426+
if err := h.applySchedulingClass(ctx, pod, nil, "tenant-acme"); err != nil {
427+
t.Fatalf("applySchedulingClass returned error: %v", err)
428+
}
429+
430+
if _, ok := pod.GetLabels()[schedulerapi.SchedulingClassLabel]; ok {
431+
t.Errorf("pod label should not be set when namespace has no scheduling-class label")
432+
}
433+
if _, ok := pod.GetAnnotations()[schedulerapi.SchedulingClassAnnotation]; ok {
434+
t.Errorf("pod annotation should not be set when namespace has no scheduling-class label")
435+
}
436+
if sched, _, _ := unstructured.NestedString(pod.Object, "spec", "schedulerName"); sched != "" {
437+
t.Errorf("pod spec.schedulerName should not be set when namespace has no scheduling-class label, got %q", sched)
438+
}
439+
}
440+
441+
func TestApplySchedulingClass_MissingClassCRSkipsInjection(t *testing.T) {
442+
ctx := context.Background()
443+
// Namespace carries the label, but the SchedulingClass CR is NOT created.
444+
h := newWebhookTestEnv(t, "tenant-acme", testSchedulingClass, false)
445+
pod := newPod("tenant-acme", "db-0")
446+
447+
if err := h.applySchedulingClass(ctx, pod, nil, "tenant-acme"); err != nil {
448+
t.Fatalf("applySchedulingClass returned error: %v", err)
449+
}
450+
451+
// Defensive behavior: pod must not be left referencing a non-existent
452+
// scheduler, otherwise it would stay Pending forever.
453+
if _, ok := pod.GetLabels()[schedulerapi.SchedulingClassLabel]; ok {
454+
t.Errorf("pod label should not be set when SchedulingClass CR is missing")
455+
}
456+
if _, ok := pod.GetAnnotations()[schedulerapi.SchedulingClassAnnotation]; ok {
457+
t.Errorf("pod annotation should not be set when SchedulingClass CR is missing")
458+
}
459+
if sched, _, _ := unstructured.NestedString(pod.Object, "spec", "schedulerName"); sched != "" {
460+
t.Errorf("pod spec.schedulerName should not be set when SchedulingClass CR is missing, got %q", sched)
461+
}
462+
}

0 commit comments

Comments
 (0)