diff --git a/pkg/controllers/v1alpha1/databackup/status_handler.go b/pkg/controllers/v1alpha1/databackup/status_handler.go index 9a1196b452f..6af0bcbf180 100644 --- a/pkg/controllers/v1alpha1/databackup/status_handler.go +++ b/pkg/controllers/v1alpha1/databackup/status_handler.go @@ -21,6 +21,7 @@ import ( "github.com/fluid-cloudnative/fluid/api/v1alpha1" "github.com/fluid-cloudnative/fluid/pkg/common" + "github.com/fluid-cloudnative/fluid/pkg/dataflow" "github.com/fluid-cloudnative/fluid/pkg/dataoperation" "github.com/fluid-cloudnative/fluid/pkg/runtime" "github.com/fluid-cloudnative/fluid/pkg/utils" @@ -52,7 +53,13 @@ func (o *OnceHandler) GetOperationStatus(ctx runtime.ReconcileRequestContext, op return } - // TODO: inject nodeaffinity like other data operations when using job instead of pod + if kubeclient.IsSucceededPod(backupPod) && result.NodeAffinity == nil { + result.NodeAffinity, err = dataflow.GenerateNodeAffinityFromPod(backupPod) + if err != nil { + ctx.Log.V(1).Info("NodeAffinity not injected", "reason", err.Error()) + err = nil + } + } var finishTime time.Time if len(backupPod.Status.Conditions) != 0 { diff --git a/pkg/controllers/v1alpha1/databackup/status_handler_test.go b/pkg/controllers/v1alpha1/databackup/status_handler_test.go index 204193fcc82..03cbbc4cd6a 100644 --- a/pkg/controllers/v1alpha1/databackup/status_handler_test.go +++ b/pkg/controllers/v1alpha1/databackup/status_handler_test.go @@ -184,5 +184,76 @@ var _ = Describe("OnceHandler", func() { Expect(opStatus.Conditions).To(HaveLen(1)) Expect(opStatus.Conditions[0].LastTransitionTime.Time).To(BeTemporally("~", conditionTime.Time, time.Second)) }) + + It("should inject NodeAffinity when backup pod succeeds with dataflow annotations", func() { + pod := &corev1.Pod{ + ObjectMeta: v1.ObjectMeta{ + Name: mockDataBackup.GetName() + "-pod", + Namespace: mockDataBackup.GetNamespace(), + Annotations: map[string]string{ + common.AnnotationDataFlowAffinityInject: "true", + common.AnnotationDataFlowCustomizedAffinityPrefix + "node-label": "node-value", + }, + }, + Status: corev1.PodStatus{ + Phase: corev1.PodSucceeded, + Conditions: []corev1.PodCondition{ + { + Type: corev1.PodReady, + Status: corev1.ConditionFalse, + LastTransitionTime: v1.Now(), + }, + }, + }, + } + c := fake.NewFakeClientWithScheme(testScheme, pod, mockDataBackup) + handler := &OnceHandler{dataBackup: mockDataBackup} + ctx := cruntime.ReconcileRequestContext{ + NamespacedName: types.NamespacedName{Namespace: "default", Name: "test"}, + Log: fake.NullLogger(), + Client: c, + } + result, err := handler.GetOperationStatus(ctx, &mockDataBackup.Status) + Expect(err).NotTo(HaveOccurred()) + Expect(result.Phase).To(Equal(common.PhaseComplete)) + Expect(result.NodeAffinity).NotTo(BeNil()) + Expect(result.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution).NotTo(BeNil()) + terms := result.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms + Expect(terms).To(HaveLen(1)) + Expect(terms[0].MatchExpressions).To(HaveLen(1)) + Expect(terms[0].MatchExpressions[0].Key).To(Equal("node-label")) + Expect(terms[0].MatchExpressions[0].Values).To(ContainElement("node-value")) + }) + + It("should not inject NodeAffinity when backup pod succeeds without dataflow annotations", func() { + pod := &corev1.Pod{ + ObjectMeta: v1.ObjectMeta{ + Name: mockDataBackup.GetName() + "-pod", + Namespace: mockDataBackup.GetNamespace(), + }, + Status: corev1.PodStatus{ + Phase: corev1.PodSucceeded, + Conditions: []corev1.PodCondition{ + { + Type: corev1.PodReady, + Status: corev1.ConditionFalse, + LastTransitionTime: v1.Now(), + }, + }, + }, + } + c := fake.NewFakeClientWithScheme(testScheme, pod, mockDataBackup) + handler := &OnceHandler{dataBackup: mockDataBackup} + ctx := cruntime.ReconcileRequestContext{ + NamespacedName: types.NamespacedName{Namespace: "default", Name: "test"}, + Log: fake.NullLogger(), + Client: c, + } + result, err := handler.GetOperationStatus(ctx, &mockDataBackup.Status) + Expect(err).NotTo(HaveOccurred()) + Expect(result.Phase).To(Equal(common.PhaseComplete)) + Expect(result.NodeAffinity).To(BeNil()) + }) + }) }) diff --git a/pkg/dataflow/helper.go b/pkg/dataflow/helper.go index a0ea084bbe4..79370746210 100644 --- a/pkg/dataflow/helper.go +++ b/pkg/dataflow/helper.go @@ -65,3 +65,47 @@ func GenerateNodeAffinity(job *batchv1.Job) (*corev1.NodeAffinity, error) { return nodeAffinity, nil } + +// GenerateNodeAffinityFromPod generates a NodeAffinity from a Pod's annotations, +// using the same annotation-based logic as GenerateNodeAffinity for Jobs. +// This is used by DataBackup which runs as a Pod rather than a Job. +func GenerateNodeAffinityFromPod(pod *corev1.Pod) (*corev1.NodeAffinity, error) { + if pod == nil { + return nil, nil + } + // not inject, i.e. feature gate not enabled + if v := pod.Annotations[common.AnnotationDataFlowAffinityInject]; v != "true" { + return nil, nil + } + + annotations := pod.Annotations + + nodeAffinity := &corev1.NodeAffinity{ + RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{ + NodeSelectorTerms: []corev1.NodeSelectorTerm{ + { + MatchExpressions: nil, + }, + }, + }, + } + + hasInjectedLabels := false + for key, value := range annotations { + if strings.HasPrefix(key, common.AnnotationDataFlowCustomizedAffinityPrefix) { + nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions = + append(nodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions, + corev1.NodeSelectorRequirement{ + Key: strings.TrimPrefix(key, common.AnnotationDataFlowCustomizedAffinityPrefix), + Operator: corev1.NodeSelectorOpIn, + Values: []string{value}, + }) + hasInjectedLabels = true + } + } + if !hasInjectedLabels { + return nil, errors.New("the affinity label is not set, wait for next reconcile") + } + + return nodeAffinity, nil +} diff --git a/pkg/dataflow/helper_test.go b/pkg/dataflow/helper_test.go index d5d0a55f4a5..e4cdf88e44e 100644 --- a/pkg/dataflow/helper_test.go +++ b/pkg/dataflow/helper_test.go @@ -128,3 +128,86 @@ func TestGenerateNodeLabels(t *testing.T) { }) } } + +func TestGenerateNodeAffinityFromPod(t *testing.T) { + tests := []struct { + name string + pod *v1.Pod + wantNil bool + wantErr bool + wantLen int + }{ + { + name: "nil pod returns nil", + pod: nil, + wantNil: true, + wantErr: false, + }, + { + name: "pod without inject annotation returns nil", + pod: &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "backup-pod", + Annotations: map[string]string{}, + }, + }, + wantNil: true, + wantErr: false, + }, + { + name: "pod with inject annotation but no affinity labels returns error", + pod: &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "backup-pod", + Annotations: map[string]string{ + common.AnnotationDataFlowAffinityInject: "true", + }, + }, + }, + wantNil: true, + wantErr: true, + }, + { + name: "pod with inject annotation and affinity labels returns NodeAffinity", + pod: &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "backup-pod", + Annotations: map[string]string{ + common.AnnotationDataFlowAffinityInject: "true", + common.AnnotationDataFlowCustomizedAffinityPrefix + common.K8sNodeNameLabelKey: "node01", + common.AnnotationDataFlowCustomizedAffinityPrefix + common.K8sZoneLabelKey: "zone01", + }, + }, + }, + wantNil: false, + wantErr: false, + wantLen: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := GenerateNodeAffinityFromPod(tt.pod) + if (err != nil) != tt.wantErr { + t.Errorf("GenerateNodeAffinityFromPod() error = %v, wantErr %v", err, tt.wantErr) + return + } + if tt.wantNil { + if got != nil { + t.Errorf("GenerateNodeAffinityFromPod() expected nil, got %v", got) + } + return + } + if got == nil { + t.Fatal("GenerateNodeAffinityFromPod() expected non-nil NodeAffinity") + } + terms := got.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms + if len(terms) != 1 { + t.Fatalf("expected 1 NodeSelectorTerm, got %d", len(terms)) + } + if len(terms[0].MatchExpressions) != tt.wantLen { + t.Errorf("expected %d MatchExpressions, got %d", tt.wantLen, len(terms[0].MatchExpressions)) + } + }) + } +}