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) }