refactor: env overwriter

.github/workflows/build.yaml

@@ -113,7 +113,7 @@ jobs:
       - name: Test common module
         working-directory: ./common
         run: |
-          go test -v ./...
+          make test
 
   test-procdiscovery:
     runs-on: ubuntu-latest

autoscaler/controllers/gateway/root.go

@@ -60,7 +60,7 @@ func Sync(ctx context.Context, client client.Client, scheme *runtime.Scheme, ima
 
 	odigosSystemNamespaceName := env.GetCurrentNamespace()
 	var odigosConfig odigosv1.OdigosConfiguration
-	if err := client.Get(ctx, types.NamespacedName{Namespace: odigosSystemNamespaceName, Name: consts.DefaultOdigosConfigurationName}, &odigosConfig); err != nil {
+	if err := client.Get(ctx, types.NamespacedName{Namespace: odigosSystemNamespaceName, Name: consts.OdigosConfigurationName}, &odigosConfig); err != nil {
 		logger.Error(err, "failed to get odigos config")
 		return err
 	}

cli/cmd/resources/applyresources.go

@@ -9,6 +9,7 @@ import (
 	"github.com/odigos-io/odigos/cli/cmd/resources/resourcemanager"
 	"github.com/odigos-io/odigos/cli/pkg/kube"
 	"github.com/odigos-io/odigos/cli/pkg/log"
+	"github.com/odigos-io/odigos/common/consts"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -41,5 +42,5 @@ func DeleteOldOdigosSystemObjects(ctx context.Context, client *kube.Client, ns s
 }
 
 func GetCurrentConfig(ctx context.Context, client *kube.Client, ns string) (*v1alpha1.OdigosConfiguration, error) {
-	return client.OdigosClient.OdigosConfigurations(ns).Get(ctx, OdigosConfigName, metav1.GetOptions{})
+	return client.OdigosClient.OdigosConfigurations(ns).Get(ctx, consts.OdigosConfigurationName, metav1.GetOptions{})
 }

cli/cmd/resources/odigosconfig.go

@@ -7,14 +7,11 @@ import (
 	"github.com/odigos-io/odigos/cli/cmd/resources/resourcemanager"
 	"github.com/odigos-io/odigos/cli/pkg/kube"
 	"github.com/odigos-io/odigos/common"
+	"github.com/odigos-io/odigos/common/consts"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
-const (
-	OdigosConfigName = "odigos-config"
-)
-
 func otelSdkConfigCommunity() (map[common.ProgrammingLanguage]common.OtelSdk, map[common.ProgrammingLanguage][]common.OtelSdk) {
 	return map[common.ProgrammingLanguage]common.OtelSdk{
 			common.JavaProgrammingLanguage:       common.OtelSdkNativeCommunity,
@@ -75,7 +72,7 @@ func NewOdigosConfiguration(ns string, config *odigosv1.OdigosConfigurationSpec)
 			APIVersion: "odigos.io/v1alpha1",
 		},
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      OdigosConfigName,
+			Name:      consts.OdigosConfigurationName,
 			Namespace: ns,
 		},
 		Spec: *config,

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"
@@ -202,6 +202,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]
@@ -226,24 +232,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/Makefile

@@ -0,0 +1,4 @@
+
+.PHONY: test
+test:
+	go test ./...

common/consts/consts.go

@@ -5,17 +5,18 @@ import (
 )
 
 const (
-	CurrentNamespaceEnvVar         = "CURRENT_NS"
-	DefaultOdigosNamespace         = "odigos-system"
-	DefaultOdigosConfigurationName = "odigos-config"
-	OTLPPort                       = 4317
-	OTLPHttpPort                   = 4318
-	PprofOdigosPort                = 6060
-	OdigosInstrumentationLabel     = "odigos-instrumentation"
-	InstrumentationEnabled         = "enabled"
-	InstrumentationDisabled        = "disabled"
-	OdigosReportedNameAnnotation   = "odigos.io/reported-name"
-	EbpfInstrumentationAnnotation  = "instrumentation.odigos.io/ebpf" // deprecated.
+	CurrentNamespaceEnvVar        = "CURRENT_NS"
+	DefaultOdigosNamespace        = "odigos-system"
+	OdigosConfigurationName       = "odigos-config"
+	OTLPPort                      = 4317
+	OTLPHttpPort                  = 4318
+	PprofOdigosPort               = 6060
+	OdigosInstrumentationLabel    = "odigos-instrumentation"
+	InstrumentationEnabled        = "enabled"
+	InstrumentationDisabled       = "disabled"
+	OdigosReportedNameAnnotation  = "odigos.io/reported-name"
+	EbpfInstrumentationAnnotation = "instrumentation.odigos.io/ebpf" // deprecated.
+
 	// Used to store the original value of the environment variable in the pod manifest.
 	// This is used to restore the original value when an instrumentation is removed
 	// or odigos is uninstalled.

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,117 +23,94 @@ 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
-	}
-
-	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
-	}
-
-	// 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
-		}
+		return nil
 	}
 
-	return true
-}
-
-func ShouldRevert(envName string, observedValue string) bool {
-	env, ok := EnvValuesMap[envName]
-	if !ok {
-		// We don't care about this environment variable
-		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 find any of the values odigos sets for any SDK in the observed value,
-	// that means that the environment variable is patched by odigos and we need to revert it
-	for _, val := range env.values {
-		if strings.Contains(observedValue, val) && observedValue != val {
-			return true
+	// 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 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)
+	// 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.
+				// this can happen when moving between SDKs.
+				patchedEvnValue := strings.ReplaceAll(observedValue, sdkEnvValue, desiredOdigosPart)
+				return &patchedEvnValue
+			}
 		}
 	}
 
-	return currentVal + env.delim + valToAppend
+	// 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 ValToAppend(envName string, sdk common.OtelSdk) (string, bool) {