1.4.5.2

EdgeDNS/internal/stats/manager.go

@@ -0,0 +1,109 @@
+// 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
+}