github.com/seeker-insurance/kit@v0.0.13/workerpool/workerpool.go (about)

     1  package workerpool
     2  
     3  import (
     4  	"sync"
     5  	"github.com/seeker-insurance/kit/brake"
     6  	"sync/atomic"
     7  	"fmt"
     8  )
     9  
    10  const (
    11  	statusMsg = "\r%v/%v (%.3f%%) | Success: %v  Errored: %v               "
    12  )
    13  
    14  type Task struct {
    15  	Err error
    16  
    17  	f func() error
    18  }
    19  
    20  func NewTask(f func() error) *Task {
    21  	return &Task{f: f}
    22  }
    23  
    24  func (t *Task) run() error {
    25  	t.Err = t.f()
    26  	return t.Err
    27  }
    28  
    29  type Pool struct {
    30  	TasksChan  chan *Task
    31  	ErrorsChan chan error
    32  
    33  	concurrency  uint
    34  	taskWg       sync.WaitGroup
    35  	errorWg      sync.WaitGroup
    36  	errorCount   uint64
    37  	successCount uint64
    38  }
    39  
    40  func NewPool(concurrency uint) *Pool {
    41  	return &Pool{
    42  		TasksChan:   make(chan *Task),
    43  		concurrency: concurrency,
    44  		ErrorsChan:  make(chan error),
    45  	}
    46  }
    47  
    48  func (p *Pool) Run() {
    49  	for i := uint(0); i < p.concurrency; i++ {
    50  		p.taskWg.Add(1)
    51  		go p.work()
    52  	}
    53  
    54  	brake.NotifyFromChan(p.ErrorsChan, &p.errorWg)
    55  
    56  	p.taskWg.Wait()
    57  
    58  	close(p.ErrorsChan)
    59  	p.errorWg.Wait()
    60  }
    61  
    62  func (p *Pool) PrintStatus(total int) {
    63  	progressCount := p.successCount + p.errorCount
    64  	percentDone := float64(progressCount) / float64(total) * 100
    65  	fmt.Printf(statusMsg, progressCount, total, percentDone, p.successCount, p.errorCount)
    66  }
    67  
    68  func (p *Pool) work() {
    69  	for task := range p.TasksChan {
    70  		if err := task.run(); err != nil {
    71  			p.ErrorsChan <- err
    72  			atomic.AddUint64(&p.errorCount, 1)
    73  			continue
    74  		}
    75  		atomic.AddUint64(&p.successCount, 1)
    76  	}
    77  	p.taskWg.Done()
    78  }