From 21dde88f3428bd4d674e495aff0e18750946a168 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Sun, 24 Mar 2024 04:56:24 +0800
Subject: [PATCH] 1

---
 funcCtrl/FuncCtrl.go      | 62 ++++++++-------------------------------
 funcCtrl/FuncCtrl_test.go | 10 ++-----
 2 files changed, 16 insertions(+), 56 deletions(-)

diff --git a/funcCtrl/FuncCtrl.go b/funcCtrl/FuncCtrl.go
index ab5cfc2..24236ce 100644
--- a/funcCtrl/FuncCtrl.go
+++ b/funcCtrl/FuncCtrl.go
@@ -2,7 +2,6 @@ package part
 
 import (
 	"context"
-	"runtime"
 	"sync"
 	"sync/atomic"
 	"unsafe"
@@ -79,62 +78,27 @@ func (t *BlockFunc) UnBlock() {
 }
 
 type BlockFuncN struct { //新的等待旧的 个数
-	n   atomic.Int64
-	Max int64
+	n chan struct{}
 }
 
-func (t *BlockFuncN) Block(failF ...func()) {
-	for {
-		now := t.n.Load()
-		if now < t.Max && now >= 0 {
-			break
-		}
-		for i := 0; i < len(failF); i++ {
-			failF[i]()
-		}
-		runtime.Gosched()
-	}
-	t.n.Add(1)
+func NewBlockFuncN(max int) *BlockFuncN {
+	return &BlockFuncN{n: make(chan struct{}, max)}
 }
 
-func (t *BlockFuncN) UnBlock(failF ...func()) {
-	for {
-		now := t.n.Load()
-		if now > 0 {
-			break
-		}
-		for i := 0; i < len(failF); i++ {
-			failF[i]()
-		}
-		runtime.Gosched()
+func (t *BlockFuncN) Block() (unBlock func()) {
+	t.n <- struct{}{}
+	return func() {
+		<-t.n
 	}
-	t.n.Add(-1)
-}
-
-func (t *BlockFuncN) BlockF(failF ...func()) (unBlock func(failF ...func())) {
-	t.Block(failF...)
-	return t.UnBlock
 }
 
-func (t *BlockFuncN) BlockAll(failF ...func()) {
-	for !t.n.CompareAndSwap(0, -1) {
-		for i := 0; i < len(failF); i++ {
-			failF[i]()
-		}
-		runtime.Gosched()
+func (t *BlockFuncN) BlockAll() (unBlock func()) {
+	for i := cap(t.n); i > 0; i-- {
+		t.n <- struct{}{}
 	}
-}
-
-func (t *BlockFuncN) UnBlockAll(failF ...func()) {
-	for !t.n.CompareAndSwap(-1, 0) {
-		for i := 0; i < len(failF); i++ {
-			failF[i]()
+	return func() {
+		for i := cap(t.n); i > 0; i-- {
+			<-t.n
 		}
-		runtime.Gosched()
 	}
 }
-
-func (t *BlockFuncN) BlockAllF(failF ...func()) (unBlock func(failF ...func())) {
-	t.BlockAll(failF...)
-	return t.UnBlockAll
-}
diff --git a/funcCtrl/FuncCtrl_test.go b/funcCtrl/FuncCtrl_test.go
index e8b220a..f54c090 100644
--- a/funcCtrl/FuncCtrl_test.go
+++ b/funcCtrl/FuncCtrl_test.go
@@ -102,12 +102,9 @@ func Test_BlockFuncN(t *testing.T) {
 	var c = make(chan string, 8)
 	var cc string
 
-	var b = &BlockFuncN{
-		Max: 2,
-	}
+	var b = NewBlockFuncN(2)
 	var a = func(i string) {
-		b.Block()
-		defer b.UnBlock()
+		defer b.Block()()
 		c <- i
 		time.Sleep(time.Second)
 		c <- i
@@ -125,8 +122,7 @@ func Test_BlockFuncN(t *testing.T) {
 		t.Fatal()
 	}
 
-	b.BlockAll()
-	b.UnBlockAll()
+	b.BlockAll()()
 
 	for len(c) > 0 {
 		cc += <-c
-- 
2.39.2