diff --git a/modules/common/affinity/affinity.go b/modules/common/affinity/affinity.go index 0e188c47..15e5f3c8 100644 --- a/modules/common/affinity/affinity.go +++ b/modules/common/affinity/affinity.go @@ -17,8 +17,12 @@ limitations under the License. package affinity import ( + "encoding/json" + "fmt" corev1 "k8s.io/api/core/v1" + v1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/util/strategicpatch" ) // DistributePods - returns rule to ensure that two replicas of the same selector @@ -27,8 +31,9 @@ func DistributePods( selectorKey string, selectorValues []string, topologyKey string, -) *corev1.Affinity { - return &corev1.Affinity{ + overrides *OverrideSpec, +) (*corev1.Affinity, error) { + defaultAffinity := &corev1.Affinity{ PodAntiAffinity: &corev1.PodAntiAffinity{ // This rule ensures that two replicas of the same selector // should not run if possible on the same worker node @@ -53,4 +58,46 @@ func DistributePods( }, }, } + // patch the default affinity Object with the data passed as input + if overrides != nil { + patchedAffinity, err := toCoreAffinity(defaultAffinity, overrides) + return patchedAffinity, err + } + return defaultAffinity, nil +} + +func toCoreAffinity( + affinity *v1.Affinity, + override *OverrideSpec, +) (*v1.Affinity, error) { + + aff := &v1.Affinity{ + PodAntiAffinity: affinity.PodAntiAffinity, + PodAffinity: affinity.PodAffinity, + } + if override != nil { + if override != nil { + origAffinit, err := json.Marshal(affinity) + if err != nil { + return aff, fmt.Errorf("error marshalling Affinity Spec: %w", err) + } + patch, err := json.Marshal(override) + if err != nil { + return aff, fmt.Errorf("error marshalling Affinity Spec: %w", err) + } + + patchedJSON, err := strategicpatch.StrategicMergePatch(origAffinit, patch, v1.Affinity{}) + if err != nil { + return aff, fmt.Errorf("error patching Affinity Spec: %w", err) + } + + patchedSpec := v1.Affinity{} + err = json.Unmarshal(patchedJSON, &patchedSpec) + if err != nil { + return aff, fmt.Errorf("error unmarshalling patched Service Spec: %w", err) + } + aff = &patchedSpec + } + } + return aff, nil } diff --git a/modules/common/affinity/affinity_test.go b/modules/common/affinity/affinity_test.go index 16a142dd..a7923d6e 100644 --- a/modules/common/affinity/affinity_test.go +++ b/modules/common/affinity/affinity_test.go @@ -47,13 +47,82 @@ var affinityObj = &corev1.Affinity{ }, } +// weightedPodAffinityTermOverride represents an Override passed to the Affinity +// tests +var weightedPodAffinityTermOverride = []corev1.WeightedPodAffinityTerm{ + { + PodAffinityTerm: corev1.PodAffinityTerm{ + LabelSelector: &metav1.LabelSelector{ + MatchExpressions: []metav1.LabelSelectorRequirement{ + { + Key: "CustomKeySelector", + Operator: metav1.LabelSelectorOpIn, + Values: []string{ + "selectorValue1", + "selectorValue2", + "selectorValue3", + }, + }, + }, + }, + TopologyKey: "CustomTopologyKey", + }, + Weight: 80, + }, +} + func TestDistributePods(t *testing.T) { t.Run("Default pod distribution", func(t *testing.T) { g := NewWithT(t) + d, _ := DistributePods("ThisSelector", []string{"selectorValue1", "selectorValue2"}, "ThisTopologyKey", nil) + g.Expect(d).To(BeEquivalentTo(affinityObj)) + }) - d := DistributePods("ThisSelector", []string{"selectorValue1", "selectorValue2"}, "ThisTopologyKey") + // Override the default AntiAffinity + t.Run("Pod distribution with overrides", func(t *testing.T) { + // The resulting affinity that should be assigned to the Pod + var expectedAffinity = &corev1.Affinity{ + PodAffinity: nil, + NodeAffinity: nil, + PodAntiAffinity: &corev1.PodAntiAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: weightedPodAffinityTermOverride, + }, + } + affinityOverride := &OverrideSpec{ + PodAffinity: nil, + PodAntiAffinity: &corev1.PodAntiAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: weightedPodAffinityTermOverride, + }, + NodeAffinity: nil, + } + g := NewWithT(t) + d, _ := DistributePods("ThisSelector", []string{"selectorValue1", "selectorValue2"}, "ThisTopologyKey", affinityOverride) + g.Expect(d).To(BeEquivalentTo(expectedAffinity)) + }) - g.Expect(d).To(BeEquivalentTo(affinityObj)) + // Override the Affinity but keep the default AntiAffinity + t.Run("Pod distribution with overrides", func(t *testing.T) { + // The resulting affinity that should be assigned to the Pod + var expectedAffinity = &corev1.Affinity{ + // the default PodAntiAffinity defined in the DistributePods function + // is applied, while PodAffinity is the result of the override passed + // as input + PodAntiAffinity: affinityObj.PodAntiAffinity, + NodeAffinity: nil, + PodAffinity: &corev1.PodAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: weightedPodAffinityTermOverride, + }, + } + affinityOverride := &OverrideSpec{ + PodAntiAffinity: nil, + PodAffinity: &corev1.PodAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: weightedPodAffinityTermOverride, + }, + NodeAffinity: nil, + } + g := NewWithT(t) + d, _ := DistributePods("ThisSelector", []string{"selectorValue1", "selectorValue2"}, "ThisTopologyKey", affinityOverride) + g.Expect(d).To(BeEquivalentTo(expectedAffinity)) }) } diff --git a/modules/common/affinity/types.go b/modules/common/affinity/types.go new file mode 100644 index 00000000..d1a763b3 --- /dev/null +++ b/modules/common/affinity/types.go @@ -0,0 +1,36 @@ +/* +Copyright 2024 Red Hat + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// +kubebuilder:object:generate:=true + +package affinity + +import ( + corev1 "k8s.io/api/core/v1" +) + +// OverrideSpec - +type OverrideSpec struct { + // Describes pod affinity scheduling rules (e.g. co-locate this pod in the same node, zone, etc. as some other pod(s)). + // +optional + PodAffinity *corev1.PodAffinity `json:"podAffinity,omitempty" protobuf:"bytes,2,opt,name=podAffinity"` + // Describes pod anti-affinity scheduling rules (e.g. avoid putting this pod in the same node, zone, etc. as some other pod(s)). + // +optional + PodAntiAffinity *corev1.PodAntiAffinity `json:"podAntiAffinity,omitempty" protobuf:"bytes,3,opt,name=podAntiAffinity"` + // Describes node affinity scheduling rules for the pod. + // +optional + NodeAffinity *corev1.NodeAffinity `json:"nodeAffinity,omitempty" protobuf:"bytes,1,opt,name=nodeAffinity"` +} diff --git a/modules/common/affinity/zz_generated.deepcopy.go b/modules/common/affinity/zz_generated.deepcopy.go new file mode 100644 index 00000000..9546ef21 --- /dev/null +++ b/modules/common/affinity/zz_generated.deepcopy.go @@ -0,0 +1,56 @@ +//go:build !ignore_autogenerated +// +build !ignore_autogenerated + +/* + + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Code generated by controller-gen. DO NOT EDIT. + +package affinity + +import ( + "k8s.io/api/core/v1" +) + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *OverrideSpec) DeepCopyInto(out *OverrideSpec) { + *out = *in + if in.PodAffinity != nil { + in, out := &in.PodAffinity, &out.PodAffinity + *out = new(v1.PodAffinity) + (*in).DeepCopyInto(*out) + } + if in.PodAntiAffinity != nil { + in, out := &in.PodAntiAffinity, &out.PodAntiAffinity + *out = new(v1.PodAntiAffinity) + (*in).DeepCopyInto(*out) + } + if in.NodeAffinity != nil { + in, out := &in.NodeAffinity, &out.NodeAffinity + *out = new(v1.NodeAffinity) + (*in).DeepCopyInto(*out) + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OverrideSpec. +func (in *OverrideSpec) DeepCopy() *OverrideSpec { + if in == nil { + return nil + } + out := new(OverrideSpec) + in.DeepCopyInto(out) + return out +}