feat: add GOMEMLIMIT to Deployment env if memory resource limits are …

…defined (#232) * add GOMEMLIMIT to Deployment env if memory resource limits are defined * lint pass'
armadaproject · May 9, 2023 · 7ce9165 · 7ce9165
1 parent f4b862a
commit 7ce9165
Show file tree

Hide file tree

Showing 8 changed files with 69 additions and 1 deletion.
diff --git a/controllers/install/armadaserver_controller.go b/controllers/install/armadaserver_controller.go
@@ -568,6 +568,7 @@ func createArmadaServerDeployment(as *installv1alpha1.ArmadaServer) (*appsv1.Dep
 	}
 	if as.Spec.Resources != nil {
 		deployment.Spec.Template.Spec.Containers[0].Resources = *as.Spec.Resources
+		deployment.Spec.Template.Spec.Containers[0].Env = addGoMemLimit(deployment.Spec.Template.Spec.Containers[0].Env, *as.Spec.Resources)
 	}
 
 	return &deployment, nil

diff --git a/controllers/install/binoculars_controller.go b/controllers/install/binoculars_controller.go
@@ -323,6 +323,7 @@ func createBinocularsDeployment(binoculars *installv1alpha1.Binoculars) (*appsv1
 	}
 	if binoculars.Spec.Resources != nil {
 		deployment.Spec.Template.Spec.Containers[0].Resources = *binoculars.Spec.Resources
+		deployment.Spec.Template.Spec.Containers[0].Env = addGoMemLimit(deployment.Spec.Template.Spec.Containers[0].Env, *binoculars.Spec.Resources)
 	}
 	return &deployment, nil
 }

diff --git a/controllers/install/common_helpers.go b/controllers/install/common_helpers.go
@@ -444,3 +444,12 @@ func createPrometheusRule(name, namespace string, scrapeInterval *metav1.Duratio
 		},
 	}
 }
+
+func addGoMemLimit(env []corev1.EnvVar, resources corev1.ResourceRequirements) []corev1.EnvVar {
+	if resources.Limits.Memory() != nil && resources.Limits.Memory().Value() != 0 {
+		val := resources.Limits.Memory().Value()
+		goMemLimit := corev1.EnvVar{Name: "GOMEMLIMIT", Value: fmt.Sprintf("%dB", val)}
+		env = append(env, goMemLimit)
+	}
+	return env
+}
diff --git a/controllers/install/common_helpers_test.go b/controllers/install/common_helpers_test.go
@@ -5,6 +5,8 @@ import (
 	"testing"
 	"time"
 
+	"sigs.k8s.io/yaml"
+
 	"context"
 
 	"github.com/stretchr/testify/assert"
@@ -720,10 +722,61 @@ func makeCommonComponents() CommonComponents {
 	secret := corev1.Secret{
 		StringData: map[string]string{"secretkey": "secretval"},
 	}
-
 	return CommonComponents{
 		Deployment:      &deployment,
 		PriorityClasses: []*schedulingv1.PriorityClass{&pc},
 		Secret:          &secret,
 	}
 }
+
+func TestAddGoMemLimit(t *testing.T) {
+	type test struct {
+		name               string
+		resourcesYaml      string
+		expectedGoMemLimit string
+	}
+
+	tests := []test{
+		{
+			name: "1Gi memory limit",
+			resourcesYaml: `limits:
+    memory: 1Gi`,
+			expectedGoMemLimit: "1073741824B",
+		},
+		{
+			name: "500Mi memory limit",
+			resourcesYaml: `limits:
+    memory: 500Mi`,
+			expectedGoMemLimit: "524288000B",
+		},
+		{
+			name:               "no memory limit",
+			resourcesYaml:      ``,
+			expectedGoMemLimit: "",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			resources := corev1.ResourceRequirements{}
+			if err := yaml.Unmarshal([]byte(tc.resourcesYaml), &resources); err != nil {
+				t.Fatalf("error unmarshalling resources yaml: %v", err)
+			}
+
+			var env []corev1.EnvVar
+			env = addGoMemLimit(env, resources)
+
+			goMemLimitFound := false
+			for _, envVar := range env {
+				if envVar.Name == "GOMEMLIMIT" {
+					goMemLimitFound = true
+					assert.Equal(t, tc.expectedGoMemLimit, envVar.Value)
+				}
+			}
+
+			if !goMemLimitFound && tc.expectedGoMemLimit != "" {
+				t.Errorf("expected GOMEMLIMIT to be set, but it was not found")
+			}
+		})
+	}
+}
diff --git a/controllers/install/eventingester_controller.go b/controllers/install/eventingester_controller.go
@@ -233,6 +233,7 @@ func (r *EventIngesterReconciler) createDeployment(eventIngester *installv1alpha
 	}
 	if eventIngester.Spec.Resources != nil {
 		deployment.Spec.Template.Spec.Containers[0].Resources = *eventIngester.Spec.Resources
+		deployment.Spec.Template.Spec.Containers[0].Env = addGoMemLimit(deployment.Spec.Template.Spec.Containers[0].Env, *eventIngester.Spec.Resources)
 	}
 
 	return &deployment, nil

diff --git a/controllers/install/executor_controller.go b/controllers/install/executor_controller.go
@@ -359,6 +359,7 @@ func (r *ExecutorReconciler) createDeployment(executor *installv1alpha1.Executor
 	}
 	if executor.Spec.Resources != nil {
 		deployment.Spec.Template.Spec.Containers[0].Resources = *executor.Spec.Resources
+		deployment.Spec.Template.Spec.Containers[0].Env = addGoMemLimit(deployment.Spec.Template.Spec.Containers[0].Env, *executor.Spec.Resources)
 	}
 	return &deployment
 }

diff --git a/controllers/install/lookout_controller.go b/controllers/install/lookout_controller.go
@@ -365,6 +365,7 @@ func createLookoutDeployment(lookout *installv1alpha1.Lookout) (*appsv1.Deployme
 	}
 	if lookout.Spec.Resources != nil {
 		deployment.Spec.Template.Spec.Containers[0].Resources = *lookout.Spec.Resources
+		deployment.Spec.Template.Spec.Containers[0].Env = addGoMemLimit(deployment.Spec.Template.Spec.Containers[0].Env, *lookout.Spec.Resources)
 	}
 	return &deployment, nil
 }

diff --git a/controllers/install/lookoutingester_controller.go b/controllers/install/lookoutingester_controller.go
@@ -215,6 +215,7 @@ func (r *LookoutIngesterReconciler) createDeployment(lookoutIngester *installv1a
 	}
 	if lookoutIngester.Spec.Resources != nil {
 		deployment.Spec.Template.Spec.Containers[0].Resources = *lookoutIngester.Spec.Resources
+		deployment.Spec.Template.Spec.Containers[0].Env = addGoMemLimit(deployment.Spec.Template.Spec.Containers[0].Env, *lookoutIngester.Spec.Resources)
 	}
 	return &deployment, nil
 }