CLOUDP-65312: arbitrary statefulset configuration (#118)

Author: Nikolas De Giorgis

.evergreen.yml

@@ -152,6 +152,8 @@ task_groups:
     - e2e_test_replica_set_tls
     - e2e_test_replica_set_tls_upgrade
     - e2e_test_replica_set_tls_rotate
+    - e2e_test_statefulset_arbitrary_config
+    - e2e_test_statefulset_arbitrary_config_update
   teardown_task:
     - func: upload_e2e_logs
 
@@ -284,6 +286,18 @@ tasks:
         vars:
           test: replica_set_tls_rotate
 
+  - name: e2e_test_statefulset_arbitrary_config
+    commands:
+      - func: run_e2e_test
+        vars:
+          test: statefulset_arbitrary_config
+
+  - name: e2e_test_statefulset_arbitrary_config_update
+    commands:
+      - func: run_e2e_test
+        vars:
+          test: statefulset_arbitrary_config_update
+
 buildvariants:
   - name: go_unit_tests
     display_name: go_unit_tests

pkg/apis/mongodb/v1/mongodb_types.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strings"
 
+	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/types"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -44,6 +45,15 @@ type MongoDBSpec struct {
 	// Users specifies the MongoDB users that should be configured in your deployment
 	// +required
 	Users []MongoDBUser `json:"users"`
+
+	// +optional
+	StatefulSetConfiguration StatefulSetConfiguration `json:"statefulset,omitempty"`
+}
+
+// StatefulSetConfiguration holds the optional custom StatefulSet
+// that should be merged into the operator created one.
+type StatefulSetConfiguration struct {
+	Spec appsv1.StatefulSetSpec `json:"spec"`
 }
 
 type MongoDBUser struct {

pkg/controller/mongodb/replica_set_controller.go

@@ -604,6 +604,7 @@ func buildStatefulSetModificationFunction(mdb mdbv1.MongoDB) statefulset.Modific
 				buildScramPodSpecModification(mdb),
 			),
 		),
+		statefulset.WithCustomSpecs(mdb.Spec.StatefulSetConfiguration.Spec),
 	)
 }
 

pkg/kube/podtemplatespec/podspec_template.go

@@ -1,6 +1,7 @@
 package podtemplatespec
 
 import (
+	"github.com/imdario/mergo"
 	"github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/container"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -227,6 +228,127 @@ func WithVolumeMounts(containerName string, volumeMounts ...corev1.VolumeMount)
 	}
 }
 
+func MergePodTemplateSpecs(defaultTemplate, overrideTemplate corev1.PodTemplateSpec) (corev1.PodTemplateSpec, error) {
+	// Containers need to be merged manually
+	mergedContainers, err := mergeContainers(defaultTemplate.Spec.Containers, overrideTemplate.Spec.Containers)
+	if err != nil {
+		return corev1.PodTemplateSpec{}, err
+	}
+
+	// InitContainers need to be merged manually
+	mergedInitContainers, err := mergeContainers(defaultTemplate.Spec.InitContainers, overrideTemplate.Spec.InitContainers)
+	if err != nil {
+		return corev1.PodTemplateSpec{}, err
+	}
+
+	// Affinity needs to be merged manually
+	mergedAffinity, err := mergeAffinity(defaultTemplate.Spec.Affinity, overrideTemplate.Spec.Affinity)
+	if err != nil {
+		return corev1.PodTemplateSpec{}, err
+	}
+
+	// Everything else can be merged with mergo
+	mergedPodTemplateSpec := *defaultTemplate.DeepCopy()
+	if err = mergo.Merge(&mergedPodTemplateSpec, overrideTemplate, mergo.WithOverride, mergo.WithAppendSlice); err != nil {
+		return corev1.PodTemplateSpec{}, err
+	}
+
+	mergedPodTemplateSpec.Spec.Containers = mergedContainers
+	mergedPodTemplateSpec.Spec.InitContainers = mergedInitContainers
+	mergedPodTemplateSpec.Spec.Affinity = mergedAffinity
+	return mergedPodTemplateSpec, nil
+}
+
+func mergeVolumeMounts(defaultMounts, overrideMounts []corev1.VolumeMount) ([]corev1.VolumeMount, error) {
+	defaultMountsMap := createMountsMap(defaultMounts)
+	overrideMountsMap := createMountsMap(overrideMounts)
+	mergedVolumeMounts := []corev1.VolumeMount{}
+	for _, defaultMount := range defaultMounts {
+		if overrideMount, ok := overrideMountsMap[defaultMount.Name]; ok {
+			// needs merge
+			if err := mergo.Merge(&defaultMount, overrideMount, mergo.WithAppendSlice); err != nil { //nolint
+				return nil, err
+			}
+		}
+		mergedVolumeMounts = append(mergedVolumeMounts, defaultMount)
+	}
+	for _, overrideMount := range overrideMounts {
+		if _, ok := defaultMountsMap[overrideMount.Name]; ok {
+			// already merged
+			continue
+		}
+		mergedVolumeMounts = append(mergedVolumeMounts, overrideMount)
+	}
+	return mergedVolumeMounts, nil
+}
+
+func createMountsMap(volumeMounts []corev1.VolumeMount) map[string]corev1.VolumeMount {
+	mountMap := make(map[string]corev1.VolumeMount)
+	for _, m := range volumeMounts {
+		mountMap[m.Name] = m
+	}
+	return mountMap
+}
+
+func mergeContainers(defaultContainers, customContainers []corev1.Container) ([]corev1.Container, error) {
+	defaultMap := createContainerMap(defaultContainers)
+	customMap := createContainerMap(customContainers)
+	mergedContainers := []corev1.Container{}
+	for _, defaultContainer := range defaultContainers {
+		if customContainer, ok := customMap[defaultContainer.Name]; ok {
+			// The container is present in both maps, so we need to merge
+			// Merge mounts
+			mergedMounts, err := mergeVolumeMounts(defaultContainer.VolumeMounts, customContainer.VolumeMounts)
+			if err != nil {
+				return nil, err
+			}
+			if err := mergo.Merge(&defaultContainer, customContainer, mergo.WithOverride); err != nil { //nolint
+				return nil, err
+			}
+			// completely override any resources that were provided
+			// this prevents issues with custom requests giving errors due
+			// to the defaulted limits
+			defaultContainer.Resources = customContainer.Resources
+			defaultContainer.VolumeMounts = mergedMounts
+		}
+		// The default container was not modified by the override, so just add it
+		mergedContainers = append(mergedContainers, defaultContainer)
+	}
+
+	// Look for customContainers that were not merged into existing ones
+	for _, customContainer := range customContainers {
+		if _, ok := defaultMap[customContainer.Name]; ok {
+			continue
+		}
+		// Need to add it
+		mergedContainers = append(mergedContainers, customContainer)
+	}
+
+	return mergedContainers, nil
+}
+
+func createContainerMap(containers []corev1.Container) map[string]corev1.Container {
+	containerMap := make(map[string]corev1.Container)
+	for _, c := range containers {
+		containerMap[c.Name] = c
+	}
+	return containerMap
+}
+
+func mergeAffinity(defaultAffinity, overrideAffinity *corev1.Affinity) (*corev1.Affinity, error) {
+	if defaultAffinity == nil {
+		return overrideAffinity, nil
+	}
+	if overrideAffinity == nil {
+		return defaultAffinity, nil
+	}
+	mergedAffinity := defaultAffinity.DeepCopy()
+	if err := mergo.Merge(mergedAffinity, *overrideAffinity, mergo.WithOverride); err != nil {
+		return nil, err
+	}
+	return mergedAffinity, nil
+}
+
 // findContainerByName will find either a container or init container by name in a pod template spec
 func findContainerByName(name string, podTemplateSpec *corev1.PodTemplateSpec) *corev1.Container {
 	containerIdx := findIndexByName(name, podTemplateSpec.Spec.Containers)