github.com/kaydxh/golang@v0.0.131/pkg/pool/task/worker.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  	"sync"
    27  
    28  	errors_ "github.com/kaydxh/golang/go/errors"
    29  	rate_ "github.com/kaydxh/golang/go/time/rate"
    30  )
    31  
    32  type WorkFunc func(work interface{}) error
    33  
    34  type Worker struct {
    35  	Burst int32
    36  
    37  	cancel context.CancelFunc
    38  	wg     sync.WaitGroup
    39  	errs   []error
    40  	errMu  sync.Mutex
    41  }
    42  
    43  func (w *Worker) Work(
    44  	ctx context.Context,
    45  	workChan <-chan interface{},
    46  	workFn WorkFunc,
    47  	errStop bool,
    48  ) (doneC <-chan struct{}) {
    49  	done := make(chan struct{})
    50  	if w.Burst == 0 {
    51  		w.Burst = 1
    52  	}
    53  
    54  	go func() {
    55  		defer close(done)
    56  
    57  		ctx, w.cancel = context.WithCancel(ctx)
    58  		limiter := rate_.NewLimiter(int(w.Burst))
    59  		for {
    60  
    61  			//util the condition is met, need one token, or will be blocked
    62  			limiter.AllowWaitUntil()
    63  
    64  			select {
    65  			case work, ok := <-workChan:
    66  				if !ok {
    67  					//workChan is closed
    68  					return
    69  				}
    70  				w.wg.Add(1)
    71  
    72  				go func(wk interface{}) {
    73  					defer limiter.Put()
    74  					defer w.wg.Done()
    75  
    76  					if err := workFn(wk); err != nil {
    77  						w.TrySetError(err)
    78  						if errStop {
    79  							w.cancel()
    80  						}
    81  						return
    82  					}
    83  
    84  				}(work)
    85  
    86  			case <-ctx.Done():
    87  				//  err: context canceled
    88  				w.TrySetError(ctx.Err())
    89  				return
    90  			}
    91  
    92  		}
    93  	}()
    94  
    95  	return done
    96  
    97  }
    98  
    99  func (w *Worker) Error() error {
   100  	w.errMu.Lock()
   101  	defer w.errMu.Unlock()
   102  
   103  	return errors_.NewAggregate(w.errs)
   104  }
   105  
   106  func (w *Worker) TrySetError(err error) {
   107  	w.errMu.Lock()
   108  	defer w.errMu.Unlock()
   109  	w.errs = append(w.errs, err)
   110  }
   111  
   112  func (w *Worker) Wait() {
   113  	w.wg.Wait()
   114  }
   115  
   116  func (w *Worker) Stop() {
   117  	w.cancel()
   118  }