Skip to content

Commit

Permalink
Safely access/mutate fargate coredns pod annotations
Browse files Browse the repository at this point in the history 
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 <hilalyamine@gmail.com>
  • Loading branch information
@a-hilaly
a-hilaly committed Jan 15, 2024
1 parent b9e936c commit 92716e1
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 12 deletions.
29 changes: 25 additions & 4 deletions pkg/fargate/coredns/coredns.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 &&
Expand Down Expand Up @@ -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 &&
Expand All @@ -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)
Expand All @@ -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)
}
Expand Down Expand Up @@ -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
}
56 changes: 48 additions & 8 deletions pkg/fargate/coredns/coredns_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,18 +122,31 @@ 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).To(Not(BeNil()))
Expect(deployment.Spec.Template.Annotations).To(HaveKeyWithValue(coredns.ComputeTypeAnnotationKey, "ec2"))
// When:
err = coredns.ScheduleOnFargate(mockClientset)
Expect(err).To(Not(HaveOccurred()))
// 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).To(Not(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(
Expand All @@ -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)},
Expand All @@ -165,6 +180,23 @@ 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(HaveOccurred())
Expect(err.Error()).To(Equal("timed out while waiting for \"coredns\" to be scheduled on Fargate"))
}
})
})

})
Expand All @@ -173,7 +205,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",
Expand All @@ -187,9 +219,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,
},
},
},
Expand All @@ -199,9 +229,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",
Expand All @@ -213,9 +249,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,
Expand All @@ -230,3 +264,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,
})
}

0 comments on commit 92716e1

Please sign in to comment.