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