You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
stario/sync.go

56 lines
1020 B
Go

3 years ago
package stario
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 (w *WaitGroup) Add(delta int) {
3 years ago
var Udelta uint32
if delta < 0 {
Udelta = uint32(-delta - 1)
} else {
Udelta = uint32(delta)
}
for {
allC := atomic.LoadUint32(&w.allCount)
if atomic.LoadUint32(&w.maxCount) == 0 || atomic.LoadUint32(&w.maxCount) >= allC+uint32(delta) {
3 years ago
if delta < 0 {
atomic.AddUint32(&w.allCount, ^uint32(Udelta))
3 years ago
} else {
atomic.AddUint32(&w.allCount, uint32(Udelta))
3 years ago
}
break
}
time.Sleep(time.Microsecond)
}
w.wg.Add(delta)
3 years ago
}
func (w *WaitGroup) Done() {
w.Add(-1)
3 years ago
}
func (w *WaitGroup) Wait() {
w.wg.Wait()
3 years ago
}
func (w *WaitGroup) GetMaxWaitNum() int {
return int(atomic.LoadUint32(&w.maxCount))
3 years ago
}
func (w *WaitGroup) SetMaxWaitNum(num int) {
atomic.AddUint32(&w.maxCount, uint32(num))
3 years ago
}