diff --git a/concurrent/gater.go b/concurrent/gater.go index 540d1ce..f9d4d21 100644 --- a/concurrent/gater.go +++ b/concurrent/gater.go @@ -1,16 +1,44 @@ package concurrent -// Gater 用于控制最大并发协程数 -type Gater chan struct{} +import ( + "sync" -func NewGater(cap int) Gater { - return make(Gater, cap) + "github.com/pyihe/go-pkg/maths" +) + +type Limiter struct { + queue chan struct{} + waiter *sync.WaitGroup } -func (g Gater) Enter() { - g <- struct{}{} +func NewLimiter(size int) *Limiter { + size = maths.MaxInt(0, size) + return &Limiter{ + queue: make(chan struct{}, size), + waiter: &sync.WaitGroup{}, + } } -func (g Gater) Leave() { - <-g +func (lim *Limiter) Add(delta int) { + switch { + case delta > 0: + for i := 0; i < delta; i++ { + lim.queue <- struct{}{} + } + default: + for i := 0; i > delta; i-- { + <-lim.queue + } + } + + lim.waiter.Add(delta) +} + +func (lim *Limiter) Done() { + <-lim.queue + lim.waiter.Done() +} + +func (lim *Limiter) Wait() { + lim.waiter.Wait() }