-
Notifications
You must be signed in to change notification settings - Fork 9
/
timeout.go
83 lines (66 loc) · 1.78 KB
/
timeout.go
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
package timeout
import (
"context"
"errors"
"net/http"
"strconv"
"time"
"github.com/gin-gonic/gin"
)
type contextKey uint64
const (
startKey contextKey = iota
)
const (
timeoutKey = "Timeout-Millisecond"
)
// TimeoutMiddleware is used to pass the remaining timeout
func TimeoutMiddleware(timeout time.Duration) func(c *gin.Context) {
return func(c *gin.Context) {
remain := timeout
headerTimeout := c.Request.Header.Get(timeoutKey)
if headerTimeout != "" {
t, _ := strconv.ParseInt(headerTimeout, 10, 64)
remain = time.Duration(t) * time.Millisecond
}
ctx, cancel := context.WithTimeout(c.Request.Context(), remain)
_ = cancel
ctx = SetStart(ctx, remain.Milliseconds())
c.Request = c.Request.WithContext(ctx)
c.Next()
}
}
// SetStart for marking current service request start
func SetStart(ctx context.Context, timeout int64) context.Context {
ctx = context.WithValue(ctx, timeoutKey, timeout)
return context.WithValue(ctx, startKey, nowMillisecond())
}
// CalcRemainTimeout is used to calculate remain timeout
func CalcRemainTimeout(ctx context.Context) (int64, error) {
timeout, ok := ctx.Value(timeoutKey).(int64)
if !ok {
return 0, nil
}
startAt, ok := ctx.Value(startKey).(int64)
if !ok {
return 0, errors.New("miss startAt")
}
remain := timeout - (nowMillisecond() - startAt)
if remain < 0 {
return 0, errors.New("timeout < diff, context deadline exceeded")
}
return remain, nil
}
// SetHeader save timeout field to http.Header
func SetHeader(ctx context.Context, header http.Header) (err error) {
remain, err := CalcRemainTimeout(ctx)
if err != nil {
return
}
header.Set(timeoutKey, strconv.FormatInt(remain, 10))
return
}
// nowMillisecond return now time nanosecond
func nowMillisecond() int64 {
return time.Now().UnixNano() / 1e6
}