feat(orchestrator): support lifecycle prestop in prod env (erda-proje…

…ct#5753)

Signed-off-by: iutx <[email protected]>

apistructs/application.go

@@ -16,6 +16,7 @@ package apistructs
 
 import (
 	"errors"
+	"strings"
 	"time"
 )
 
@@ -333,6 +334,10 @@ const (
 	ProdWorkspace DiceWorkspace = "PROD"
 )
 
+const (
+	DiceWorkspaceEnvKey = "DICE_WORKSPACE"
+)
+
 var DiceWorkspaceSlice = []DiceWorkspace{DevWorkspace, TestWorkspace, StagingWorkspace, ProdWorkspace}
 
 // Deployable 是否可部署的合法环境
@@ -373,3 +378,7 @@ func (mode ApplicationMode) CheckAppMode() error {
 func (w DiceWorkspace) String() string {
 	return string(w)
 }
+
+func (w DiceWorkspace) Equal(t DiceWorkspace) bool {
+	return strings.ToUpper(w.String()) == strings.ToUpper(t.String())
+}

internal/tools/orchestrator/scheduler/executor/plugins/k8s/deployment.go

@@ -628,6 +628,9 @@ func (k *Kubernetes) newDeployment(service *apistructs.Service, serviceGroup *ap
 	// Configure health check
 	SetHealthCheck(&deployment.Spec.Template.Spec.Containers[0], service)
 
+	// Add default lifecycle
+	k.AddLifeCycle(service, &deployment.Spec.Template.Spec)
+
 	if err := k.AddContainersEnv(deployment.Spec.Template.Spec.Containers /*containers*/, service, serviceGroup); err != nil {
 		return nil, err
 	}

internal/tools/orchestrator/scheduler/executor/plugins/k8s/lifecycle.go

@@ -0,0 +1,53 @@
+// Copyright (c) 2021 Terminus, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package k8s
+
+import (
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/utils/pointer"
+
+	"github.com/erda-project/erda/apistructs"
+	"github.com/erda-project/erda/internal/tools/orchestrator/scheduler/executor/util"
+)
+
+const (
+	DefaultProdTerminationGracePeriodSeconds = 45
+)
+
+var (
+	DefaultProdLifecyclePreStopHandler = &corev1.Handler{
+		Exec: &corev1.ExecAction{
+			Command: []string{"sh", "-c", "sleep 10"},
+		},
+	}
+)
+
+func (k *Kubernetes) AddLifeCycle(service *apistructs.Service, podSpec *corev1.PodSpec) {
+	if podSpec == nil {
+		return
+	}
+
+	workspace, _ := util.GetDiceWorkspaceFromEnvs(service.Env)
+	if workspace.Equal(apistructs.ProdWorkspace) {
+		if len(podSpec.Containers) == 0 {
+			return
+		}
+
+		podSpec.TerminationGracePeriodSeconds = pointer.Int64(DefaultProdTerminationGracePeriodSeconds)
+		podSpec.Containers[0].Lifecycle = &corev1.Lifecycle{
+			PreStop: DefaultProdLifecyclePreStopHandler,
+		}
+	}
+}

internal/tools/orchestrator/scheduler/executor/plugins/k8s/lifecycle_test.go

@@ -0,0 +1,89 @@
+// Copyright (c) 2021 Terminus, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package k8s
+
+import (
+	"testing"
+
+	corev1 "k8s.io/api/core/v1"
+
+	"github.com/erda-project/erda/apistructs"
+)
+
+func TestAddLifecycle(t *testing.T) {
+	k := Kubernetes{}
+
+	type args struct {
+		service *apistructs.Service
+		podSpec *corev1.PodSpec
+	}
+
+	tests := []struct {
+		name string
+		args args
+		want bool
+	}{
+		{
+			name: "production environment",
+			args: args{
+				service: &apistructs.Service{
+					Env: map[string]string{
+						apistructs.DiceWorkspaceEnvKey: apistructs.ProdWorkspace.String(),
+					},
+				},
+				podSpec: &corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name:  "busybox",
+							Image: "busybox",
+						},
+					},
+				},
+			},
+			want: true,
+		},
+		{
+			name: "none production environment",
+			args: args{
+				service: &apistructs.Service{
+					Env: map[string]string{
+						apistructs.DiceWorkspaceEnvKey: apistructs.TestWorkspace.String(),
+					},
+				},
+				podSpec: &corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name:  "busybox",
+							Image: "busybox",
+						},
+					},
+				},
+			},
+			want: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			k.AddLifeCycle(tt.args.service, tt.args.podSpec)
+			got := tt.args.podSpec.Containers[0].Lifecycle != nil &&
+				tt.args.podSpec.TerminationGracePeriodSeconds != nil
+			if got != tt.want {
+				t.Fatalf("add lifecycle fail, got: lifecycle: %+v, termination grace period seconds: %v",
+					tt.args.podSpec.Containers[0].Lifecycle, tt.args.podSpec.TerminationGracePeriodSeconds)
+			}
+		})
+	}
+}

internal/tools/orchestrator/scheduler/executor/util/util.go

@@ -354,3 +354,11 @@ func ParseAnnotationFromEnv(k string) string {
 	}
 	return ""
 }
+
+func GetDiceWorkspaceFromEnvs(envs map[string]string) (apistructs.DiceWorkspace, error) {
+	if workspace, ok := envs[apistructs.DiceWorkspaceEnvKey]; ok {
+		return apistructs.DiceWorkspace(strings.ToUpper(workspace)), nil
+	} else {
+		return "", errors.New("not found workspace env")
+	}
+}