110 lines
2.3 KiB
Go
110 lines
2.3 KiB
Go
// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
|
|
|
|
package stats
|
|
|
|
import (
|
|
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
|
|
teaconst "github.com/TeaOSLab/EdgeDNS/internal/const"
|
|
"github.com/TeaOSLab/EdgeDNS/internal/events"
|
|
"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
|
|
"github.com/TeaOSLab/EdgeDNS/internal/rpc"
|
|
"github.com/iwind/TeaGo/Tea"
|
|
timeutil "github.com/iwind/TeaGo/utils/time"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var SharedManager = NewStatManager()
|
|
|
|
func init() {
|
|
if !teaconst.IsMain {
|
|
return
|
|
}
|
|
|
|
events.On(events.EventStart, func() {
|
|
SharedManager.Start()
|
|
})
|
|
}
|
|
|
|
type StatManager struct {
|
|
statMap map[string]*pb.NSRecordHourlyStat // record@hour => *NSRecordHourlyStat
|
|
locker sync.Mutex
|
|
ticker *time.Ticker
|
|
}
|
|
|
|
func NewStatManager() *StatManager {
|
|
return &StatManager{
|
|
statMap: map[string]*pb.NSRecordHourlyStat{},
|
|
}
|
|
}
|
|
|
|
func (this *StatManager) Start() {
|
|
this.ticker = time.NewTicker(5 * time.Minute)
|
|
if Tea.IsTesting() {
|
|
this.ticker = time.NewTicker(1 * time.Minute)
|
|
}
|
|
go func() {
|
|
for range this.ticker.C {
|
|
err := this.Loop()
|
|
if err != nil {
|
|
if rpc.IsConnError(err) {
|
|
remotelogs.Debug("STAT", "upload failed: "+err.Error())
|
|
} else {
|
|
remotelogs.Error("STAT", "upload failed: "+err.Error())
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (this *StatManager) Add(domainId int64, recordId int64, bytes int64) {
|
|
this.locker.Lock()
|
|
defer this.locker.Unlock()
|
|
|
|
var hour = timeutil.Format("H")
|
|
var key = strconv.FormatInt(recordId, 10) + "@" + hour
|
|
stat, ok := this.statMap[key]
|
|
if ok {
|
|
stat.CountRequests++
|
|
stat.Bytes += bytes
|
|
} else {
|
|
this.statMap[key] = &pb.NSRecordHourlyStat{
|
|
NsDomainId: domainId,
|
|
NsRecordId: recordId,
|
|
Bytes: bytes,
|
|
CountRequests: 1,
|
|
CreatedAt: time.Now().Unix(),
|
|
}
|
|
}
|
|
}
|
|
|
|
func (this *StatManager) Loop() error {
|
|
rpcClient, err := rpc.SharedRPC()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
this.locker.Lock()
|
|
var m = this.statMap
|
|
this.statMap = map[string]*pb.NSRecordHourlyStat{}
|
|
this.locker.Unlock()
|
|
|
|
if len(m) == 0 {
|
|
return nil
|
|
}
|
|
|
|
var stats = []*pb.NSRecordHourlyStat{}
|
|
for _, stat := range m {
|
|
stats = append(stats, stat)
|
|
}
|
|
|
|
_, err = rpcClient.NSRecordHourlyStatRPC.UploadNSRecordHourlyStats(rpcClient.Context(), &pb.UploadNSRecordHourlyStatsRequest{
|
|
Stats: stats,
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|