53 lines
962 B
Go
53 lines
962 B
Go
// Filename: internal/middleware/timeout.go
|
|
|
|
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func TimeoutMiddleware(timeout time.Duration) gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
// 创建带超时的 context
|
|
ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
|
|
defer cancel()
|
|
|
|
// 替换 request context
|
|
c.Request = c.Request.WithContext(ctx)
|
|
|
|
// 使用 channel 等待请求完成
|
|
finished := make(chan struct{})
|
|
|
|
go func() {
|
|
c.Next()
|
|
close(finished)
|
|
}()
|
|
|
|
select {
|
|
case <-finished:
|
|
// 请求正常完成
|
|
return
|
|
case <-ctx.Done():
|
|
// 超时
|
|
c.AbortWithStatusJSON(http.StatusGatewayTimeout, gin.H{
|
|
"error": "Request timeout",
|
|
"code": "TIMEOUT",
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
// 使用示例
|
|
func SetupTimeout(r *gin.Engine) {
|
|
// 对 API 路由设置 30 秒超时
|
|
api := r.Group("/api")
|
|
api.Use(TimeoutMiddleware(30 * time.Second))
|
|
{
|
|
// ... API routes
|
|
}
|
|
}
|