Repository: zeromicro/go-zero
Pull request #1203 (2021-11-03 13:23:48)

feat: support customizing timeout for specific route
Usage:

```go
	srv.AddRoute(rest.Route{
		Method:  http.MethodGet,
		Path:    "/hello",
		Handler: handleHomeHello,
	}, rest.WithTimeout(time.Second))
	srv.AddRoutes([]rest.Route{
		{
			Method:  http.MethodGet,
			Path:    "/hello",
			Handler: handleApiHello,
		},
		{
			Method:  http.MethodGet,
			Path:    "/hi",
			Handler: handleApiHi,
		},
	}, rest.WithPrefix("/api"), rest.WithTimeout(time.Second*5))
```


Patch:
diff --git a/rest/engine.go b/rest/engine.go
index 410f5acf..e19a56b1 100644
--- a/rest/engine.go
+++ b/rest/engine.go
@@ -120,7 +120,7 @@ func (ng *engine) bindRoute(fr featuredRoutes, router httpx.Router, metrics *sta
 		handler.MaxConns(ng.conf.MaxConns),
 		handler.BreakerHandler(route.Method, route.Path, metrics),
 		handler.SheddingHandler(ng.getShedder(fr.priority), metrics),
-		handler.TimeoutHandler(time.Duration(ng.conf.Timeout)*time.Millisecond),
+		handler.TimeoutHandler(ng.checkedTimeout(fr.timeout)),
 		handler.RecoverHandler,
 		handler.MetricHandler(metrics),
 		handler.MaxBytesHandler(ng.conf.MaxBytes),
@@ -148,6 +148,14 @@ func (ng *engine) bindRoutes(router httpx.Router) error {
 	return nil
 }
 
+func (ng *engine) checkedTimeout(timeout time.Duration) time.Duration {
+	if timeout > 0 {
+		return timeout
+	}
+
+	return time.Duration(ng.conf.Timeout) * time.Millisecond
+}
+
 func (ng *engine) createMetrics() *stat.Metrics {
 	var metrics *stat.Metrics
 
diff --git a/rest/server.go b/rest/server.go
index 57d6434e..18b41fe2 100644
--- a/rest/server.go
+++ b/rest/server.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"net/http"
 	"path"
+	"time"
 
 	"github.com/tal-tech/go-zero/core/logx"
 	"github.com/tal-tech/go-zero/rest/handler"
@@ -204,6 +205,13 @@ func WithSignature(signature SignatureConf) RouteOption {
 	}
 }
 
+// WithTimeout returns a RouteOption to set timeout with given value.
+func WithTimeout(timeout time.Duration) RouteOption {
+	return func(r *featuredRoutes) {
+		r.timeout = timeout
+	}
+}
+
 // WithTLSConfig returns a RunOption that with given tls config.
 func WithTLSConfig(cfg *tls.Config) RunOption {
 	return func(srv *Server) {
diff --git a/rest/types.go b/rest/types.go
index 102ce2b3..cc4636c5 100644
--- a/rest/types.go
+++ b/rest/types.go
@@ -1,6 +1,9 @@
 package rest
 
-import "net/http"
+import (
+	"net/http"
+	"time"
+)
 
 type (
 	// Middleware defines the middleware method.
@@ -28,6 +31,7 @@ type (
 	}
 
 	featuredRoutes struct {
+		timeout   time.Duration
 		priority  bool
 		jwt       jwtSetting
 		signature signatureSetting


Test patch:
diff --git a/rest/engine_test.go b/rest/engine_test.go
index f5eac8f1..dfd084c2 100644
--- a/rest/engine_test.go
+++ b/rest/engine_test.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/conf"
@@ -154,6 +155,41 @@ Verbose: true
 	}
 }
 
+func TestEngine_checkedTimeout(t *testing.T) {
+	tests := []struct {
+		name    string
+		timeout time.Duration
+		expect  time.Duration
+	}{
+		{
+			name:   "not set",
+			expect: time.Second,
+		},
+		{
+			name:    "less",
+			timeout: time.Millisecond * 500,
+			expect:  time.Millisecond * 500,
+		},
+		{
+			name:    "equal",
+			timeout: time.Second,
+			expect:  time.Second,
+		},
+		{
+			name:    "more",
+			timeout: time.Millisecond * 1500,
+			expect:  time.Millisecond * 1500,
+		},
+	}
+
+	ng := newEngine(RestConf{
+		Timeout: 1000,
+	})
+	for _, test := range tests {
+		assert.Equal(t, test.expect, ng.checkedTimeout(test.timeout))
+	}
+}
+
 type mockedRouter struct{}
 
 func (m mockedRouter) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
diff --git a/rest/server_test.go b/rest/server_test.go
index 30c2383b..3208048c 100644
--- a/rest/server_test.go
+++ b/rest/server_test.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/conf"
@@ -238,6 +239,12 @@ func TestWithPriority(t *testing.T) {
 	assert.True(t, fr.priority)
 }
 
+func TestWithTimeout(t *testing.T) {
+	var fr featuredRoutes
+	WithTimeout(time.Hour)(&fr)
+	assert.Equal(t, time.Hour, fr.timeout)
+}
+
 func TestWithTLSConfig(t *testing.T) {
 	const configYaml = `
 Name: foo