Add handling of context on more places

Signed-off-by: Michel Hollands <michel.hollands@gmail.com>
grafana · Aug 3, 2023 · 5e60a3c · 5e60a3c
1 parent cc09f23
commit 5e60a3c
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/pkg/admission/prep_downscale.go b/pkg/admission/prep_downscale.go
@@ -277,6 +277,11 @@ func getResourceAnnotations(ctx context.Context, ar v1.AdmissionReview, api kube
 }
 
 func addDownscaledAnnotationToStatefulSet(ctx context.Context, api kubernetes.Interface, namespace, stsName string) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
 	client := api.AppsV1().StatefulSets(namespace)
 	patch := fmt.Sprintf(`{"metadata":{"annotations":{"%v":"%v"}}}`, config.LastDownscaleAnnotationKey, time.Now().UTC().Format(time.RFC3339))
 	_, err := client.Patch(ctx, stsName, types.StrategicMergePatchType, []byte(patch), metav1.PatchOptions{})
@@ -290,6 +295,11 @@ type statefulSetDownscale struct {
 }
 
 func findDownscalesDoneMinTimeAgo(ctx context.Context, api kubernetes.Interface, namespace, stsName, rolloutGroup string) (*statefulSetDownscale, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
 	client := api.AppsV1().StatefulSets(namespace)
 	groupReq, err := labels.NewRequirement(config.RolloutGroupLabelKey, selection.Equals, []string{rolloutGroup})
 	if err != nil {