Rework MergePodSpecs logic

flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go

@@ -670,8 +670,9 @@
 		return podSpec, objectMeta, nil
 	}
 
-	// merge podSpec with podTemplate
-	mergedPodSpec, err := MergePodSpecs(&podTemplate.Template.Spec, podSpec, primaryContainerName, primaryInitContainerName)
+	// merge podTemplate onto podSpec
+	templateSpec := &podTemplate.Template.Spec
+	mergedPodSpec, err := MergeBasePodSpecOntoTemplate(templateSpec, podSpec, primaryContainerName, primaryInitContainerName)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -685,50 +686,54 @@
 	return mergedPodSpec, mergedObjectMeta, nil
 }
 
-// MergePodSpecs merges the two provided PodSpecs. This process uses the first as the base configuration, where values
-// set by the first PodSpec are overwritten by the second in the return value. Additionally, this function applies
-// container-level configuration from the basePodSpec.
-func MergePodSpecs(basePodSpec *v1.PodSpec, podSpec *v1.PodSpec, primaryContainerName string, primaryInitContainerName string) (*v1.PodSpec, error) {
-	if basePodSpec == nil || podSpec == nil {
-		return nil, errors.New("neither the basePodSpec or the podSpec can be nil")
+// MergeBasePodSpecOntoTemplate merges a base pod spec onto a template pod spec. The template pod spec has some
+// magic values that allow users to specify templates that target all containers and primary containers. Aside from
+// magic values this method will merge containers that have matching names.
+func MergeBasePodSpecOntoTemplate(templatePodSpec *v1.PodSpec, basePodSpec *v1.PodSpec, primaryContainerName string, primaryInitContainerName string) (*v1.PodSpec, error) {
+	if templatePodSpec == nil || basePodSpec == nil {
+		return nil, errors.New("neither the templatePodSpec or the basePodSpec can be nil")
 	}
 
-	// extract defaultContainerTemplate and primaryContainerTemplate
+	// extract primaryContainerTemplate. The base should always contain the primary container.
 	var defaultContainerTemplate, primaryContainerTemplate *v1.Container
-	for i := 0; i < len(basePodSpec.Containers); i++ {
-		if basePodSpec.Containers[i].Name == defaultContainerTemplateName {
-			defaultContainerTemplate = &basePodSpec.Containers[i]
-		} else if basePodSpec.Containers[i].Name == primaryContainerName {
-			primaryContainerTemplate = &basePodSpec.Containers[i]
+
+	// extract default container template
+	for i := 0; i < len(templatePodSpec.Containers); i++ {
+		if templatePodSpec.Containers[i].Name == defaultContainerTemplateName {
+			defaultContainerTemplate = &templatePodSpec.Containers[i]
+		} else if templatePodSpec.Containers[i].Name == primaryContainerTemplateName {
+			primaryContainerTemplate = &templatePodSpec.Containers[i]
 		}
 	}
 
-	// extract defaultInitContainerTemplate and primaryInitContainerTemplate
+	// extract primaryInitContainerTemplate. The base should always contain the primary container.
 	var defaultInitContainerTemplate, primaryInitContainerTemplate *v1.Container
-	for i := 0; i < len(basePodSpec.InitContainers); i++ {
-		if basePodSpec.InitContainers[i].Name == defaultInitContainerTemplateName {
-			defaultInitContainerTemplate = &basePodSpec.InitContainers[i]
-		} else if basePodSpec.InitContainers[i].Name == primaryInitContainerName {
-			primaryInitContainerTemplate = &basePodSpec.InitContainers[i]
+
+	// extract defaultInitContainerTemplate
+	for i := 0; i < len(templatePodSpec.InitContainers); i++ {
+		if templatePodSpec.InitContainers[i].Name == defaultInitContainerTemplateName {
+			defaultInitContainerTemplate = &templatePodSpec.InitContainers[i]
+		} else if templatePodSpec.InitContainers[i].Name == primaryInitContainerTemplateName {
+			primaryInitContainerTemplate = &templatePodSpec.InitContainers[i]
 		}
 	}
 
-	// merge PodTemplate PodSpec with podSpec
-	var mergedPodSpec *v1.PodSpec = basePodSpec.DeepCopy()
-	if err := mergo.Merge(mergedPodSpec, podSpec, mergo.WithOverride, mergo.WithAppendSlice); err != nil {
+	// Merge base into template
+	mergedPodSpec := templatePodSpec.DeepCopy()
+	if err := mergo.Merge(mergedPodSpec, basePodSpec, mergo.WithOverride, mergo.WithAppendSlice); err != nil {
 		return nil, err
 	}
 
 	// merge PodTemplate containers
 	var mergedContainers []v1.Container
-	for _, container := range podSpec.Containers {
+	for _, container := range basePodSpec.Containers {
 		// if applicable start with defaultContainerTemplate
 		var mergedContainer *v1.Container
 		if defaultContainerTemplate != nil {
 			mergedContainer = defaultContainerTemplate.DeepCopy()
 		}
 
-		// if applicable merge with primaryContainerTemplate
+		// If this is a primary container handle the template
 		if container.Name == primaryContainerName && primaryContainerTemplate != nil {
 			if mergedContainer == nil {
 				mergedContainer = primaryContainerTemplate.DeepCopy()
@@ -740,35 +745,48 @@
 			}
 		}
 
-		// if applicable merge with existing container
+		// Check for any name matching template containers
+		for _, templateContainer := range templatePodSpec.Containers {
+			if templateContainer.Name != container.Name {
+				continue
+			}
+
+			if mergedContainer == nil {
+				mergedContainer = &templateContainer
+			} else {
+				err := mergo.Merge(mergedContainer, templateContainer, mergo.WithOverride, mergo.WithAppendSlice)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		// Merge in the base container
 		if mergedContainer == nil {
-			mergedContainers = append(mergedContainers, container)
+			mergedContainer = container.DeepCopy()
 		} else {
 			err := mergo.Merge(mergedContainer, container, mergo.WithOverride, mergo.WithAppendSlice)
 			if err != nil {
 				return nil, err
 			}
-
-			mergedContainers = append(mergedContainers, *mergedContainer)
 		}
-	}
 
-	if mergedContainers == nil {
-		mergedContainers = basePodSpec.Containers
+		mergedContainers = append(mergedContainers, *mergedContainer)
+
 	}
 
 	mergedPodSpec.Containers = mergedContainers
 
 	// merge PodTemplate init containers
 	var mergedInitContainers []v1.Container
-	for _, initContainer := range podSpec.InitContainers {
+	for _, initContainer := range basePodSpec.InitContainers {
 		// if applicable start with defaultContainerTemplate
 		var mergedInitContainer *v1.Container
 		if defaultInitContainerTemplate != nil {
 			mergedInitContainer = defaultInitContainerTemplate.DeepCopy()
 		}
 
-		// if applicable merge with primaryInitContainerTemplate
+		// If this is a primary init container handle the template
 		if initContainer.Name == primaryInitContainerName && primaryInitContainerTemplate != nil {
 			if mergedInitContainer == nil {
 				mergedInitContainer = primaryInitContainerTemplate.DeepCopy()
@@ -780,21 +798,86 @@
 			}
 		}
 
-		// if applicable merge with existing init initContainer
+		// Check for any name matching template containers
+		for _, templateInitContainer := range templatePodSpec.InitContainers {
+			if templateInitContainer.Name != initContainer.Name {
+				continue
+			}
+
+			if mergedInitContainer == nil {
+				mergedInitContainer = &templateInitContainer
+			} else {
+				err := mergo.Merge(mergedInitContainer, templateInitContainer, mergo.WithOverride, mergo.WithAppendSlice)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		// Merge in the base init container
 		if mergedInitContainer == nil {
-			mergedInitContainers = append(mergedInitContainers, initContainer)
+			mergedInitContainer = initContainer.DeepCopy()
 		} else {
 			err := mergo.Merge(mergedInitContainer, initContainer, mergo.WithOverride, mergo.WithAppendSlice)
 			if err != nil {
 				return nil, err
 			}
+		}
+
+		mergedInitContainers = append(mergedInitContainers, *mergedInitContainer)
+	}
 
-			mergedInitContainers = append(mergedInitContainers, *mergedInitContainer)
+	mergedPodSpec.InitContainers = mergedInitContainers
+
+	return mergedPodSpec, nil
+}
+
+// MergeOverlayPodSpecOntoBase merges a customized pod spec onto a base pod spec. At a container level it will
+// merge containers that have matching names.
+func MergeOverlayPodSpecOntoBase(basePodSpec *v1.PodSpec, overlayPodSpec *v1.PodSpec) (*v1.PodSpec, error) {
+	if basePodSpec == nil || overlayPodSpec == nil {
+		return nil, errors.New("basePodSpec and overlayPodSpec must not be nil")
+	}
+
+	mergedPodSpec := basePodSpec.DeepCopy()
+	if err := mergo.Merge(mergedPodSpec, overlayPodSpec, mergo.WithOverride, mergo.WithAppendSlice); err != nil {
+		return nil, err
+	}
+
+	// merge PodTemplate containers
+	var mergedContainers []v1.Container
+	for _, container := range basePodSpec.Containers {
+
+		mergedContainer := container.DeepCopy()
+
+		for _, overlayContainer := range overlayPodSpec.Containers {
+			if mergedContainer.Name == overlayContainer.Name {
+				err := mergo.Merge(mergedContainer, overlayContainer, mergo.WithOverride, mergo.WithAppendSlice)
+				if err != nil {
+					return nil, err
+				}
+			}
 		}
+		mergedContainers = append(mergedContainers, *mergedContainer)
 	}
 
-	if mergedInitContainers == nil {
-		mergedInitContainers = basePodSpec.InitContainers
+	mergedPodSpec.Containers = mergedContainers
+
+	// merge PodTemplate init containers
+	var mergedInitContainers []v1.Container
+	for _, initContainer := range basePodSpec.InitContainers {
+
+		mergedInitContainer := initContainer.DeepCopy()
+
+		for _, overlayInitContainer := range overlayPodSpec.InitContainers {
+			if mergedInitContainer.Name == overlayInitContainer.Name {
+				err := mergo.Merge(mergedInitContainer, overlayInitContainer, mergo.WithOverride, mergo.WithAppendSlice)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+		mergedInitContainers = append(mergedInitContainers, *mergedInitContainer)
 	}
 
 	mergedPodSpec.InitContainers = mergedInitContainers