增加服务带宽统计

internal/rpc/services/service_server_bandwidth_stat.go

@@ -0,0 +1,80 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package services
+
+import (
+	"context"
+	"github.com/TeaOSLab/EdgeAPI/internal/db/models/stats"
+	"github.com/TeaOSLab/EdgeAPI/internal/goman"
+	"github.com/TeaOSLab/EdgeAPI/internal/remotelogs"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/iwind/TeaGo/Tea"
+	"github.com/iwind/TeaGo/dbs"
+	"github.com/iwind/TeaGo/types"
+	"sync"
+	"time"
+)
+
+var serverBandwidthStatsMap = map[string]*pb.ServerBandwidthStat{} // key => bandwidth
+var serverBandwidthStatsLocker = &sync.Mutex{}
+
+func init() {
+	var ticker = time.NewTicker(5 * time.Minute)
+	if Tea.IsTesting() {
+		ticker = time.NewTicker(1 * time.Minute)
+	}
+
+	dbs.OnReadyDone(func() {
+		goman.New(func() {
+			for range ticker.C {
+				func() {
+					var tx *dbs.Tx
+
+					serverBandwidthStatsLocker.Lock()
+					var m = serverBandwidthStatsMap
+					serverBandwidthStatsMap = map[string]*pb.ServerBandwidthStat{}
+					serverBandwidthStatsLocker.Unlock()
+
+					for _, stat := range m {
+						err := stats.SharedServerBandwidthStatDAO.UpdateServerBandwidth(tx, stat.ServerId, stat.Day, stat.TimeAt, stat.Bytes)
+						if err != nil {
+							remotelogs.Error("ServerBandwidthStatService", "dump bandwidth stats failed: "+err.Error())
+						}
+					}
+				}()
+			}
+		})
+	})
+}
+
+type ServerBandwidthStatService struct {
+	BaseService
+}
+
+// UploadServerBandwidthStats 上传带宽统计
+func (this *ServerBandwidthStatService) UploadServerBandwidthStats(ctx context.Context, req *pb.UploadServerBandwidthStatsRequest) (*pb.RPCSuccess, error) {
+	_, err := this.ValidateNode(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, stat := range req.ServerBandwidthStats {
+		var key = types.String(stat.ServerId) + "@" + stat.Day + "@" + stat.TimeAt
+		serverBandwidthStatsLocker.Lock()
+		oldStat, ok := serverBandwidthStatsMap[key]
+		if ok {
+			oldStat.Bytes += stat.Bytes
+		} else {
+			serverBandwidthStatsMap[key] = &pb.ServerBandwidthStat{
+				Id:       0,
+				ServerId: stat.ServerId,
+				Day:      stat.Day,
+				TimeAt:   stat.TimeAt,
+				Bytes:    stat.Bytes,
+			}
+		}
+		serverBandwidthStatsLocker.Unlock()
+	}
+
+	return this.Success()
+}