feat: refactor pod patch configuration and patching logic (#10)

- Renamed WorkerTemplate to Worker in config
- Updated patching logic to use PatchToPod and PatchToContainer
- Refined container environment variable patching in pod webhook

Author: 0x5457

cmd/main.go

@@ -167,7 +167,7 @@ func main() {
 		Scheme:    mgr.GetScheme(),
 		Scheduler: scheduler,
 		WorkerGenerator: &worker.WorkerGenerator{
-			PodTemplate: &config.WorkerTemplate,
+			PodTemplate: &config.Worker,
 		},
 	}).SetupWithManager(mgr); err != nil {
 		setupLog.Error(err, "unable to create controller", "controller", "TensorFusionConnection")

internal/config/config.go

@@ -8,24 +8,14 @@ import (
 	"k8s.io/utils/ptr"
 )
 
-type Pod struct {
-	Spec PodSpec `json:"spec,omitempty"`
-}
-
-type PodSpec struct {
-	InitContainers   []corev1.Container `json:"initContainers,omitempty" patchStrategy:"merge" patchMergeKey:"name" protobuf:"bytes,20,rep,name=initContainers"`
-	Containers       []corev1.Container `json:"containers,omitempty" patchStrategy:"merge" patchMergeKey:"name" protobuf:"bytes,2,rep,name=containers"`
-	RuntimeClassName *string            `json:"runtimeClassName,omitempty" protobuf:"bytes,29,opt,name=runtimeClassName"`
-}
-
 type Config struct {
-	WorkerTemplate corev1.PodTemplate `json:"workerTemplate"`
-	PodMutator     PodMutator         `json:"podMutator"`
+	Worker     corev1.PodTemplate `json:"worker"`
+	PodMutator PodMutator         `json:"podMutator"`
 }
 
 type PodMutator struct {
-	PatchStrategicMerge Pod             `json:"patchStrategicMerge"`
-	PatchEnvVars        []corev1.EnvVar `json:"envVars"`
+	PatchToPod       any `json:"patchToPod"`
+	PatchToContainer any `json:"patchToContainer"`
 }
 
 func LoadConfig(filename string) (*Config, error) {
@@ -43,7 +33,7 @@ func LoadConfig(filename string) (*Config, error) {
 
 func NewDefaultConfig() *Config {
 	return &Config{
-		WorkerTemplate: corev1.PodTemplate{
+		Worker: corev1.PodTemplate{
 			Template: corev1.PodTemplateSpec{
 				Spec: corev1.PodSpec{
 					TerminationGracePeriodSeconds: ptr.To[int64](0),
@@ -58,20 +48,22 @@ func NewDefaultConfig() *Config {
 			},
 		},
 		PodMutator: PodMutator{
-			PatchStrategicMerge: Pod{
-				Spec: PodSpec{
-					InitContainers: []corev1.Container{
+			PatchToPod: map[string]any{
+				"spec": map[string]any{
+					"initContainers": []corev1.Container{
 						{
 							Name:  "inject-lib",
 							Image: "busybox:stable-glibc",
 						},
 					},
 				},
 			},
-			PatchEnvVars: []corev1.EnvVar{
-				{
-					Name:  "LD_PRELOAD",
-					Value: "tensorfusion.so",
+			PatchToContainer: map[string]any{
+				"env": []corev1.EnvVar{
+					{
+						Name:  "LD_PRELOAD",
+						Value: "tensorfusion.so",
+					},
 				},
 			},
 		},

internal/controller/tensorfusionconnection_controller_test.go

@@ -75,7 +75,7 @@ var _ = Describe("TensorFusionConnection Controller", func() {
 				Client: k8sClient,
 				Scheme: k8sClient.Scheme(),
 				WorkerGenerator: &worker.WorkerGenerator{
-					PodTemplate: &config.WorkerTemplate,
+					PodTemplate: &config.Worker,
 				},
 			}
 			_, err := controllerReconciler.Reconcile(ctx, reconcile.Request{

internal/webhook/v1/pod_webhook.go

@@ -22,7 +22,6 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/NexusGPU/tensor-fusion-operator/internal/config"
 	"gomodules.xyz/jsonpatch/v2"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -33,6 +32,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 
+	"github.com/NexusGPU/tensor-fusion-operator/internal/config"
 	"github.com/NexusGPU/tensor-fusion-operator/internal/constants"
 )
 
@@ -147,30 +147,43 @@ func (m *TensorFusionPodMutator) patchTFClient(pod *corev1.Pod, tfReq []TFReq) (
 		return nil, fmt.Errorf("marshal current pod: %v", err)
 	}
 
-	// Patch env vars
+	// Patch to Container
 	for _, req := range tfReq {
 		for i := range pod.Spec.Containers {
 			container := &pod.Spec.Containers[i]
 			if container.Name == req.ContainerName {
-				container.Env = append(container.Env, m.Config.PatchEnvVars...)
+				containerJSON, err := json.Marshal(container)
+				if err != nil {
+					return nil, fmt.Errorf("marshal container: %v", err)
+				}
+				patchJSON, err := json.Marshal(m.Config.PatchToContainer)
+				if err != nil {
+					return nil, fmt.Errorf("marshal patchToContainer: %v", err)
+				}
+
+				patchedJSON, err := strategicpatch.StrategicMergePatch(containerJSON, patchJSON, corev1.Container{})
+				if err != nil {
+					return nil, fmt.Errorf("apply strategic merge patch to container: %v", err)
+				}
+
+				if err := json.Unmarshal(patchedJSON, container); err != nil {
+					return nil, fmt.Errorf("unmarshal patched container: %v", err)
+				}
 			}
 		}
 	}
-	envPatchedStatus, err := json.Marshal(pod)
+
+	containerPatchedJSON, err := json.Marshal(pod)
 	if err != nil {
 		return nil, fmt.Errorf("marshal current pod: %v", err)
 	}
-	patches, err := jsonpatch.CreatePatch(currentBytes, envPatchedStatus)
+	patches, err := jsonpatch.CreatePatch(currentBytes, containerPatchedJSON)
 	if err != nil {
-		return nil, fmt.Errorf("patch env: %v", err)
+		return nil, fmt.Errorf("patch to container: %v", err)
 	}
 
-	podPatch := m.Config.PatchStrategicMerge
-	// Copy containers
-	podPatch.Spec.Containers = append([]corev1.Container{}, podPatch.Spec.Containers...)
-
 	// Convert the strategic merge patch to JSON
-	patchBytes, err := json.Marshal(m.Config.PatchStrategicMerge)
+	patchBytes, err := json.Marshal(m.Config.PatchToPod)
 
 	if err != nil {
 		return nil, fmt.Errorf("marshal patch: %v", err)