From: qydysky <32743305+qydysky@users.noreply.github.com>
Date: Sun, 22 Jan 2023 16:51:21 +0000 (+0800)
Subject: Impove BlockN
X-Git-Tag: v0.21.11
X-Git-Url: http://127.0.0.1:8081/?a=commitdiff_plain;h=f22a590d55277b4cf83a54d6b1fc9621427c0232;p=part%2F.git

Impove BlockN
---

diff --git a/funcCtrl/FuncCtrl.go b/funcCtrl/FuncCtrl.go
index 9ba7211..77756eb 100644
--- a/funcCtrl/FuncCtrl.go
+++ b/funcCtrl/FuncCtrl.go
@@ -63,59 +63,49 @@ func (t *BlockFunc) UnBlock() {
 }
 
 type BlockFuncN struct { //新的等待旧的 个数
-	plan int64
-	n    int64
-	Max  int64
+	n   atomic.Int64
+	Max int64
 }
 
-func (t *BlockFuncN) Block() {
+func (t *BlockFuncN) Block(failF ...func()) {
 	for {
-		now := atomic.LoadInt64(&t.n)
+		now := t.n.Load()
 		if now < t.Max && now >= 0 {
 			break
 		}
+		for i := 0; i < len(failF); i++ {
+			failF[i]()
+		}
 		runtime.Gosched()
 	}
-	atomic.AddInt64(&t.n, 1)
+	t.n.Add(1)
 }
 
-func (t *BlockFuncN) UnBlock() {
+func (t *BlockFuncN) UnBlock(failF ...func()) {
 	for {
-		now := atomic.LoadInt64(&t.n)
+		now := t.n.Load()
 		if now > 0 {
 			break
 		}
+		for i := 0; i < len(failF); i++ {
+			failF[i]()
+		}
 		runtime.Gosched()
 	}
-	atomic.AddInt64(&t.n, -1)
-	if atomic.LoadInt64(&t.plan) > 0 {
-		atomic.AddInt64(&t.plan, -1)
-	}
-}
-
-func (t *BlockFuncN) None() {
-	for !atomic.CompareAndSwapInt64(&t.n, 0, -1) {
-		runtime.Gosched()
-	}
-}
-
-func (t *BlockFuncN) UnNone() {
-	for !atomic.CompareAndSwapInt64(&t.n, -1, 0) {
-		runtime.Gosched()
-	}
+	t.n.Add(-1)
 }
 
-func (t *BlockFuncN) Plan(n int64) {
-	for !atomic.CompareAndSwapInt64(&t.plan, 0, n) {
+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) PlanDone(switchFuncs ...func()) {
-	for atomic.LoadInt64(&t.plan) > 0 {
-		for i := 0; i < len(switchFuncs); i++ {
-			switchFuncs[i]()
-		}
-		runtime.Gosched()
+func (t *BlockFuncN) UnBlockAll() {
+	if !t.n.CompareAndSwap(-1, 0) {
+		panic("must BlockAll First")
 	}
 }
diff --git a/funcCtrl/FuncCtrl_test.go b/funcCtrl/FuncCtrl_test.go
index 365d90b..7b7acc8 100644
--- a/funcCtrl/FuncCtrl_test.go
+++ b/funcCtrl/FuncCtrl_test.go
@@ -72,30 +72,7 @@ func Test_BlockFuncN(t *testing.T) {
 	go a(1)
 	go a(2)
 	go a(3)
-	b.None()
-	b.UnNone()
-	t.Log(`fin`)
-}
-
-func Test_BlockFuncNPlan(t *testing.T) {
-	var b = &BlockFuncN{
-		Max: 2,
-	}
-	var a = func(i int) {
-		defer b.UnBlock()
-		b.Block()
-		t.Log(i, `.`)
-		time.Sleep(time.Second)
-		t.Log(i, `.`)
-	}
-	t.Log(`show two . at one time`)
-	b.Plan(4)
-	go a(0)
-	go a(1)
-	go a(2)
-	go a(3)
-	b.PlanDone(func() {
-		time.Sleep(time.Microsecond * 10)
-	})
+	b.BlockAll()
+	b.UnBlockAll()
 	t.Log(`fin`)
 }