github.com/iikira/iikira-go-utils@v0.0.0-20230610031953-f2cb11cde33a/utils/waitgroup/wait_group.go (about)

     1  // Package waitgroup sync.WaitGroup extension
     2  package waitgroup
     3  
     4  import "sync"
     5  
     6  // WaitGroup 在 sync.WaitGroup 的基础上, 新增线程控制功能
     7  type WaitGroup struct {
     8  	wg sync.WaitGroup
     9  	p  chan struct{}
    10  
    11  	sync.RWMutex
    12  }
    13  
    14  // NewWaitGroup returns a pointer to a new `WaitGroup` object.
    15  // parallel 为最大并发数, 0 代表无限制
    16  func NewWaitGroup(parallel int) (w *WaitGroup) {
    17  	w = &WaitGroup{
    18  		wg: sync.WaitGroup{},
    19  	}
    20  
    21  	if parallel <= 0 {
    22  		return
    23  	}
    24  
    25  	w.p = make(chan struct{}, parallel)
    26  	return
    27  }
    28  
    29  // AddDelta sync.WaitGroup.Add(1)
    30  func (w *WaitGroup) AddDelta() {
    31  	if w.p != nil {
    32  		w.p <- struct{}{}
    33  	}
    34  
    35  	w.wg.Add(1)
    36  }
    37  
    38  // Done sync.WaitGroup.Done()
    39  func (w *WaitGroup) Done() {
    40  	w.wg.Done()
    41  
    42  	if w.p != nil {
    43  		<-w.p
    44  	}
    45  }
    46  
    47  // Wait 参照 sync.WaitGroup 的 Wait 方法
    48  func (w *WaitGroup) Wait() {
    49  	w.wg.Wait()
    50  	if w.p != nil {
    51  		close(w.p)
    52  	}
    53  }
    54  
    55  // Parallel 返回当前正在进行的任务数量
    56  func (w *WaitGroup) Parallel() int {
    57  	return len(w.p)
    58  }