|
1 | 1 | package podtemplatespec
|
2 | 2 |
|
3 | 3 | import (
|
| 4 | + "github.com/imdario/mergo" |
4 | 5 | "github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/container"
|
5 | 6 | corev1 "k8s.io/api/core/v1"
|
6 | 7 | metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
@@ -227,6 +228,127 @@ func WithVolumeMounts(containerName string, volumeMounts ...corev1.VolumeMount)
|
227 | 228 | }
|
228 | 229 | }
|
229 | 230 |
|
| 231 | +func MergePodTemplateSpecs(defaultTemplate, overrideTemplate corev1.PodTemplateSpec) (corev1.PodTemplateSpec, error) { |
| 232 | + // Containers need to be merged manually |
| 233 | + mergedContainers, err := mergeContainers(defaultTemplate.Spec.Containers, overrideTemplate.Spec.Containers) |
| 234 | + if err != nil { |
| 235 | + return corev1.PodTemplateSpec{}, err |
| 236 | + } |
| 237 | + |
| 238 | + // InitContainers need to be merged manually |
| 239 | + mergedInitContainers, err := mergeContainers(defaultTemplate.Spec.InitContainers, overrideTemplate.Spec.InitContainers) |
| 240 | + if err != nil { |
| 241 | + return corev1.PodTemplateSpec{}, err |
| 242 | + } |
| 243 | + |
| 244 | + // Affinity needs to be merged manually |
| 245 | + mergedAffinity, err := mergeAffinity(defaultTemplate.Spec.Affinity, overrideTemplate.Spec.Affinity) |
| 246 | + if err != nil { |
| 247 | + return corev1.PodTemplateSpec{}, err |
| 248 | + } |
| 249 | + |
| 250 | + // Everything else can be merged with mergo |
| 251 | + mergedPodTemplateSpec := *defaultTemplate.DeepCopy() |
| 252 | + if err = mergo.Merge(&mergedPodTemplateSpec, overrideTemplate, mergo.WithOverride, mergo.WithAppendSlice); err != nil { |
| 253 | + return corev1.PodTemplateSpec{}, err |
| 254 | + } |
| 255 | + |
| 256 | + mergedPodTemplateSpec.Spec.Containers = mergedContainers |
| 257 | + mergedPodTemplateSpec.Spec.InitContainers = mergedInitContainers |
| 258 | + mergedPodTemplateSpec.Spec.Affinity = mergedAffinity |
| 259 | + return mergedPodTemplateSpec, nil |
| 260 | +} |
| 261 | + |
| 262 | +func mergeVolumeMounts(defaultMounts, overrideMounts []corev1.VolumeMount) ([]corev1.VolumeMount, error) { |
| 263 | + defaultMountsMap := createMountsMap(defaultMounts) |
| 264 | + overrideMountsMap := createMountsMap(overrideMounts) |
| 265 | + mergedVolumeMounts := []corev1.VolumeMount{} |
| 266 | + for _, defaultMount := range defaultMounts { |
| 267 | + if overrideMount, ok := overrideMountsMap[defaultMount.Name]; ok { |
| 268 | + // needs merge |
| 269 | + if err := mergo.Merge(&defaultMount, overrideMount, mergo.WithAppendSlice); err != nil { //nolint |
| 270 | + return nil, err |
| 271 | + } |
| 272 | + } |
| 273 | + mergedVolumeMounts = append(mergedVolumeMounts, defaultMount) |
| 274 | + } |
| 275 | + for _, overrideMount := range overrideMounts { |
| 276 | + if _, ok := defaultMountsMap[overrideMount.Name]; ok { |
| 277 | + // already merged |
| 278 | + continue |
| 279 | + } |
| 280 | + mergedVolumeMounts = append(mergedVolumeMounts, overrideMount) |
| 281 | + } |
| 282 | + return mergedVolumeMounts, nil |
| 283 | +} |
| 284 | + |
| 285 | +func createMountsMap(volumeMounts []corev1.VolumeMount) map[string]corev1.VolumeMount { |
| 286 | + mountMap := make(map[string]corev1.VolumeMount) |
| 287 | + for _, m := range volumeMounts { |
| 288 | + mountMap[m.Name] = m |
| 289 | + } |
| 290 | + return mountMap |
| 291 | +} |
| 292 | + |
| 293 | +func mergeContainers(defaultContainers, customContainers []corev1.Container) ([]corev1.Container, error) { |
| 294 | + defaultMap := createContainerMap(defaultContainers) |
| 295 | + customMap := createContainerMap(customContainers) |
| 296 | + mergedContainers := []corev1.Container{} |
| 297 | + for _, defaultContainer := range defaultContainers { |
| 298 | + if customContainer, ok := customMap[defaultContainer.Name]; ok { |
| 299 | + // The container is present in both maps, so we need to merge |
| 300 | + // Merge mounts |
| 301 | + mergedMounts, err := mergeVolumeMounts(defaultContainer.VolumeMounts, customContainer.VolumeMounts) |
| 302 | + if err != nil { |
| 303 | + return nil, err |
| 304 | + } |
| 305 | + if err := mergo.Merge(&defaultContainer, customContainer, mergo.WithOverride); err != nil { //nolint |
| 306 | + return nil, err |
| 307 | + } |
| 308 | + // completely override any resources that were provided |
| 309 | + // this prevents issues with custom requests giving errors due |
| 310 | + // to the defaulted limits |
| 311 | + defaultContainer.Resources = customContainer.Resources |
| 312 | + defaultContainer.VolumeMounts = mergedMounts |
| 313 | + } |
| 314 | + // The default container was not modified by the override, so just add it |
| 315 | + mergedContainers = append(mergedContainers, defaultContainer) |
| 316 | + } |
| 317 | + |
| 318 | + // Look for customContainers that were not merged into existing ones |
| 319 | + for _, customContainer := range customContainers { |
| 320 | + if _, ok := defaultMap[customContainer.Name]; ok { |
| 321 | + continue |
| 322 | + } |
| 323 | + // Need to add it |
| 324 | + mergedContainers = append(mergedContainers, customContainer) |
| 325 | + } |
| 326 | + |
| 327 | + return mergedContainers, nil |
| 328 | +} |
| 329 | + |
| 330 | +func createContainerMap(containers []corev1.Container) map[string]corev1.Container { |
| 331 | + containerMap := make(map[string]corev1.Container) |
| 332 | + for _, c := range containers { |
| 333 | + containerMap[c.Name] = c |
| 334 | + } |
| 335 | + return containerMap |
| 336 | +} |
| 337 | + |
| 338 | +func mergeAffinity(defaultAffinity, overrideAffinity *corev1.Affinity) (*corev1.Affinity, error) { |
| 339 | + if defaultAffinity == nil { |
| 340 | + return overrideAffinity, nil |
| 341 | + } |
| 342 | + if overrideAffinity == nil { |
| 343 | + return defaultAffinity, nil |
| 344 | + } |
| 345 | + mergedAffinity := defaultAffinity.DeepCopy() |
| 346 | + if err := mergo.Merge(mergedAffinity, *overrideAffinity, mergo.WithOverride); err != nil { |
| 347 | + return nil, err |
| 348 | + } |
| 349 | + return mergedAffinity, nil |
| 350 | +} |
| 351 | + |
230 | 352 | // findContainerByName will find either a container or init container by name in a pod template spec
|
231 | 353 | func findContainerByName(name string, podTemplateSpec *corev1.PodTemplateSpec) *corev1.Container {
|
232 | 354 | containerIdx := findIndexByName(name, podTemplateSpec.Spec.Containers)
|
|
0 commit comments