Merge pull request #547 from whitebear009/hpa-controlled-by

make function GetPodTemplate and IsHPAControlledByEHPA more precise

pkg/controller/ehpa/hpa_event_handler.go

@@ -1,15 +1,14 @@
 package ehpa
 
 import (
-	"strings"
-
 	autoscalingv2 "k8s.io/api/autoscaling/v2beta2"
 	"k8s.io/client-go/util/workqueue"
 	"k8s.io/klog/v2"
 	"sigs.k8s.io/controller-runtime/pkg/event"
 	"sigs.k8s.io/controller-runtime/pkg/handler"
 
 	"github.com/gocrane/crane/pkg/metrics"
+	"github.com/gocrane/crane/pkg/utils"
 )
 
 type hpaEventHandler struct {
@@ -38,7 +37,7 @@ func (h *hpaEventHandler) Update(evt event.UpdateEvent, q workqueue.RateLimiting
 		for _, cond := range newHpa.Status.Conditions {
 			if cond.Reason == "SucceededRescale" || cond.Reason == "SucceededOverloadRescale" {
 				scaleType := "hpa"
-				if strings.HasPrefix("ehpa-", newHpa.Name) {
+				if utils.IsHPAControlledByEHPA(newHpa) {
 					scaleType = "ehpa"
 				}
 

pkg/utils/hpa.go

@@ -6,10 +6,10 @@ import (
 
 	autoscalingv2 "k8s.io/api/autoscaling/v2beta2"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	autoscalingapi "github.com/gocrane/api/autoscaling/v1alpha1"
-
 	"github.com/gocrane/crane/pkg/known"
 )
 
@@ -59,3 +59,16 @@ func GetEHPAFromScaleTarget(context context.Context, kubeClient client.Client, n
 
 	return nil, nil
 }
+
+func IsHPAControlledByEHPA(hpa *autoscalingv2.HorizontalPodAutoscaler) bool {
+	for _, ownerReference := range hpa.OwnerReferences {
+		gv, err := schema.ParseGroupVersion(ownerReference.APIVersion)
+		if err != nil {
+			return false
+		}
+		if gv.Group == autoscalingapi.GroupName && ownerReference.Kind == "EffectiveHorizontalPodAutoscaler" {
+			return true
+		}
+	}
+	return false
+}

pkg/utils/pod_template.go

@@ -4,11 +4,11 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"strings"
 
 	appsv1 "k8s.io/api/apps/v1"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/klog/v2"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
@@ -20,17 +20,20 @@ func GetPodTemplate(context context.Context, namespace string, name string, kind
 		Name:      name,
 		Namespace: namespace,
 	}
-
-	if kind == "Deployment" && strings.HasPrefix(apiVersion, "apps") {
+	gv, err := schema.ParseGroupVersion(apiVersion)
+	if err != nil {
+		return nil, err
+	}
+	if kind == "Deployment" && gv.Group == appsv1.GroupName {
 		deployment := &appsv1.Deployment{}
-		err := kubeClient.Get(context, key, deployment)
+		err = kubeClient.Get(context, key, deployment)
 		if err != nil {
 			return nil, err
 		}
 		templateSpec = &deployment.Spec.Template
-	} else if kind == "StatefulSet" && strings.HasPrefix(apiVersion, "apps") {
+	} else if kind == "StatefulSet" && gv.Group == appsv1.GroupName {
 		statefulSet := &appsv1.StatefulSet{}
-		err := kubeClient.Get(context, key, statefulSet)
+		err = kubeClient.Get(context, key, statefulSet)
 		if err != nil {
 			return nil, err
 		}
@@ -39,7 +42,7 @@ func GetPodTemplate(context context.Context, namespace string, name string, kind
 		unstructed := &unstructured.Unstructured{}
 		unstructed.SetAPIVersion(apiVersion)
 		unstructed.SetKind(kind)
-		err := kubeClient.Get(context, key, unstructed)
+		err = kubeClient.Get(context, key, unstructed)
 		if err != nil {
 			return nil, err
 		}

pkg/utils/pod_template_test.go

@@ -0,0 +1,99 @@
+package utils
+
+import (
+	"context"
+	"testing"
+
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	fakeClient "sigs.k8s.io/controller-runtime/pkg/client/fake"
+)
+
+func TestGetPodTemplate(t *testing.T) {
+	ctx := context.TODO()
+	testNamespace := "test-namespace"
+	testName := "test-name"
+	testImage := "test-image"
+	testCases := []struct {
+		object     client.Object
+		kind       string
+		apiVersion string
+	}{
+		{
+			object: &appsv1.Deployment{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: testNamespace,
+					Name:      testName,
+				},
+				Spec: appsv1.DeploymentSpec{
+					Template: corev1.PodTemplateSpec{
+						Spec: corev1.PodSpec{
+							Containers: []corev1.Container{
+								{
+									Image: testImage,
+								},
+							},
+						},
+					},
+				},
+			},
+			kind:       "Deployment",
+			apiVersion: "apps/v1",
+		},
+		{
+			object: &appsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: testNamespace,
+					Name:      testName,
+				},
+				Spec: appsv1.StatefulSetSpec{
+					Template: corev1.PodTemplateSpec{
+						Spec: corev1.PodSpec{
+							Containers: []corev1.Container{
+								{
+									Image: testImage,
+								},
+							},
+						},
+					},
+				},
+			},
+			kind:       "StatefulSet",
+			apiVersion: "apps/v1",
+		},
+		{
+			object: &appsv1.DaemonSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: testNamespace,
+					Name:      testName,
+				},
+				Spec: appsv1.DaemonSetSpec{
+					Template: corev1.PodTemplateSpec{
+						Spec: corev1.PodSpec{
+							Containers: []corev1.Container{
+								{
+									Image: testImage,
+								},
+							},
+						},
+					},
+				},
+			},
+			kind:       "DaemonSet",
+			apiVersion: "apps/v1",
+		},
+	}
+
+	for _, tc := range testCases {
+		fakeClient := fakeClient.NewClientBuilder().WithObjects(tc.object).Build()
+		podTemplate, err := GetPodTemplate(ctx, testNamespace, testName, tc.kind, tc.apiVersion, fakeClient)
+		if err != nil {
+			t.Errorf("get pod template error: %v", err)
+		}
+		if len(podTemplate.Spec.Containers) == 1 && podTemplate.Spec.Containers[0].Image != testImage {
+			t.Errorf("the container image of pod template is inconsistent: expect %s, actual is %s", testImage, podTemplate.Spec.Containers[0].Image)
+		}
+	}
+}