56 lines
1.1 KiB
Go
56 lines
1.1 KiB
Go
package middleware
|
||
|
||
import (
|
||
"context"
|
||
"net/http"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
// TimeoutMiddleware 创建超时中间件
|
||
// timeout: 超时时间（例如：time.Second * 30）
|
||
// 使用示例: router.GET("path", middleware.TimeoutMiddleware(30*time.Second), HandleFunc)
|
||
func TimeoutMiddleware(timeout time.Duration) gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
|
||
defer cancel()
|
||
|
||
c.Request = c.Request.WithContext(ctx)
|
||
|
||
// 使用 buffered channel 避免 goroutine 泄漏
|
||
done := make(chan struct{}, 1)
|
||
panicChan := make(chan interface{}, 1)
|
||
|
||
go func() {
|
||
defer func() {
|
||
if p := recover(); p != nil {
|
||
select {
|
||
case panicChan <- p:
|
||
default:
|
||
}
|
||
}
|
||
select {
|
||
case done <- struct{}{}:
|
||
default:
|
||
}
|
||
}()
|
||
c.Next()
|
||
}()
|
||
|
||
select {
|
||
case p := <-panicChan:
|
||
panic(p)
|
||
case <-done:
|
||
return
|
||
case <-ctx.Done():
|
||
c.Header("Connection", "close")
|
||
c.AbortWithStatusJSON(http.StatusGatewayTimeout, gin.H{
|
||
"code": 504,
|
||
"msg": "请求超时",
|
||
})
|
||
return
|
||
}
|
||
}
|
||
}
|