parent
4bf0637e04
commit
c8e513063c
@ -0,0 +1,55 @@
|
||||
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) >= 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))
|
||||
}
|
Loading…
Reference in New Issue