Impove BlockN

diff --git a/funcCtrl/FuncCtrl.go b/funcCtrl/FuncCtrl.go
index 9ba7211b7b24a41f5e074a3330130ca886ddac3c..77756eb20d0977f2c4a19483dff5b57a7ad970a5 100644 (file)
--- 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 365d90be61bb92ada2cac1c5e4937883378da1ef..7b7acc8a5b0620a2d3d099d36d8caa6f45ab84b7 100644 (file)
--- 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`)
 }