github.com/icyphox/x@v0.0.355-0.20220311094250-029bd783e8b8/httpx/resilient_client.go (about)

     1  package httpx
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"log"
     7  	"net/http"
     8  	"time"
     9  
    10  	"github.com/opentracing/opentracing-go"
    11  
    12  	"github.com/ory/x/tracing"
    13  
    14  	"github.com/hashicorp/go-retryablehttp"
    15  
    16  	"github.com/ory/x/logrusx"
    17  )
    18  
    19  type resilientOptions struct {
    20  	ctx           context.Context
    21  	c             *http.Client
    22  	l             interface{}
    23  	retryWaitMin  time.Duration
    24  	retryWaitMax  time.Duration
    25  	retryMax      int
    26  	noInternalIPs bool
    27  	tracer        opentracing.Tracer
    28  }
    29  
    30  func newResilientOptions() *resilientOptions {
    31  	connTimeout := time.Minute
    32  	return &resilientOptions{
    33  		c:            &http.Client{Timeout: connTimeout},
    34  		retryWaitMin: 1 * time.Second,
    35  		retryWaitMax: 30 * time.Second,
    36  		retryMax:     4,
    37  		l:            log.New(io.Discard, "", log.LstdFlags),
    38  	}
    39  }
    40  
    41  // ResilientOptions is a set of options for the ResilientClient.
    42  type ResilientOptions func(o *resilientOptions)
    43  
    44  // ResilientClientWithClient sets the underlying http client to use.
    45  func ResilientClientWithClient(c *http.Client) ResilientOptions {
    46  	return func(o *resilientOptions) {
    47  		o.c = c
    48  	}
    49  }
    50  
    51  // ResilientClientWithTracer wraps the http clients transport with a tracing instrumentation
    52  func ResilientClientWithTracer(tracer opentracing.Tracer) ResilientOptions {
    53  	return func(o *resilientOptions) {
    54  		o.tracer = tracer
    55  	}
    56  }
    57  
    58  // ResilientClientWithMaxRetry sets the maximum number of retries.
    59  func ResilientClientWithMaxRetry(retryMax int) ResilientOptions {
    60  	return func(o *resilientOptions) {
    61  		o.retryMax = retryMax
    62  	}
    63  }
    64  
    65  // ResilientClientWithMinxRetryWait sets the minimum wait time between retries.
    66  func ResilientClientWithMinxRetryWait(retryWaitMin time.Duration) ResilientOptions {
    67  	return func(o *resilientOptions) {
    68  		o.retryWaitMin = retryWaitMin
    69  	}
    70  }
    71  
    72  // ResilientClientWithMaxRetryWait sets the maximum wait time for a retry.
    73  func ResilientClientWithMaxRetryWait(retryWaitMax time.Duration) ResilientOptions {
    74  	return func(o *resilientOptions) {
    75  		o.retryWaitMax = retryWaitMax
    76  	}
    77  }
    78  
    79  // ResilientClientWithConnectionTimeout sets the connection timeout for the client.
    80  func ResilientClientWithConnectionTimeout(connTimeout time.Duration) ResilientOptions {
    81  	return func(o *resilientOptions) {
    82  		o.c.Timeout = connTimeout
    83  	}
    84  }
    85  
    86  // ResilientClientWithLogger sets the logger to be used by the client.
    87  func ResilientClientWithLogger(l *logrusx.Logger) ResilientOptions {
    88  	return func(o *resilientOptions) {
    89  		o.l = l
    90  	}
    91  }
    92  
    93  // ResilientClientDisallowInternalIPs disallows internal IPs from being used.
    94  func ResilientClientDisallowInternalIPs() ResilientOptions {
    95  	return func(o *resilientOptions) {
    96  		o.noInternalIPs = true
    97  	}
    98  }
    99  
   100  // NewResilientClient creates a new ResilientClient.
   101  func NewResilientClient(opts ...ResilientOptions) *retryablehttp.Client {
   102  	o := newResilientOptions()
   103  	for _, f := range opts {
   104  		f(o)
   105  	}
   106  
   107  	if o.noInternalIPs == true {
   108  		o.c.Transport = &NoInternalIPRoundTripper{RoundTripper: o.c.Transport}
   109  	}
   110  
   111  	if o.tracer != nil {
   112  		o.c.Transport = tracing.RoundTripper(o.tracer, o.c.Transport)
   113  	}
   114  
   115  	return &retryablehttp.Client{
   116  		HTTPClient:   o.c,
   117  		Logger:       o.l,
   118  		RetryWaitMin: o.retryWaitMin,
   119  		RetryWaitMax: o.retryWaitMax,
   120  		RetryMax:     o.retryMax,
   121  		CheckRetry:   retryablehttp.DefaultRetryPolicy,
   122  		Backoff:      retryablehttp.DefaultBackoff,
   123  	}
   124  }