From 3ede597a15581f957296938c0a9f835e209cd24b Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Wed, 3 Nov 2021 22:20:32 +0800
Subject: [PATCH] feat: support customizing timeout for specific route (#1203)

* feat: support customizing timeout for specific route

* test: add more tests
---
 rest/engine.go      | 10 +++++++++-
 rest/engine_test.go | 36 ++++++++++++++++++++++++++++++++++++
 rest/server.go      |  8 ++++++++
 rest/server_test.go |  7 +++++++
 rest/types.go       |  6 +++++-
 5 files changed, 65 insertions(+), 2 deletions(-)

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/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.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/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
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