1 package middleware 2 3 import ( 4 "context" 5 "net/http" 6 "time" 7 ) 8 9 // Timeout is a middleware that cancels ctx after a given timeout and return 10 // a 504 Gateway Timeout error to the client. 11 // 12 // It's required that you select the ctx.Done() channel to check for the signal 13 // if the context has reached its deadline and return, otherwise the timeout 14 // signal will be just ignored. 15 // 16 // ie. a route/handler may look like: 17 // 18 // r.Get("/long", func(w http.ResponseWriter, r *http.Request) { 19 // ctx := r.Context() 20 // processTime := time.Duration(rand.Intn(4)+1) * time.Second 21 // 22 // select { 23 // case <-ctx.Done(): 24 // return 25 // 26 // case <-time.After(processTime): 27 // // The above channel simulates some hard work. 28 // } 29 // 30 // w.Write([]byte("done")) 31 // }) 32 // 33 func Timeout(timeout time.Duration) func(next http.Handler) http.Handler { 34 return func(next http.Handler) http.Handler { 35 fn := func(w http.ResponseWriter, r *http.Request) { 36 ctx, cancel := context.WithTimeout(r.Context(), timeout) 37 defer func() { 38 cancel() 39 if ctx.Err() == context.DeadlineExceeded { 40 w.WriteHeader(http.StatusGatewayTimeout) 41 } 42 }() 43 44 r = r.WithContext(ctx) 45 next.ServeHTTP(w, r) 46 } 47 return http.HandlerFunc(fn) 48 } 49 } 50