From 087d5d783d0f65541bfe6ae71621aba3806d29fe Mon Sep 17 00:00:00 2001 From: Amine Hilaly Date: Mon, 15 Jan 2024 15:04:28 -0600 Subject: [PATCH] Safely access/mutate fargate coredns pod annotations Prior to this patch, the `pkg/fargate/coredns` package had some bits of code that accessed/mutated pod annotations assuming that they'll always be instantiated correctly. This patch adds utility functions to safely mutate and access fargate pod annotations. Signed-off-by: Amine Hilaly --- pkg/fargate/coredns/coredns.go | 29 ++++++++++++--- pkg/fargate/coredns/coredns_test.go | 55 ++++++++++++++++++++++++----- 2 files changed, 72 insertions(+), 12 deletions(-) diff --git a/pkg/fargate/coredns/coredns.go b/pkg/fargate/coredns/coredns.go index b451161102..8c4b89fa4d 100644 --- a/pkg/fargate/coredns/coredns.go +++ b/pkg/fargate/coredns/coredns.go @@ -67,7 +67,7 @@ func isDeploymentScheduledOnFargate(clientSet kubeclient.Interface) (bool, error if coredns.Spec.Replicas == nil { return false, errors.New("nil spec.replicas in coredns deployment") } - computeType, exists := coredns.Spec.Template.Annotations[ComputeTypeAnnotationKey] + computeType, exists := safeGetAnnotationValue(coredns.Spec.Template.Annotations, ComputeTypeAnnotationKey) logger.Debug("deployment %q with compute type %q currently has %v/%v replicas running", Name, computeType, coredns.Status.ReadyReplicas, *coredns.Spec.Replicas) scheduled := exists && computeType == computeTypeFargate && @@ -95,7 +95,7 @@ func arePodsScheduledOnFargate(clientSet kubeclient.Interface) (bool, error) { } func isRunningOnFargate(pod *v1.Pod) bool { - computeType, exists := pod.Annotations[ComputeTypeAnnotationKey] + computeType, exists := safeGetAnnotationValue(pod.Annotations, ComputeTypeAnnotationKey) logger.Debug("pod %q with compute type %q and status %q is scheduled on %q", pod.Name, computeType, pod.Status.Phase, pod.Spec.NodeName) return exists && computeType == computeTypeFargate && @@ -119,7 +119,7 @@ func scheduleOnFargate(clientSet kubeclient.Interface) error { if err != nil { return err } - coredns.Spec.Template.Annotations[ComputeTypeAnnotationKey] = computeTypeFargate + coredns.Spec.Template.Annotations = safeSetAnnotation(coredns.Spec.Template.Annotations, ComputeTypeAnnotationKey, computeTypeFargate) bytes, err := runtime.Encode(unstructured.UnstructuredJSONScheme, coredns) if err != nil { return errors.Wrapf(err, "failed to marshal %q deployment", Name) @@ -128,7 +128,8 @@ func scheduleOnFargate(clientSet kubeclient.Interface) error { if err != nil { return errors.Wrap(err, "failed to patch deployment") } - value, exists := patched.Spec.Template.Annotations[ComputeTypeAnnotationKey] + + value, exists := safeGetAnnotationValue(patched.Spec.Template.Annotations, ComputeTypeAnnotationKey) if !exists { return fmt.Errorf("could not find annotation %q on patched deployment %q: patching must have failed", ComputeTypeAnnotationKey, Name) } @@ -156,3 +157,23 @@ func WaitForScheduleOnFargate(clientSet kubeclient.Interface, retryPolicy retry. } return fmt.Errorf("timed out while waiting for %q to be scheduled on Fargate", Name) } + +// safeGetAnnotationValue safely gets the value of an annotation from a map. It +// returns the value and a boolean indicating whether the key was found. +func safeGetAnnotationValue(annotations map[string]string, key string) (string, bool) { + if annotations == nil { + return "", false + } + value, exist := annotations[key] + return value, exist +} + +// safeSetAnnotation safely sets the value of an annotation in a map. It will +// initialize the annotaions map if it is nil. +func safeSetAnnotation(annotations map[string]string, key, value string) map[string]string { + if annotations == nil { + annotations = make(map[string]string) + } + annotations[key] = value + return annotations +} diff --git a/pkg/fargate/coredns/coredns_test.go b/pkg/fargate/coredns/coredns_test.go index 4b5c7e2d2e..f2a93df1a0 100644 --- a/pkg/fargate/coredns/coredns_test.go +++ b/pkg/fargate/coredns/coredns_test.go @@ -122,6 +122,7 @@ var _ = Describe("coredns", func() { mockClientset := mockClientsetWith(deployment("ec2", 0, 2)) deployment, err := mockClientset.AppsV1().Deployments(coredns.Namespace).Get(context.Background(), coredns.Name, metav1.GetOptions{}) Expect(err).To(Not(HaveOccurred())) + Expect(deployment.Spec.Template.Annotations).NotTo(BeNil()) Expect(deployment.Spec.Template.Annotations).To(HaveKeyWithValue(coredns.ComputeTypeAnnotationKey, "ec2")) // When: err = coredns.ScheduleOnFargate(mockClientset) @@ -129,11 +130,23 @@ var _ = Describe("coredns", func() { // Then: deployment, err = mockClientset.AppsV1().Deployments(coredns.Namespace).Get(context.Background(), coredns.Name, metav1.GetOptions{}) Expect(err).To(Not(HaveOccurred())) + Expect(deployment.Spec.Template.Annotations).NotTo(BeNil()) Expect(deployment.Spec.Template.Annotations).To(HaveKeyWithValue(coredns.ComputeTypeAnnotationKey, "fargate")) }) }) Describe("WaitForScheduleOnFargate", func() { + It("should error if the annotations are not set", func() { + // Given: + mockClientset := mockClientsetWith( + deployment("fargate", 2, 2), podWithAnnotations("fargate", v1.PodRunning, nil), pod("fargate", v1.PodRunning), + ) + // When: + err := coredns.WaitForScheduleOnFargate(mockClientset, retryPolicy) + // Then: + Expect(err).To(HaveOccurred()) + }) + It("should wait for coredns to be scheduled on Fargate and return w/o any error", func() { // Given: mockClientset := mockClientsetWith( @@ -149,6 +162,8 @@ var _ = Describe("coredns", func() { failureCases := [][]runtime.Object{ {deployment("ec2", 2, 2), pod("ec2", v1.PodRunning), pod("ec2", v1.PodRunning)}, {deployment("ec2", 0, 2), pod("ec2", v1.PodPending), pod("ec2", v1.PodPending)}, + {deploymentWithAnnotations("ec2", 0, 2, nil), podWithAnnotations("ec2", v1.PodFailed, nil), podWithAnnotations("ec2", v1.PodFailed, nil)}, + {deploymentWithAnnotations("ec2", 0, 2, nil), podWithAnnotations("ec2", v1.PodFailed, map[string]string{}), podWithAnnotations("ec2", v1.PodFailed, map[string]string{})}, {deployment("fargate", 0, 2), pod("fargate", v1.PodPending), pod("fargate", v1.PodPending)}, {deployment("fargate", 0, 2), pod("fargate", v1.PodFailed), pod("fargate", v1.PodFailed)}, {deployment("fargate", 0, 2), pod("fargate", v1.PodPending), pod("fargate", v1.PodFailed)}, @@ -165,6 +180,22 @@ var _ = Describe("coredns", func() { Expect(err.Error()).To(Equal("timed out while waiting for \"coredns\" to be scheduled on Fargate")) } }) + + It("Should timeout if coredns pods do not have the correct annotations", func() { + failureCases := [][]runtime.Object{ + {deploymentWithAnnotations("fargate", 0, 2, nil), podWithAnnotations("fargate", v1.PodPending, nil), podWithAnnotations("fargate", v1.PodRunning, nil)}, + {deploymentWithAnnotations("ec2", 0, 2, nil), podWithAnnotations("ec2", v1.PodFailed, nil), podWithAnnotations("ec2", v1.PodRunning, nil)}, + {deploymentWithAnnotations("ec2", 0, 2, map[string]string{}), podWithAnnotations("ec2", v1.PodRunning, map[string]string{}), podWithAnnotations("ec2", v1.PodRunning, map[string]string{})}, + } + for _, failureCase := range failureCases { + // Given: + mockClientset := mockClientsetWith(failureCase...) + // When: + err := coredns.WaitForScheduleOnFargate(mockClientset, retryPolicy) + // Then: + Expect(err).To(MatchError("timed out while waiting for \"coredns\" to be scheduled on Fargate")) + } + }) }) }) @@ -173,7 +204,7 @@ func mockClientsetWith(objects ...runtime.Object) kubeclient.Interface { return fake.NewSimpleClientset(objects...) } -func deployment(computeType string, numReady, numReplicas int32) *appsv1.Deployment { +func deploymentWithAnnotations(computeType string, numReady, numReplicas int32, annotations map[string]string) *appsv1.Deployment { return &appsv1.Deployment{ TypeMeta: metav1.TypeMeta{ Kind: "Deployment", @@ -187,9 +218,7 @@ func deployment(computeType string, numReady, numReplicas int32) *appsv1.Deploym Replicas: &numReplicas, Template: v1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ - Annotations: map[string]string{ - coredns.ComputeTypeAnnotationKey: computeType, - }, + Annotations: annotations, }, }, }, @@ -199,9 +228,15 @@ func deployment(computeType string, numReady, numReplicas int32) *appsv1.Deploym } } +func deployment(computeType string, numReady, numReplicas int32) *appsv1.Deployment { + return deploymentWithAnnotations(computeType, numReady, numReplicas, map[string]string{ + coredns.ComputeTypeAnnotationKey: computeType, + }) +} + const chars = "abcdef0123456789" -func pod(computeType string, phase v1.PodPhase) *v1.Pod { +func podWithAnnotations(computeType string, phase v1.PodPhase, annotatations map[string]string) *v1.Pod { pod := &v1.Pod{ TypeMeta: metav1.TypeMeta{ Kind: "Pod", @@ -213,9 +248,7 @@ func pod(computeType string, phase v1.PodPhase) *v1.Pod { Labels: map[string]string{ "eks.amazonaws.com/component": coredns.Name, }, - Annotations: map[string]string{ - coredns.ComputeTypeAnnotationKey: computeType, - }, + Annotations: annotatations, }, Status: v1.PodStatus{ Phase: phase, @@ -230,3 +263,9 @@ func pod(computeType string, phase v1.PodPhase) *v1.Pod { } return pod } + +func pod(computeType string, phase v1.PodPhase) *v1.Pod { + return podWithAnnotations(computeType, phase, map[string]string{ + coredns.ComputeTypeAnnotationKey: computeType, + }) +}