diff --git a/server/middleware/timeout.go b/server/middleware/timeout.go new file mode 100644 index 00000000..473abf6c --- /dev/null +++ b/server/middleware/timeout.go @@ -0,0 +1,55 @@ +package middleware + +import ( + "context" + "github.com/gin-gonic/gin" + "net/http" + "time" +) + +// TimeoutMiddleware 创建超时中间件 +// 入参 timeout 设置超时时间(例如:time.Second * 5) +// 使用示例 xxx.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 + } + } +}