Expose PodAffinity and PodAntiAffinity struct and build overrides

This patch introduces a first, basic implementation that allows to
expose the Pod Affinity/Antiaffinity interface as part of the CR
spec and allows to override it.

Signed-off-by: Francesco Pantano <[email protected]>

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,
+	overrides *AffinityOverrideSpec,
 ) *corev1.Affinity {
-	return &corev1.Affinity{
+	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, _ := toCoreAffinity(defaultAffinity, overrides)
+		return patchedAffinity
+	}
+	return defaultAffinity
+}
+
+func toCoreAffinity(
+	affinity *v1.Affinity,
+	override *AffinityOverrideSpec,
+) (*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
 }

modules/common/affinity/types.go

@@ -0,0 +1,38 @@
+/*
+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 - service override configuration for the Affinity propagated to the Pods
+// Allows for the manifest of the created StatefulSet to be overwritten with custom Pod affinity configuration.
+type OverrideSpec struct {
+	Spec                       *AffinityOverrideSpec `json:"spec,omitempty" protobuf:"bytes,2,opt,name=spec"`
+}
+
+type AffinityOverrideSpec 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"`
+}