From 250f89dc72bfeb5cae354f9a79b12187d093a522 Mon Sep 17 00:00:00 2001 From: mattia-eleuteri Date: Tue, 12 May 2026 14:29:28 +0200 Subject: [PATCH] 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 --- internal/lineagecontrollerwebhook/webhook.go | 17 +- .../lineagecontrollerwebhook/webhook_test.go | 187 ++++++++++++++++++ 2 files changed, 202 insertions(+), 2 deletions(-) diff --git a/internal/lineagecontrollerwebhook/webhook.go b/internal/lineagecontrollerwebhook/webhook.go index 7633eb53b..ea8b0d01b 100644 --- a/internal/lineagecontrollerwebhook/webhook.go +++ b/internal/lineagecontrollerwebhook/webhook.go @@ -230,8 +230,14 @@ func hasAllLineageLabels(o *unstructured.Unstructured) bool { return true } -// applySchedulingClass injects schedulerName and scheduling-class annotation -// into Pods whose namespace carries the scheduler.cozystack.io/scheduling-class label. +// applySchedulingClass injects schedulerName, scheduling-class annotation, and +// scheduling-class label into Pods whose namespace carries the +// scheduler.cozystack.io/scheduling-class label. +// +// The label (in addition to the annotation) lets SchedulingClass authors write +// podAffinity rules whose labelSelector matches pods across applications that +// share the same class. +// // If the referenced SchedulingClass CR does not exist (e.g. the scheduler // package is not installed), the injection is silently skipped so that pods // are not left Pending. @@ -283,6 +289,13 @@ func (h *LineageControllerWebhook) applySchedulingClass(ctx context.Context, obj annotations[schedulerapi.SchedulingClassAnnotation] = schedulingClass obj.SetAnnotations(annotations) + podLabels := obj.GetLabels() + if podLabels == nil { + podLabels = make(map[string]string) + } + podLabels[schedulerapi.SchedulingClassLabel] = schedulingClass + obj.SetLabels(podLabels) + return nil } diff --git a/internal/lineagecontrollerwebhook/webhook_test.go b/internal/lineagecontrollerwebhook/webhook_test.go index ea128370e..d0d0fb791 100644 --- a/internal/lineagecontrollerwebhook/webhook_test.go +++ b/internal/lineagecontrollerwebhook/webhook_test.go @@ -120,6 +120,7 @@ func TestHandle_PatchTouchesOnlyLineageAndScheduling(t *testing.T) { "add /metadata/labels/apps.cozystack.io~1application.name": "demo", "add /metadata/labels/internal.cozystack.io~1tenantresource": "false", "add /metadata/annotations/scheduler.cozystack.io~1scheduling-class": "gpu", + "add /metadata/labels/scheduler.cozystack.io~1scheduling-class": "gpu", "add /spec/schedulerName": schedulerapi.SchedulerName, } @@ -273,3 +274,189 @@ func TestHandle_ReadmissionShortCircuit(t *testing.T) { }) } } + +const testSchedulingClass = "co-region-test" + +// newSchedulingClassCR builds an unstructured SchedulingClass CR the fake +// dynamic client can serve. +func newSchedulingClassCR(name string) *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(schema.GroupVersionKind{ + Group: schedulerapi.Group, + Version: schedulerapi.Version, + Kind: "SchedulingClass", + }) + u.SetName(name) + return u +} + +// newWebhookTestEnv builds a LineageControllerWebhook wired with fakes, +// pre-populated with a Namespace carrying the scheduling-class label and a +// SchedulingClass CR by that name. +func newWebhookTestEnv(t *testing.T, namespace string, classOnNamespace string, classExists bool) *LineageControllerWebhook { + t.Helper() + + testScheme := runtime.NewScheme() + _ = clientgoscheme.AddToScheme(testScheme) + + ns := &corev1.Namespace{ + ObjectMeta: metav1.ObjectMeta{ + Name: namespace, + }, + } + if classOnNamespace != "" { + ns.Labels = map[string]string{ + schedulerapi.SchedulingClassLabel: classOnNamespace, + } + } + + c := fake.NewClientBuilder(). + WithScheme(testScheme). + WithObjects(ns). + Build() + + gvr := schema.GroupVersionResource{ + Group: schedulerapi.Group, + Version: schedulerapi.Version, + Resource: schedulerapi.Resource, + } + + dynObjs := []runtime.Object{} + if classExists { + dynObjs = append(dynObjs, newSchedulingClassCR(classOnNamespace)) + } + dyn := dynamicfake.NewSimpleDynamicClientWithCustomListKinds( + testScheme, + map[schema.GroupVersionResource]string{gvr: "SchedulingClassList"}, + dynObjs..., + ) + + return &LineageControllerWebhook{ + Client: c, + Scheme: testScheme, + dynClient: dyn, + } +} + +// newPod returns a minimal unstructured Pod for mutation tests. +func newPod(namespace, name string) *unstructured.Unstructured { + u := &unstructured.Unstructured{} + u.SetGroupVersionKind(schema.GroupVersionKind{Version: "v1", Kind: "Pod"}) + u.SetNamespace(namespace) + u.SetName(name) + return u +} + +func TestApplySchedulingClass_SetsLabelAnnotationAndSchedulerName(t *testing.T) { + ctx := context.Background() + h := newWebhookTestEnv(t, "tenant-acme", testSchedulingClass, true) + pod := newPod("tenant-acme", "db-0") + + if err := h.applySchedulingClass(ctx, pod, nil, "tenant-acme"); err != nil { + t.Fatalf("applySchedulingClass returned error: %v", err) + } + + gotLabel := pod.GetLabels()[schedulerapi.SchedulingClassLabel] + if gotLabel != testSchedulingClass { + t.Errorf("expected pod label %s=%s, got %q", schedulerapi.SchedulingClassLabel, testSchedulingClass, gotLabel) + } + + gotAnnotation := pod.GetAnnotations()[schedulerapi.SchedulingClassAnnotation] + if gotAnnotation != testSchedulingClass { + t.Errorf("expected pod annotation %s=%s, got %q", schedulerapi.SchedulingClassAnnotation, testSchedulingClass, gotAnnotation) + } + + gotSched, _, _ := unstructured.NestedString(pod.Object, "spec", "schedulerName") + if gotSched != schedulerapi.SchedulerName { + t.Errorf("expected pod spec.schedulerName=%s, got %q", schedulerapi.SchedulerName, gotSched) + } +} + +func TestApplySchedulingClass_PreservesExistingLabels(t *testing.T) { + ctx := context.Background() + h := newWebhookTestEnv(t, "tenant-acme", testSchedulingClass, true) + pod := newPod("tenant-acme", "db-0") + pod.SetLabels(map[string]string{ + "app.kubernetes.io/name": "postgres", + "role": "primary", + }) + + if err := h.applySchedulingClass(ctx, pod, nil, "tenant-acme"); err != nil { + t.Fatalf("applySchedulingClass returned error: %v", err) + } + + labels := pod.GetLabels() + if labels[schedulerapi.SchedulingClassLabel] != testSchedulingClass { + t.Errorf("scheduling-class label not set") + } + if labels["app.kubernetes.io/name"] != "postgres" { + t.Errorf("existing label app.kubernetes.io/name was clobbered: %q", labels["app.kubernetes.io/name"]) + } + if labels["role"] != "primary" { + t.Errorf("existing label role was clobbered: %q", labels["role"]) + } +} + +func TestApplySchedulingClass_NonPodIsNoOp(t *testing.T) { + ctx := context.Background() + h := newWebhookTestEnv(t, "tenant-acme", testSchedulingClass, true) + + svc := &unstructured.Unstructured{} + svc.SetGroupVersionKind(schema.GroupVersionKind{Version: "v1", Kind: "Service"}) + svc.SetNamespace("tenant-acme") + svc.SetName("svc-0") + + if err := h.applySchedulingClass(ctx, svc, nil, "tenant-acme"); err != nil { + t.Fatalf("applySchedulingClass returned error: %v", err) + } + + if _, ok := svc.GetLabels()[schedulerapi.SchedulingClassLabel]; ok { + t.Errorf("non-Pod object should not be mutated") + } + if _, ok := svc.GetAnnotations()[schedulerapi.SchedulingClassAnnotation]; ok { + t.Errorf("non-Pod object should not be mutated") + } +} + +func TestApplySchedulingClass_NoClassOnNamespaceIsNoOp(t *testing.T) { + ctx := context.Background() + h := newWebhookTestEnv(t, "tenant-acme", "", false) + pod := newPod("tenant-acme", "db-0") + + if err := h.applySchedulingClass(ctx, pod, nil, "tenant-acme"); err != nil { + t.Fatalf("applySchedulingClass returned error: %v", err) + } + + if _, ok := pod.GetLabels()[schedulerapi.SchedulingClassLabel]; ok { + t.Errorf("pod label should not be set when namespace has no scheduling-class label") + } + if _, ok := pod.GetAnnotations()[schedulerapi.SchedulingClassAnnotation]; ok { + t.Errorf("pod annotation should not be set when namespace has no scheduling-class label") + } + if sched, _, _ := unstructured.NestedString(pod.Object, "spec", "schedulerName"); sched != "" { + t.Errorf("pod spec.schedulerName should not be set when namespace has no scheduling-class label, got %q", sched) + } +} + +func TestApplySchedulingClass_MissingClassCRSkipsInjection(t *testing.T) { + ctx := context.Background() + // Namespace carries the label, but the SchedulingClass CR is NOT created. + h := newWebhookTestEnv(t, "tenant-acme", testSchedulingClass, false) + pod := newPod("tenant-acme", "db-0") + + if err := h.applySchedulingClass(ctx, pod, nil, "tenant-acme"); err != nil { + t.Fatalf("applySchedulingClass returned error: %v", err) + } + + // Defensive behavior: pod must not be left referencing a non-existent + // scheduler, otherwise it would stay Pending forever. + if _, ok := pod.GetLabels()[schedulerapi.SchedulingClassLabel]; ok { + t.Errorf("pod label should not be set when SchedulingClass CR is missing") + } + if _, ok := pod.GetAnnotations()[schedulerapi.SchedulingClassAnnotation]; ok { + t.Errorf("pod annotation should not be set when SchedulingClass CR is missing") + } + if sched, _, _ := unstructured.NestedString(pod.Object, "spec", "schedulerName"); sched != "" { + t.Errorf("pod spec.schedulerName should not be set when SchedulingClass CR is missing, got %q", sched) + } +}