TRKJ/EdgeAPI
2
0
Fork 0
mirror of https://github.com/TeaOSLab/EdgeAPI.git synced 2025-11-04 16:00:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

改进流量限制

Browse Source
This commit is contained in:
GoEdgeLab
2021-11-10 14:54:27 +08:00
parent ab1686bca5
commit 4227bfd7fc
9 changed files with 253 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
internal/db/models/node_dao.go
View File

@@ -729,7 +729,7 @@ func (this *NodeDAO) ComposeNodeConfig(tx *dbs.Tx, nodeId int64, cacheMap maps.M
}
for _, server := range servers {
serverConfig, err := SharedServerDAO.ComposeServerConfig(tx, server, cacheMap)
serverConfig, err := SharedServerDAO.ComposeServerConfig(tx, server, cacheMap, true)
if err != nil {
return nil, err
}

37
internal/db/models/plan_dao.go
View File

@@ -1,11 +1,14 @@
package models
import (
"encoding/json"
"github.com/TeaOSLab/EdgeAPI/internal/errors"
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
_ "github.com/go-sql-driver/mysql"
"github.com/iwind/TeaGo/Tea"
"github.com/iwind/TeaGo/dbs"
"github.com/iwind/TeaGo/maps"
"github.com/iwind/TeaGo/types"
)
const (
@@ -210,6 +213,40 @@ func (this *PlanDAO) SortPlans(tx *dbs.Tx, planIds []int64) error {
return nil
}
// FindEnabledPlanTrafficLimit 获取套餐的流量限制
func (this *PlanDAO) FindEnabledPlanTrafficLimit(tx *dbs.Tx, planId int64, cacheMap maps.Map) (*serverconfigs.TrafficLimitConfig, error) {
var cacheKey = this.Table + ":FindEnabledPlanTrafficLimit:" + types.String(planId)
if cacheMap != nil {
cache, ok := cacheMap[cacheKey]
if ok {
return cache.(*serverconfigs.TrafficLimitConfig), nil
}
}
trafficLimit, err := this.Query(tx).
Pk(planId).
State(PlanStateEnabled).
Result("trafficLimit").
FindStringCol("")
if err != nil {
return nil, err
}
if len(trafficLimit) == 0 {
return nil, nil
}
var config = &serverconfigs.TrafficLimitConfig{}
err = json.Unmarshal([]byte(trafficLimit), config)
if err != nil {
return nil, err
}
if cacheMap != nil {
cacheMap[cacheKey] = config
}
return config, nil
}
// NotifyUpdate 通知变更
func (this *PlanDAO) NotifyUpdate(tx *dbs.Tx, planId int64) error {
// 这里不要加入状态参数，因为需要适应删除后的更新

19
internal/db/models/server_daily_stat_dao.go
View File

@@ -105,15 +105,22 @@ func (this *ServerDailyStatDAO) SaveStats(tx *dbs.Tx, stats []*pb.ServerDailySta
}
// 更新流量限制状态
trafficLimit, err := SharedServerDAO.FindServerTrafficLimitConfig(tx, stat.ServerId, cacheMap)
if err != nil {
return err
}
if trafficLimit != nil && trafficLimit.IsOn && !trafficLimit.IsEmpty() {
err = SharedServerDAO.UpdateServerTrafficLimitStatus(tx, trafficLimit, stat.ServerId, false)
if stat.CheckTrafficLimiting {
trafficLimitConfig, err := SharedServerDAO.CalculateServerTrafficLimitConfig(tx, stat.ServerId, cacheMap)
if err != nil {
return err
}
if trafficLimitConfig != nil && trafficLimitConfig.IsOn && !trafficLimitConfig.IsEmpty() {
err = SharedServerDAO.IncreaseServerTotalTraffic(tx, stat.ServerId, stat.Bytes)
if err != nil {
return err
}
err = SharedServerDAO.UpdateServerTrafficLimitStatus(tx, trafficLimitConfig, stat.ServerId, false)
if err != nil {
return err
}
}
}
}

181
internal/db/models/server_dao.go
View File

@@ -3,7 +3,6 @@ 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"
@@ -827,7 +826,7 @@ func (this *ServerDAO) FindServerNodeFilters(tx *dbs.Tx, serverId int64) (isOk b
}
// ComposeServerConfigWithServerId 构造服务的Config
func (this *ServerDAO) ComposeServerConfigWithServerId(tx *dbs.Tx, serverId int64) (*serverconfigs.ServerConfig, error) {
func (this *ServerDAO) ComposeServerConfigWithServerId(tx *dbs.Tx, serverId int64, forNode bool) (*serverconfigs.ServerConfig, error) {
server, err := this.FindEnabledServer(tx, serverId)
if err != nil {
return nil, err
@@ -835,11 +834,12 @@ func (this *ServerDAO) ComposeServerConfigWithServerId(tx *dbs.Tx, serverId int6
if server == nil {
return nil, ErrNotFound
}
return this.ComposeServerConfig(tx, server, nil)
return this.ComposeServerConfig(tx, server, nil, forNode)
}
// ComposeServerConfig 构造服务的Config
func (this *ServerDAO) ComposeServerConfig(tx *dbs.Tx, server *Server, cacheMap maps.Map) (*serverconfigs.ServerConfig, error) {
// forNode 是否是节点请求
func (this *ServerDAO) ComposeServerConfig(tx *dbs.Tx, server *Server, cacheMap maps.Map, forNode bool) (*serverconfigs.ServerConfig, error) {
if server == nil {
return nil, ErrNotFound
}
@@ -1039,24 +1039,11 @@ func (this *ServerDAO) ComposeServerConfig(tx *dbs.Tx, server *Server, cacheMap
return nil, err
}
config.TrafficLimit = trafficLimitConfig
if trafficLimitConfig.IsOn && !trafficLimitConfig.IsEmpty() {
if len(server.TrafficLimitStatus) > 0 {
var status = &serverconfigs.TrafficLimitStatus{}
err = json.Unmarshal([]byte(server.TrafficLimitStatus), status)
if err != nil {
return nil, err
}
if status.IsValid() {
config.TrafficLimitStatus = status
}
}
}
}
// 用户套餐
if server.UserPlanId > 0 {
userPlan, err := SharedUserPlanDAO.FindEnabledUserPlan(tx, int64(server.UserPlanId))
if forNode && server.UserPlanId > 0 {
userPlan, err := SharedUserPlanDAO.FindEnabledUserPlan(tx, int64(server.UserPlanId), cacheMap)
if err != nil {
return nil, err
}
@@ -1087,6 +1074,19 @@ func (this *ServerDAO) ComposeServerConfig(tx *dbs.Tx, server *Server, cacheMap
}
}
if config.TrafficLimit != nil && config.TrafficLimit.IsOn && !config.TrafficLimit.IsEmpty() {
if len(server.TrafficLimitStatus) > 0 {
var status = &serverconfigs.TrafficLimitStatus{}
err = json.Unmarshal([]byte(server.TrafficLimitStatus), status)
if err != nil {
return nil, err
}
if status.IsValid() {
config.TrafficLimitStatus = status
}
}
}
return config, nil
}
@@ -1819,30 +1819,106 @@ func (this *ServerDAO) FindServerTrafficLimitConfig(tx *dbs.Tx, serverId int64,
return result.(*serverconfigs.TrafficLimitConfig), nil
}
trafficLimit, err := this.Query(tx).
serverOne, err := this.Query(tx).
Pk(serverId).
Result("trafficLimit").
FindStringCol("")
Find()
if err != nil {
return nil, err
}
var limit = &serverconfigs.TrafficLimitConfig{}
if len(trafficLimit) == 0 {
if serverOne == nil {
return limit, nil
}
var trafficLimit = serverOne.(*Server).TrafficLimit
err = json.Unmarshal([]byte(trafficLimit), limit)
if err != nil {
return nil, err
}
// TODO 套餐
cacheMap[cacheKey] = limit
return limit, nil
}
// CalculateServerTrafficLimitConfig 计算服务的流量限制
// TODO 优化性能
func (this *ServerDAO) CalculateServerTrafficLimitConfig(tx *dbs.Tx, serverId int64, cacheMap maps.Map) (*serverconfigs.TrafficLimitConfig, error) {
if cacheMap == nil {
cacheMap = maps.Map{}
}
var cacheKey = this.Table + ":FindServerTrafficLimitConfig:" + types.String(serverId)
result, ok := cacheMap[cacheKey]
if ok {
return result.(*serverconfigs.TrafficLimitConfig), nil
}
serverOne, err := this.Query(tx).
Pk(serverId).
Result("trafficLimit", "userPlanId").
Find()
if err != nil {
return nil, err
}
var limitConfig = &serverconfigs.TrafficLimitConfig{}
if serverOne == nil {
return limitConfig, nil
}
var trafficLimit = serverOne.(*Server).TrafficLimit
var userPlanId = int64(serverOne.(*Server).UserPlanId)
if len(trafficLimit) == 0 {
if userPlanId > 0 {
userPlan, err := SharedUserPlanDAO.FindEnabledUserPlan(tx, userPlanId, cacheMap)
if err != nil {
return nil, err
}
if userPlan != nil {
planLimit, err := SharedPlanDAO.FindEnabledPlanTrafficLimit(tx, int64(userPlan.PlanId), cacheMap)
if err != nil {
return nil, err
}
if planLimit != nil {
return planLimit, nil
}
}
}
return limitConfig, nil
}
err = json.Unmarshal([]byte(trafficLimit), limitConfig)
if err != nil {
return nil, err
}
if !limitConfig.IsOn {
if userPlanId > 0 {
userPlan, err := SharedUserPlanDAO.FindEnabledUserPlan(tx, userPlanId, cacheMap)
if err != nil {
return nil, err
}
if userPlan != nil {
planLimit, err := SharedPlanDAO.FindEnabledPlanTrafficLimit(tx, int64(userPlan.PlanId), cacheMap)
if err != nil {
return nil, err
}
if planLimit != nil {
return planLimit, nil
}
}
}
}
cacheMap[cacheKey] = limitConfig
return limitConfig, nil
}
// UpdateServerTrafficLimitConfig 修改服务的流量限制
func (this *ServerDAO) UpdateServerTrafficLimitConfig(tx *dbs.Tx, serverId int64, trafficLimitConfig *serverconfigs.TrafficLimitConfig) error {
if serverId <= 0 {
@@ -1873,16 +1949,22 @@ func (this *ServerDAO) UpdateServerTrafficLimitStatus(tx *dbs.Tx, trafficLimitCo
return nil
}
oldStatusString, err := this.Query(tx).
serverOne, err := this.Query(tx).
Pk(serverId).
Result("trafficLimitStatus").
FindStringCol("")
Result("trafficLimitStatus", "totalTraffic", "totalDailyTraffic", "totalMonthlyTraffic", "trafficDay", "trafficMonth").
Find()
if err != nil {
return err
}
if serverOne == nil {
return nil
}
var server = serverOne.(*Server)
var oldStatus = &serverconfigs.TrafficLimitStatus{}
if len(oldStatusString) > 0 {
err = json.Unmarshal([]byte(oldStatusString), oldStatus)
if len(server.TrafficLimitStatus) > 0 {
err = json.Unmarshal([]byte(server.TrafficLimitStatus), oldStatus)
if err != nil {
return err
}
@@ -1897,38 +1979,22 @@ func (this *ServerDAO) UpdateServerTrafficLimitStatus(tx *dbs.Tx, trafficLimitCo
// daily
if trafficLimitConfig.DailyBytes() > 0 {
stat, err := SharedServerDailyStatDAO.SumDailyStat(tx, serverId, timeutil.Format("Ymd"))
if err != nil {
return err
}
if stat != nil && stat.Bytes >= trafficLimitConfig.DailyBytes() {
if server.TrafficDay == timeutil.Format("Ymd") && server.TotalDailyTraffic >= float64(trafficLimitConfig.DailyBytes())/1024/1024/1024 {
untilDay = timeutil.Format("Ymd")
}
}
// monthly
if trafficLimitConfig.MonthlyBytes() > 0 {
stat, err := SharedServerDailyStatDAO.SumMonthlyStat(tx, serverId, timeutil.Format("Ym"))
if err != nil {
return err
}
if stat != nil && stat.Bytes >= trafficLimitConfig.MonthlyBytes() {
untilDay = timeutil.Format("Ym") + fmt.Sprintf("%02d", types.Int(timeutil.Format("t")))
if server.TrafficMonth == timeutil.Format("Ym") && trafficLimitConfig.MonthlyBytes() > 0 {
if server.TotalMonthlyTraffic >= float64(trafficLimitConfig.MonthlyBytes())/1024/1024/1024 {
untilDay = timeutil.Format("Ym32")
}
}
// totally
if trafficLimitConfig.TotalBytes() > 0 {
totalTraffic, err := this.Query(tx).
Pk(serverId).
Result("totalTraffic").
FindFloat64Col(0)
if err != nil {
return err
}
if totalTraffic >= float64(trafficLimitConfig.TotalBytes()) {
untilDay = "20990101"
if server.TotalTraffic >= float64(trafficLimitConfig.TotalBytes())/1024/1024/1024 {
untilDay = "30000101"
}
}
@@ -1958,9 +2024,17 @@ func (this *ServerDAO) UpdateServerTrafficLimitStatus(tx *dbs.Tx, trafficLimitCo
// IncreaseServerTotalTraffic 增加服务的总流量
func (this *ServerDAO) IncreaseServerTotalTraffic(tx *dbs.Tx, serverId int64, bytes int64) error {
var gb = float64(bytes) / 1024 / 1024 / 1024
var day = timeutil.Format("Ymd")
var month = timeutil.Format("Ym")
return this.Query(tx).
Pk(serverId).
Set("totalDailyTraffic", dbs.SQL("IF(trafficDay=:day, totalDailyTraffic, 0)+:trafficGB")).
Set("totalMonthlyTraffic", dbs.SQL("IF(trafficMonth=:month, totalMonthlyTraffic, 0)+:trafficGB")).
Set("totalTraffic", dbs.SQL("totalTraffic+:trafficGB")).
Set("trafficDay", day).
Set("trafficMonth", month).
Param("day", day).
Param("month", month).
Param("trafficGB", gb).
UpdateQuickly()
@@ -1970,7 +2044,8 @@ func (this *ServerDAO) IncreaseServerTotalTraffic(tx *dbs.Tx, serverId int64, by
func (this *ServerDAO) ResetServerTotalTraffic(tx *dbs.Tx, serverId int64) error {
return this.Query(tx).
Pk(serverId).
Set("totalTraffic", 0).
Set("totalDailyTraffic", 0).
Set("totalMonthlyTraffic", 0).
UpdateQuickly()
}
@@ -1994,7 +2069,7 @@ func (this *ServerDAO) UpdateServersClusterIdWithPlanId(tx *dbs.Tx, planId int64
// UpdateServerUserPlanId 设置服务所属套餐
func (this *ServerDAO) UpdateServerUserPlanId(tx *dbs.Tx, serverId int64, userPlanId int64) error {
userPlan, err := SharedUserPlanDAO.FindEnabledUserPlan(tx, userPlanId)
userPlan, err := SharedUserPlanDAO.FindEnabledUserPlan(tx, userPlanId, nil)
if err != nil {
return err
}

45
internal/db/models/server_dao_test.go
View File

@@ -8,6 +8,7 @@ import (
_ "github.com/go-sql-driver/mysql"
"github.com/iwind/TeaGo/dbs"
"github.com/iwind/TeaGo/logs"
"github.com/iwind/TeaGo/maps"
"testing"
"time"
)
@@ -15,7 +16,7 @@ import (
func TestServerDAO_ComposeServerConfig(t *testing.T) {
dbs.NotifyReady()
var tx *dbs.Tx
config, err := SharedServerDAO.ComposeServerConfigWithServerId(tx, 1)
config, err := SharedServerDAO.ComposeServerConfigWithServerId(tx, 1, false)
if err != nil {
t.Fatal(err)
}
@@ -25,7 +26,7 @@ func TestServerDAO_ComposeServerConfig(t *testing.T) {
func TestServerDAO_ComposeServerConfig_AliasServerNames(t *testing.T) {
dbs.NotifyReady()
var tx *dbs.Tx
config, err := SharedServerDAO.ComposeServerConfigWithServerId(tx, 14)
config, err := SharedServerDAO.ComposeServerConfigWithServerId(tx, 14, false)
if err != nil {
t.Fatal(err)
}
@@ -35,7 +36,7 @@ func TestServerDAO_ComposeServerConfig_AliasServerNames(t *testing.T) {
func TestServerDAO_UpdateServerConfig(t *testing.T) {
dbs.NotifyReady()
var tx *dbs.Tx
config, err := SharedServerDAO.ComposeServerConfigWithServerId(tx, 1)
config, err := SharedServerDAO.ComposeServerConfigWithServerId(tx, 1, false)
if err != nil {
t.Fatal(err)
}
@@ -176,7 +177,7 @@ func TestServerDAO_UpdateServerTrafficLimitStatus(t *testing.T) {
IsOn: true,
DailySize: &shared.SizeCapacity{Count: 1, Unit: "mb"},
MonthlySize: &shared.SizeCapacity{Count: 10, Unit: "mb"},
TotalSize: nil,
TotalSize: &shared.SizeCapacity{Count: 100, Unit: "gb"},
NoticePageBody: "",
}, 23, false)
if err != nil {
@@ -185,6 +186,42 @@ func TestServerDAO_UpdateServerTrafficLimitStatus(t *testing.T) {
t.Log("ok")
}
func TestServerDAO_CalculateServerTrafficLimitConfig(t *testing.T) {
dbs.NotifyReady()
var tx *dbs.Tx
before := time.Now()
defer func() {
t.Log(time.Since(before).Seconds()*1000, "ms")
}()
var cacheMap = maps.Map{}
config, err := SharedServerDAO.CalculateServerTrafficLimitConfig(tx, 23, cacheMap)
if err != nil {
t.Fatal(err)
}
logs.PrintAsJSON(config, t)
}
func TestServerDAO_CalculateServerTrafficLimitConfig_Cache(t *testing.T) {
dbs.NotifyReady()
var tx *dbs.Tx
before := time.Now()
defer func() {
t.Log(time.Since(before).Seconds()*1000, "ms")
}()
var cacheMap = maps.Map{}
for i := 0; i < 10; i++ {
config, err := SharedServerDAO.CalculateServerTrafficLimitConfig(tx, 23, cacheMap)
if err != nil {
t.Fatal(err)
}
_ = config
}
}
func BenchmarkServerDAO_CountAllEnabledServers(b *testing.B) {
SharedServerDAO = NewServerDAO()

12
internal/db/models/server_model.go
View File

@@ -35,8 +35,12 @@ type Server struct {
UdpPorts string `field:"udpPorts"` // 所包含UDP端口
SupportCNAME uint8 `field:"supportCNAME"` // 允许CNAME不在域名名单
TrafficLimit string `field:"trafficLimit"` // 流量限制
TotalTraffic float64 `field:"totalTraffic"` // 总流量用量（单位GB）
TrafficDay string `field:"trafficDay"` // YYYYMMDD
TrafficMonth string `field:"trafficMonth"` // YYYYMM
TotalDailyTraffic float64 `field:"totalDailyTraffic"` // 日流量
TotalMonthlyTraffic float64 `field:"totalMonthlyTraffic"` // 月流量
TrafficLimitStatus string `field:"trafficLimitStatus"` // 流量限制状态
TotalTraffic float64 `field:"totalTraffic"` // 总流量
UserPlanId uint32 `field:"userPlanId"` // 所属套餐ID
}
@@ -74,8 +78,12 @@ type ServerOperator struct {
UdpPorts interface{} // 所包含UDP端口
SupportCNAME interface{} // 允许CNAME不在域名名单
TrafficLimit interface{} // 流量限制
TotalTraffic interface{} // 总流量用量（单位GB）
TrafficDay interface{} // YYYYMMDD
TrafficMonth interface{} // YYYYMM
TotalDailyTraffic interface{} // 日流量
TotalMonthlyTraffic interface{} // 月流量
TrafficLimitStatus interface{} // 流量限制状态
TotalTraffic interface{} // 总流量
UserPlanId interface{} // 所属套餐ID
}

17
internal/db/models/user_plan_dao.go
View File

@@ -5,6 +5,8 @@ import (
_ "github.com/go-sql-driver/mysql"
"github.com/iwind/TeaGo/Tea"
"github.com/iwind/TeaGo/dbs"
"github.com/iwind/TeaGo/maps"
"github.com/iwind/TeaGo/types"
timeutil "github.com/iwind/TeaGo/utils/time"
"time"
)
@@ -59,7 +61,15 @@ func (this *UserPlanDAO) DisableUserPlan(tx *dbs.Tx, id int64) error {
}
// FindEnabledUserPlan 查找启用中的条目
func (this *UserPlanDAO) FindEnabledUserPlan(tx *dbs.Tx, userPlanId int64) (*UserPlan, error) {
func (this *UserPlanDAO) FindEnabledUserPlan(tx *dbs.Tx, userPlanId int64, cacheMap maps.Map) (*UserPlan, error) {
var cacheKey = this.Table + ":FindEnabledUserPlan:" + types.String(userPlanId)
if cacheMap != nil {
cache, ok := cacheMap[cacheKey]
if ok {
return cache.(*UserPlan), nil
}
}
result, err := this.Query(tx).
Pk(userPlanId).
Attr("state", UserPlanStateEnabled).
@@ -67,6 +77,11 @@ func (this *UserPlanDAO) FindEnabledUserPlan(tx *dbs.Tx, userPlanId int64) (*Use
if result == nil {
return nil, err
}
if cacheMap != nil {
cacheMap[cacheKey] = result
}
return result.(*UserPlan), err
}

12
internal/rpc/services/service_server.go
View File

@@ -100,9 +100,9 @@ func (this *ServerService) CreateServer(ctx context.Context, req *pb.CreateServe
}
}
// 校验用户套餐
// 检查用户套餐
if req.UserPlanId > 0 {
userPlan, err := models.SharedUserPlanDAO.FindEnabledUserPlan(tx, req.UserPlanId)
userPlan, err := models.SharedUserPlanDAO.FindEnabledUserPlan(tx, req.UserPlanId, nil)
if err != nil {
return nil, err
}
@@ -643,7 +643,7 @@ func (this *ServerService) ListEnabledServersMatch(ctx context.Context, req *pb.
}
// 配置
config, err := models.SharedServerDAO.ComposeServerConfig(tx, server, nil)
config, err := models.SharedServerDAO.ComposeServerConfig(tx, server, nil, false)
if err != nil {
return nil, err
}
@@ -784,7 +784,7 @@ func (this *ServerService) FindEnabledServer(ctx context.Context, req *pb.FindEn
}
// 配置
config, err := models.SharedServerDAO.ComposeServerConfig(tx, server, nil)
config, err := models.SharedServerDAO.ComposeServerConfig(tx, server, nil, false)
if err != nil {
return nil, err
}
@@ -844,7 +844,7 @@ func (this *ServerService) FindEnabledServerConfig(ctx context.Context, req *pb.
}
}
config, err := models.SharedServerDAO.ComposeServerConfigWithServerId(tx, req.ServerId)
config, err := models.SharedServerDAO.ComposeServerConfigWithServerId(tx, req.ServerId, false)
if err != nil {
return nil, err
}
@@ -1779,7 +1779,7 @@ func (this *ServerService) UpdateServerUserPlan(ctx context.Context, req *pb.Upd
return nil, errors.New("the server is not belong to any user")
}
userPlan, err := models.SharedUserPlanDAO.FindEnabledUserPlan(tx, req.UserPlanId)
userPlan, err := models.SharedUserPlanDAO.FindEnabledUserPlan(tx, req.UserPlanId, nil)
if err != nil {
return nil, err
}

2
internal/setup/sql.go
View File

File diff suppressed because one or more lines are too long