1.4.5.2

EdgeNode/internal/nodes/plan_bandwidth_limiter_plus.go

@@ -0,0 +1,82 @@
+// 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() {
+		if sharedNodeConfig == nil {
+			return
+		}
+		sharedPlanBandwidthLimiter.UpdatePlans(sharedNodeConfig.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)
+}