Skip to content

Commit 34cd968

Browse files
committed
fix: resource revert paniced when no limits were defined at all
1 parent bee8920 commit 34cd968

File tree

3 files changed

+107
-62
lines changed

3 files changed

+107
-62
lines changed

internal/boost/pod/pod.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,16 +72,23 @@ func RevertResourceBoost(pod *corev1.Pod) error {
7272
}
7373
delete(pod.Labels, BoostLabelKey)
7474
delete(pod.Annotations, BoostAnnotationKey)
75-
for _, container := range pod.Spec.Containers {
75+
for i := range pod.Spec.Containers {
76+
container := &pod.Spec.Containers[i]
7677
if request, ok := annotation.InitCPURequests[container.Name]; ok {
7778
if reqQuantity, err := apiResource.ParseQuantity(request); err == nil {
79+
if container.Resources.Requests == nil {
80+
container.Resources.Requests = corev1.ResourceList{}
81+
}
7882
container.Resources.Requests[corev1.ResourceCPU] = reqQuantity
7983
} else {
8084
return fmt.Errorf("failed to parse CPU request: %s", err)
8185
}
8286
}
8387
if limit, ok := annotation.InitCPULimits[container.Name]; ok {
8488
if limitQuantity, err := apiResource.ParseQuantity(limit); err == nil {
89+
if container.Resources.Limits == nil {
90+
container.Resources.Limits = corev1.ResourceList{}
91+
}
8592
container.Resources.Limits[corev1.ResourceCPU] = limitQuantity
8693
} else {
8794
return fmt.Errorf("failed to parse CPU limit: %s", err)

internal/boost/pod/pod_suite_test.go

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,64 @@ package pod_test
1717
import (
1818
"testing"
1919

20+
bpod "github.com/google/kube-startup-cpu-boost/internal/boost/pod"
2021
. "github.com/onsi/ginkgo/v2"
2122
. "github.com/onsi/gomega"
23+
corev1 "k8s.io/api/core/v1"
24+
apiResource "k8s.io/apimachinery/pkg/api/resource"
25+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2226
)
2327

2428
func TestPod(t *testing.T) {
2529
RegisterFailHandler(Fail)
2630
RunSpecs(t, "Pod Suite")
2731
}
32+
33+
var (
34+
podTemplate *corev1.Pod
35+
containerOneName string
36+
containerTwoName string
37+
)
38+
39+
var _ = BeforeSuite(func() {
40+
containerOneName = "container-one"
41+
containerTwoName = "container-two"
42+
reqQuantity, err := apiResource.ParseQuantity("1")
43+
Expect(err).ShouldNot(HaveOccurred())
44+
limitQuantity, err := apiResource.ParseQuantity("2")
45+
Expect(err).ShouldNot(HaveOccurred())
46+
podTemplate = &corev1.Pod{
47+
ObjectMeta: metav1.ObjectMeta{
48+
Name: "test-pod",
49+
Labels: map[string]string{
50+
bpod.BoostLabelKey: "boost-001",
51+
},
52+
},
53+
Spec: corev1.PodSpec{
54+
Containers: []corev1.Container{
55+
{
56+
Name: containerOneName,
57+
Resources: corev1.ResourceRequirements{
58+
Requests: corev1.ResourceList{
59+
corev1.ResourceCPU: reqQuantity,
60+
},
61+
Limits: corev1.ResourceList{
62+
corev1.ResourceCPU: limitQuantity,
63+
},
64+
},
65+
},
66+
{
67+
Name: containerTwoName,
68+
Resources: corev1.ResourceRequirements{
69+
Requests: corev1.ResourceList{
70+
corev1.ResourceCPU: reqQuantity,
71+
},
72+
Limits: corev1.ResourceList{
73+
corev1.ResourceCPU: limitQuantity,
74+
},
75+
},
76+
},
77+
},
78+
},
79+
}
80+
})

internal/boost/pod/pod_test.go

Lines changed: 46 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -21,88 +21,46 @@ import (
2121
. "github.com/onsi/ginkgo/v2"
2222
. "github.com/onsi/gomega"
2323
corev1 "k8s.io/api/core/v1"
24-
apiResource "k8s.io/apimachinery/pkg/api/resource"
25-
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2624
)
2725

2826
var _ = Describe("Pod", func() {
29-
var pod *corev1.Pod
30-
var annot *bpod.BoostPodAnnotation
31-
var containerOne, containerTwo string
32-
var reqQuantity, limitQuantity apiResource.Quantity
33-
var err error
27+
var (
28+
annot *bpod.BoostPodAnnotation
29+
pod *corev1.Pod
30+
err error
31+
)
3432

3533
BeforeEach(func() {
36-
containerOne = "container-one"
37-
containerTwo = "container-one"
38-
reqQuantity, err = apiResource.ParseQuantity("1")
39-
Expect(err).ShouldNot(HaveOccurred())
40-
limitQuantity, err = apiResource.ParseQuantity("2")
41-
Expect(err).ShouldNot(HaveOccurred())
4234
annot = &bpod.BoostPodAnnotation{
4335
BoostTimestamp: time.Now(),
4436
InitCPURequests: map[string]string{
45-
containerOne: "500m",
46-
containerTwo: "500m",
37+
containerOneName: "500m",
38+
containerTwoName: "500m",
4739
},
4840
InitCPULimits: map[string]string{
49-
containerOne: "1",
50-
containerTwo: "1",
41+
containerOneName: "1",
42+
containerTwoName: "1",
5143
},
5244
}
53-
pod = &corev1.Pod{
54-
ObjectMeta: metav1.ObjectMeta{
55-
Name: "test-pod",
56-
Labels: map[string]string{
57-
bpod.BoostLabelKey: "boost-001",
58-
},
59-
Annotations: map[string]string{
60-
bpod.BoostAnnotationKey: annot.ToJSON(),
61-
},
62-
},
63-
Spec: corev1.PodSpec{
64-
Containers: []corev1.Container{
65-
{
66-
Name: containerOne,
67-
Resources: corev1.ResourceRequirements{
68-
Requests: corev1.ResourceList{
69-
corev1.ResourceCPU: reqQuantity,
70-
},
71-
Limits: corev1.ResourceList{
72-
corev1.ResourceCPU: limitQuantity,
73-
},
74-
},
75-
},
76-
{
77-
Name: containerTwo,
78-
Resources: corev1.ResourceRequirements{
79-
Requests: corev1.ResourceList{
80-
corev1.ResourceCPU: reqQuantity,
81-
},
82-
Limits: corev1.ResourceList{
83-
corev1.ResourceCPU: limitQuantity,
84-
},
85-
},
86-
},
87-
},
88-
},
45+
pod = podTemplate.DeepCopy()
46+
pod.Annotations = map[string]string{
47+
bpod.BoostAnnotationKey: annot.ToJSON(),
8948
}
9049
})
9150

9251
Describe("Reverts the POD container resources to original values", func() {
52+
JustBeforeEach(func() {
53+
err = bpod.RevertResourceBoost(pod)
54+
})
9355
When("POD is missing startup-cpu-boost annotation", func() {
9456
BeforeEach(func() {
9557
delete(pod.ObjectMeta.Annotations, bpod.BoostAnnotationKey)
96-
err = bpod.RevertResourceBoost(pod)
9758
})
9859
It("errors", func() {
9960
Expect(err).Should(HaveOccurred())
10061
})
10162
})
10263
When("POD has valid startup-cpu-boost annotation", func() {
103-
BeforeEach(func() {
104-
err = bpod.RevertResourceBoost(pod)
105-
})
10664
It("does not error", func() {
10765
Expect(err).ShouldNot(HaveOccurred())
10866
})
@@ -115,14 +73,41 @@ var _ = Describe("Pod", func() {
11573
It("reverts CPU requests to initial values", func() {
11674
cpuReqOne := pod.Spec.Containers[0].Resources.Requests[corev1.ResourceCPU]
11775
cpuReqTwo := pod.Spec.Containers[1].Resources.Requests[corev1.ResourceCPU]
118-
Expect(cpuReqOne.String()).Should(Equal(annot.InitCPURequests[containerOne]))
119-
Expect(cpuReqTwo.String()).Should(Equal(annot.InitCPURequests[containerTwo]))
76+
Expect(cpuReqOne.String()).Should(Equal(annot.InitCPURequests[containerOneName]))
77+
Expect(cpuReqTwo.String()).Should(Equal(annot.InitCPURequests[containerTwoName]))
12078
})
12179
It("reverts CPU limits to initial values", func() {
12280
cpuReqOne := pod.Spec.Containers[0].Resources.Limits[corev1.ResourceCPU]
12381
cpuReqTwo := pod.Spec.Containers[1].Resources.Limits[corev1.ResourceCPU]
124-
Expect(cpuReqOne.String()).Should(Equal(annot.InitCPULimits[containerOne]))
125-
Expect(cpuReqTwo.String()).Should(Equal(annot.InitCPULimits[containerTwo]))
82+
Expect(cpuReqOne.String()).Should(Equal(annot.InitCPULimits[containerOneName]))
83+
Expect(cpuReqTwo.String()).Should(Equal(annot.InitCPULimits[containerTwoName]))
84+
})
85+
When("Limits were removed during boost", func() {
86+
BeforeEach(func() {
87+
pod.Spec.Containers[0].Resources.Limits = nil
88+
pod.Spec.Containers[1].Resources.Limits = nil
89+
})
90+
It("does not error", func() {
91+
Expect(err).ShouldNot(HaveOccurred())
92+
})
93+
It("removes startup-cpu-boost label", func() {
94+
Expect(pod.Labels).NotTo(HaveKey(bpod.BoostLabelKey))
95+
})
96+
It("removes startup-cpu-boost annotation", func() {
97+
Expect(pod.Annotations).NotTo(HaveKey(bpod.BoostAnnotationKey))
98+
})
99+
It("reverts CPU requests to initial values", func() {
100+
cpuReqOne := pod.Spec.Containers[0].Resources.Requests[corev1.ResourceCPU]
101+
cpuReqTwo := pod.Spec.Containers[1].Resources.Requests[corev1.ResourceCPU]
102+
Expect(cpuReqOne.String()).Should(Equal(annot.InitCPURequests[containerOneName]))
103+
Expect(cpuReqTwo.String()).Should(Equal(annot.InitCPURequests[containerTwoName]))
104+
})
105+
It("reverts CPU limits to initial values", func() {
106+
cpuReqOne := pod.Spec.Containers[0].Resources.Limits[corev1.ResourceCPU]
107+
cpuReqTwo := pod.Spec.Containers[1].Resources.Limits[corev1.ResourceCPU]
108+
Expect(cpuReqOne.String()).Should(Equal(annot.InitCPULimits[containerOneName]))
109+
Expect(cpuReqTwo.String()).Should(Equal(annot.InitCPULimits[containerTwoName]))
110+
})
126111
})
127112
})
128113
})

0 commit comments

Comments
 (0)