timeout.go 1.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. package middleware
  2. import (
  3. "context"
  4. "github.com/gin-gonic/gin"
  5. "net/http"
  6. "time"
  7. )
  8. // TimeoutMiddleware 创建超时中间件
  9. // 入参 timeout 设置超时时间(例如:time.Second * 5)
  10. // 使用示例 xxx.Get("path",middleware.TimeoutMiddleware(30*time.Second),HandleFunc)
  11. func TimeoutMiddleware(timeout time.Duration) gin.HandlerFunc {
  12. return func(c *gin.Context) {
  13. ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
  14. defer cancel()
  15. c.Request = c.Request.WithContext(ctx)
  16. // 使用 buffered channel 避免 goroutine 泄漏
  17. done := make(chan struct{}, 1)
  18. panicChan := make(chan interface{}, 1)
  19. go func() {
  20. defer func() {
  21. if p := recover(); p != nil {
  22. select {
  23. case panicChan <- p:
  24. default:
  25. }
  26. }
  27. select {
  28. case done <- struct{}{}:
  29. default:
  30. }
  31. }()
  32. c.Next()
  33. }()
  34. select {
  35. case p := <-panicChan:
  36. panic(p)
  37. case <-done:
  38. return
  39. case <-ctx.Done():
  40. // 确保服务器超时设置足够长
  41. c.Header("Connection", "close")
  42. c.AbortWithStatusJSON(http.StatusGatewayTimeout, gin.H{
  43. "code": 504,
  44. "msg": "请求超时",
  45. })
  46. return
  47. }
  48. }
  49. }