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"))
}