Gin增加超时控制
背景
- 有时候很多API有业务超时要求,就是数据必须在几秒内返回,超时的话立刻结束,不用卡死在某一个页面上。
目标
- 对这一类API特定处理,增加超时控制,不要影响别的API
方案
- 方案很简单,就是利用go提供的context技术对gin框架的request对象进行改造,使其具备超时控制机能。
代码:
package main import ( "context" "fmt" "log" "net/http" "time" "github.com/gin-gonic/gin" ) // 超时控制中间件 func timeoutMiddleware(timeout time.Duration) func(c *gin.Context) { return func(c *gin.Context) { // 用超时context wrap request的context ctx, cancel := context.WithTimeout(c.Request.Context(), timeout) defer func() { // 检查是否超时 if ctx.Err() == context.DeadlineExceeded { c.Writer.WriteHeader(http.StatusGatewayTimeout) c.Abort() } //清理资源 cancel() }() // 替换 c.Request = c.Request.WithContext(ctx) c.Next() } } func timedHandler(duration time.Duration) func(c *gin.Context) { return func(c *gin.Context) { // 获取替换之后的context 它具备了超时控制 ctx := c.Request.Context() // 定义响应struct type responseData struct { status int body map[string]interface{} } // 创建一个done chan表明request要完成了 doneChan := make(chan responseData) // 模拟API耗时的处理 go func() { time.Sleep(duration) doneChan <- responseData{ status: 200, body: gin.H{"hello": "world"}, } }() // 监听两个chan谁先到达 select { // 超时 case <-ctx.Done(): return // 请求完成 case res := <-doneChan: c.JSON(res.status, res.body) } } } func main() { engine := gin.New() // 不用超时控制 group := engine.Group("/normal") { group.GET("/short", timedHandler(time.Second)) } // 需要超时控制 时间控制在2s timeoutGroup := engine.Group("/timeout") timeoutGroup.Use(timeoutMiddleware(time.Second * 2)) { // 延迟5s timeoutGroup.GET("/long", timedHandler(time.Second*5)) } // run the server log.Fatal(engine.Run(":5090")) }
验证
- 正常请求
- 超时请求