github.com/dtroyer-salad/og2/v2@v2.0.0-20240412154159-c47231610877/registry/remote/retry/client.go (about)

     1  /*
     2  Copyright The ORAS Authors.
     3  Licensed under the Apache License, Version 2.0 (the "License");
     4  you may not use this file except in compliance with the License.
     5  You may obtain a copy of the License at
     6  
     7  http://www.apache.org/licenses/LICENSE-2.0
     8  
     9  Unless required by applicable law or agreed to in writing, software
    10  distributed under the License is distributed on an "AS IS" BASIS,
    11  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  See the License for the specific language governing permissions and
    13  limitations under the License.
    14  */
    15  
    16  package retry
    17  
    18  import (
    19  	"net/http"
    20  	"time"
    21  )
    22  
    23  // DefaultClient is a client with the default retry policy.
    24  var DefaultClient = NewClient()
    25  
    26  // NewClient creates an HTTP client with the default retry policy.
    27  func NewClient() *http.Client {
    28  	return &http.Client{
    29  		Transport: NewTransport(nil),
    30  	}
    31  }
    32  
    33  // Transport is an HTTP transport with retry policy.
    34  type Transport struct {
    35  	// Base is the underlying HTTP transport to use.
    36  	// If nil, http.DefaultTransport is used for round trips.
    37  	Base http.RoundTripper
    38  
    39  	// Policy returns a retry Policy to use for the request.
    40  	// If nil, DefaultPolicy is used to determine if the request should be retried.
    41  	Policy func() Policy
    42  }
    43  
    44  // NewTransport creates an HTTP Transport with the default retry policy.
    45  func NewTransport(base http.RoundTripper) *Transport {
    46  	return &Transport{
    47  		Base: base,
    48  	}
    49  }
    50  
    51  // RoundTrip executes a single HTTP transaction, returning a Response for the
    52  // provided Request.
    53  // It relies on the configured Policy to determine if the request should be
    54  // retried and to backoff.
    55  func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
    56  	ctx := req.Context()
    57  	policy := t.policy()
    58  	attempt := 0
    59  	for {
    60  		resp, respErr := t.roundTrip(req)
    61  		duration, err := policy.Retry(attempt, resp, respErr)
    62  		if err != nil {
    63  			if respErr == nil {
    64  				resp.Body.Close()
    65  			}
    66  			return nil, err
    67  		}
    68  		if duration < 0 {
    69  			return resp, respErr
    70  		}
    71  
    72  		// rewind the body if possible
    73  		if req.Body != nil {
    74  			if req.GetBody == nil {
    75  				// body can't be rewound, so we can't retry
    76  				return resp, respErr
    77  			}
    78  			body, err := req.GetBody()
    79  			if err != nil {
    80  				// failed to rewind the body, so we can't retry
    81  				return resp, respErr
    82  			}
    83  			req.Body = body
    84  		}
    85  
    86  		// close the response body if needed
    87  		if respErr == nil {
    88  			resp.Body.Close()
    89  		}
    90  
    91  		timer := time.NewTimer(duration)
    92  		select {
    93  		case <-ctx.Done():
    94  			timer.Stop()
    95  			return nil, ctx.Err()
    96  		case <-timer.C:
    97  		}
    98  		attempt++
    99  	}
   100  }
   101  
   102  func (t *Transport) roundTrip(req *http.Request) (*http.Response, error) {
   103  	if t.Base == nil {
   104  		return http.DefaultTransport.RoundTrip(req)
   105  	}
   106  	return t.Base.RoundTrip(req)
   107  }
   108  
   109  func (t *Transport) policy() Policy {
   110  	if t.Policy == nil {
   111  		return DefaultPolicy
   112  	}
   113  	return t.Policy()
   114  }