github.com/kaydxh/golang@v0.0.131/pkg/pool/task/pool.go (about)

     1  /*
     2   *Copyright (c) 2022, kaydxh
     3   *
     4   *Permission is hereby granted, free of charge, to any person obtaining a copy
     5   *of this software and associated documentation files (the "Software"), to deal
     6   *in the Software without restriction, including without limitation the rights
     7   *to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     8   *copies of the Software, and to permit persons to whom the Software is
     9   *furnished to do so, subject to the following conditions:
    10   *
    11   *The above copyright notice and this permission notice shall be included in all
    12   *copies or substantial portions of the Software.
    13   *
    14   *THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    15   *IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    16   *FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    17   *AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    18   *LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    19   *OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    20   *SOFTWARE.
    21   */
    22  package task
    23  
    24  import (
    25  	"context"
    26  	"fmt"
    27  	"sync"
    28  
    29  	errors_ "github.com/kaydxh/golang/go/errors"
    30  	rate_ "github.com/kaydxh/golang/go/time/rate"
    31  )
    32  
    33  type TaskHandler func(task interface{}) error
    34  
    35  type PoolConfig struct {
    36  	burst int
    37  	// errstop, only take effect to the tasks that excessed the number burst,
    38  	errStop bool
    39  }
    40  
    41  func defaultPoolConfig() PoolConfig {
    42  	return PoolConfig{
    43  		burst:   1,
    44  		errStop: false,
    45  	}
    46  }
    47  
    48  type Pool struct {
    49  	TaskFunc TaskHandler
    50  	opts     PoolConfig
    51  	ctx      context.Context
    52  
    53  	taskChan     chan interface{}
    54  	errs         []error
    55  	cancel       context.CancelFunc
    56  	workDoneChan chan struct{}
    57  
    58  	wg    sync.WaitGroup
    59  	errMu sync.Mutex
    60  }
    61  
    62  func New(taskFunc TaskHandler, opts ...PoolOptions) *Pool {
    63  	p := &Pool{
    64  		TaskFunc:     taskFunc,
    65  		taskChan:     make(chan interface{}),
    66  		workDoneChan: make(chan struct{}),
    67  		opts:         defaultPoolConfig(),
    68  	}
    69  	p.ctx, p.cancel = context.WithCancel(context.Background())
    70  	p.ApplyOptions(opts...)
    71  
    72  	go p.run()
    73  	return p
    74  }
    75  
    76  func (p *Pool) Put(task interface{}) error {
    77  	select {
    78  	case p.taskChan <- task:
    79  	case <-p.workDoneChan:
    80  		return fmt.Errorf("workdone channel is closed, work goroutine is exit")
    81  	case <-p.ctx.Done():
    82  		return fmt.Errorf("work is stoped, work goroutine is exit")
    83  
    84  	}
    85  
    86  	return nil
    87  }
    88  
    89  func (p *Pool) Wait() {
    90  	p.wg.Wait()
    91  }
    92  
    93  func (p *Pool) Stop() {
    94  	p.cancel()
    95  }
    96  
    97  func (p *Pool) Error() error {
    98  	p.errMu.Lock()
    99  	defer p.errMu.Unlock()
   100  
   101  	return errors_.NewAggregate(p.errs)
   102  }
   103  
   104  func (p *Pool) trySetError(err error) {
   105  	p.errMu.Lock()
   106  	defer p.errMu.Unlock()
   107  	p.errs = append(p.errs, err)
   108  }
   109  
   110  func (p *Pool) run() (doneC <-chan struct{}) {
   111  
   112  	go func() {
   113  		defer close(p.workDoneChan)
   114  
   115  		limiter := rate_.NewLimiter(int(p.opts.burst))
   116  		for {
   117  			//util the condition is met, need one token, or will be blocked
   118  			limiter.AllowWaitUntil()
   119  
   120  			select {
   121  			case task, ok := <-p.taskChan:
   122  				if !ok {
   123  					return
   124  				}
   125  				p.wg.Add(1)
   126  
   127  				go func(t interface{}) {
   128  					defer limiter.Put()
   129  					defer p.wg.Done()
   130  
   131  					if err := p.TaskFunc(t); err != nil {
   132  						p.trySetError(err)
   133  						if p.opts.errStop {
   134  							p.cancel()
   135  						}
   136  						return
   137  					}
   138  
   139  				}(task)
   140  
   141  			case <-p.ctx.Done():
   142  				//  err: context canceled
   143  				p.trySetError(p.ctx.Err())
   144  				return
   145  			}
   146  
   147  		}
   148  	}()
   149  
   150  	return p.workDoneChan
   151  }