如何实现gin的timeout-Middleware

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

package buffpool

import (
	"bytes"
	"sync"
)

const BuffSize = 10 * 1024

var buffPool sync.Pool

func GetBuff() *bytes.Buffer {
	var buffer *bytes.Buffer
	item := buffPool.Get()
	if item == nil {
		var byteSlice []byte
		byteSlice = make([]byte, 0, BuffSize)
		buffer = bytes.NewBuffer(byteSlice)

	} else {
		buffer = item.(*bytes.Buffer)
	}
	return buffer
}

func PutBuff(buffer *bytes.Buffer) {
	buffer.Reset()
	buffPool.Put(buffer)
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

package timeout

import (
	"bytes"
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
	"sync"
)

type TimeoutWriter struct {
	gin.ResponseWriter
	// body
	body *bytes.Buffer
    // header
    // 让子协程和父协程分别写不同的header
	h http.Header

	mu          sync.Mutex
	timedOut    bool
    wroteHeader bool
    //让子协程和父协程分别写不同的code
	code        int
}

func (tw *TimeoutWriter) Write(b []byte) (int, error) {
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut {
        // 已经超时了，就不再写数据
		return 0, nil
	}

	return tw.body.Write(b)
}

func (tw *TimeoutWriter) WriteHeader(code int) {
	checkWriteHeaderCode(code)
	tw.mu.Lock()
	defer tw.mu.Unlock()
	if tw.timedOut || tw.wroteHeader {
		return
	}
	tw.writeHeader(code)
}

func (tw *TimeoutWriter) writeHeader(code int) {
	tw.wroteHeader = true
	tw.code = code
}

func (tw *TimeoutWriter) WriteHeaderNow() {}

func (tw *TimeoutWriter) Header() http.Header {
	return tw.h
}

func checkWriteHeaderCode(code int) {
	if code < 100 || code > 999 {
		panic(fmt.Sprintf("invalid WriteHeader code %v", code))
	}
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

package timeout

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/vearne/gin-timeout/buffpool"
	"net/http"
	"time"
)

const (
	HandlerFuncTimeout = "E509"
	ErrUnknowError     = "E003"
)

func Timeout(t time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		// sync.Pool
		buffer := buffpool.GetBuff()

		tw := &TimeoutWriter{body: buffer, ResponseWriter: c.Writer, h: make(http.Header)}
		c.Writer = tw

		// wrap the request context with a timeout
		ctx, cancel := context.WithTimeout(c.Request.Context(), t)
		defer cancel()
        //我们需要在超时发生时，通知子协程，让其也尽快退出:将context传入到Request中.
		c.Request = c.Request.WithContext(ctx)

		// Channel capacity must be greater than 0.
		// Otherwise, if the parent coroutine quit due to timeout,
        // the child coroutine may never be able to quit.
        // 具体的处理逻辑在子协程中完成
		finish := make(chan struct{}, 1)
		panicChan := make(chan interface{}, 1)
        go func() { 
        // 子协程只会将返回数据写入到内存buff中
        // 创建的子协程没有recover，存在程序崩溃的风险
			defer func() {
				if p := recover(); p != nil {
					panicChan <- p
				}
			}()
			c.Next()
			finish <- struct{}{}
		}()
        //主协程阻塞
        //1）或者等待deadline到达
        //2）或者接收到子协程的通知(通过channel)
		select {
		case p := <-panicChan:
			c.Abort()
			tw.ResponseWriter.WriteHeader(http.StatusInternalServerError)
			bt, _ := json.Marshal(errResponse{Code: ErrUnknowError,
				Msg: fmt.Sprintf("unknow internal error, %v", p)})
			tw.ResponseWriter.Write(bt)
        //如果dealline到达，返回给调用方HTTP Code 504

        //子协程的所有的所有输出,都只暂存在内存buff中。主协程退出后，这些数据也就没有机会被返回给调用方
		case <-ctx.Done():
			tw.mu.Lock()
			defer tw.mu.Unlock()
            // ****************注意***************
            // 子协程和父协程存在同时修改Header的风险
            // 由于Header是个map，可能诱发 
            // fatal error: concurrent map read and map write
			tw.ResponseWriter.WriteHeader(http.StatusServiceUnavailable)
			bt, _ := json.Marshal(errResponse{Code: HandlerFuncTimeout,
				Msg: http.ErrHandlerTimeout.Error()})
			tw.ResponseWriter.Write(bt)
			c.Abort()
            tw.timedOut = true
            // 如果超时的话，buffer无法主动清除，只能等待GC回收
			// If timeout happen, the buffer cannot be cleared actively,
            // but wait for the GC to recycle.
        //主协程收到子协程的消息, 将header头和内存buff中的数据，写入到流中，发送给调用方
		case <-finish:
			tw.mu.Lock()
			defer tw.mu.Unlock()
			dst := tw.ResponseWriter.Header()
			for k, vv := range tw.Header() {
				dst[k] = vv
			}
			tw.ResponseWriter.WriteHeader(tw.code)
			tw.ResponseWriter.Write(buffer.Bytes())
			buffpool.PutBuff(buffer)
		}
	}
}

type errResponse struct {
	Code string `json:"code"`
	Msg  string `json:"msg"`
}