starmap/sync.go

package starmap

import (
	"sync"
	"sync/atomic"
	"time"
)

type WaitGroup struct {
	wg       *sync.WaitGroup
	maxCount uint32
	allCount uint32
}

func NewWaitGroup(maxCount int) WaitGroup {
	return WaitGroup{wg: &sync.WaitGroup{}, maxCount: uint32(maxCount)}
}

func (swg *WaitGroup) Add(delta int) {
	var Udelta uint32
	if delta < 0 {
		Udelta = uint32(-delta - 1)
	} else {
		Udelta = uint32(delta)
	}
	for {
		allC := atomic.LoadUint32(&swg.allCount)
		if atomic.LoadUint32(&swg.maxCount) == 0 || atomic.LoadUint32(&swg.maxCount) >= allC+uint32(delta) {
			if delta < 0 {
				atomic.AddUint32(&swg.allCount, ^uint32(Udelta))
			} else {
				atomic.AddUint32(&swg.allCount, uint32(Udelta))
			}
			break
		}
		time.Sleep(time.Microsecond)
	}
	swg.wg.Add(delta)
}

func (swg *WaitGroup) Done() {
	swg.Add(-1)
}

func (swg *WaitGroup) Wait() {
	swg.wg.Wait()
}

func (swg *WaitGroup) GetMaxWaitNum() int {
	return int(atomic.LoadUint32(&swg.maxCount))
}

func (swg *WaitGroup) SetMaxWaitNum(num int) {
	atomic.AddUint32(&swg.maxCount, uint32(num))
}