Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Safely access/mutate fargate coredns pod annotations #7480

Merged
merged 1 commit into from
Jan 17, 2024
Merged

Safely access/mutate fargate coredns pod annotations #7480

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 25 additions & 4 deletions pkg/fargate/coredns/coredns.go
View file
Open in desktop
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
}
55 changes: 47 additions & 8 deletions pkg/fargate/coredns/coredns_test.go
View file
Open in desktop
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).NotTo(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).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(
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,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"))
}
})
})

})
Expand All @@ -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",
Expand All @@ -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,
},
},
},
Expand All @@ -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",
Expand All @@ -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,
Expand All @@ -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,
})
}