84 lines
1.8 KiB
Go
84 lines
1.8 KiB
Go
// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
|
|
//go:build plus
|
|
|
|
package nodes
|
|
|
|
import (
|
|
"context"
|
|
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
|
|
teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
|
|
"github.com/TeaOSLab/EdgeNode/internal/events"
|
|
"github.com/TeaOSLab/EdgeNode/internal/utils/ratelimit"
|
|
"sync"
|
|
)
|
|
|
|
func init() {
|
|
if !teaconst.IsMain {
|
|
return
|
|
}
|
|
|
|
events.On(events.EventLoaded, func() {
|
|
var cfg = nodeConfig()
|
|
if cfg == nil {
|
|
return
|
|
}
|
|
sharedPlanBandwidthLimiter.UpdatePlans(cfg.FindAllPlans())
|
|
})
|
|
}
|
|
|
|
var sharedPlanBandwidthLimiter = NewPlanBandwidthLimiter()
|
|
|
|
// PlanBandwidthLimiter 套餐带宽限制
|
|
type PlanBandwidthLimiter struct {
|
|
m map[int64]*ratelimit.Bandwidth // id => *Bandwidth
|
|
|
|
mu *sync.RWMutex
|
|
}
|
|
|
|
// NewPlanBandwidthLimiter 获取新的Limiter
|
|
func NewPlanBandwidthLimiter() *PlanBandwidthLimiter {
|
|
return &PlanBandwidthLimiter{
|
|
mu: &sync.RWMutex{},
|
|
m: make(map[int64]*ratelimit.Bandwidth),
|
|
}
|
|
}
|
|
|
|
// UpdatePlans 更新套餐信息
|
|
func (this *PlanBandwidthLimiter) UpdatePlans(planMap map[int64]*serverconfigs.PlanConfig) {
|
|
var bandwidthMap = map[int64]*ratelimit.Bandwidth{}
|
|
for _, plan := range planMap {
|
|
if plan.BandwidthLimitPerNode == nil || plan.BandwidthLimitPerNode.Count <= 0 {
|
|
continue
|
|
}
|
|
var bits = plan.BandwidthLimitPerNode.Bits()
|
|
if bits <= 0 {
|
|
continue
|
|
}
|
|
|
|
var bytes = bits / 8
|
|
if bytes == 0 {
|
|
continue
|
|
}
|
|
bandwidthMap[plan.Id] = ratelimit.NewBandwidth(bytes)
|
|
}
|
|
|
|
this.mu.Lock()
|
|
this.m = bandwidthMap
|
|
this.mu.Unlock()
|
|
}
|
|
|
|
// Ack 获取下一次发送数据的机会
|
|
func (this *PlanBandwidthLimiter) Ack(ctx context.Context, planId int64, newBytes int) {
|
|
if planId <= 0 {
|
|
return
|
|
}
|
|
|
|
this.mu.RLock()
|
|
defer this.mu.RUnlock()
|
|
bandwidth, ok := this.m[planId]
|
|
if !ok {
|
|
return
|
|
}
|
|
bandwidth.Ack(ctx, newBytes)
|
|
}
|