github.com/pubgo/xprocess@v0.1.11/xprocess_waitgroup/waitgroup.go (about) 1 package xprocess_waitgroup 2 3 import ( 4 "runtime" 5 "sync" 6 "sync/atomic" 7 _ "unsafe" 8 ) 9 10 //go:linkname state sync.(*WaitGroup).state 11 func state(*sync.WaitGroup) (*uint64, *uint32) 12 13 func New(check bool, c ...uint16) WaitGroup { 14 cc := uint16(runtime.NumCPU() * 2) 15 if len(c) > 0 { 16 cc = c[0] 17 } 18 19 return WaitGroup{Check: check, Concurrent: cc} 20 } 21 22 type WaitGroup struct { 23 _ int8 24 Check bool 25 Concurrent uint16 26 wg sync.WaitGroup 27 } 28 29 func (t *WaitGroup) Count() uint16 { 30 count, _ := state(&t.wg) 31 return uint16(atomic.LoadUint64(count) >> 32) 32 } 33 34 func (t *WaitGroup) check() { 35 if !t.Check { 36 return 37 } 38 39 if t.Concurrent == 0 { 40 t.Concurrent = uint16(runtime.NumCPU() * 2) 41 } 42 43 if t.Count() >= t.Concurrent { 44 t.wg.Wait() 45 } 46 } 47 48 func (t *WaitGroup) EnableCheck() { t.Check = true } 49 func (t *WaitGroup) Inc() { t.check(); t.wg.Add(1) } 50 func (t *WaitGroup) Dec() { t.wg.Done() } 51 func (t *WaitGroup) Done() { t.wg.Done() } 52 func (t *WaitGroup) Wait() { t.wg.Wait() } 53 func (t *WaitGroup) Add(delta int) { t.check(); t.wg.Add(delta) }