fix(orchestrator): eci fluent-bit sidecar log collect error (#6464)

Signed-off-by: Ash <[email protected]>
erda-project · Nov 1, 2024 · 3e3030f · 3e3030f
1 parent 8dc57ee
commit 3e3030f
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 28 deletions.
diff --git a/internal/tools/orchestrator/scheduler/executor/plugins/k8s/lifecycle.go b/internal/tools/orchestrator/scheduler/executor/plugins/k8s/lifecycle.go
@@ -35,19 +35,23 @@ var (
 )
 
 func (k *Kubernetes) AddLifeCycle(service *apistructs.Service, podSpec *corev1.PodSpec) {
-	if podSpec == nil {
+	if podSpec == nil || len(podSpec.Containers) == 0 {
 		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{
+		setPreStopHandler(&podSpec.Containers[0])
+	}
+}
+
+func setPreStopHandler(container *corev1.Container) {
+	if container.Lifecycle == nil {
+		container.Lifecycle = &corev1.Lifecycle{
 			PreStop: DefaultProdLifecyclePreStopHandler,
 		}
+		return
 	}
+	container.Lifecycle.PreStop = DefaultProdLifecyclePreStopHandler
 }
diff --git a/internal/tools/orchestrator/scheduler/executor/plugins/k8s/lifecycle_test.go b/internal/tools/orchestrator/scheduler/executor/plugins/k8s/lifecycle_test.go
@@ -15,9 +15,11 @@
 package k8s
 
 import (
+	"reflect"
 	"testing"
 
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/utils/pointer"
 
 	"github.com/erda-project/erda/apistructs"
 )
@@ -29,60 +31,117 @@ func TestAddLifecycle(t *testing.T) {
 		service *apistructs.Service
 		podSpec *corev1.PodSpec
 	}
-
 	tests := []struct {
 		name string
 		args args
-		want bool
+		want *corev1.PodSpec
 	}{
 		{
-			name: "production environment",
+			name: "nil podSpec",
+			args: args{
+				service: &apistructs.Service{
+					Env: map[string]string{
+						apistructs.DiceWorkspaceEnvKey: apistructs.ProdWorkspace.String(),
+					},
+				},
+				podSpec: nil,
+			},
+			want: nil,
+		},
+		{
+			name: "empty containers",
 			args: args{
 				service: &apistructs.Service{
 					Env: map[string]string{
 						apistructs.DiceWorkspaceEnvKey: apistructs.ProdWorkspace.String(),
 					},
 				},
 				podSpec: &corev1.PodSpec{
-					Containers: []corev1.Container{
-						{
-							Name:  "busybox",
-							Image: "busybox",
-						},
+					Containers: []corev1.Container{},
+				},
+			},
+			want: &corev1.PodSpec{
+				Containers: []corev1.Container{},
+			},
+		},
+		{
+			name: "non-PROD workspace",
+			args: args{
+				service: &apistructs.Service{
+					Env: map[string]string{
+						apistructs.DiceWorkspaceEnvKey: "DEV",
+					},
+				},
+				podSpec: &corev1.PodSpec{
+					Containers: []corev1.Container{{}},
+				},
+			},
+			want: &corev1.PodSpec{
+				Containers: []corev1.Container{{}},
+			},
+		},
+		{
+			name: "PROD workspace with containers",
+			args: args{
+				service: &apistructs.Service{
+					Env: map[string]string{
+						apistructs.DiceWorkspaceEnvKey: apistructs.ProdWorkspace.String(),
 					},
 				},
+				podSpec: &corev1.PodSpec{
+					Containers: []corev1.Container{{}},
+				},
+			},
+			want: &corev1.PodSpec{
+				TerminationGracePeriodSeconds: pointer.Int64(DefaultProdTerminationGracePeriodSeconds),
+				Containers: []corev1.Container{{
+					Lifecycle: &corev1.Lifecycle{
+						PreStop: DefaultProdLifecyclePreStopHandler,
+					},
+				}},
 			},
-			want: true,
 		},
 		{
-			name: "none production environment",
+			name: "container with existing Lifecycle and PostStart",
 			args: args{
 				service: &apistructs.Service{
 					Env: map[string]string{
-						apistructs.DiceWorkspaceEnvKey: apistructs.TestWorkspace.String(),
+						apistructs.DiceWorkspaceEnvKey: apistructs.ProdWorkspace.String(),
 					},
 				},
 				podSpec: &corev1.PodSpec{
-					Containers: []corev1.Container{
-						{
-							Name:  "busybox",
-							Image: "busybox",
+					Containers: []corev1.Container{{
+						Lifecycle: &corev1.Lifecycle{
+							PostStart: &corev1.LifecycleHandler{
+								Exec: &corev1.ExecAction{
+									Command: []string{"echo", "postStart"},
+								},
+							},
 						},
-					},
+					}},
 				},
 			},
-			want: false,
+			want: &corev1.PodSpec{
+				TerminationGracePeriodSeconds: pointer.Int64(DefaultProdTerminationGracePeriodSeconds),
+				Containers: []corev1.Container{{
+					Lifecycle: &corev1.Lifecycle{
+						PostStart: &corev1.LifecycleHandler{
+							Exec: &corev1.ExecAction{
+								Command: []string{"echo", "postStart"},
+							},
+						},
+						PreStop: DefaultProdLifecyclePreStopHandler,
+					},
+				}},
+			},
 		},
 	}
 
 	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)
+			if !reflect.DeepEqual(tt.args.podSpec, tt.want) {
+				t.Errorf("AddLifeCycle() = %v, want %v", tt.args.podSpec, tt.want)
 			}
 		})
 	}