套餐可以设置带宽限制

internal/utils/ratelimit/bandwidth.go

@@ -0,0 +1,61 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package ratelimit
+
+import (
+	"context"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
+	"sync/atomic"
+	"time"
+)
+
+// Bandwidth lossy bandwidth limiter
+type Bandwidth struct {
+	totalBytes int64
+
+	currentTimestamp int64
+	currentBytes     int64
+}
+
+// NewBandwidth create new bandwidth limiter
+func NewBandwidth(totalBytes int64) *Bandwidth {
+	return &Bandwidth{totalBytes: totalBytes}
+}
+
+// Ack acquire next chance to send data
+func (this *Bandwidth) Ack(ctx context.Context, newBytes int) {
+	if newBytes <= 0 {
+		return
+	}
+	if this.totalBytes <= 0 {
+		return
+	}
+
+	var timestamp = fasttime.Now().Unix()
+	if this.currentTimestamp != 0 && this.currentTimestamp != timestamp {
+		this.currentTimestamp = timestamp
+		this.currentBytes = int64(newBytes)
+
+		// 第一次发送直接放行，不需要判断
+		return
+	}
+
+	if this.currentTimestamp == 0 {
+		this.currentTimestamp = timestamp
+	}
+	if atomic.AddInt64(&this.currentBytes, int64(newBytes)) <= this.totalBytes {
+		return
+	}
+
+	var timeout = time.NewTimer(1 * time.Second)
+	if ctx != nil {
+		select {
+		case <-timeout.C:
+		case <-ctx.Done():
+		}
+	} else {
+		select {
+		case <-timeout.C:
+		}
+	}
+}

internal/utils/ratelimit/bandwidth_test.go

@@ -0,0 +1,27 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package ratelimit_test
+
+import (
+	"context"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/ratelimit"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/testutils"
+	"testing"
+)
+
+func TestBandwidth(t *testing.T) {
+	if !testutils.IsSingleTesting() {
+		return
+	}
+
+	var bandwidth = ratelimit.NewBandwidth(32 << 10)
+	bandwidth.Ack(context.Background(), 123)
+	bandwidth.Ack(context.Background(), 16 << 10)
+	bandwidth.Ack(context.Background(), 32 << 10)
+}
+
+func TestBandwidth_0(t *testing.T) {
+	var bandwidth = ratelimit.NewBandwidth(0)
+	bandwidth.Ack(context.Background(), 123)
+	bandwidth.Ack(context.Background(), 123456)
+}

internal/utils/ratelimit/counter.go

@@ -0,0 +1,55 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package ratelimit
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/zero"
+	"sync"
+)
+
+type Counter struct {
+	count     int
+	sem       chan zero.Zero
+	done      chan zero.Zero
+	closeOnce sync.Once
+}
+
+func NewCounter(count int) *Counter {
+	return &Counter{
+		count: count,
+		sem:   make(chan zero.Zero, count),
+		done:  make(chan zero.Zero),
+	}
+}
+
+func (this *Counter) Count() int {
+	return this.count
+}
+
+// Len 已占用数量
+func (this *Counter) Len() int {
+	return len(this.sem)
+}
+
+func (this *Counter) Ack() bool {
+	select {
+	case this.sem <- zero.New():
+		return true
+	case <-this.done:
+		return false
+	}
+}
+
+func (this *Counter) Release() {
+	select {
+	case <-this.sem:
+	default:
+		// 总是能Release成功
+	}
+}
+
+func (this *Counter) Close() {
+	this.closeOnce.Do(func() {
+		close(this.done)
+	})
+}

internal/utils/ratelimit/counter_test.go

@@ -0,0 +1,44 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package ratelimit_test
+
+import (
+	"github.com/TeaOSLab/EdgeNode/internal/utils/ratelimit"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/testutils"
+	"testing"
+	"time"
+)
+
+func TestCounter_ACK(t *testing.T) {
+	if !testutils.IsSingleTesting() {
+		return
+	}
+
+	var counter = ratelimit.NewCounter(10)
+
+	go func() {
+		for i := 0; i < 10; i++ {
+			counter.Ack()
+		}
+		//counter.Release()
+		t.Log("waiting", time.Now().Unix())
+		counter.Ack()
+		t.Log("done", time.Now().Unix())
+	}()
+
+	time.Sleep(1 * time.Second)
+	counter.Close()
+	time.Sleep(1 * time.Second)
+}
+
+func TestCounter_Release(t *testing.T) {
+	var counter = ratelimit.NewCounter(10)
+
+	for i := 0; i < 10; i++ {
+		counter.Ack()
+	}
+	for i := 0; i < 10; i++ {
+		counter.Release()
+	}
+	t.Log(counter.Len())
+}