1 /* 2 Copyright 2023 The Kubernetes Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package framework 18 19 import ( 20 "context" 21 "errors" 22 "fmt" 23 "time" 24 25 "github.com/onsi/gomega" 26 27 apierrors "k8s.io/apimachinery/pkg/api/errors" 28 metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" 29 ) 30 31 // GetFunc is a function which retrieves a certain object. 32 type GetFunc[T any] func(ctx context.Context) (T, error) 33 34 // APIGetFunc is a get functions as used in client-go. 35 type APIGetFunc[T any] func(ctx context.Context, name string, getOptions metav1.GetOptions) (T, error) 36 37 // APIListFunc is a list functions as used in client-go. 38 type APIListFunc[T any] func(ctx context.Context, listOptions metav1.ListOptions) (T, error) 39 40 // GetObject takes a get function like clientset.CoreV1().Pods(ns).Get 41 // and the parameters for it and returns a function that executes that get 42 // operation in a [gomega.Eventually] or [gomega.Consistently]. 43 // 44 // Delays and retries are handled by [HandleRetry]. A "not found" error is 45 // a fatal error that causes polling to stop immediately. If that is not 46 // desired, then wrap the result with [IgnoreNotFound]. 47 func GetObject[T any](get APIGetFunc[T], name string, getOptions metav1.GetOptions) GetFunc[T] { 48 return HandleRetry(func(ctx context.Context) (T, error) { 49 return get(ctx, name, getOptions) 50 }) 51 } 52 53 // ListObjects takes a list function like clientset.CoreV1().Pods(ns).List 54 // and the parameters for it and returns a function that executes that list 55 // operation in a [gomega.Eventually] or [gomega.Consistently]. 56 // 57 // Delays and retries are handled by [HandleRetry]. 58 func ListObjects[T any](list APIListFunc[T], listOptions metav1.ListOptions) GetFunc[T] { 59 return HandleRetry(func(ctx context.Context) (T, error) { 60 return list(ctx, listOptions) 61 }) 62 } 63 64 // HandleRetry wraps an arbitrary get function. When the wrapped function 65 // returns an error, HandleGetError will decide whether the call should be 66 // retried and if requested, will sleep before doing so. 67 // 68 // This is meant to be used inside [gomega.Eventually] or [gomega.Consistently]. 69 func HandleRetry[T any](get GetFunc[T]) GetFunc[T] { 70 return func(ctx context.Context) (T, error) { 71 t, err := get(ctx) 72 if err != nil { 73 if retry, delay := ShouldRetry(err); retry { 74 if delay > 0 { 75 // We could return 76 // gomega.TryAgainAfter(delay) here, 77 // but then we need to funnel that 78 // error through any other 79 // wrappers. Waiting directly is simpler. 80 ctx, cancel := context.WithTimeout(ctx, delay) 81 defer cancel() 82 <-ctx.Done() 83 } 84 return t, err 85 } 86 // Give up polling immediately. 87 var null T 88 return t, gomega.StopTrying(fmt.Sprintf("Unexpected final error while getting %T", null)).Wrap(err) 89 } 90 return t, nil 91 } 92 } 93 94 // ShouldRetry decides whether to retry an API request. Optionally returns a 95 // delay to retry after. 96 func ShouldRetry(err error) (retry bool, retryAfter time.Duration) { 97 // if the error sends the Retry-After header, we respect it as an explicit confirmation we should retry. 98 if delay, shouldRetry := apierrors.SuggestsClientDelay(err); shouldRetry { 99 return shouldRetry, time.Duration(delay) * time.Second 100 } 101 102 // these errors indicate a transient error that should be retried. 103 if apierrors.IsTimeout(err) || 104 apierrors.IsTooManyRequests(err) || 105 apierrors.IsServiceUnavailable(err) || 106 errors.As(err, &transientError{}) { 107 return true, 0 108 } 109 110 return false, 0 111 } 112 113 // RetryNotFound wraps an arbitrary get function. When the wrapped function 114 // encounters a "not found" error, that error is treated as a transient problem 115 // and polling continues. 116 // 117 // This is meant to be used inside [gomega.Eventually] or [gomega.Consistently]. 118 func RetryNotFound[T any](get GetFunc[T]) GetFunc[T] { 119 return func(ctx context.Context) (T, error) { 120 t, err := get(ctx) 121 if apierrors.IsNotFound(err) { 122 // If we are wrapping HandleRetry, then the error will 123 // be gomega.StopTrying. We need to get rid of that, 124 // otherwise gomega.Eventually will stop. 125 var stopTryingErr gomega.PollingSignalError 126 if errors.As(err, &stopTryingErr) { 127 if wrappedErr := errors.Unwrap(stopTryingErr); wrappedErr != nil { 128 err = wrappedErr 129 } 130 } 131 132 // Mark the error as transient in case that we get 133 // wrapped by HandleRetry. 134 err = transientError{error: err} 135 } 136 return t, err 137 } 138 } 139 140 // transientError wraps some other error and indicates that the 141 // wrapper error is something that may go away. 142 type transientError struct { 143 error 144 } 145 146 func (err transientError) Unwrap() error { 147 return err.error 148 } 149