12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849 |
- package middleware
-
- import (
- "context"
- "net/http"
- "time"
- )
-
- // Timeout is a middleware that cancels ctx after a given timeout and return
- // a 504 Gateway Timeout error to the client.
- //
- // It's required that you select the ctx.Done() channel to check for the signal
- // if the context has reached its deadline and return, otherwise the timeout
- // signal will be just ignored.
- //
- // ie. a route/handler may look like:
- //
- // r.Get("/long", func(w http.ResponseWriter, r *http.Request) {
- // ctx := r.Context()
- // processTime := time.Duration(rand.Intn(4)+1) * time.Second
- //
- // select {
- // case <-ctx.Done():
- // return
- //
- // case <-time.After(processTime):
- // // The above channel simulates some hard work.
- // }
- //
- // w.Write([]byte("done"))
- // })
- //
- func Timeout(timeout time.Duration) func(next http.Handler) http.Handler {
- return func(next http.Handler) http.Handler {
- fn := func(w http.ResponseWriter, r *http.Request) {
- ctx, cancel := context.WithTimeout(r.Context(), timeout)
- defer func() {
- cancel()
- if ctx.Err() == context.DeadlineExceeded {
- w.WriteHeader(http.StatusGatewayTimeout)
- }
- }()
-
- r = r.WithContext(ctx)
- next.ServeHTTP(w, r)
- }
- return http.HandlerFunc(fn)
- }
- }
|