实现单个服务的带宽限制（商业版）

internal/db/models/server_daily_stat_dao.go

@@ -1,6 +1,7 @@
 package models
 
 import (
+	"github.com/TeaOSLab/EdgeAPI/internal/errors"
 	"github.com/TeaOSLab/EdgeAPI/internal/utils"
 	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
 	_ "github.com/go-sql-driver/mysql"
@@ -53,12 +54,14 @@ func init() {
 // SaveStats 提交数据
 func (this *ServerDailyStatDAO) SaveStats(tx *dbs.Tx, stats []*pb.ServerDailyStat) error {
 	var serverUserMap = map[int64]int64{} // serverId => userId
+	var cacheMap = maps.Map{}
 	for _, stat := range stats {
 		day := timeutil.FormatTime("Ymd", stat.CreatedAt)
 		hour := timeutil.FormatTime("YmdH", stat.CreatedAt)
 		timeFrom := timeutil.FormatTime("His", stat.CreatedAt)
 		timeTo := timeutil.FormatTime("His", stat.CreatedAt+5*60-1) // 5分钟
 
+		// 所属用户
 		serverUserId, ok := serverUserMap[stat.ServerId]
 		if !ok {
 			userId, err := SharedServerDAO.FindServerUserId(tx, stat.ServerId)
@@ -100,7 +103,20 @@ func (this *ServerDailyStatDAO) SaveStats(tx *dbs.Tx, stats []*pb.ServerDailySta
 		if err != nil {
 			return err
 		}
+
+		// 更新带宽限制状态
+		bandwidthLimit, err := SharedServerDAO.FindServerBandwidthLimitConfig(tx, stat.ServerId, cacheMap)
+		if err != nil {
+			return err
 		}
+		if bandwidthLimit != nil && bandwidthLimit.IsOn && !bandwidthLimit.IsEmpty() {
+			err = SharedServerDAO.UpdateServerBandwidthLimitStatus(tx, bandwidthLimit, stat.ServerId, false)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -233,7 +249,7 @@ func (this *ServerDailyStatDAO) SumDailyStat(tx *dbs.Tx, serverId int64, day str
 	stat = &pb.ServerDailyStat{}
 
 	if !regexp.MustCompile(`^\d{8}$`).MatchString(day) {
-		return
+		return nil, errors.New("invalid day '" + day + "'")
 	}
 
 	one, _, err := this.Query(tx).
@@ -258,6 +274,37 @@ func (this *ServerDailyStatDAO) SumDailyStat(tx *dbs.Tx, serverId int64, day str
 	return
 }
 
+// SumMonthlyStat 获取某月内的流量
+// month 格式为YYYYMM
+func (this *ServerDailyStatDAO) SumMonthlyStat(tx *dbs.Tx, serverId int64, month string) (stat *pb.ServerDailyStat, err error) {
+	stat = &pb.ServerDailyStat{}
+
+	if !regexp.MustCompile(`^\d{6}$`).MatchString(month) {
+		return
+	}
+
+	one, _, err := this.Query(tx).
+		Result("SUM(bytes) AS bytes, SUM(cachedBytes) AS cachedBytes, SUM(countRequests) AS countRequests, SUM(countCachedRequests) AS countCachedRequests, SUM(countAttackRequests) AS countAttackRequests, SUM(attackBytes) AS attackBytes").
+		Attr("serverId", serverId).
+		Between("day", month+"01", month+"31").
+		FindOne()
+	if err != nil {
+		return nil, err
+	}
+
+	if one == nil {
+		return
+	}
+
+	stat.Bytes = one.GetInt64("bytes")
+	stat.CachedBytes = one.GetInt64("cachedBytes")
+	stat.CountRequests = one.GetInt64("countRequests")
+	stat.CountCachedRequests = one.GetInt64("countCachedRequests")
+	stat.CountAttackRequests = one.GetInt64("countAttackRequests")
+	stat.AttackBytes = one.GetInt64("attackBytes")
+	return
+}
+
 // FindDailyStats 按天统计
 func (this *ServerDailyStatDAO) FindDailyStats(tx *dbs.Tx, serverId int64, dayFrom string, dayTo string) (result []*ServerDailyStat, err error) {
 	ones, err := this.Query(tx).

internal/db/models/server_dao.go

@@ -3,6 +3,7 @@ package models
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"github.com/TeaOSLab/EdgeAPI/internal/db/models/dns"
 	"github.com/TeaOSLab/EdgeAPI/internal/utils/numberutils"
 	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
@@ -16,6 +17,7 @@ import (
 	"github.com/iwind/TeaGo/maps"
 	"github.com/iwind/TeaGo/rands"
 	"github.com/iwind/TeaGo/types"
+	timeutil "github.com/iwind/TeaGo/utils/time"
 	"regexp"
 	"strconv"
 	"strings"
@@ -50,7 +52,7 @@ func init() {
 
 // Init 初始化
 func (this *ServerDAO) Init() {
-	this.DAOObject.Init()
+	_ = this.DAOObject.Init()
 
 	// 这里不处理增删改事件，是为了避免Server修改本身的时候，也要触发别的Server变更
 }
@@ -1025,6 +1027,29 @@ func (this *ServerDAO) ComposeServerConfig(tx *dbs.Tx, server *Server, cacheMap
 		config.HTTPCachePolicyId = httpCachePolicyId
 	}
 
+	// bandwidth limit
+	if len(server.BandwidthLimit) > 0 {
+		var bandwidthLimitConfig = &serverconfigs.BandwidthLimitConfig{}
+		err = json.Unmarshal([]byte(server.BandwidthLimit), bandwidthLimitConfig)
+		if err != nil {
+			return nil, err
+		}
+		config.BandwidthLimit = bandwidthLimitConfig
+
+		if bandwidthLimitConfig.IsOn && !bandwidthLimitConfig.IsEmpty() {
+			if len(server.BandwidthLimitStatus) > 0 {
+				var status = &serverconfigs.BandwidthLimitStatus{}
+				err = json.Unmarshal([]byte(server.BandwidthLimitStatus), status)
+				if err != nil {
+					return nil, err
+				}
+				if status.IsValid() {
+					config.BandwidthLimitStatus = status
+				}
+			}
+		}
+	}
+
 	return config, nil
 }
 
@@ -1738,6 +1763,170 @@ func (this *ServerDAO) NotifyServerPortsUpdate(tx *dbs.Tx, serverId int64) error
 		UpdateQuickly()
 }
 
+// FindServerBandwidthLimitConfig 查找服务的带宽限制
+func (this *ServerDAO) FindServerBandwidthLimitConfig(tx *dbs.Tx, serverId int64, cacheMap maps.Map) (*serverconfigs.BandwidthLimitConfig, error) {
+	if cacheMap == nil {
+		cacheMap = maps.Map{}
+	}
+	var cacheKey = this.Table + ":FindServerBandwidthLimitConfig:" + types.String(serverId)
+	result, ok := cacheMap[cacheKey]
+	if ok {
+		return result.(*serverconfigs.BandwidthLimitConfig), nil
+	}
+
+	bandwidthLimit, err := this.Query(tx).
+		Pk(serverId).
+		Result("bandwidthLimit").
+		FindStringCol("")
+	if err != nil {
+		return nil, err
+	}
+
+	var limit = &serverconfigs.BandwidthLimitConfig{}
+	if len(bandwidthLimit) == 0 {
+		return limit, nil
+	}
+	err = json.Unmarshal([]byte(bandwidthLimit), limit)
+	if err != nil {
+		return nil, err
+	}
+
+	cacheMap[cacheKey] = limit
+
+	return limit, nil
+}
+
+// UpdateServerBandwidthLimitConfig 修改服务的带宽限制
+func (this *ServerDAO) UpdateServerBandwidthLimitConfig(tx *dbs.Tx, serverId int64, bandwidthLimitConfig *serverconfigs.BandwidthLimitConfig) error {
+	if serverId <= 0 {
+		return errors.New("invalid serverId")
+	}
+	limitJSON, err := json.Marshal(bandwidthLimitConfig)
+	if err != nil {
+		return err
+	}
+
+	err = this.Query(tx).
+		Pk(serverId).
+		Set("bandwidthLimit", limitJSON).
+		UpdateQuickly()
+	if err != nil {
+		return err
+	}
+
+	// 更新状态
+	return this.UpdateServerBandwidthLimitStatus(tx, bandwidthLimitConfig, serverId, true)
+}
+
+func (this *ServerDAO) UpdateServerBandwidthLimitStatus(tx *dbs.Tx, bandwidthLimitConfig *serverconfigs.BandwidthLimitConfig, serverId int64, isUpdatingConfig bool) error {
+	if !bandwidthLimitConfig.IsOn {
+		if isUpdatingConfig {
+			return this.NotifyUpdate(tx, serverId)
+		}
+		return nil
+	}
+
+	oldStatusString, err := this.Query(tx).
+		Pk(serverId).
+		Result("bandwidthLimitStatus").
+		FindStringCol("")
+	if err != nil {
+		return err
+	}
+	var oldStatus = &serverconfigs.BandwidthLimitStatus{}
+	if len(oldStatusString) > 0 {
+		err = json.Unmarshal([]byte(oldStatusString), oldStatus)
+		if err != nil {
+			return err
+		}
+
+		// 如果已经达到限制了，而且还在有效期，那就没必要再更新
+		if !isUpdatingConfig && oldStatus.IsValid() {
+			return nil
+		}
+	}
+
+	var untilDay = ""
+
+	// daily
+	if bandwidthLimitConfig.DailyBytes() > 0 {
+		stat, err := SharedServerDailyStatDAO.SumDailyStat(tx, serverId, timeutil.Format("Ymd"))
+		if err != nil {
+			return err
+		}
+		if stat != nil && stat.Bytes >= bandwidthLimitConfig.DailyBytes() {
+			untilDay = timeutil.Format("Ymd")
+		}
+	}
+
+	// monthly
+	if bandwidthLimitConfig.MonthlyBytes() > 0 {
+		stat, err := SharedServerDailyStatDAO.SumMonthlyStat(tx, serverId, timeutil.Format("Ym"))
+		if err != nil {
+			return err
+		}
+		if stat != nil && stat.Bytes >= bandwidthLimitConfig.MonthlyBytes() {
+			untilDay = timeutil.Format("Ym") + fmt.Sprintf("%02d", types.Int(timeutil.Format("t")))
+		}
+	}
+
+	// totally
+	if bandwidthLimitConfig.TotalBytes() > 0 {
+		totalBandwidth, err := this.Query(tx).
+			Pk(serverId).
+			Result("totalBandwidth").
+			FindFloat64Col(0)
+		if err != nil {
+			return err
+		}
+
+		if totalBandwidth >= float64(bandwidthLimitConfig.TotalBytes()) {
+			untilDay = "20990101"
+		}
+	}
+
+	var isChanged = oldStatus.UntilDay != untilDay
+	if isChanged {
+		statusJSON, err := json.Marshal(&serverconfigs.BandwidthLimitStatus{UntilDay: untilDay})
+		if err != nil {
+			return err
+		}
+
+		err = this.Query(tx).
+			Pk(serverId).
+			Set("bandwidthLimitStatus", statusJSON).
+			UpdateQuickly()
+		if err != nil {
+			return err
+		}
+		return this.NotifyUpdate(tx, serverId)
+	}
+
+	if isUpdatingConfig {
+		return this.NotifyUpdate(tx, serverId)
+	}
+	return nil
+}
+
+// IncreaseServerTotalBandwidth 增加服务的总带宽
+func (this *ServerDAO) IncreaseServerTotalBandwidth(tx *dbs.Tx, serverId int64, bytes int64) error {
+	var gb = float64(bytes) / 1024 / 1024 / 1024
+	return this.Query(tx).
+		Pk(serverId).
+		Set("totalBandwidth", dbs.SQL("totalBandwidth+:bandwidthGB")).
+		Param("bandwidthGB", gb).
+		UpdateQuickly()
+
+}
+
+// ResetServerTotalBandwidth 重置服务总带宽
+func (this *ServerDAO) ResetServerTotalBandwidth(tx *dbs.Tx, serverId int64) error {
+	return this.Query(tx).
+		Pk(serverId).
+		Set("totalBandwidth", 0).
+		UpdateQuickly()
+}
+
 // NotifyUpdate 同步集群
 func (this *ServerDAO) NotifyUpdate(tx *dbs.Tx, serverId int64) error {
 	// 创建任务

internal/db/models/server_dao_test.go

@@ -3,10 +3,13 @@ package models
 import (
 	"crypto/md5"
 	"encoding/json"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
+	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs/shared"
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/iwind/TeaGo/dbs"
 	"github.com/iwind/TeaGo/logs"
 	"testing"
+	"time"
 )
 
 func TestServerDAO_ComposeServerConfig(t *testing.T) {
@@ -161,6 +164,27 @@ func TestServerDAO_FindAllEnabledServersWithDomain(t *testing.T) {
 	}
 }
 
+func TestServerDAO_UpdateServerBandwidthLimitStatus(t *testing.T) {
+	dbs.NotifyReady()
+
+	var tx *dbs.Tx
+	before := time.Now()
+	defer func() {
+		t.Log(time.Since(before).Seconds()*1000, "ms")
+	}()
+	err := NewServerDAO().UpdateServerBandwidthLimitStatus(tx, &serverconfigs.BandwidthLimitConfig{
+		IsOn:           true,
+		DailySize:      &shared.SizeCapacity{Count: 1, Unit: "mb"},
+		MonthlySize:    &shared.SizeCapacity{Count: 10, Unit: "mb"},
+		TotalSize:      nil,
+		NoticePageBody: "",
+	}, 23)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log("ok")
+}
+
 func BenchmarkServerDAO_CountAllEnabledServers(b *testing.B) {
 	SharedServerDAO = NewServerDAO()
 

internal/db/models/server_model.go

@@ -34,6 +34,9 @@ type Server struct {
 	TcpPorts             string  `field:"tcpPorts"`             // 所包含TCP端口
 	UdpPorts             string  `field:"udpPorts"`             // 所包含UDP端口
 	SupportCNAME         uint8   `field:"supportCNAME"`         // 允许CNAME不在域名名单
+	BandwidthLimit       string  `field:"bandwidthLimit"`       // 带宽限制
+	TotalBandwidth       float64 `field:"totalBandwidth"`       // 总带宽用量（单位GB）
+	BandwidthLimitStatus string  `field:"bandwidthLimitStatus"` // 带宽限制状态
 }
 
 type ServerOperator struct {
@@ -69,6 +72,9 @@ type ServerOperator struct {
 	TcpPorts             interface{} // 所包含TCP端口
 	UdpPorts             interface{} // 所包含UDP端口
 	SupportCNAME         interface{} // 允许CNAME不在域名名单
+	BandwidthLimit       interface{} // 带宽限制
+	TotalBandwidth       interface{} // 总带宽用量（单位GB）
+	BandwidthLimitStatus interface{} // 带宽限制状态
 }
 
 func NewServerOperator() *ServerOperator {

internal/rpc/services/service_server.go

@@ -1652,3 +1652,45 @@ func (this *ServerService) PurgeServerCache(ctx context.Context, req *pb.PurgeSe
 
 	return purgeResponse, nil
 }
+
+// FindEnabledServerBandwidthLimit 查找带宽限制
+func (this *ServerService) FindEnabledServerBandwidthLimit(ctx context.Context, req *pb.FindEnabledServerBandwidthLimitRequest) (*pb.FindEnabledServerBandwidthLimitResponse, error) {
+	_, _, err := this.ValidateAdminAndUser(ctx, 0, 0)
+	if err != nil {
+		return nil, err
+	}
+
+	var tx = this.NullTx()
+	limitConfig, err := models.SharedServerDAO.FindServerBandwidthLimitConfig(tx, req.ServerId, nil)
+	if err != nil {
+		return nil, err
+	}
+	limitConfigJSON, err := json.Marshal(limitConfig)
+	if err != nil {
+		return nil, err
+	}
+	return &pb.FindEnabledServerBandwidthLimitResponse{
+		BandwidthLimitJSON: limitConfigJSON,
+	}, nil
+}
+
+// UpdateServerBandwidthLimit 设置带宽限制
+func (this *ServerService) UpdateServerBandwidthLimit(ctx context.Context, req *pb.UpdateServerBandwidthLimitRequest) (*pb.RPCSuccess, error) {
+	_, _, err := this.ValidateAdminAndUser(ctx, 0, 0)
+	if err != nil {
+		return nil, err
+	}
+
+	var tx = this.NullTx()
+	var config = &serverconfigs.BandwidthLimitConfig{}
+	err = json.Unmarshal(req.BandwidthLimitJSON, config)
+	if err != nil {
+		return nil, err
+	}
+
+	err = models.SharedServerDAO.UpdateServerBandwidthLimitConfig(tx, req.ServerId, config)
+	if err != nil {
+		return nil, err
+	}
+	return this.Success()
+}

internal/rpc/services/service_server_daily_stat.go

@@ -25,6 +25,7 @@ func (this *ServerDailyStatService) UploadServerDailyStats(ctx context.Context,
 
 	tx := this.NullTx()
 
+	// 保存统计数据
 	err = models.SharedServerDailyStatDAO.SaveStats(tx, req.Stats)
 	if err != nil {
 		return nil, err