github.com/amp-space/amp-sdk-go@v0.7.6/stdlib/task/pool_worker.go (about)

     1  package task
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  )
     8  
     9  type PoolWorker interface {
    10  	Context
    11  	Add(item PoolWorkerItem)
    12  	ForceRetry(uniqueID PoolUniqueID)
    13  }
    14  
    15  type PoolWorkerItem interface {
    16  	PoolUniqueIDer
    17  	Work(ctx context.Context) (retry bool)
    18  }
    19  
    20  type PoolWorkerScheduler interface {
    21  	CheckForRetriesInterval() time.Duration
    22  	RetryWhen(item PoolWorkerItem) time.Time
    23  }
    24  
    25  type poolWorker struct {
    26  	Context
    27  	name        string
    28  	concurrency int
    29  	pool        *Pool
    30  	scheduler   PoolWorkerScheduler
    31  }
    32  
    33  func StartNewPoolWorker(name string, concurrency int, scheduler PoolWorkerScheduler) (*poolWorker, error) {
    34  	w := &poolWorker{
    35  		name:        name,
    36  		concurrency: concurrency,
    37  		scheduler:   scheduler,
    38  	}
    39  	
    40  	w.pool = StartNewPool(name, concurrency, scheduler.CheckForRetriesInterval())
    41  	var err error
    42  	w.Context, err = Start(&Task{
    43  		Label: "tester",
    44  		OnStart: w.OnContextStarted,
    45  	})
    46  	return w, err
    47  }
    48  
    49  func (w *poolWorker) OnContextStarted(ctx Context) error {
    50  	_, err := w.StartChild(&Task{
    51  		Label: "poolWorker",
    52  		OnStart: w.pool.OnContextStarted,
    53  	})
    54  	if err != nil {
    55  		return err
    56  	}
    57  
    58  	for i := 0; i < w.concurrency; i++ {
    59  		w.Context.Go(fmt.Sprintf("worker %v", i), func(ctx Context) {
    60  			for {
    61  				select {
    62  				case <-ctx.Done():
    63  					return
    64  				default:
    65  				}
    66  
    67  				x, err := w.pool.Get(ctx)
    68  				if err != nil {
    69  					return
    70  				}
    71  				item := x.(PoolWorkerItem)
    72  
    73  				retry := item.Work(ctx)
    74  				if retry {
    75  					w.pool.RetryLater(item.ID(), w.scheduler.RetryWhen(item))
    76  				} else {
    77  					w.pool.Complete(item.ID())
    78  				}
    79  			}
    80  		})
    81  	}
    82  	return nil
    83  }
    84  
    85  func (w *poolWorker) Add(item PoolWorkerItem) {
    86  	w.pool.Add(item)
    87  }
    88  
    89  func (w *poolWorker) ForceRetry(id PoolUniqueID) {
    90  	w.pool.ForceRetry(id)
    91  }
    92  
    93  type StaticScheduler struct {
    94  	checkForRetriesInterval time.Duration
    95  	retryAfter              time.Duration
    96  }
    97  
    98  var _ PoolWorkerScheduler = StaticScheduler{}
    99  
   100  func NewStaticScheduler(checkForRetriesInterval time.Duration, retryAfter time.Duration) StaticScheduler {
   101  	return StaticScheduler{checkForRetriesInterval, retryAfter}
   102  }
   103  
   104  func (s StaticScheduler) CheckForRetriesInterval() time.Duration { return s.checkForRetriesInterval }
   105  func (s StaticScheduler) RetryWhen(item PoolWorkerItem) time.Time {
   106  	return time.Now().Add(s.retryAfter)
   107  }