github.com/juju/juju@v0.0.0-20240327075706-a90865de2538/worker/periodicworker.go (about)

     1  // Copyright 2014 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package worker
     5  
     6  import (
     7  	"errors"
     8  	"math/rand"
     9  	"time"
    10  
    11  	"github.com/juju/worker/v3"
    12  	"gopkg.in/tomb.v2"
    13  )
    14  
    15  // ErrKilled can be returned by the PeriodicWorkerCall to signify that
    16  // the function has returned as a result of a Stop() or Kill() signal
    17  // and that the function was able to stop cleanly
    18  var ErrKilled = errors.New("worker killed")
    19  
    20  // PeriodicWorkerOption is an optional parameter of the NewPeriodicWorker function and can be
    21  // used to set optional parameters of the new periodic worker.
    22  type PeriodicWorkerOption func(w *periodicWorker)
    23  
    24  // Jitter will introduce a jitter in the worker's period by the specified amount (as percents - i.e. between 0 and 1).
    25  func Jitter(amount float64) PeriodicWorkerOption {
    26  	return func(w *periodicWorker) {
    27  		w.jitter = amount
    28  	}
    29  }
    30  
    31  // periodicWorker implements the worker returned by NewPeriodicWorker.
    32  type periodicWorker struct {
    33  	jitter   float64
    34  	tomb     tomb.Tomb
    35  	newTimer NewTimerFunc
    36  }
    37  
    38  // PeriodicWorkerCall represents the callable to be passed
    39  // to the periodic worker to be run every elapsed period.
    40  type PeriodicWorkerCall func(stop <-chan struct{}) error
    41  
    42  // PeriodicTimer is an interface for the timer that periodicworker
    43  // will use to handle the calls.
    44  type PeriodicTimer interface {
    45  	// Reset changes the timer to expire after duration d.
    46  	// It returns true if the timer had been active, false
    47  	// if the timer had expired or been stopped.
    48  	Reset(time.Duration) bool
    49  	// CountDown returns the channel used to signal expiration of
    50  	// the timer duration. The channel is called C in the base
    51  	// implementation of timer but the name is confusing.
    52  	CountDown() <-chan time.Time
    53  }
    54  
    55  // NewTimerFunc is a constructor used to obtain the instance
    56  // of PeriodicTimer periodicWorker uses on its loop.
    57  // TODO(fwereade): 2016-03-17 lp:1558657
    58  type NewTimerFunc func(time.Duration) PeriodicTimer
    59  
    60  // Timer implements PeriodicTimer.
    61  type Timer struct {
    62  	timer *time.Timer
    63  }
    64  
    65  // Reset implements PeriodicTimer.
    66  func (t *Timer) Reset(d time.Duration) bool {
    67  	return t.timer.Reset(d)
    68  }
    69  
    70  // CountDown implements PeriodicTimer.
    71  func (t *Timer) CountDown() <-chan time.Time {
    72  	return t.timer.C
    73  }
    74  
    75  // NewTimer is the default implementation of NewTimerFunc.
    76  func NewTimer(d time.Duration) PeriodicTimer {
    77  	return &Timer{time.NewTimer(d)}
    78  }
    79  
    80  // NewPeriodicWorker returns a worker that runs the given function continually
    81  // sleeping for sleepDuration in between each call, until Kill() is called
    82  // The stopCh argument will be closed when the worker is killed. The error returned
    83  // by the doWork function will be returned by the worker's Wait function.
    84  func NewPeriodicWorker(call PeriodicWorkerCall, period time.Duration, timerFunc NewTimerFunc, options ...PeriodicWorkerOption) worker.Worker {
    85  	w := &periodicWorker{newTimer: timerFunc}
    86  	for _, option := range options {
    87  		option(w)
    88  	}
    89  	w.tomb.Go(func() error {
    90  		return w.run(call, period)
    91  	})
    92  	return w
    93  }
    94  
    95  func (w *periodicWorker) run(call PeriodicWorkerCall, period time.Duration) error {
    96  	timer := w.newTimer(0)
    97  	stop := w.tomb.Dying()
    98  	for {
    99  		select {
   100  		case <-stop:
   101  			return tomb.ErrDying
   102  		case <-timer.CountDown():
   103  			if err := call(stop); err != nil {
   104  				if err == ErrKilled {
   105  					return tomb.ErrDying
   106  				}
   107  				return err
   108  			}
   109  		}
   110  		timer.Reset(nextPeriod(period, w.jitter))
   111  	}
   112  }
   113  
   114  var nextPeriod = func(period time.Duration, jitter float64) time.Duration {
   115  	r := rand.New(rand.NewSource(time.Now().Unix()))
   116  	p := period
   117  	if jitter != 0 {
   118  		lower := (1.0 - jitter) * float64(period)
   119  		window := (2.0 * jitter) * float64(period)
   120  		offset := float64(r.Int63n(int64(window)))
   121  		p = time.Duration(lower + offset)
   122  	}
   123  	return p
   124  }
   125  
   126  // Kill implements Worker.Kill() and will close the channel given to the doWork
   127  // function.
   128  func (w *periodicWorker) Kill() {
   129  	w.tomb.Kill(nil)
   130  }
   131  
   132  // Wait implements Worker.Wait(), and will return the error returned by
   133  // the doWork function.
   134  func (w *periodicWorker) Wait() error {
   135  	return w.tomb.Wait()
   136  }