limit.go

// LimitWaitGroup - is the special kind of Wait Group.
// It works similar to sync.WaitGroup with Add and Done methods,
// but instead of Wait method it brings InsureLimit method.
// InsureLimit - will block the thread if pending
// proccesses counter breaks to provided threshold (limit) count
// Also SetLimit method allows dynamic change of group limit
type LimitWaitGroup struct {
	limit int
	count int
	cond  *sync.Cond
}

func NewLimitWaitGroup(limit int) *LimitWaitGroup {
	return &LimitWaitGroup{
		limit: limit,
		count: 0,
		cond:  sync.NewCond(&sync.RWMutex{}),
	}
}

func (t *LimitWaitGroup) Add() {
	t.cond.L.Lock()
	defer t.cond.L.Unlock()
	if t.count == t.limit {
		return
	}
	t.count++
}

func (t *LimitWaitGroup) Done() {
	t.cond.L.Lock()
	defer t.cond.L.Unlock()
	if t.count == 0 {
		return
	}
	if t.count == t.limit {
		t.cond.Broadcast()
	}
	t.count--
}

func (t *LimitWaitGroup) InsureLimit() {
	t.cond.L.Lock()
	defer t.cond.L.Unlock()
	if t.count < t.limit {
		return
	}
	t.cond.Wait()
}

func (t *LimitWaitGroup) SetLimit(limit int) {
	t.cond.L.Lock()
	defer t.cond.L.Unlock()
	t.limit = limit
	if t.count < t.limit {
		t.cond.Broadcast()
	}
}