-
Notifications
You must be signed in to change notification settings - Fork 3
/
bandwidth_stat_manager.go
338 lines (287 loc) · 9.04 KB
/
bandwidth_stat_manager.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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
package stats
import (
"encoding/json"
"github.com/TeaOSLab/EdgeCommon/pkg/nodeconfigs"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
"github.com/TeaOSLab/EdgeNode/internal/events"
"github.com/TeaOSLab/EdgeNode/internal/goman"
"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
"github.com/TeaOSLab/EdgeNode/internal/rpc"
"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
"github.com/iwind/TeaGo/Tea"
"github.com/iwind/TeaGo/logs"
"github.com/iwind/TeaGo/types"
timeutil "github.com/iwind/TeaGo/utils/time"
"os"
"sync"
"time"
)
var SharedBandwidthStatManager = NewBandwidthStatManager()
const bandwidthTimestampDelim = 2 // N秒平均,更为精确
func init() {
if !teaconst.IsMain {
return
}
events.On(events.EventLoaded, func() {
goman.New(func() {
SharedBandwidthStatManager.Start()
})
})
events.OnClose(func() {
SharedBandwidthStatManager.Cancel()
err := SharedBandwidthStatManager.Save()
if err != nil {
remotelogs.Error("STAT", "save bandwidth stats failed: "+err.Error())
}
})
}
type BandwidthStat struct {
Day string `json:"day"`
TimeAt string `json:"timeAt"`
UserId int64 `json:"userId"`
ServerId int64 `json:"serverId"`
CurrentBytes int64 `json:"currentBytes"`
CurrentTimestamp int64 `json:"currentTimestamp"`
MaxBytes int64 `json:"maxBytes"`
TotalBytes int64 `json:"totalBytes"`
CachedBytes int64 `json:"cachedBytes"`
AttackBytes int64 `json:"attackBytes"`
CountRequests int64 `json:"countRequests"`
CountCachedRequests int64 `json:"countCachedRequests"`
CountAttackRequests int64 `json:"countAttackRequests"`
CountWebsocketConnections int64 `json:"countWebsocketConnections"`
UserPlanId int64 `json:"userPlanId"`
}
// BandwidthStatManager 服务带宽统计
type BandwidthStatManager struct {
m map[string]*BandwidthStat // serverId@day@time => *BandwidthStat
pbStats []*pb.ServerBandwidthStat
lastTime string // 上一次执行的时间
ticker *time.Ticker
locker sync.Mutex
cacheFile string
}
func NewBandwidthStatManager() *BandwidthStatManager {
return &BandwidthStatManager{
m: map[string]*BandwidthStat{},
ticker: time.NewTicker(1 * time.Minute), // 时间小于1分钟是为了更快速地上传结果
cacheFile: Tea.Root + "/data/stat_bandwidth.cache",
}
}
func (this *BandwidthStatManager) Start() {
// 初始化DAU统计
{
err := SharedDAUManager.Init()
if err != nil {
remotelogs.Error("DAU_MANAGER", "initialize DAU manager failed: "+err.Error())
}
}
// 从上次数据中恢复
this.locker.Lock()
this.recover()
this.locker.Unlock()
// 循环上报数据
for range this.ticker.C {
err := this.Loop()
if err != nil && !rpc.IsConnError(err) {
remotelogs.Error("BANDWIDTH_STAT_MANAGER", err.Error())
}
}
}
func (this *BandwidthStatManager) Loop() error {
var regionId int64
nodeConfig, _ := nodeconfigs.SharedNodeConfig()
if nodeConfig != nil {
regionId = nodeConfig.RegionId
}
var now = time.Now()
var day = timeutil.Format("Ymd", now)
var currentTime = timeutil.FormatTime("Hi", now.Unix()/300*300) // 300s = 5 minutes
if this.lastTime == currentTime {
return nil
}
this.lastTime = currentTime
var pbStats = []*pb.ServerBandwidthStat{}
// 历史未提交记录
if len(this.pbStats) > 0 {
var expiredTime = timeutil.FormatTime("Hi", time.Now().Unix()-1200) // 只保留20分钟
for _, stat := range this.pbStats {
if stat.TimeAt > expiredTime {
pbStats = append(pbStats, stat)
}
}
this.pbStats = nil
}
var ipStatMap = SharedDAUManager.ReadStatMap()
this.locker.Lock()
for key, stat := range this.m {
if stat.Day < day || stat.TimeAt < currentTime {
// 防止数据出现错误
if stat.CachedBytes > stat.TotalBytes || stat.CountCachedRequests == stat.CountRequests {
stat.CachedBytes = stat.TotalBytes
}
if stat.AttackBytes > stat.TotalBytes {
stat.AttackBytes = stat.TotalBytes
}
var ipKey = "server_" + stat.Day + "_" + types.String(stat.ServerId)
pbStats = append(pbStats, &pb.ServerBandwidthStat{
Id: 0,
UserId: stat.UserId,
ServerId: stat.ServerId,
Day: stat.Day,
TimeAt: stat.TimeAt,
Bytes: stat.MaxBytes / bandwidthTimestampDelim,
TotalBytes: stat.TotalBytes,
CachedBytes: stat.CachedBytes,
AttackBytes: stat.AttackBytes,
CountRequests: stat.CountRequests,
CountCachedRequests: stat.CountCachedRequests,
CountAttackRequests: stat.CountAttackRequests,
CountWebsocketConnections: stat.CountWebsocketConnections,
CountIPs: ipStatMap[ipKey],
UserPlanId: stat.UserPlanId,
NodeRegionId: regionId,
})
delete(this.m, key)
}
}
this.locker.Unlock()
if len(pbStats) > 0 {
// 上传
rpcClient, err := rpc.SharedRPC()
if err != nil {
return err
}
_, err = rpcClient.ServerBandwidthStatRPC.UploadServerBandwidthStats(rpcClient.Context(), &pb.UploadServerBandwidthStatsRequest{ServerBandwidthStats: pbStats})
if err != nil {
this.pbStats = pbStats
return err
}
}
return nil
}
// AddBandwidth 添加带宽数据
func (this *BandwidthStatManager) AddBandwidth(userId int64, userPlanId int64, serverId int64, peekBytes int64, totalBytes int64) {
if serverId <= 0 || (peekBytes == 0 && totalBytes == 0) {
return
}
var now = fasttime.Now()
var timestamp = now.Unix() / bandwidthTimestampDelim * bandwidthTimestampDelim // 将时间戳均分成N等份
var day = now.Ymd()
var timeAt = now.Round5Hi()
var key = types.String(serverId) + "@" + day + "@" + timeAt
// 增加TCP Header尺寸,这里默认MTU为1500,且默认为IPv4
const mtu = 1500
const tcpHeaderSize = 20
if peekBytes > mtu {
peekBytes += peekBytes * tcpHeaderSize / mtu
}
this.locker.Lock()
stat, ok := this.m[key]
if ok {
// 此刻如果发生用户ID(userId)的变化也忽略,等N分钟后有新记录后再换
if stat.CurrentTimestamp == timestamp {
stat.CurrentBytes += peekBytes
} else {
stat.CurrentBytes = peekBytes
stat.CurrentTimestamp = timestamp
}
if stat.CurrentBytes > stat.MaxBytes {
stat.MaxBytes = stat.CurrentBytes
}
stat.TotalBytes += totalBytes
} else {
this.m[key] = &BandwidthStat{
Day: day,
TimeAt: timeAt,
UserId: userId,
UserPlanId: userPlanId,
ServerId: serverId,
CurrentBytes: peekBytes,
MaxBytes: peekBytes,
TotalBytes: totalBytes,
CurrentTimestamp: timestamp,
}
}
this.locker.Unlock()
}
// AddTraffic 添加请求数据
func (this *BandwidthStatManager) AddTraffic(serverId int64, cachedBytes int64, countRequests int64, countCachedRequests int64, countAttacks int64, attackBytes int64, countWebsocketConnections int64) {
var now = fasttime.Now()
var day = now.Ymd()
var timeAt = now.Round5Hi()
var key = types.String(serverId) + "@" + day + "@" + timeAt
this.locker.Lock()
// 只有有记录了才会添加
stat, ok := this.m[key]
if ok {
stat.CachedBytes += cachedBytes
stat.CountRequests += countRequests
stat.CountCachedRequests += countCachedRequests
stat.CountAttackRequests += countAttacks
stat.AttackBytes += attackBytes
stat.CountWebsocketConnections += countWebsocketConnections
}
this.locker.Unlock()
}
func (this *BandwidthStatManager) Inspect() {
this.locker.Lock()
logs.PrintAsJSON(this.m)
this.locker.Unlock()
}
func (this *BandwidthStatManager) Map() map[int64]int64 /** serverId => max bytes **/ {
this.locker.Lock()
defer this.locker.Unlock()
var m = map[int64]int64{}
for _, v := range this.m {
m[v.ServerId] = v.MaxBytes / bandwidthTimestampDelim
}
return m
}
// Save 保存到本地磁盘
func (this *BandwidthStatManager) Save() error {
this.locker.Lock()
defer this.locker.Unlock()
if len(this.m) == 0 {
return nil
}
data, err := json.Marshal(this.m)
if err != nil {
return err
}
_ = os.Remove(this.cacheFile)
return os.WriteFile(this.cacheFile, data, 0666)
}
// Cancel 取消上传
func (this *BandwidthStatManager) Cancel() {
this.ticker.Stop()
}
// 从本地缓存文件中恢复数据
func (this *BandwidthStatManager) recover() {
cacheData, err := os.ReadFile(this.cacheFile)
if err == nil {
var m = map[string]*BandwidthStat{}
err = json.Unmarshal(cacheData, &m)
if err == nil && len(m) > 0 {
var lastTime = ""
for _, stat := range m {
if stat.Day != fasttime.Now().Ymd() {
continue
}
if len(lastTime) == 0 || stat.TimeAt > lastTime {
lastTime = stat.TimeAt
}
}
if len(lastTime) > 0 {
var availableTime = timeutil.FormatTime("Hi", (time.Now().Unix()-300) /** 只保留5分钟的 **/ /300*300) // 300s = 5 minutes
if lastTime >= availableTime {
this.m = m
this.lastTime = lastTime
}
}
}
_ = os.Remove(this.cacheFile)
}
}