From 65dcdaaf8da039353f71a293d31936028a4a0a7d Mon Sep 17 00:00:00 2001
From: Muzry <muzrry@gmail.com>
Date: Fri, 29 Oct 2021 17:40:10 +0800
Subject: [PATCH] fix: update the container env when scale the service group
 (#2672)

---
 .../executor/plugins/k8s/deployment.go        | 31 +++++++++++
 .../executor/plugins/k8s/deployment_test.go   | 55 +++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/modules/scheduler/executor/plugins/k8s/deployment.go b/modules/scheduler/executor/plugins/k8s/deployment.go
index e458995039d..a88d80d3782 100644
--- a/modules/scheduler/executor/plugins/k8s/deployment.go
+++ b/modules/scheduler/executor/plugins/k8s/deployment.go
@@ -905,6 +905,8 @@ func (k *Kubernetes) scaleDeployment(ctx context.Context, sg *apistructs.Service
 		return setContainerErr
 	}
 
+	k.UpdateContainerResourceEnv(scalingService.Resources, &container)
+
 	deploy.Spec.Template.Spec.Containers[0] = container
 
 	newCPU, newMem := getRequestsResources(deploy.Spec.Template.Spec.Containers)
@@ -977,3 +979,32 @@ func (k *Kubernetes) setContainerResources(service apistructs.Service, container
 
 	return nil
 }
+func (k *Kubernetes) UpdateContainerResourceEnv(originResource apistructs.Resources, container *apiv1.Container) {
+	for index, env := range container.Env {
+		var needToUpdate = false
+		switch env.Name {
+		case "DICE_CPU_ORIGIN":
+			needToUpdate = true
+			env.Value = fmt.Sprintf("%f", originResource.Cpu)
+		case "DICE_CPU_REQUEST":
+			needToUpdate = true
+			env.Value = container.Resources.Requests.Cpu().AsDec().String()
+		case "DICE_CPU_LIMIT":
+			needToUpdate = true
+			env.Value = container.Resources.Limits.Cpu().AsDec().String()
+		case "DICE_MEM_ORIGIN":
+			needToUpdate = true
+			env.Value = fmt.Sprintf("%f", originResource.Mem)
+		case "DICE_MEM_REQUEST":
+			needToUpdate = true
+			env.Value = fmt.Sprintf("%d", container.Resources.Requests.Memory().Value()/1024/1024)
+		case "DICE_MEM_LIMIT":
+			needToUpdate = true
+			env.Value = fmt.Sprintf("%d", container.Resources.Limits.Memory().Value()/1024/1024)
+		}
+		if needToUpdate {
+			container.Env[index] = env
+		}
+	}
+	return
+}
diff --git a/modules/scheduler/executor/plugins/k8s/deployment_test.go b/modules/scheduler/executor/plugins/k8s/deployment_test.go
index b6324fe9ae5..802a7858f42 100644
--- a/modules/scheduler/executor/plugins/k8s/deployment_test.go
+++ b/modules/scheduler/executor/plugins/k8s/deployment_test.go
@@ -21,6 +21,7 @@ import (
 	"bou.ke/monkey"
 	"github.com/stretchr/testify/assert"
 	apiv1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 
 	"github.com/erda-project/erda/apistructs"
 	"github.com/erda-project/erda/modules/scheduler/executor/plugins/k8s/secret"
@@ -96,3 +97,57 @@ func TestNewDeployment(t *testing.T) {
 	_, err := k.newDeployment(service, servicegroup)
 	assert.Equal(t, err, nil)
 }
+
+func TestUpdateContainerResourceEnv(t *testing.T) {
+	k := Kubernetes{}
+
+	container := apiv1.Container{
+		Name:       "",
+		Image:      "",
+		Command:    nil,
+		Args:       nil,
+		WorkingDir: "",
+		Ports:      nil,
+		EnvFrom:    nil,
+		Env: []apiv1.EnvVar{
+			{
+				Name:  "DICE_CPU_ORIGIN",
+				Value: "0.100000",
+			},
+			{
+				Name:  "DICE_CPU_REQUEST",
+				Value: "0.010",
+			},
+			{
+				Name:  "DICE_CPU_LIMIT",
+				Value: "0.100",
+			},
+			{
+				Name:  "DICE_MEM_ORIGIN",
+				Value: "1024.000000",
+			},
+			{
+				Name:  "DICE_MEM_REQUEST",
+				Value: "1024",
+			},
+			{
+				Name:  "DICE_MEM_LIMIT",
+				Value: "1024",
+			},
+		},
+		Resources: apiv1.ResourceRequirements{
+			Limits: apiv1.ResourceList{
+				"cpu":    resource.MustParse("100m"),
+				"memory": resource.MustParse("1024Mi"),
+			},
+			Requests: apiv1.ResourceList{
+				"cpu":    resource.MustParse("10m"),
+				"memory": resource.MustParse("1024Mi"),
+			},
+		},
+	}
+	originResource := apistructs.Resources{Cpu: 0.1, Mem: 1024}
+	k.UpdateContainerResourceEnv(originResource, &container)
+	assert.Equal(t, container.Env[0].Value, "0.100000")
+	assert.Equal(t, container.Env[5].Value, "1024")
+}