github.com/TeaOSLab/EdgeNode@v1.3.8/internal/goman/task_group.go (about)

     1  // Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package goman
     4  
     5  import (
     6  	"github.com/TeaOSLab/EdgeNode/internal/zero"
     7  	"runtime"
     8  	"sync"
     9  )
    10  
    11  type TaskGroup struct {
    12  	semi   chan zero.Zero
    13  	wg     *sync.WaitGroup
    14  	locker *sync.RWMutex
    15  }
    16  
    17  func NewTaskGroup() *TaskGroup {
    18  	var concurrent = runtime.NumCPU()
    19  	if concurrent <= 1 {
    20  		concurrent = 2
    21  	}
    22  	return &TaskGroup{
    23  		semi:   make(chan zero.Zero, concurrent),
    24  		wg:     &sync.WaitGroup{},
    25  		locker: &sync.RWMutex{},
    26  	}
    27  }
    28  
    29  func (this *TaskGroup) Run(f func()) {
    30  	this.wg.Add(1)
    31  	go func() {
    32  		defer this.wg.Done()
    33  
    34  		this.semi <- zero.Zero{}
    35  
    36  		f()
    37  
    38  		<-this.semi
    39  	}()
    40  }
    41  
    42  func (this *TaskGroup) Wait() {
    43  	this.wg.Wait()
    44  }
    45  
    46  func (this *TaskGroup) Lock() {
    47  	this.locker.Lock()
    48  }
    49  
    50  func (this *TaskGroup) Unlock() {
    51  	this.locker.Unlock()
    52  }