Retry requests support implemented for client.

client/client.go

@@ -10,6 +10,7 @@ import (
 	"golang.org/x/text/transform"
 	"io"
 	"io/ioutil"
+	"log"
 	"net"
 	"net/http"
 	"net/url"
@@ -25,14 +26,25 @@ var (
 // Client is a small wrapper around *http.Client to provide new methods.
 type Client struct {
 	*http.Client
+	maxBodySize           int64
+	charsetDetectDisabled bool
+	retryTimes            int
+	retryHTTPCodes        []int
 }
 
-const DefaultUserAgent = "Geziyor 1.0"
-const DefaultMaxBody int64 = 1024 * 1024 * 1024 // 1GB
+const (
+	DefaultUserAgent        = "Geziyor 1.0"
+	DefaultMaxBody    int64 = 1024 * 1024 * 1024 // 1GB
+	DefaultRetryTimes       = 2
+)
+
+var (
+	DefaultRetryHTTPCodes = []int{500, 502, 503, 504, 522, 524, 408}
+)
 
 // NewClient creates http.Client with modified values for typical web scraper
-func NewClient() *Client {
-	client := &http.Client{
+func NewClient(maxBodySize int64, charsetDetectDisabled bool, retryTimes int, retryHTTPCodes []int) *Client {
+	httpClient := &http.Client{
 		Transport: &http.Transport{
 			Proxy: http.ProxyFromEnvironment,
 			DialContext: (&net.Dialer{
@@ -48,31 +60,57 @@ func NewClient() *Client {
 		},
 		Timeout: time.Second * 180, // Google's timeout
 	}
-	return &Client{Client: client}
+
+	client := Client{
+		Client:                httpClient,
+		maxBodySize:           maxBodySize,
+		charsetDetectDisabled: charsetDetectDisabled,
+		retryTimes:            retryTimes,
+		retryHTTPCodes:        retryHTTPCodes,
+	}
+
+	return &client
 }
 
 // DoRequest selects appropriate request handler, client or Chrome
-func (c *Client) DoRequest(req *Request, maxBodySize int64, charsetDetectDisabled bool) (*Response, error) {
+func (c *Client) DoRequest(req *Request) (*Response, error) {
 	if !req.Rendered {
-		return c.DoRequestClient(req, maxBodySize, charsetDetectDisabled)
+		return c.DoRequestClient(req)
 	} else {
 		return c.DoRequestChrome(req)
 	}
 }
 
 // DoRequestClient is a simple wrapper to read response according to options.
-func (c *Client) DoRequestClient(req *Request, maxBodySize int64, charsetDetectDisabled bool) (*Response, error) {
+func (c *Client) DoRequestClient(req *Request) (*Response, error) {
 	// Do request
 	resp, err := c.Do(req.Request)
 	if resp != nil {
 		defer resp.Body.Close()
 	}
 	if err != nil {
+		// Retry on Error
+		if req.retryCounter < c.retryTimes {
+			req.retryCounter++
+			log.Println("Retrying:", req.URL.String())
+			return c.DoRequestClient(req)
+		}
 		return nil, errors.Wrap(err, "Response error")
 	}
 
+	// Checks status code to retry
+	if req.retryCounter < c.retryTimes {
+		for _, statusCode := range c.retryHTTPCodes {
+			if resp.StatusCode == statusCode {
+				req.retryCounter++
+				log.Println("Retrying:", req.URL.String(), resp.StatusCode)
+				return c.DoRequestClient(req)
+			}
+		}
+	}
+
 	// Limit response body reading
-	bodyReader := io.LimitReader(resp.Body, maxBodySize)
+	bodyReader := io.LimitReader(resp.Body, c.maxBodySize)
 
 	// Decode response
 	if resp.Request.Method != "HEAD" {
@@ -81,7 +119,7 @@ func (c *Client) DoRequestClient(req *Request, maxBodySize int64, charsetDetectD
 				bodyReader = transform.NewReader(bodyReader, enc.NewDecoder())
 			}
 		} else {
-			if !charsetDetectDisabled {
+			if !c.charsetDetectDisabled {
 				bodyReader, err = charset.NewReader(bodyReader, req.Header.Get("Content-Type"))
 				if err != nil {
 					return nil, errors.Wrap(err, "Reading determined encoding error")

client/client_test.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"fmt"
+	"github.com/stretchr/testify/assert"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
@@ -100,7 +101,7 @@ func TestCharsetFromHeaders(t *testing.T) {
 	defer ts.Close()
 
 	req, _ := NewRequest("GET", ts.URL, nil)
-	res, _ := NewClient().DoRequestClient(req, DefaultMaxBody, false)
+	res, _ := NewClient(DefaultMaxBody, false, DefaultRetryTimes, DefaultRetryHTTPCodes).DoRequestClient(req)
 
 	if string(res.Body) != "Gültekin" {
 		t.Fatal(string(res.Body))
@@ -115,7 +116,7 @@ func TestCharsetFromBody(t *testing.T) {
 	defer ts.Close()
 
 	req, _ := NewRequest("GET", ts.URL, nil)
-	res, _ := NewClient().DoRequestClient(req, DefaultMaxBody, false)
+	res, _ := NewClient(DefaultMaxBody, false, DefaultRetryTimes, DefaultRetryHTTPCodes).DoRequestClient(req)
 
 	if string(res.Body) != "Gültekin" {
 		t.Fatal(string(res.Body))
@@ -131,9 +132,16 @@ func TestCharsetProvidedWithRequest(t *testing.T) {
 
 	req, _ := NewRequest("GET", ts.URL, nil)
 	req.Encoding = "windows-1254"
-	res, _ := NewClient().DoRequestClient(req, DefaultMaxBody, false)
+	res, _ := NewClient(DefaultMaxBody, false, DefaultRetryTimes, DefaultRetryHTTPCodes).DoRequestClient(req)
 
 	if string(res.Body) != "Gültekin" {
 		t.Fatal(string(res.Body))
 	}
 }
+
+func TestRetry(t *testing.T) {
+	req, _ := NewRequest("GET", "https://httpbin.org/status/500", nil)
+	res, err := NewClient(DefaultMaxBody, false, DefaultRetryTimes, DefaultRetryHTTPCodes).DoRequestClient(req)
+	assert.Nil(t, res)
+	assert.Error(t, err)
+}

client/request.go

@@ -26,6 +26,8 @@ type Request struct {
 
 	// Set this true to cancel requests. Should be used on middlewares.
 	Cancelled bool
+
+	retryCounter int
 }
 
 // Cancel request