实现全局的TCP最大连接数

internal/ratelimit/counter.go

@@ -0,0 +1,56 @@
+// 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/ratelimit/counter_test.go

@@ -0,0 +1,38 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package ratelimit
+
+import (
+	"testing"
+	"time"
+)
+
+func TestCounter_ACK(t *testing.T) {
+	var counter = 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 = NewCounter(10)
+
+	for i := 0; i < 10; i++ {
+		counter.Ack()
+	}
+	for i := 0; i < 10; i++ {
+		counter.Release()
+	}
+	t.Log(len(counter.sem))
+}