一个较为完善的服务，必须要考虑到业务处理超时导致的出现雪崩现象的可能。为此，开发了一个简单实用的超时中间件，项目地址：https://github.com/cyj19/gin-timeout

1. 思路

在网上搜索没有发现现成的可用代码，但是超时控制明显是一个常见的问题，不可能只有我遇到。后来换了思路，搜索Golang如何进行超时控制处理，发现标准库net/http下的server.go中有一个TimeoutHandler，模仿该Handler实现超时中间件。
核心思路：把业务处理的响应结果存放到缓冲区而不是直接写入响应流，在中间件判断是否超时来选择把哪些内容写入响应流

2. 实现

有了思路，效仿之前的访问日志中间件，重写ResponseWriter涉及的方法，把结果存储到缓冲区。需要特别注意的是要把业务处理返回的响应头也是存储到缓冲区，否则会出现超时返回和正常业务返回同时写响应头的问题
writer.go

  1. // 这部分可以参照net/http server.go中的timeoutWriter来写，稍微修改一下就行
  3. // 在不入侵业务代码的情况下，我们是无法阻止业务结果写入响应的，因此需要自定义Writer，
  4. // 把响应头和响应内容都只存放到缓冲区，不直接写入响应流
  6. type timeoutWriter struct {
  7.     gin.ResponseWriter
  8.     h           http.Header   // response header
  9.     wbuf        *bytes.Buffer // response content
  10.     mu          sync.Mutex
  11.     timedOut    bool
  12.     wroteHeader bool
  13.     code        int // response code
  14. }
  16. func (tw *timeoutWriter) Header() http.Header {
  17.     return tw.h
  18. }
  20. func checkWriteHeaderCode(code int) {
  21.     if code < 100 || code > 999 {
  22.         panic(fmt.Sprintf("invalid WriteHeader code %v", code))
  23.     }
  24. }
  26. func (tw *timeoutWriter) writeHeaderLocked(code int) {
  27.     checkWriteHeaderCode(code)
  29.     switch {
  30.     case tw.timedOut:
  31.         return
  32.     case tw.wroteHeader:
  33.         return
  34.     default:
  35.         tw.wroteHeader = true
  36.         tw.code = code
  37.         return
  38.     }
  39. }
  41. func (tw *timeoutWriter) WriteHeader(code int) {
  42.     tw.mu.Lock()
  43.     defer tw.mu.Unlock()
  44.     tw.writeHeaderLocked(code)
  45. }
  47. func (tw *timeoutWriter) Write(p []byte) (int, error) {
  48.     tw.mu.Lock()
  49.     defer tw.mu.Unlock()
  50.     if tw.timedOut {
  51.         return 0, nil
  52.     }
  53.     if !tw.wroteHeader {
  54.         tw.WriteHeader(http.StatusOK)
  55.     }
  56.     // normal response content is written to wbuf
  57.     return tw.wbuf.Write(p)
  58. }

option.go

  1. // Option timeout configuration
  2. type Option struct {
  3.     Timeout *time.Duration // The timeout time is generally configured in the configuration file. In order to facilitate hot update, use the pointer
  4.     Code    int            // response code
  5.     Msg     string         // response message
  6. }

timeout.go

  1. // 这部分可以参考net/http server.go的TimeHandler来写
  3. func ContextTimeout(opt Option) gin.HandlerFunc {
  4.     return func(c *gin.Context) {
  5.         var tw = &timeoutWriter{
  6.             wbuf:           bytes.NewBufferString(""),
  7.             ResponseWriter: c.Writer,
  8.             h:              make(http.Header),
  9.         }
  10.         c.Writer = tw
  11.         ctx, cancel := context.WithTimeout(c.Request.Context(), *opt.Timeout)
  12.         defer cancel()
  13.         // 通知业务处理结束
  14.         done := make(chan struct{})
  15.         // 通知发生恐慌
  16.         panicChan := make(chan interface{}, 1)
  17.         go func() {
  18.             // 恐慌恢复
  19.             defer func() {
  20.                 if p := recover(); p != nil {
  21.                     panicChan <- p
  22.                 }
  23.             }()
  24.             c.Next()
  25.             close(done)
  26.         }()
  28.         select {
  29.         case p := <-panicChan:
  30.             panic(p)
  31.         case <-ctx.Done():
  32.             tw.mu.Lock()
  33.             defer tw.mu.Unlock()
  34.             tw.ResponseWriter.WriteHeader(opt.Code)
  35.             _, _ = tw.ResponseWriter.Write([]byte(opt.Msg))
  36.             tw.timedOut = true
  37.             c.Abort()
  38.         case <-done:
  39.             tw.mu.Lock()
  40.             defer tw.mu.Unlock()
  41.             dst := tw.ResponseWriter.Header()
  42.             // add the header of timeoutWriter to the original header
  43.             for k, vv := range tw.h {
  44.                 dst[k] = vv
  45.             }
  46.             if !tw.wroteHeader {
  47.                 tw.code = http.StatusOK
  48.             }
  49.             tw.ResponseWriter.WriteHeader(tw.code)
  50.             _, _ = tw.ResponseWriter.Write(tw.wbuf.Bytes())
  51.         }
  52.     }
  53. }

3. 总结

习惯了遇到问题就想找现成的代码，其实超时中间件和之前的访问日志中间件的实现原理都是大同小异的。多思考，多思考，多思考，不要吝啬这些时间。