// Package pester provides additional resiliency over the standard http client methods by // allowing you to control concurrency, retries, and a backoff strategy. package pester import ( "bytes" "errors" "fmt" "io" "io/ioutil" "math/rand" "net/http" "net/url" "sync" "time" ) //ErrUnexpectedMethod occurs when an http.Client method is unable to be mapped from a calling method in the pester client var ErrUnexpectedMethod = errors.New("unexpected client method, must be one of Do, Get, Head, Post, or PostFrom") // ErrReadingBody happens when we cannot read the body bytes var ErrReadingBody = errors.New("error reading body") // ErrReadingRequestBody happens when we cannot read the request body bytes var ErrReadingRequestBody = errors.New("error reading request body") // Client wraps the http client and exposes all the functionality of the http.Client. // Additionally, Client provides pester specific values for handling resiliency. type Client struct { // wrap it to provide access to http built ins hc *http.Client Transport http.RoundTripper CheckRedirect func(req *http.Request, via []*http.Request) error Jar http.CookieJar Timeout time.Duration // pester specific Concurrency int MaxRetries int Backoff BackoffStrategy KeepLog bool LogHook LogHook SuccessReqNum int SuccessRetryNum int wg *sync.WaitGroup sync.Mutex ErrLog []ErrEntry } // ErrEntry is used to provide the LogString() data and is populated // each time an error happens if KeepLog is set. // ErrEntry.Retry is deprecated in favor of ErrEntry.Attempt type ErrEntry struct { Time time.Time Method string URL string Verb string Request int Retry int Attempt int Err error } // result simplifies the channel communication for concurrent request handling type result struct { resp *http.Response err error req int retry int } // params represents all the params needed to run http client calls and pester errors type params struct { method string verb string req *http.Request url string bodyType string body io.Reader data url.Values } var random *rand.Rand func init() { random = rand.New(rand.NewSource(time.Now().UnixNano())) } // New constructs a new DefaultClient with sensible default values func New() *Client { return &Client{ Concurrency: DefaultClient.Concurrency, MaxRetries: DefaultClient.MaxRetries, Backoff: DefaultClient.Backoff, ErrLog: DefaultClient.ErrLog, wg: &sync.WaitGroup{}, } } // NewExtendedClient allows you to pass in an http.Client that is previously set up // and extends it to have Pester's features of concurrency and retries. func NewExtendedClient(hc *http.Client) *Client { c := New() c.hc = hc return c } // PrintErrStrategy is used to log attempts as they happen. // You know, more visible type LogHook func(e ErrEntry) // BackoffStrategy is used to determine how long a retry request should wait until attempted type BackoffStrategy func(retry int) time.Duration // DefaultClient provides sensible defaults var DefaultClient = &Client{Concurrency: 1, MaxRetries: 3, Backoff: DefaultBackoff, ErrLog: []ErrEntry{}} // DefaultBackoff always returns 1 second func DefaultBackoff(_ int) time.Duration { return 1 * time.Second } // ExponentialBackoff returns ever increasing backoffs by a power of 2 func ExponentialBackoff(i int) time.Duration { return time.Duration(1< 0 { p.req.Body = ioutil.NopCloser(bytes.NewBuffer(originalRequestBody)) } if len(originalBody) > 0 { p.body = bytes.NewBuffer(originalBody) } var resp *http.Response // route the calls switch p.method { case "Do": resp, err = httpClient.Do(p.req) case "Get": resp, err = httpClient.Get(p.url) case "Head": resp, err = httpClient.Head(p.url) case "Post": resp, err = httpClient.Post(p.url, p.bodyType, p.body) case "PostForm": resp, err = httpClient.PostForm(p.url, p.data) default: err = ErrUnexpectedMethod } // Early return if we have a valid result // Only retry (ie, continue the loop) on 5xx status codes if err == nil && resp.StatusCode < 500 { multiplexCh <- result{resp: resp, err: err, req: n, retry: i} return } c.log(ErrEntry{ Time: time.Now(), Method: p.method, Verb: p.verb, URL: p.url, Request: n, Retry: i + 1, // would remove, but would break backward compatibility Attempt: i, Err: err, }) // if it is the last iteration, grab the result (which is an error at this point) if i == AttemptLimit { multiplexCh <- result{resp: resp, err: err} return } // if we are retrying, we should close this response body to free the fd if resp != nil { resp.Body.Close() } // prevent a 0 from causing the tick to block, pass additional microsecond <-time.After(c.Backoff(i) + 1*time.Microsecond) } }(req, p) } // spin off the go routine so it can continually listen in on late results and close the response bodies go func() { gotFirstResult := false for { select { case res := <-multiplexCh: if !gotFirstResult { gotFirstResult = true close(finishCh) resultCh <- res } else if res.resp != nil { // we only return one result to the caller; close all other response bodies that come back // drain the body before close as to not prevent keepalive. see https://gist.github.com/mholt/eba0f2cc96658be0f717 io.Copy(ioutil.Discard, res.resp.Body) res.resp.Body.Close() } case <-allRequestsBackCh: // don't leave this goroutine running return } } }() res := <-resultCh c.Lock() defer c.Unlock() c.SuccessReqNum = res.req c.SuccessRetryNum = res.retry return res.resp, res.err } // LogString provides a string representation of the errors the client has seen func (c *Client) LogString() string { c.Lock() defer c.Unlock() var res string for _, e := range c.ErrLog { res += c.FormatError(e) } return res } // Format the Error to human readable string func (c *Client) FormatError(e ErrEntry) string { return fmt.Sprintf("%d %s [%s] %s request-%d retry-%d error: %s\n", e.Time.Unix(), e.Method, e.Verb, e.URL, e.Request, e.Retry, e.Err) } // LogErrCount is a helper method used primarily for test validation func (c *Client) LogErrCount() int { c.Lock() defer c.Unlock() return len(c.ErrLog) } // EmbedHTTPClient allows you to extend an existing Pester client with an // underlying http.Client, such as https://godoc.org/golang.org/x/oauth2/google#DefaultClient func (c *Client) EmbedHTTPClient(hc *http.Client) { c.hc = hc } func (c *Client) log(e ErrEntry) { if c.KeepLog { c.Lock() defer c.Unlock() c.ErrLog = append(c.ErrLog, e) } else if c.LogHook != nil { // NOTE: There is a possibility that Log Printing hook slows it down. // but the consumer can always do the Job in a go-routine. c.LogHook(e) } } // Do provides the same functionality as http.Client.Do func (c *Client) Do(req *http.Request) (resp *http.Response, err error) { return c.pester(params{method: "Do", req: req, verb: req.Method, url: req.URL.String()}) } // Get provides the same functionality as http.Client.Get func (c *Client) Get(url string) (resp *http.Response, err error) { return c.pester(params{method: "Get", url: url, verb: "GET"}) } // Head provides the same functionality as http.Client.Head func (c *Client) Head(url string) (resp *http.Response, err error) { return c.pester(params{method: "Head", url: url, verb: "HEAD"}) } // Post provides the same functionality as http.Client.Post func (c *Client) Post(url string, bodyType string, body io.Reader) (resp *http.Response, err error) { return c.pester(params{method: "Post", url: url, bodyType: bodyType, body: body, verb: "POST"}) } // PostForm provides the same functionality as http.Client.PostForm func (c *Client) PostForm(url string, data url.Values) (resp *http.Response, err error) { return c.pester(params{method: "PostForm", url: url, data: data, verb: "POST"}) } //////////////////////////////////////// // Provide self-constructing variants // //////////////////////////////////////// // Do provides the same functionality as http.Client.Do and creates its own constructor func Do(req *http.Request) (resp *http.Response, err error) { c := New() return c.Do(req) } // Get provides the same functionality as http.Client.Get and creates its own constructor func Get(url string) (resp *http.Response, err error) { c := New() return c.Get(url) } // Head provides the same functionality as http.Client.Head and creates its own constructor func Head(url string) (resp *http.Response, err error) { c := New() return c.Head(url) } // Post provides the same functionality as http.Client.Post and creates its own constructor func Post(url string, bodyType string, body io.Reader) (resp *http.Response, err error) { c := New() return c.Post(url, bodyType, body) } // PostForm provides the same functionality as http.Client.PostForm and creates its own constructor func PostForm(url string, data url.Values) (resp *http.Response, err error) { c := New() return c.PostForm(url, data) }