refactor: env overwriter

cli/cmd/uninstall.go

@@ -8,7 +8,7 @@ import (
 	"strings"
 	"time"
 
-	"github.com/odigos-io/odigos/common/envOverwrite"
+	"github.com/odigos-io/odigos/k8sutils/pkg/envoverwrite"
 
 	"github.com/odigos-io/odigos/cli/cmd/resources"
 	"github.com/odigos-io/odigos/cli/pkg/confirm"
@@ -201,6 +201,12 @@ func getWorkloadRolloutJsonPatch(obj client.Object, pts *v1.PodTemplateSpec) ([]
 	}
 
 	// read the original env vars (of the manifest) from the annotation
+	manifestEnvOriginal, err := envoverwrite.NewOrigWorkloadEnvValues(obj)
+	if err != nil {
+		fmt.Println("Failed to get original env vars from annotation: ", err)
+		manifestEnvOriginal = &envoverwrite.OrigWorkloadEnvValues{}
+	}
+
 	var origManifestEnv map[string]map[string]string
 	if obj.GetAnnotations() != nil {
 		manifestEnvAnnotation, ok := obj.GetAnnotations()[consts.ManifestEnvOriginalValAnnotation]
@@ -225,24 +231,20 @@ func getWorkloadRolloutJsonPatch(obj client.Object, pts *v1.PodTemplateSpec) ([]
 			}
 		}
 
-		containerOriginalEnv := origManifestEnv[c.Name]
-
-		for iEnv, envVar := range c.Env {
-			if envOverwrite.ShouldRevert(envVar.Name, envVar.Value) {
-				if origVal, ok := containerOriginalEnv[envVar.Name]; ok {
-					// revert the env var to its original value if we have it
-					patchOperations = append(patchOperations, map[string]interface{}{
-						"op":    "replace",
-						"path":  fmt.Sprintf("/spec/template/spec/containers/%d/env/%d/value", iContainer, iEnv),
-						"value": origVal,
-					})
-				} else {
-					// remove the env var
-					patchOperations = append(patchOperations, map[string]interface{}{
-						"op":   "remove",
-						"path": fmt.Sprintf("/spec/template/spec/containers/%d/env/%d", iContainer, iEnv),
-					})
-				}
+		for envName, originalEnvValue := range manifestEnvOriginal.GetContainerStoredEnvs(c.Name) {
+			if origManifestEnv == nil {
+				// originally the value was absent, so we remove it
+				patchOperations = append(patchOperations, map[string]interface{}{
+					"op":   "remove",
+					"path": fmt.Sprintf("/spec/template/spec/containers/%d/env/%d", iContainer, envName),
+				})
+			} else {
+				// revert the env var to its original value
+				patchOperations = append(patchOperations, map[string]interface{}{
+					"op":    "replace",
+					"path":  fmt.Sprintf("/spec/template/spec/containers/%d/env/%d/value", iContainer, envName),
+					"value": *originalEnvValue,
+				})
 			}
 		}
 	}

common/envOverwrite/originalenv.go

@@ -0,0 +1,14 @@
+package envOverwrite
+
+// this type is used to store the original values of the environment variables the user has set
+// for a specific deployment manifest.
+// In k8s we need to store such map for each container in the pod.
+// In systemd, this can store the original values of the systemd service.
+//
+// When we want to rollback the changes we made to the environment variables, we can fetch the original
+// values from this map and set them back to the manifest.
+//
+// The key is the environment variable name.
+// The value is either the original value of the environment variable or nil if the environment variable
+// was not set by the user.
+type OriginalEnv map[string]*string

common/envOverwrite/overwriter.go

@@ -7,9 +7,10 @@ import (
 )
 
 type envValues struct {
-	delim string
+	delim  string
 	values map[common.OtelSdk]string
 }
+
 // EnvValuesMap is a map of environment variables odigos uses for various languages and goals.
 // The key is the environment variable name and the value is the value to be set or appended
 // to the environment variable. We need to make sure that in case any of these environment
@@ -22,67 +23,93 @@ var EnvValuesMap = map[string]envValues{
 		delim: " ",
 		values: map[common.OtelSdk]string{
 			common.OtelSdkNativeCommunity: "--require /var/odigos/nodejs/autoinstrumentation.js",
-			common.OtelSdkEbpfEnterprise: "--require /var/odigos/nodejs-ebpf/autoinstrumentation.js",
+			common.OtelSdkEbpfEnterprise:  "--require /var/odigos/nodejs-ebpf/autoinstrumentation.js",
 		},
 	},
 	"PYTHONPATH": {
 		delim: ":",
 		values: map[common.OtelSdk]string{
 			common.OtelSdkNativeCommunity: "/var/odigos/python:/var/odigos/python/opentelemetry/instrumentation/auto_instrumentation",
-			common.OtelSdkEbpfEnterprise: "/var/odigos/python-ebpf:/var/odigos/python/opentelemetry/instrumentation/auto_instrumentation:/var/odigos/python",
+			common.OtelSdkEbpfEnterprise:  "/var/odigos/python-ebpf:/var/odigos/python/opentelemetry/instrumentation/auto_instrumentation:/var/odigos/python",
 		},
 	},
 	"JAVA_OPTS": {
 		delim: " ",
 		values: map[common.OtelSdk]string{
 			common.OtelSdkNativeCommunity: "-javaagent:/var/odigos/java/javaagent.jar",
-			common.OtelSdkEbpfEnterprise: "-javaagent:/var/odigos/java-ebpf/dtrace-injector.jar",
+			common.OtelSdkEbpfEnterprise:  "-javaagent:/var/odigos/java-ebpf/dtrace-injector.jar",
 			common.OtelSdkNativeEnterprise: "-javaagent:/var/odigos/java-ext-ebpf/javaagent.jar " +
-			"-Dotel.javaagent.extensions=/var/odigos/java-ext-ebpf/otel_agent_extension.jar",
+				"-Dotel.javaagent.extensions=/var/odigos/java-ext-ebpf/otel_agent_extension.jar",
 		},
 	},
 	"JAVA_TOOL_OPTIONS": {
 		delim: " ",
 		values: map[common.OtelSdk]string{
 			common.OtelSdkNativeCommunity: "-javaagent:/var/odigos/java/javaagent.jar",
-			common.OtelSdkEbpfEnterprise: "-javaagent:/var/odigos/java-ebpf/dtrace-injector.jar",
+			common.OtelSdkEbpfEnterprise:  "-javaagent:/var/odigos/java-ebpf/dtrace-injector.jar",
 			common.OtelSdkNativeEnterprise: "-javaagent:/var/odigos/java-ext-ebpf/javaagent.jar " +
-			"-Dotel.javaagent.extensions=/var/odigos/java-ext-ebpf/otel_agent_extension.jar",
+				"-Dotel.javaagent.extensions=/var/odigos/java-ext-ebpf/otel_agent_extension.jar",
 		},
 	},
 }
 
-func ShouldPatch(envName string, observedValue string, sdk common.OtelSdk) bool {
-	env, ok := EnvValuesMap[envName]
+// returns the current value that should be populated in a specific environment variable.
+// if we should not patch the value, returns nil.
+// the are 2 parts to the environment value: odigos part and user part.
+// either one can be set or empty.
+// so we have 4 cases to handle:
+func GetPatchedEnvValue(envName string, observedValue string, currentSdk common.OtelSdk) *string {
+	envMetadata, ok := EnvValuesMap[envName]
 	if !ok {
 		// Odigos does not manipulate this environment variable, so ignore it
-		return false
+		return nil
 	}
 
-	desiredValue, ok := env.values[sdk]
+	desiredOdigosPart, ok := envMetadata.values[currentSdk]
 	if !ok {
 		// No specific overwrite is required for this SDK
-		return false
+		return nil
 	}
 
-	if strings.Contains(observedValue, desiredValue) {
-		// if the observed value contains the value odigos sets for this SDK,
-		// that means that either:
-		//    1. the user does not add any additional values and we see here our own value only,
-		//    2. we already patched the value in the deployment manifest and we see here the patched value,
-		// so we should not patch it
-		return false
+	// scenario 1: no user defined values and no odigos value
+	// happens: might be the case right after the source is instrumented, and before the instrumentation is applied.
+	// action: there are no user defined values, so no need to make any changes.
+	if observedValue == "" {
+		return nil
 	}
 
-	// if we are moving from one SDK to another, avoid patching the value
-	// it there is no user defined value (observedValue is equal to the value odigos sets for the previous SDK)
-	for _, v := range env.values {
-		if v == observedValue {
-			return false
+	// scenario 2: no user defined values, only odigos value
+	// happens: when the user did not set any value to this env (either via manifest or dockerfile)
+	// action: we don't need to overwrite the value, just let odigos handle it
+	for _, sdkEnvValue := range envMetadata.values {
+		if sdkEnvValue == observedValue {
+			return nil
 		}
 	}
 
-	return true
+	// Scenario 3: both odigos and user defined values are present
+	// happens: when the user set some values to this env (either via manifest or dockerfile) and odigos instrumentation is applied.
+	// action: we want to keep the user defined values and upsert the odigos value.
+	for _, sdkEnvValue := range envMetadata.values {
+		if strings.Contains(observedValue, sdkEnvValue) {
+			if sdkEnvValue == desiredOdigosPart {
+				// shortcut, the value is already patched
+				// both the odigos part equals to the new value, and the user part we want to keep
+				return &observedValue
+			} else {
+				// The environment variable is patched by some other odigos sdk.
+				// replace just the odigos part with the new desired value.
+				patchedEvnValue := strings.ReplaceAll(observedValue, sdkEnvValue, desiredOdigosPart)
+				return &patchedEvnValue
+			}
+		}
+	}
+
+	// Scenario 4: only user defined values are present
+	// happens: when the user set some values to this env (either via manifest or dockerfile) and odigos instrumentation not yet applied.
+	// action: we want to keep the user defined values and append the odigos value.
+	mergedEnvValue := observedValue + envMetadata.delim + desiredOdigosPart
+	return &mergedEnvValue
 }
 
 func ShouldRevert(envName string, observedValue string) bool {
@@ -103,38 +130,6 @@ func ShouldRevert(envName string, observedValue string) bool {
 	return false
 }
 
-func Patch(envName string, currentVal string, sdk common.OtelSdk) string {
-	env, exists := EnvValuesMap[envName]
-	if !exists {
-		return ""
-	}
-
-	valToAppend, ok := env.values[sdk]
-	if !ok {
-		return ""
-	}
-
-	if currentVal == "" {
-		return valToAppend
-	}
-
-	if strings.Contains(currentVal, valToAppend) {
-		// The environment variable is already patched with the correct value for this SDK
-		return currentVal
-	}
-
-	for sdkKey, val := range env.values {
-		if sdkKey != sdk && strings.Contains(currentVal, val){
-			// The environment variable is patched by another SDK
-			// but we need to append our value to it, so replace the
-			// value of the other SDK with ours
-			return strings.ReplaceAll(currentVal, val, valToAppend)
-		}
-	}
-
-	return currentVal + env.delim + valToAppend
-}
-
 func ValToAppend(envName string, sdk common.OtelSdk) (string, bool) {
 	env, exists := EnvValuesMap[envName]
 	if !exists {

instrumentor/Makefile

@@ -36,7 +36,8 @@ vet: ## Run go vet against code.
 
 .PHONY: test
 test: fmt vet envtest ## Run tests.
-	KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) -p path)" go test ./... -coverprofile cover.out
+	KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) -p path)" go test ./... -p 1 -coverprofile cover.out
+	# KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) -p path)" go test ./controllers/instrumentationdevice -p 1 -coverprofile cover.out
 
 ##@ Build
 

instrumentor/controllers/instrumentationdevice/common.go

@@ -128,7 +128,11 @@ func uninstrument(logger logr.Logger, ctx context.Context, kubeClient client.Cli
 			return err
 		}
 
-		instrumentation.Revert(podSpec, obj)
+		instrumentation.RevertInstrumentationDevices(podSpec)
+		err = instrumentation.RevertEnvOverwrites(obj, podSpec)
+		if err != nil {
+			return err
+		}
 		return nil
 	})