feat: add timeout callback

go.sum

@@ -7,10 +7,13 @@ github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302/go.mod h1:SGn
 github.com/alicebob/miniredis/v2 v2.33.0 h1:uvTF0EDeu9RLnUEG27Db5I68ESoIxTiXbNUiji6lZrA=
 github.com/alicebob/miniredis/v2 v2.33.0/go.mod h1:MhP4a3EU7aENRi9aO+tHfTBZicLqQevyi/DJpoj6mi0=
 github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
+github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
+github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
 github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
+github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
 github.com/bufbuild/protocompile v0.10.0 h1:+jW/wnLMLxaCEG8AX9lD0bQ5v9h1RUiMKOBOT5ll9dM=
 github.com/bufbuild/protocompile v0.10.0/go.mod h1:G9qQIQo0xZ6Uyj6CMNz0saGmx2so+KONo8/KrELABiY=
 github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
@@ -55,6 +58,7 @@ github.com/go-openapi/swag v0.22.4/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+
 github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
 github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
 github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI=
+github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls=
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
@@ -75,6 +79,7 @@ github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/
 github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0=
 github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1 h1:K6RDEckDVWvDI9JAJYCmNdQXq6neHJOYx3V6jnqNEec=
+github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 h1:bkypFPDjIYGfCYD5mRBvpqxfYX1YCS1PXdKYWi8FsN0=
@@ -102,6 +107,7 @@ github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0N
 github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
+github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
@@ -130,15 +136,19 @@ github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32/go.mod h1:9wM+0iRr9ahx58uY
 github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
 github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
 github.com/onsi/ginkgo/v2 v2.13.0 h1:0jY9lJquiL8fcf3M4LAXN5aMlS/b2BV86HFFPCPMgE4=
+github.com/onsi/ginkgo/v2 v2.13.0/go.mod h1:TE309ZR8s5FsKKpuB1YAQYBzCaAfUgatB/xlT/ETL/o=
 github.com/onsi/gomega v1.29.0 h1:KIA/t2t5UBzoirT4H9tsML45GEbo3ouUnBHsCfD2tVg=
+github.com/onsi/gomega v1.29.0/go.mod h1:9sxs+SwGrKI0+PWe4Fxa9tFQQBG5xSsSbMXOI8PPpoQ=
 github.com/openzipkin/zipkin-go v0.4.3 h1:9EGwpqkgnwdEIJ+Od7QVSEIH+ocmm5nPat0G7sjsSdg=
 github.com/openzipkin/zipkin-go v0.4.3/go.mod h1:M9wCJZFWCo2RiY+o1eBCEMe0Dp2S5LDHcMZmk3RmK7c=
 github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM=
 github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
+github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
 github.com/prometheus/client_golang v1.19.1 h1:wZWJDwK+NameRJuPGDhlnFgx8e8HN3XHQeLaYJFJBOE=
 github.com/prometheus/client_golang v1.19.1/go.mod h1:mP78NwGzrVks5S2H6ab8+ZZGJLZUq1hoULYBAYBw1Ho=
 github.com/prometheus/client_model v0.5.0 h1:VQw1hfvPvk3Uv6Qf29VrPF32JB6rtbgI6cYPYQjL0Qw=
@@ -152,9 +162,11 @@ github.com/redis/go-redis/v9 v9.6.1/go.mod h1:0C0c6ycQsdpVNQpxb1njEQIqkx5UcsM8FJ
 github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
+github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog=
 github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI=
 github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
@@ -283,6 +295,7 @@ golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4f
 golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
 golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d h1:vU5i/LfpvrRCpgM/VPfJLg5KjxD3E+hfT1SH+d9zLwg=
+golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

rest/engine.go

@@ -31,6 +31,7 @@ type engine struct {
 	timeout              time.Duration
 	unauthorizedCallback handler.UnauthorizedCallback
 	unsignedCallback     handler.UnsignedCallback
+	timeoutCallback      handler.TimeoutCallback
 	chain                chain.Chain
 	middlewares          []Middleware
 	shedder              load.Shedder
@@ -64,7 +65,8 @@ func (ng *engine) addRoutes(r featuredRoutes) {
 }
 
 func (ng *engine) appendAuthHandler(fr featuredRoutes, chn chain.Chain,
-	verifier func(chain.Chain) chain.Chain) chain.Chain {
+	verifier func(chain.Chain) chain.Chain,
+) chain.Chain {
 	if fr.jwt.enabled {
 		if len(fr.jwt.prevSecret) == 0 {
 			chn = chn.Append(handler.Authorize(fr.jwt.secret,
@@ -95,7 +97,8 @@ func (ng *engine) bindFeaturedRoutes(router httpx.Router, fr featuredRoutes, met
 }
 
 func (ng *engine) bindRoute(fr featuredRoutes, router httpx.Router, metrics *stat.Metrics,
-	route Route, verifier func(chain.Chain) chain.Chain) error {
+	route Route, verifier func(chain.Chain) chain.Chain,
+) error {
 	chn := ng.chain
 	if chn == nil {
 		chn = ng.buildChainWithNativeMiddlewares(fr, route, metrics)
@@ -124,7 +127,8 @@ func (ng *engine) bindRoutes(router httpx.Router) error {
 }
 
 func (ng *engine) buildChainWithNativeMiddlewares(fr featuredRoutes, route Route,
-	metrics *stat.Metrics) chain.Chain {
+	metrics *stat.Metrics,
+) chain.Chain {
 	chn := chain.New()
 
 	if ng.conf.Middlewares.Trace {
@@ -148,7 +152,7 @@ func (ng *engine) buildChainWithNativeMiddlewares(fr featuredRoutes, route Route
 		chn = chn.Append(handler.SheddingHandler(ng.getShedder(fr.priority), metrics))
 	}
 	if ng.conf.Middlewares.Timeout {
-		chn = chn.Append(handler.TimeoutHandler(ng.checkedTimeout(fr.timeout)))
+		chn = chn.Append(handler.TimeoutHandler(ng.checkedTimeout(fr.timeout), handler.WithTimeoutCallback(ng.timeoutCallback)))
 	}
 	if ng.conf.Middlewares.Recover {
 		chn = chn.Append(handler.RecoverHandler)
@@ -265,6 +269,10 @@ func (ng *engine) setUnsignedCallback(callback handler.UnsignedCallback) {
 	ng.unsignedCallback = callback
 }
 
+func (ng *engine) setTimeoutCallback(callback handler.TimeoutCallback) {
+	ng.timeoutCallback = callback
+}
+
 func (ng *engine) signatureVerifier(signature signatureSetting) (func(chain.Chain) chain.Chain, error) {
 	if !signature.enabled {
 		return func(chn chain.Chain) chain.Chain {

rest/engine_test.go

@@ -215,6 +215,8 @@ Verbose: true
 					ng.setUnsignedCallback(func(w http.ResponseWriter, r *http.Request,
 						next http.Handler, strict bool, code int) {
 					})
+					ng.setTimeoutCallback(func(w http.ResponseWriter, r *http.Request, err error) {
+					})
 				}
 				ng.addRoutes(route)
 				ng.use(func(next http.HandlerFunc) http.HandlerFunc {
@@ -429,8 +431,7 @@ func TestEngine_start(t *testing.T) {
 	})
 }
 
-type mockedRouter struct {
-}
+type mockedRouter struct{}
 
 func (m mockedRouter) ServeHTTP(_ http.ResponseWriter, _ *http.Request) {
 }

rest/handler/timeouthandler.go

@@ -28,10 +28,40 @@ const (
 	valueSSE                  = "text/event-stream"
 )
 
+type (
+	// TimeoutCallback defines the method of timeout callback.
+	TimeoutCallback = func(w http.ResponseWriter, r *http.Request, err error)
+	// TimeoutOption defines the method to customize a TimeoutOptions.
+	TimeoutOptions struct {
+		Callback TimeoutCallback
+	}
+
+	// TimeoutOption defines the method to customize an TimeoutOptions.
+	TimeoutOption func(opts *TimeoutOptions)
+)
+
+var defaultTimeoutCallback TimeoutCallback = func(w http.ResponseWriter, r *http.Request, err error) {
+	if errors.Is(err, context.Canceled) {
+		w.WriteHeader(statusClientClosedRequest)
+	} else {
+		w.WriteHeader(http.StatusServiceUnavailable)
+	}
+	_, _ = io.WriteString(w, reason)
+}
+
 // TimeoutHandler returns the handler with given timeout.
 // If client closed request, code 499 will be logged.
 // Notice: even if canceled in server side, 499 will be logged as well.
-func TimeoutHandler(duration time.Duration) func(http.Handler) http.Handler {
+func TimeoutHandler(duration time.Duration, opt ...TimeoutOption) func(http.Handler) http.Handler {
+	var opts TimeoutOptions
+	for _, o := range opt {
+		o(&opts)
+	}
+
+	if opts.Callback == nil {
+		opts.Callback = defaultTimeoutCallback
+	}
+
 	return func(next http.Handler) http.Handler {
 		if duration <= 0 {
 			return next
@@ -40,6 +70,7 @@ func TimeoutHandler(duration time.Duration) func(http.Handler) http.Handler {
 		return &timeoutHandler{
 			handler: next,
 			dt:      duration,
+			cb:      opts.Callback,
 		}
 	}
 }
@@ -51,10 +82,7 @@ func TimeoutHandler(duration time.Duration) func(http.Handler) http.Handler {
 type timeoutHandler struct {
 	handler http.Handler
 	dt      time.Duration
-}
-
-func (h *timeoutHandler) errorBody() string {
-	return reason
+	cb      TimeoutCallback
 }
 
 func (h *timeoutHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -106,15 +134,10 @@ func (h *timeoutHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	case <-ctx.Done():
 		tw.mu.Lock()
 		defer tw.mu.Unlock()
-		// there isn't any user-defined middleware before TimoutHandler,
+		// there isn't any user-defined middleware before TimeoutHandler,
 		// so we can guarantee that cancelation in biz related code won't come here.
 		httpx.ErrorCtx(r.Context(), w, ctx.Err(), func(w http.ResponseWriter, err error) {
-			if errors.Is(err, context.Canceled) {
-				w.WriteHeader(statusClientClosedRequest)
-			} else {
-				w.WriteHeader(http.StatusServiceUnavailable)
-			}
-			_, _ = io.WriteString(w, h.errorBody())
+			h.cb(w, r, err)
 		})
 		tw.timedOut = true
 	}
@@ -244,3 +267,10 @@ func relevantCaller() runtime.Frame {
 
 	return frame
 }
+
+// WithTimeoutCallback returns an AuthorizeOption with setting Timeout callback.
+func WithTimeoutCallback(callback TimeoutCallback) TimeoutOption {
+	return func(opts *TimeoutOptions) {
+		opts.Callback = callback
+	}
+}

rest/handler/timeouthandler_test.go

@@ -88,6 +88,20 @@ func TestTimeout(t *testing.T) {
 	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
 }
 
+func TestTimeHandlerCallback(t *testing.T) {
+	timeoutHandler := TimeoutHandler(time.Millisecond, WithTimeoutCallback(func(w http.ResponseWriter, r *http.Request, err error) {
+		w.WriteHeader(486)
+	}))
+	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(time.Minute)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, req)
+	assert.Equal(t, 486, resp.Code)
+}
+
 func TestWithinTimeout(t *testing.T) {
 	timeoutHandler := TimeoutHandler(time.Second)
 	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -112,7 +126,7 @@ func TestWithinTimeoutBadCode(t *testing.T) {
 	assert.Equal(t, http.StatusInternalServerError, resp.Code)
 }
 
-func TestWithTimeoutTimedout(t *testing.T) {
+func TestWithTimeoutTimeout(t *testing.T) {
 	timeoutHandler := TimeoutHandler(time.Millisecond)
 	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		time.Sleep(time.Millisecond * 10)

rest/server.go

@@ -176,7 +176,8 @@ func WithCorsHeaders(headers ...string) RunOption {
 // WithCustomCors returns a func to enable CORS for given origin, or default to all origins (*),
 // fn lets caller customizing the response.
 func WithCustomCors(middlewareFn func(header http.Header), notAllowedFn func(http.ResponseWriter),
-	origin ...string) RunOption {
+	origin ...string,
+) RunOption {
 	return func(server *Server) {
 		server.router.SetNotAllowedHandler(cors.NotAllowedHandler(notAllowedFn, origin...))
 		server.router = newCorsRouter(server.router, middlewareFn, origin...)
@@ -326,6 +327,13 @@ func WithUnsignedCallback(callback handler.UnsignedCallback) RunOption {
 	}
 }
 
+// WithTimeoutCallback returns a RunOption that with given timeout callback set.
+func WithTimeoutCallback(callback handler.TimeoutCallback) RunOption {
+	return func(svr *Server) {
+		svr.ngin.setTimeoutCallback(callback)
+	}
+}
+
 func handleError(err error) {
 	// ErrServerClosed means the server is closed manually
 	if err == nil || errors.Is(err, http.ErrServerClosed) {

rest/server_test.go

@@ -68,6 +68,10 @@ Port: 0
 			c:    cnf,
 			opts: []RunOption{WithUnsignedCallback(nil), WithRouter(mockedRouter{})},
 		},
+		{
+			c:    cnf,
+			opts: []RunOption{WithTimeoutCallback(nil), WithRouter(mockedRouter{})},
+		},
 	}
 
 	for _, test := range tests {