代码:
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"))
}