/
initialization.go
86 lines (73 loc) · 1.73 KB
/
initialization.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
84
85
86
package middleware
import (
"strings"
"github.com/vicanso/pike/performance"
"github.com/vicanso/pike/pike"
"github.com/vicanso/pike/util"
)
const (
defaultConcurrency = 256 * 1000
)
type (
// InitializationConfig 初始化配置
InitializationConfig struct {
Header []string
RequestHeader []string
Concurrency int
}
)
func genHeader(header []string) map[string]string {
m := make(map[string]string)
// 将自定义的http response header格式化
for _, v := range header {
arr := strings.Split(v, ":")
if len(arr) != 2 {
continue
}
value := arr[1]
v := util.CheckAndGetValueFromEnv(value)
if len(v) != 0 {
value = v
}
m[arr[0]] = value
}
if len(m) == 0 {
return nil
}
return m
}
// Initialization 相关一些初始化的操作
func Initialization(config InitializationConfig) pike.Middleware {
customHeader := genHeader(config.Header)
customReqHeader := genHeader(config.RequestHeader)
// 获取限制并发请求数
concurrency := uint32(defaultConcurrency)
if config.Concurrency != 0 {
concurrency = uint32(config.Concurrency)
}
return func(c *pike.Context, next pike.Next) error {
done := c.ServerTiming.Start(pike.ServerTimingInitialization)
performance.IncreaseRequestCount()
defer func() {
performance.DecreaseConcurrency()
status := c.Response.Status()
use := util.GetTimeConsuming(c.CreatedAt)
performance.AddRequestStats(status, use)
}()
resHeader := c.Response.Header()
for k, v := range customHeader {
resHeader.Add(k, v)
}
reqHeader := c.Request.Header
for k, v := range customReqHeader {
reqHeader.Add(k, v)
}
v := performance.IncreaseConcurrency()
if v > concurrency {
done()
return ErrTooManyRequest
}
done()
return next()
}
}