1

diff --git a/funcCtrl/FuncCtrl.go b/funcCtrl/FuncCtrl.go
index ab5cfc26c39c26c9f81cffa965dcd7d689160c9f..24236ce9eed4a08cb3ab15890539b7b0925507a2 100644 (file)
--- 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 e8b220a0095f31e084bfdaefc40c43654630bb0f..f54c090477c471566e0bf7a1b55ff840d2a610d1 100644 (file)
--- 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