Handle the context in webhook

Signed-off-by: Michel Hollands <[email protected]>

Author: MichelHollands

pkg/admission/annotations.go

@@ -18,6 +18,12 @@ const (
 )
 
 func addDownscaledAnnotationToStatefulSet(ctx context.Context, api kubernetes.Interface, namespace, stsName string) error {
+	select {
+	case <-ctx.Done():
+		return fmt.Errorf("context cancelled")
+	default:
+	}
+
 	client := api.AppsV1().StatefulSets(namespace)
 	patch := fmt.Sprintf(`{"metadata":{"annotations":{"%v":"%v"}}}`, LastDownscaleAnnotationKey, time.Now().UTC().Format(time.RFC3339))
 	_, err := client.Patch(ctx, stsName, types.StrategicMergePatchType, []byte(patch), v1.PatchOptions{})
@@ -31,6 +37,11 @@ type statefulSet struct {
 }
 
 func findDownscalesDoneMinTimeAgo(ctx context.Context, api kubernetes.Interface, namespace, stsName, rolloutGroup string) (*statefulSet, error) {
+	select {
+	case <-ctx.Done():
+		return nil, fmt.Errorf("context cancelled")
+	default:
+	}
 	client := api.AppsV1().StatefulSets(namespace)
 	groupReq, err := labels.NewRequirement(RolloutGroupLabelKey, selection.Equals, []string{rolloutGroup})
 	if err != nil {

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, fmt.Errorf("context cancelled")
+	default:
+	}
+
 	switch ar.Request.Resource.Resource {
 	case "statefulsets":
 		obj, err := api.AppsV1().StatefulSets(ar.Request.Namespace).Get(ctx, ar.Request.Name, metav1.GetOptions{})

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("context cancelled")
+	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,11 +259,17 @@ 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, fmt.Errorf("context cancelled")
+	default:
+	}
+
 	switch ar.Request.Resource.Resource {
 	case "statefulsets":
 		obj, err := api.AppsV1().StatefulSets(ar.Request.Namespace).Get(ctx, ar.Request.Name, metav1.GetOptions{})
 		if err != nil {
-				return nil, err
+			return nil, err
 		}
 		return obj.Annotations, nil
 	}