From 6212bdef7a0c2476b682e1fe3e18e5d13e088f04 Mon Sep 17 00:00:00 2001
From: Romain Dauby <romain.dauby@gmail.com>
Date: Fri, 27 Dec 2024 11:48:05 -0500
Subject: [PATCH] Fix integration test for ECS

---
 .../integration_test.go                       | 106 ++++++++++++++----
 1 file changed, 82 insertions(+), 24 deletions(-)

diff --git a/extension/cgroupruntimeextension/integration_test.go b/extension/cgroupruntimeextension/integration_test.go
index a9084821b71e..8440a515d855 100644
--- a/extension/cgroupruntimeextension/integration_test.go
+++ b/extension/cgroupruntimeextension/integration_test.go
@@ -9,7 +9,6 @@ package cgroupruntimeextension // import "github.com/open-telemetry/opentelemetr
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"math"
 	"net/http"
@@ -33,6 +32,7 @@ import (
 
 const (
 	defaultCgroup2Path = "/sys/fs/cgroup"
+	ecsMetadataUri     = "ECS_CONTAINER_METADATA_URI_V4"
 )
 
 // checkCgroupSystem skips the test if is not run in a cgroupv2 system
@@ -66,19 +66,24 @@ func cgroupMaxCpu(filename string) (quota int64, period uint64, err error) {
 	return quota, period, err
 }
 
-func startMockECSServer() *httptest.Server {
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		response := map[string]interface{}{
-			"DockerID": "container-id",
-			"Limits": map[string]interface{}{
-				"CPU": 2.0,
-			},
-		}
-		w.Header().Set("Content-Type", "application/json")
-		json.NewEncoder(w).Encode(response)
+func testServerECSMetadata(t *testing.T, containerCPU, taskCPU int) *httptest.Server {
+	t.Helper()
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
+		_, err := w.Write([]byte(fmt.Sprintf(`{"Limits":{"CPU":%d},"DockerId":"container-id"}`, containerCPU)))
+		assert.NoError(t, err)
+	})
+	mux.HandleFunc("/task", func(w http.ResponseWriter, _ *http.Request) {
+		_, err := w.Write([]byte(fmt.Sprintf(
+			`{"Containers":[{"DockerId":"container-id","Limits":{"CPU":%d}}],"Limits":{"CPU":%d}}`,
+			containerCPU,
+			taskCPU,
+		)))
+		assert.NoError(t, err)
 	})
 
-	return httptest.NewServer(handler)
+	return httptest.NewServer(mux)
 }
 
 func TestCgroupV2SudoIntegration(t *testing.T) {
@@ -164,10 +169,11 @@ func TestCgroupV2SudoIntegration(t *testing.T) {
 			expectedGoMemLimit: 13421772,
 		},
 		{
-			name:            "running on AWS ECS with 90% of max cgroup memory and 2 GOMAXPROCS",
-			cgroupCpuQuota:  pointerInt64(-1),
+			name:            "AWS ECS 90% the max cgroup memory and 12 GOMAXPROCS",
+			cgroupCpuQuota:  pointerInt64(100000),
 			cgroupCpuPeriod: 8000,
-			cgroupMaxMemory: 134217728, // 128 MB
+			// 128 Mb
+			cgroupMaxMemory: 134217728,
 			config: &Config{
 				GoMaxProcs: GoMaxProcsConfig{
 					Enabled: true,
@@ -177,8 +183,54 @@ func TestCgroupV2SudoIntegration(t *testing.T) {
 					Ratio:   0.9,
 				},
 			},
-			expectedGoMaxProcs: 22,
-			expectedGoMemLimit: 120795955, // 134217728 * 0.9
+			// 100000 / 8000
+			expectedGoMaxProcs: 12,
+			// 134217728 * 0.9
+			expectedGoMemLimit: 120795955,
+			setECSMetadataURI:  true,
+		},
+		{
+			name:            "AWS ECS 50% of the max cgroup memory and 1 GOMAXPROCS",
+			cgroupCpuQuota:  pointerInt64(100000),
+			cgroupCpuPeriod: 100000,
+			// 128 Mb
+			cgroupMaxMemory: 134217728,
+			config: &Config{
+				GoMaxProcs: GoMaxProcsConfig{
+					Enabled: true,
+				},
+				GoMemLimit: GoMemLimitConfig{
+					Enabled: true,
+					Ratio:   0.5,
+				},
+			},
+			// 100000 / 100000
+			expectedGoMaxProcs: 1,
+			// 134217728 * 0.5
+			expectedGoMemLimit: 67108864,
+			setECSMetadataURI:  true,
+		},
+		{
+			name:            "AWS ECS 10% of the max cgroup memory, max cpu, default GOMAXPROCS",
+			cgroupCpuQuota:  nil,
+			cgroupCpuPeriod: 100000,
+			// 128 Mb
+			cgroupMaxMemory: 134217728,
+			config: &Config{
+				GoMaxProcs: GoMaxProcsConfig{
+					Enabled: true,
+				},
+				GoMemLimit: GoMemLimitConfig{
+					Enabled: true,
+					Ratio:   0.1,
+				},
+			},
+			// GOMAXPROCS is set to the value of  `cpu.max / cpu.period`
+			// If cpu.max is set to max, GOMAXPROCS should not be
+			// modified
+			expectedGoMaxProcs: runtime.GOMAXPROCS(-1),
+			// 134217728 * 0.1
+			expectedGoMemLimit: 13421772,
 			setECSMetadataURI:  true,
 		},
 	}
@@ -235,12 +287,25 @@ func TestCgroupV2SudoIntegration(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
+			// if running in ECS environment, set the ECS metedata URI environment variable
+			// to get the Cgroup CPU quota from the httptest server
+			cleanECS := func() {}
+			if test.setECSMetadataURI {
+				server := testServerECSMetadata(t, test.expectedGoMaxProcs*1024, test.expectedGoMaxProcs*1024)
+				os.Setenv(ecsMetadataUri, server.URL)
+				cleanECS = func() {
+					server.Close()
+					os.Unsetenv(ecsMetadataUri)
+				}
+			}
+
 			// restore startup cgroup initial resource values
 			t.Cleanup(func() {
 				debug.SetMemoryLimit(initialGoMem)
 				runtime.GOMAXPROCS(initialGoProcs)
 				memoryCgroupCleanUp()
 				cpuCgroupCleanUp()
+				cleanECS()
 			})
 
 			err = manager.Update(&cgroup2.Resources{
@@ -257,13 +322,6 @@ func TestCgroupV2SudoIntegration(t *testing.T) {
 			})
 			require.NoError(t, err)
 
-			if test.setECSMetadataURI {
-				server := startMockECSServer()
-				defer server.Close()
-				os.Setenv("ECS_CONTAINER_METADATA_URI_V4", server.URL)
-				defer os.Unsetenv("ECS_CONTAINER_METADATA_URI_V4")
-			}
-
 			factory := NewFactory()
 			ctx := context.Background()
 			extension, err := factory.Create(ctx, extensiontest.NewNopSettings(), test.config)