github.com/unionj-cloud/go-doudou/v2@v2.3.5/toolkit/gormgen/internal/utils/pools/pool.go (about)

     1  package pools
     2  
     3  import "sync"
     4  
     5  // Pool goroutine pool
     6  type Pool interface {
     7  	// Wait 等待令牌
     8  	Wait()
     9  	// Done 归还令牌
    10  	Done()
    11  	// Num 当前发放的令牌书
    12  	Num() int
    13  	// Size 总令牌数
    14  	Size() int
    15  
    16  	// WaitAll 同步等待令牌全部归还
    17  	WaitAll()
    18  	// AsyncWaitAll 异步等待令牌全部归还
    19  	AsyncWaitAll() <-chan struct{}
    20  }
    21  
    22  type pool struct {
    23  	pool chan struct{}
    24  
    25  	wg sync.WaitGroup
    26  }
    27  
    28  func (p *pool) Init(size int) {
    29  	if size >= 0 {
    30  		p.pool = make(chan struct{}, size)
    31  	}
    32  }
    33  
    34  func (p *pool) Wait() {
    35  	if p.pool != nil {
    36  		p.wg.Add(1)
    37  		p.pool <- struct{}{}
    38  	}
    39  }
    40  
    41  func (p *pool) Done() {
    42  	if p.pool != nil {
    43  		<-p.pool
    44  		p.wg.Done()
    45  	}
    46  }
    47  
    48  func (p *pool) Num() int {
    49  	if p.pool != nil {
    50  		return len(p.pool)
    51  	}
    52  	return 0
    53  }
    54  
    55  func (p *pool) Size() int {
    56  	if p.pool != nil {
    57  		return cap(p.pool)
    58  	}
    59  	return 0
    60  }
    61  
    62  func (p *pool) WaitAll() {
    63  	p.wg.Wait()
    64  }
    65  
    66  func (p *pool) AsyncWaitAll() <-chan struct{} {
    67  	sig := make(chan struct{})
    68  	go func() {
    69  		p.WaitAll()
    70  		sig <- struct{}{}
    71  	}()
    72  	return sig
    73  }