Support retries of failed proxy requests

middleware/proxy.go

@@ -29,6 +29,19 @@ type (
 		// Required.
 		Balancer ProxyBalancer
 
+		// RetryCount defines the number of times a proxied request to an unavailable
+		// ProxyTarget should be retried using the next available ProxyTarget. Defaults
+		// to 0, meaning requests are never retried.
+		RetryCount int
+
+		// RetryFilter defines a function used to determine if a failed request to an
+		// unavailable ProxyTarget should be retried. The RetryFilter will only be called
+		// when the number of previous retries is less than RetryCount. If the function returns
+		// true, the request will be retried. When not specified, DefaultProxyRetryFilter
+		// will be used, which will always retry requests. A user defined ProxyRetryFilter
+		// can be provided to only retry specific requests, for example only retry GET requests.
+		RetryFilter ProxyRetryFilter
+
 		// Rewrite defines URL path rewrite rules. The values captured in asterisk can be
 		// retrieved by index e.g. $1, $2 and so on.
 		// Examples:
@@ -71,11 +84,17 @@ type (
 		Next(echo.Context) *ProxyTarget
 	}
 
-	// TargetProvider defines an interface that gives the opportunity for balancer to return custom errors when selecting target.
+	// TargetProvider defines an interface that gives the opportunity for balancer
+	// to return custom errors when selecting target.
 	TargetProvider interface {
 		NextTarget(echo.Context) (*ProxyTarget, error)
 	}
 
+	// ProxyRetryFilter defines a function that determines if a failed request to
+	// an unavailable ProxyTarget should be retried using the next available ProxyTarget.
+	// When the function returns true, the request will be retried.
+	ProxyRetryFilter func(c echo.Context) bool
+
 	commonBalancer struct {
 		targets []*ProxyTarget
 		mutex   sync.Mutex
@@ -98,8 +117,9 @@ type (
 var (
 	// DefaultProxyConfig is the default Proxy middleware config.
 	DefaultProxyConfig = ProxyConfig{
-		Skipper:    DefaultSkipper,
-		ContextKey: "target",
+		Skipper:     DefaultSkipper,
+		ContextKey:  "target",
+		RetryFilter: DefaultProxyRetryFilter,
 	}
 )
 
@@ -141,6 +161,11 @@ func proxyRaw(t *ProxyTarget, c echo.Context) http.Handler {
 	})
 }
 
+// DefaultProxyRetryFilter is a ProxyRetryFilter that always retries requests
+func DefaultProxyRetryFilter(c echo.Context) bool {
+	return true
+}
+
 // NewRandomBalancer returns a random proxy balancer.
 func NewRandomBalancer(targets []*ProxyTarget) ProxyBalancer {
 	b := randomBalancer{}
@@ -232,14 +257,16 @@ func Proxy(balancer ProxyBalancer) echo.MiddlewareFunc {
 // ProxyWithConfig returns a Proxy middleware with config.
 // See: `Proxy()`
 func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc {
+	if config.Balancer == nil {
+		panic("echo: proxy middleware requires balancer")
+	}
 	// Defaults
 	if config.Skipper == nil {
 		config.Skipper = DefaultProxyConfig.Skipper
 	}
-	if config.Balancer == nil {
-		panic("echo: proxy middleware requires balancer")
+	if config.RetryFilter == nil {
+		config.RetryFilter = DefaultProxyConfig.RetryFilter
 	}
-
 	if config.Rewrite != nil {
 		if config.RegexRewrite == nil {
 			config.RegexRewrite = make(map[*regexp.Regexp]string)
@@ -251,25 +278,13 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc {
 
 	provider, isTargetProvider := config.Balancer.(TargetProvider)
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
-		return func(c echo.Context) (err error) {
+		return func(c echo.Context) error {
 			if config.Skipper(c) {
 				return next(c)
 			}
 
 			req := c.Request()
 			res := c.Response()
-
-			var tgt *ProxyTarget
-			if isTargetProvider {
-				tgt, err = provider.NextTarget(c)
-				if err != nil {
-					return err
-				}
-			} else {
-				tgt = config.Balancer.Next(c)
-			}
-			c.Set(config.ContextKey, tgt)
-
 			if err := rewriteURL(config.RegexRewrite, req); err != nil {
 				return err
 			}
@@ -287,19 +302,52 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc {
 				req.Header.Set(echo.HeaderXForwardedFor, c.RealIP())
 			}
 
-			// Proxy
-			switch {
-			case c.IsWebSocket():
-				proxyRaw(tgt, c).ServeHTTP(res, req)
-			case req.Header.Get(echo.HeaderAccept) == "text/event-stream":
-			default:
-				proxyHTTP(tgt, c, config).ServeHTTP(res, req)
-			}
-			if e, ok := c.Get("_error").(error); ok {
-				err = e
-			}
+			retries := config.RetryCount
+			for {
+				var tgt *ProxyTarget
+				var err error
+				if isTargetProvider {
+					tgt, err = provider.NextTarget(c)
+					if err != nil {
+						return err
+					}
+				} else {
+					tgt = config.Balancer.Next(c)
+				}
+				c.Set(config.ContextKey, tgt)
+
+				// Proxy
+				switch {
+				case c.IsWebSocket():
+					proxyRaw(tgt, c).ServeHTTP(res, req)
+				case req.Header.Get(echo.HeaderAccept) == "text/event-stream":
+				default:
+					proxyHTTP(tgt, c, config).ServeHTTP(res, req)
+				}
 
-			return
+				e, hasError := c.Get("_error").(error)
+				if !hasError {
+					return nil
+				}
+
+				retry := false
+				if httpErr, ok := e.(*echo.HTTPError); ok {
+					if httpErr.Code == http.StatusBadGateway {
+						retry = retries > 0 && config.RetryFilter(c)
+					}
+				}
+
+				if !retry {
+					return e
+				}
+
+				retries--
+
+				//Try request again. Clear the previous error and target
+				//server from context.
+				c.Set("_error", nil)
+				c.Set(config.ContextKey, nil)
+			}
 		}
 	}
 }

middleware/proxy_test.go

@@ -393,6 +393,189 @@ func TestProxyError(t *testing.T) {
 	assert.Equal(t, http.StatusBadGateway, rec.Code)
 }
 
+func TestProxyRetries(t *testing.T) {
+
+	newServer := func(res int) (*url.URL, *httptest.Server) {
+		server := httptest.NewServer(
+			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(res)
+			}),
+		)
+		targetUrl, _ := url.Parse(server.URL)
+		return targetUrl, server
+	}
+
+	targetUrl, server := newServer(http.StatusOK)
+	defer server.Close()
+	goodTarget := &ProxyTarget{
+		Name: "Good",
+		URL:  targetUrl,
+	}
+
+	targetUrl, server = newServer(http.StatusBadRequest)
+	defer server.Close()
+	goodTargetWith40X := &ProxyTarget{
+		Name: "Good with 40X",
+		URL:  targetUrl,
+	}
+
+	targetUrl, _ = url.Parse("http://127.0.0.1:27121")
+	badTarget := &ProxyTarget{
+		Name: "Bad",
+		URL:  targetUrl,
+	}
+
+	var alwaysRetryFilter ProxyRetryFilter = func(c echo.Context) bool { return true }
+	var neverRetryFilter ProxyRetryFilter = func(c echo.Context) bool { return false }
+
+	testCases := []struct {
+		name             string
+		retryCount       int
+		retryFilters     []ProxyRetryFilter
+		targets          []*ProxyTarget
+		expectedResponse int
+	}{
+		{
+			"retry count 0 does not attempt retry on fail",
+			0,
+			nil,
+			[]*ProxyTarget{
+				badTarget,
+				goodTarget,
+			},
+			http.StatusBadGateway,
+		},
+		{
+			"retry count 1 does not attempt retry on success",
+			1,
+			nil,
+			[]*ProxyTarget{
+				goodTarget,
+			},
+			http.StatusOK,
+		},
+		{
+			"retry count 1 does retry on handler return true",
+			1,
+			[]ProxyRetryFilter{
+				alwaysRetryFilter,
+			},
+			[]*ProxyTarget{
+				badTarget,
+				goodTarget,
+			},
+			http.StatusOK,
+		},
+		{
+			"retry count 1 does not retry on handler return false",
+			1,
+			[]ProxyRetryFilter{
+				neverRetryFilter,
+			},
+			[]*ProxyTarget{
+				badTarget,
+				goodTarget,
+			},
+			http.StatusBadGateway,
+		},
+		{
+			"retry count 2 returns error when no more retries left",
+			2,
+			[]ProxyRetryFilter{
+				alwaysRetryFilter,
+				alwaysRetryFilter,
+			},
+			[]*ProxyTarget{
+				badTarget,
+				badTarget,
+				badTarget,
+				goodTarget, //Should never be reached as only 2 retries
+			},
+			http.StatusBadGateway,
+		},
+		{
+			"retry count 2 returns error when retries left but handler returns false",
+			3,
+			[]ProxyRetryFilter{
+				alwaysRetryFilter,
+				alwaysRetryFilter,
+				neverRetryFilter,
+			},
+			[]*ProxyTarget{
+				badTarget,
+				badTarget,
+				badTarget,
+				goodTarget, //Should never be reached as retry handler returns false on 2nd check
+			},
+			http.StatusBadGateway,
+		},
+		{
+			"retry count 3 succeeds",
+			3,
+			[]ProxyRetryFilter{
+				alwaysRetryFilter,
+				alwaysRetryFilter,
+				alwaysRetryFilter,
+			},
+			[]*ProxyTarget{
+				badTarget,
+				badTarget,
+				badTarget,
+				goodTarget,
+			},
+			http.StatusOK,
+		},
+		{
+			"40x responses are not retried",
+			1,
+			nil,
+			[]*ProxyTarget{
+				goodTargetWith40X,
+				goodTarget,
+			},
+			http.StatusBadRequest,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+
+			retryFilterCall := 0
+			retryFilter := func(c echo.Context) bool {
+				if len(tc.retryFilters) == 0 {
+					assert.FailNow(t, fmt.Sprintf("unexpected calls, %d, to retry handler", retryFilterCall))
+				}
+
+				retryFilterCall++
+
+				nextRetryFilter := tc.retryFilters[0]
+				tc.retryFilters = tc.retryFilters[1:]
+
+				return nextRetryFilter(c)
+			}
+
+			e := echo.New()
+			e.Use(ProxyWithConfig(
+				ProxyConfig{
+					Balancer:    NewRoundRobinBalancer(tc.targets),
+					RetryCount:  tc.retryCount,
+					RetryFilter: retryFilter,
+				},
+			))
+
+			req := httptest.NewRequest(http.MethodGet, "/", nil)
+			rec := httptest.NewRecorder()
+
+			e.ServeHTTP(rec, req)
+
+			assert.Equal(t, tc.expectedResponse, rec.Code)
+			if len(tc.retryFilters) > 0 {
+				assert.FailNow(t, fmt.Sprintf("expected %d more retry handler calls", len(tc.retryFilters)))
+			}
+		})
+	}
+}
+
 func TestClientCancelConnectionResultsHTTPCode499(t *testing.T) {
 	var timeoutStop sync.WaitGroup
 	timeoutStop.Add(1)