🚀 Configure proxy.Balancer with custom timeout #1362 (#1364)

* 🚀 new possibility to escape special routing parameters, which gives the possibility to follow the google api design guide https://cloud.google.com/apis/design/custom_methods

* 🚀 new possibility to escape special routing parameters, which gives the possibility to follow the google api design guide https://cloud.google.com/apis/design/custom_methods

* 🚀 Configure proxy.Balancer with custom timeout or client #1362

* 🚀 Configure proxy.Balancer with custom timeout or client #1362

middleware/proxy/README.md

@@ -99,6 +99,20 @@ type Config struct {
 	//
 	// Optional. Default: nil
 	ModifyResponse fiber.Handler
+	
+	// Timeout is the request timeout used when calling the proxy client
+	//
+	// Optional. Default: 1 second
+	Timeout time.Duration
+
+	// Per-connection buffer size for requests' reading.
+	// This also limits the maximum header size.
+	// Increase this buffer if your clients send multi-KB RequestURIs
+	// and/or multi-KB headers (for example, BIG cookies).
+	ReadBufferSize int
+    
+	// Per-connection buffer size for responses' writing.
+	WriteBufferSize int
 }
 ```
 

middleware/proxy/config.go

@@ -1,7 +1,10 @@
 package proxy
 
 import (
+	"time"
+
 	"github.com/gofiber/fiber/v2"
+	"github.com/valyala/fasthttp"
 )
 
 // Config defines the config for middleware.
@@ -29,6 +32,11 @@ type Config struct {
 	// Optional. Default: nil
 	ModifyResponse fiber.Handler
 
+	// Timeout is the request timeout used when calling the proxy client
+	//
+	// Optional. Default: 1 second
+	Timeout time.Duration
+
 	// Per-connection buffer size for requests' reading.
 	// This also limits the maximum header size.
 	// Increase this buffer if your clients send multi-KB RequestURIs
@@ -44,6 +52,7 @@ var ConfigDefault = Config{
 	Next:           nil,
 	ModifyRequest:  nil,
 	ModifyResponse: nil,
+	Timeout:        fasthttp.DefaultLBClientTimeout,
 }
 
 // configDefault function to set default values
@@ -56,6 +65,11 @@ func configDefault(config ...Config) Config {
 	// Override default config
 	cfg := config[0]
 
+	// Set default values
+	if cfg.Timeout <= 0 {
+		cfg.Timeout = ConfigDefault.Timeout
+	}
+
 	// Set default values
 	if len(cfg.Servers) == 0 {
 		panic("Servers cannot be empty")

middleware/proxy/proxy.go

@@ -23,6 +23,8 @@ func Balancer(config Config) fiber.Handler {
 
 	// Load balanced client
 	var lbc fasthttp.LBClient
+	// Set timeout
+	lbc.Timeout = cfg.Timeout
 
 	// Scheme must be provided, falls back to http
 	// TODO add https support

middleware/proxy/proxy_test.go

@@ -12,6 +12,23 @@ import (
 	"github.com/gofiber/fiber/v2/utils"
 )
 
+func createProxyTestServer(handler fiber.Handler, t *testing.T) (*fiber.App, string) {
+	target := fiber.New(fiber.Config{DisableStartupMessage: true})
+	target.Get("/", handler)
+
+	ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0")
+	utils.AssertEqual(t, nil, err)
+
+	go func() {
+		utils.AssertEqual(t, nil, target.Listener(ln))
+	}()
+
+	time.Sleep(2 * time.Second)
+	addr := ln.Addr().String()
+
+	return target, addr
+}
+
 // go test -run Test_Proxy_Empty_Host
 func Test_Proxy_Empty_Upstream_Servers(t *testing.T) {
 	t.Parallel()
@@ -46,20 +63,9 @@ func Test_Proxy_Next(t *testing.T) {
 func Test_Proxy(t *testing.T) {
 	t.Parallel()
 
-	target := fiber.New(fiber.Config{DisableStartupMessage: true})
+	target, addr := createProxyTestServer(
-	target.Get("/", func(c *fiber.Ctx) error {
+		func(c *fiber.Ctx) error { return c.SendStatus(fiber.StatusTeapot) }, t,
-		return c.SendStatus(fiber.StatusTeapot)
+	)
-	})
-
-	ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0")
-	utils.AssertEqual(t, nil, err)
-
-	go func() {
-		utils.AssertEqual(t, nil, target.Listener(ln))
-	}()
-
-	time.Sleep(2 * time.Second)
-	addr := ln.Addr().String()
 
 	resp, err := target.Test(httptest.NewRequest("GET", "/", nil), 2000)
 	utils.AssertEqual(t, nil, err)
@@ -76,25 +82,15 @@ func Test_Proxy(t *testing.T) {
 	utils.AssertEqual(t, fiber.StatusTeapot, resp.StatusCode)
 }
 
+// go test -run Test_Proxy_Forward
 func Test_Proxy_Forward(t *testing.T) {
 	t.Parallel()
 
 	app := fiber.New()
 
-	target := fiber.New(fiber.Config{DisableStartupMessage: true})
+	_, addr := createProxyTestServer(
-	target.Get("/", func(c *fiber.Ctx) error {
+		func(c *fiber.Ctx) error { return c.SendString("forwarded") }, t,
-		return c.SendString("forwarded")
+	)
-	})
-
-	ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0")
-	utils.AssertEqual(t, nil, err)
-
-	go func() {
-		utils.AssertEqual(t, nil, target.Listener(ln))
-	}()
-
-	time.Sleep(2 * time.Second)
-	addr := ln.Addr().String()
 
 	app.Use(Forward("http://" + addr))
 
@@ -107,23 +103,13 @@ func Test_Proxy_Forward(t *testing.T) {
 	utils.AssertEqual(t, "forwarded", string(b))
 }
 
+// go test -run Test_Proxy_Modify_Response
 func Test_Proxy_Modify_Response(t *testing.T) {
 	t.Parallel()
 
-	target := fiber.New(fiber.Config{DisableStartupMessage: true})
+	_, addr := createProxyTestServer(func(c *fiber.Ctx) error {
-	target.Get("/", func(c *fiber.Ctx) error {
 		return c.Status(500).SendString("not modified")
-	})
+	}, t)
-
-	ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0")
-	utils.AssertEqual(t, nil, err)
-
-	go func() {
-		utils.AssertEqual(t, nil, target.Listener(ln))
-	}()
-
-	time.Sleep(2 * time.Second)
-	addr := ln.Addr().String()
 
 	app := fiber.New()
 	app.Use(Balancer(Config{
@@ -143,24 +129,14 @@ func Test_Proxy_Modify_Response(t *testing.T) {
 	utils.AssertEqual(t, "modified response", string(b))
 }
 
+// go test -run Test_Proxy_Modify_Request
 func Test_Proxy_Modify_Request(t *testing.T) {
 	t.Parallel()
 
-	target := fiber.New(fiber.Config{DisableStartupMessage: true})
+	_, addr := createProxyTestServer(func(c *fiber.Ctx) error {
-	target.Get("/", func(c *fiber.Ctx) error {
 		b := c.Request().Body()
 		return c.SendString(string(b))
-	})
+	}, t)
-
-	ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0")
-	utils.AssertEqual(t, nil, err)
-
-	go func() {
-		utils.AssertEqual(t, nil, target.Listener(ln))
-	}()
-
-	time.Sleep(2 * time.Second)
-	addr := ln.Addr().String()
 
 	app := fiber.New()
 	app.Use(Balancer(Config{
@@ -180,25 +156,63 @@ func Test_Proxy_Modify_Request(t *testing.T) {
 	utils.AssertEqual(t, "modified request", string(b))
 }
 
+// go test -run Test_Proxy_Timeout_Slow_Server
+func Test_Proxy_Timeout_Slow_Server(t *testing.T) {
+	t.Parallel()
+
+	_, addr := createProxyTestServer(func(c *fiber.Ctx) error {
+		time.Sleep(2 * time.Second)
+		return c.SendString("fiber is awesome")
+	}, t)
+
+	app := fiber.New()
+	app.Use(Balancer(Config{
+		Servers: []string{addr},
+		Timeout: 3 * time.Second,
+	}))
+
+	resp, err := app.Test(httptest.NewRequest("GET", "/", nil), 5000)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, fiber.StatusOK, resp.StatusCode)
+
+	b, err := ioutil.ReadAll(resp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "fiber is awesome", string(b))
+}
+
+// go test -run Test_Proxy_With_Timeout
+func Test_Proxy_With_Timeout(t *testing.T) {
+	t.Parallel()
+
+	_, addr := createProxyTestServer(func(c *fiber.Ctx) error {
+		time.Sleep(1 * time.Second)
+		return c.SendString("fiber is awesome")
+	}, t)
+
+	app := fiber.New()
+	app.Use(Balancer(Config{
+		Servers: []string{addr},
+		Timeout: 100 * time.Millisecond,
+	}))
+
+	resp, err := app.Test(httptest.NewRequest("GET", "/", nil), 2000)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, fiber.StatusInternalServerError, resp.StatusCode)
+
+	b, err := ioutil.ReadAll(resp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "timeout", string(b))
+}
+
+// go test -run Test_Proxy_Buffer_Size_Response
 func Test_Proxy_Buffer_Size_Response(t *testing.T) {
 	t.Parallel()
 
-	target := fiber.New(fiber.Config{DisableStartupMessage: true})
+	_, addr := createProxyTestServer(func(c *fiber.Ctx) error {
-	target.Get("/", func(c *fiber.Ctx) error {
 		long := strings.Join(make([]string, 5000), "-")
-		c.Response().Header.Set("Very-Long-Header", long)
+		c.Set("Very-Long-Header", long)
 		return c.SendString("ok")
-	})
+	}, t)
-
-	ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0")
-	utils.AssertEqual(t, nil, err)
-
-	go func() {
-		utils.AssertEqual(t, nil, target.Listener(ln))
-	}()
-
-	time.Sleep(2 * time.Second)
-	addr := ln.Addr().String()
 
 	app := fiber.New()
 	app.Use(Balancer(Config{Servers: []string{addr}}))