(fix) internal/civisibility/utils: improve environment tags mutability (

#3076)
DataDog · Jan 13, 2025 · 42997e1 · 42997e1
1 parent 27624f6
commit 42997e1
Show file tree

Hide file tree

Showing 2 changed files with 202 additions and 28 deletions.
diff --git a/internal/civisibility/utils/environmentTags.go b/internal/civisibility/utils/environmentTags.go
@@ -7,6 +7,7 @@ package utils
 
 import (
 	"fmt"
+	"maps"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -21,13 +22,16 @@ import (
 
 var (
 	// ciTags holds the CI/CD environment variable information.
-	ciTags      map[string]string
-	addedTags   map[string]string
-	ciTagsMutex sync.Mutex
+	currentCiTags  map[string]string // currentCiTags holds the CI/CD tags after originalCiTags + addedTags
+	originalCiTags map[string]string // originalCiTags holds the original CI/CD tags after all the CMDs
+	addedTags      map[string]string // addedTags holds the tags added by the user
+	ciTagsMutex    sync.Mutex
 
 	// ciMetrics holds the CI/CD environment numeric variable information
-	ciMetrics      map[string]float64
-	ciMetricsMutex sync.Mutex
+	currentCiMetrics  map[string]float64 // currentCiMetrics holds the CI/CD metrics after originalCiMetrics + addedMetrics
+	originalCiMetrics map[string]float64 // originalCiMetrics holds the original CI/CD metrics after all the CMDs
+	addedMetrics      map[string]float64 // addedMetrics holds the metrics added by the user
+	ciMetricsMutex    sync.Mutex
 )
 
 // GetCITags retrieves and caches the CI/CD tags from environment variables.
@@ -41,11 +45,71 @@ func GetCITags() map[string]string {
 	ciTagsMutex.Lock()
 	defer ciTagsMutex.Unlock()
 
-	if ciTags == nil {
-		ciTags = createCITagsMap()
+	// Return the current tags if they are already initialized
+	if currentCiTags != nil {
+		return currentCiTags
 	}
 
-	return ciTags
+	if originalCiTags == nil {
+		// If the original tags are not initialized, create them
+		originalCiTags = createCITagsMap()
+	}
+
+	// Create a new map with the added tags
+	newTags := maps.Clone(originalCiTags)
+	for k, v := range addedTags {
+		newTags[k] = v
+	}
+
+	// Update the current tags
+	currentCiTags = newTags
+	return currentCiTags
+}
+
+// AddCITags adds a new tag to the CI/CD tags map.
+func AddCITags(tagName, tagValue string) {
+	ciTagsMutex.Lock()
+	defer ciTagsMutex.Unlock()
+
+	// Add the tag to the added tags dictionary
+	if addedTags == nil {
+		addedTags = make(map[string]string)
+	}
+	addedTags[tagName] = tagValue
+
+	// Reset the current tags
+	currentCiTags = nil
+}
+
+// AddCITagsMap adds a new map of tags to the CI/CD tags map.
+func AddCITagsMap(tags map[string]string) {
+	if tags == nil {
+		return
+	}
+
+	ciTagsMutex.Lock()
+	defer ciTagsMutex.Unlock()
+
+	// Add the tag to the added tags dictionary
+	if addedTags == nil {
+		addedTags = make(map[string]string)
+	}
+	for k, v := range tags {
+		addedTags[k] = v
+	}
+
+	// Reset the current tags
+	currentCiTags = nil
+}
+
+// ResetCITags resets the CI/CD tags to their original values.
+func ResetCITags() {
+	ciTagsMutex.Lock()
+	defer ciTagsMutex.Unlock()
+
+	originalCiTags = nil
+	currentCiTags = nil
+	addedTags = nil
 }
 
 // GetCIMetrics retrieves and caches the CI/CD metrics from environment variables.
@@ -59,30 +123,71 @@ func GetCIMetrics() map[string]float64 {
 	ciMetricsMutex.Lock()
 	defer ciMetricsMutex.Unlock()
 
-	if ciMetrics == nil {
-		ciMetrics = createCIMetricsMap()
+	// Return the current metrics if they are already initialized
+	if currentCiMetrics != nil {
+		return currentCiMetrics
+	}
+
+	if originalCiMetrics == nil {
+		// If the original metrics are not initialized, create them
+		originalCiMetrics = createCIMetricsMap()
 	}
 
-	return ciMetrics
+	// Create a new map with the added metrics
+	newMetrics := maps.Clone(originalCiMetrics)
+	for k, v := range addedMetrics {
+		newMetrics[k] = v
+	}
+
+	// Update the current metrics
+	currentCiMetrics = newMetrics
+	return currentCiMetrics
 }
 
-// AddCITags adds a new tag to the CI/CD tags map.
-func AddCITags(tagName, tagValue string) {
-	ciTagsMutex.Lock()
-	defer ciTagsMutex.Unlock()
+// AddCIMetrics adds a new metric to the CI/CD metrics map.
+func AddCIMetrics(metricName string, metricValue float64) {
+	ciMetricsMutex.Lock()
+	defer ciMetricsMutex.Unlock()
 
-	// Add the tag to the added tags dictionary
-	if addedTags == nil {
-		addedTags = make(map[string]string)
+	// Add the metric to the added metrics dictionary
+	if addedMetrics == nil {
+		addedMetrics = make(map[string]float64)
 	}
-	addedTags[tagName] = tagValue
+	addedMetrics[metricName] = metricValue
 
-	// Create a new map with the added tags
-	newTags := createCITagsMap()
-	for k, v := range addedTags {
-		newTags[k] = v
+	// Reset the current metrics
+	currentCiMetrics = nil
+}
+
+// AddCIMetricsMap adds a new map of metrics to the CI/CD metrics map.
+func AddCIMetricsMap(metrics map[string]float64) {
+	if metrics == nil {
+		return
 	}
-	ciTags = newTags
+
+	ciMetricsMutex.Lock()
+	defer ciMetricsMutex.Unlock()
+
+	// Add the metric to the added metrics dictionary
+	if addedMetrics == nil {
+		addedMetrics = make(map[string]float64)
+	}
+	for k, v := range metrics {
+		addedMetrics[k] = v
+	}
+
+	// Reset the current metrics
+	currentCiMetrics = nil
+}
+
+// ResetCIMetrics resets the CI/CD metrics to their original values.
+func ResetCIMetrics() {
+	ciMetricsMutex.Lock()
+	defer ciMetricsMutex.Unlock()
+
+	originalCiMetrics = nil
+	currentCiMetrics = nil
+	addedMetrics = nil
 }
 
 // GetRelativePathFromCITagsSourceRoot calculates the relative path from the CI workspace root to the specified path.

diff --git a/internal/civisibility/utils/environmentTags_test.go b/internal/civisibility/utils/environmentTags_test.go
@@ -14,7 +14,8 @@ import (
 )
 
 func TestGetCITagsCache(t *testing.T) {
-	ciTags = map[string]string{"key": "value"}
+	ResetCITags()
+	originalCiTags = map[string]string{"key": "value"}
 
 	// First call to initialize ciTags
 	tags := GetCITags()
@@ -25,8 +26,41 @@ func TestGetCITagsCache(t *testing.T) {
 	assert.Equal(t, "newvalue", tags["key"])
 }
 
+func TestAddCITags(t *testing.T) {
+	ResetCITags()
+	originalCiTags = map[string]string{"key": "value"}
+
+	// First call to initialize ciTags
+	tags := GetCITags()
+	assert.Equal(t, "value", tags["key"])
+
+	AddCITags("key", "newvalue")
+	AddCITags("key2", "value2")
+	tags = GetCITags()
+	assert.Equal(t, "newvalue", tags["key"])
+	assert.Equal(t, "value2", tags["key2"])
+}
+
+func TestAddCITagsMap(t *testing.T) {
+	ResetCITags()
+	originalCiTags = map[string]string{"key": "value"}
+
+	// First call to initialize ciTags
+	tags := GetCITags()
+	assert.Equal(t, "value", tags["key"])
+
+	nmap := map[string]string{}
+	nmap["key"] = "newvalue"
+	nmap["key2"] = "value2"
+	AddCITagsMap(nmap)
+	tags = GetCITags()
+	assert.Equal(t, "newvalue", tags["key"])
+	assert.Equal(t, "value2", tags["key2"])
+}
+
 func TestGetCIMetricsCache(t *testing.T) {
-	ciMetrics = map[string]float64{"key": float64(1)}
+	ResetCIMetrics()
+	originalCiMetrics = map[string]float64{"key": float64(1)}
 
 	// First call to initialize ciMetrics
 	tags := GetCIMetrics()
@@ -37,16 +71,51 @@ func TestGetCIMetricsCache(t *testing.T) {
 	assert.Equal(t, float64(42), tags["key"])
 }
 
+func TestAddCIMetrics(t *testing.T) {
+	ResetCIMetrics()
+	originalCiMetrics = map[string]float64{"key": float64(1)}
+
+	// First call to initialize ciMetrics
+	tags := GetCIMetrics()
+	assert.Equal(t, float64(1), tags["key"])
+
+	AddCIMetrics("key", float64(42))
+	AddCIMetrics("key2", float64(2))
+	tags = GetCIMetrics()
+	assert.Equal(t, float64(42), tags["key"])
+	assert.Equal(t, float64(2), tags["key2"])
+}
+
+func TestAddCIMetricsMap(t *testing.T) {
+	ResetCIMetrics()
+	originalCiMetrics = map[string]float64{"key": float64(1)}
+
+	// First call to initialize ciMetrics
+	tags := GetCIMetrics()
+	assert.Equal(t, float64(1), tags["key"])
+
+	nmap := map[string]float64{}
+	nmap["key"] = float64(42)
+	nmap["key2"] = float64(2)
+	AddCIMetricsMap(nmap)
+	tags = GetCIMetrics()
+	assert.Equal(t, float64(42), tags["key"])
+	assert.Equal(t, float64(2), tags["key2"])
+}
+
 func TestGetRelativePathFromCITagsSourceRoot(t *testing.T) {
-	ciTags = map[string]string{constants.CIWorkspacePath: "/ci/workspace"}
+	ResetCITags()
+	originalCiTags = map[string]string{constants.CIWorkspacePath: "/ci/workspace"}
+
 	absPath := "/ci/workspace/subdir/file.txt"
 	expectedRelPath := "subdir/file.txt"
 
 	relPath := GetRelativePathFromCITagsSourceRoot(absPath)
 	assert.Equal(t, expectedRelPath, relPath)
 
 	// Test case when CIWorkspacePath is not set in ciTags
-	ciTags = map[string]string{}
+	originalCiTags = map[string]string{}
+	currentCiTags = nil
 	relPath = GetRelativePathFromCITagsSourceRoot(absPath)
 	assert.Equal(t, absPath, relPath)
 }