diff --git a/pkg/admission/no_downscale.go b/pkg/admission/no_downscale.go index 06b7b7493..c53cf137b 100644 --- a/pkg/admission/no_downscale.go +++ b/pkg/admission/no_downscale.go @@ -140,6 +140,12 @@ func allowErr(logger log.Logger, msg string, err error) *v1.AdmissionResponse { } func getResourceLabels(ctx context.Context, ar v1.AdmissionReview, api kubernetes.Interface) (map[string]string, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + switch ar.Request.Resource.Resource { case "statefulsets": obj, err := api.AppsV1().StatefulSets(ar.Request.Namespace).Get(ctx, ar.Request.Name, metav1.GetOptions{}) diff --git a/pkg/admission/prep_downscale.go b/pkg/admission/prep_downscale.go index 7d543650e..713022cc3 100644 --- a/pkg/admission/prep_downscale.go +++ b/pkg/admission/prep_downscale.go @@ -39,6 +39,12 @@ type httpClient interface { } func prepareDownscale(ctx context.Context, logger log.Logger, ar v1.AdmissionReview, api kubernetes.Interface, client httpClient) *v1.AdmissionResponse { + select { + case <-ctx.Done(): + return deny(ctx.Err().Error()) + default: + } + logger = log.With(logger, "name", ar.Request.Name, "resource", ar.Request.Resource.Resource, "namespace", ar.Request.Namespace) oldObj, oldGVK, err := codecs.UniversalDeserializer().Decode(ar.Request.OldObject.Raw, nil, nil) @@ -253,6 +259,12 @@ func deny(msg string, args ...any) *v1.AdmissionResponse { } func getResourceAnnotations(ctx context.Context, ar v1.AdmissionReview, api kubernetes.Interface) (map[string]string, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + switch ar.Request.Resource.Resource { case "statefulsets": obj, err := api.AppsV1().StatefulSets(ar.Request.Namespace).Get(ctx, ar.Request.Name, metav1.GetOptions{}) @@ -265,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{}) @@ -278,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 {