From 0fd2ea51d5eaf9c04ec1b52adf5694826be2f501 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Mon, 29 Apr 2024 23:10:02 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E8=AF=BB=E5=86=99=E7=BA=BF?=
 =?UTF-8?q?=E7=A8=8B=E9=99=90=E5=88=B6=E7=9B=B8=E5=85=B3=E6=B5=8B=E8=AF=95?=
 =?UTF-8?q?=E7=94=A8=E4=BE=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/utils/fs/limiter.go      | 11 ++++++-
 internal/utils/fs/limiter_test.go | 54 +++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 internal/utils/fs/limiter_test.go

diff --git a/internal/utils/fs/limiter.go b/internal/utils/fs/limiter.go
index 040bb8b..c08bfd6 100644
--- a/internal/utils/fs/limiter.go
+++ b/internal/utils/fs/limiter.go
@@ -3,6 +3,7 @@
 package fsutils
 
 import (
+	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
 	"runtime"
 	"time"
 )
@@ -68,5 +69,13 @@ func (this *Limiter) TryAck() bool {
 }
 
 func (this *Limiter) Release() {
-	this.threads <- struct{}{}
+	select {
+	case this.threads <- struct{}{}:
+	default:
+		remotelogs.Error("FS_LIMITER", "Limiter Ack()/Release() should appeared as a pair")
+	}
+}
+
+func (this *Limiter) FreeThreads() int {
+	return len(this.threads)
 }
diff --git a/internal/utils/fs/limiter_test.go b/internal/utils/fs/limiter_test.go
new file mode 100644
index 0000000..f54fd4d
--- /dev/null
+++ b/internal/utils/fs/limiter_test.go
@@ -0,0 +1,54 @@
+// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package fsutils_test
+
+import (
+	fsutils "github.com/TeaOSLab/EdgeNode/internal/utils/fs"
+	"github.com/iwind/TeaGo/assert"
+	"testing"
+)
+
+func TestLimiter_Ack(t *testing.T) {
+	var a = assert.NewAssertion(t)
+
+	{
+		var limiter = fsutils.NewLimiter(4)
+		a.IsTrue(limiter.FreeThreads() == 4)
+		limiter.Ack()
+		a.IsTrue(limiter.FreeThreads() == 3)
+		limiter.Ack()
+		a.IsTrue(limiter.FreeThreads() == 2)
+		limiter.Release()
+		a.IsTrue(limiter.FreeThreads() == 3)
+		limiter.Release()
+		a.IsTrue(limiter.FreeThreads() == 4)
+	}
+}
+
+func TestLimiter_TryAck(t *testing.T) {
+	var a = assert.NewAssertion(t)
+
+	{
+		var limiter = fsutils.NewLimiter(4)
+		var count = limiter.FreeThreads()
+		a.IsTrue(count == 4)
+		for i := 0; i < count; i++ {
+			limiter.Ack()
+		}
+		a.IsTrue(limiter.FreeThreads() == 0)
+		a.IsFalse(limiter.TryAck())
+		a.IsTrue(limiter.FreeThreads() == 0)
+	}
+
+	{
+		var limiter = fsutils.NewLimiter(4)
+		var count = limiter.FreeThreads()
+		a.IsTrue(count == 4)
+		for i := 0; i < count-1; i++ {
+			limiter.Ack()
+		}
+		a.IsTrue(limiter.FreeThreads() == 1)
+		a.IsTrue(limiter.TryAck())
+		a.IsTrue(limiter.FreeThreads() == 0)
+	}
+}