github.com/mhilton/juju-juju@v0.0.0-20150901100907-a94dd2c73455/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  	"time"
     9  
    10  	"launchpad.net/tomb"
    11  )
    12  
    13  // ErrKilled can be returned by the PeriodicWorkerCall to signify that
    14  // the function has returned as a result of a Stop() or Kill() signal
    15  // and that the function was able to stop cleanly
    16  var ErrKilled = errors.New("worker killed")
    17  
    18  // periodicWorker implements the worker returned by NewPeriodicWorker.
    19  type periodicWorker struct {
    20  	tomb     tomb.Tomb
    21  	newTimer NewTimerFunc
    22  }
    23  
    24  // PeriodicWorkerCall represents the callable to be passed
    25  // to the periodic worker to be run every elapsed period.
    26  type PeriodicWorkerCall func(stop <-chan struct{}) error
    27  
    28  // PeriodicTimer is an interface for the timer that periodicworker
    29  // will use to handle the calls.
    30  type PeriodicTimer interface {
    31  	// Reset changes the timer to expire after duration d.
    32  	// It returns true if the timer had been active, false
    33  	// if the timer had expired or been stopped.
    34  	Reset(time.Duration) bool
    35  	// CountDown returns the channel used to signal expiration of
    36  	// the timer duration. The channel is called C in the base
    37  	// implementation of timer but the name is confusing.
    38  	CountDown() <-chan time.Time
    39  }
    40  
    41  // NewTimerFunc is a constructor used to obtain the instance
    42  // of PeriodicTimer periodicWorker uses on its loop.
    43  type NewTimerFunc func(time.Duration) PeriodicTimer
    44  
    45  // Timer implements PeriodicTimer.
    46  type Timer struct {
    47  	timer *time.Timer
    48  }
    49  
    50  // Reset implements PeriodicTimer.
    51  func (t *Timer) Reset(d time.Duration) bool {
    52  	return t.timer.Reset(d)
    53  }
    54  
    55  // CountDown implements PeriodicTimer.
    56  func (t *Timer) CountDown() <-chan time.Time {
    57  	return t.timer.C
    58  }
    59  
    60  // NewTimer is the default implementation of NewTimerFunc.
    61  func NewTimer(d time.Duration) PeriodicTimer {
    62  	return &Timer{time.NewTimer(d)}
    63  }
    64  
    65  // NewPeriodicWorker returns a worker that runs the given function continually
    66  // sleeping for sleepDuration in between each call, until Kill() is called
    67  // The stopCh argument will be closed when the worker is killed. The error returned
    68  // by the doWork function will be returned by the worker's Wait function.
    69  func NewPeriodicWorker(call PeriodicWorkerCall, period time.Duration, timerFunc NewTimerFunc) Worker {
    70  	w := &periodicWorker{newTimer: timerFunc}
    71  	go func() {
    72  		defer w.tomb.Done()
    73  		w.tomb.Kill(w.run(call, period))
    74  	}()
    75  	return w
    76  }
    77  
    78  func (w *periodicWorker) run(call PeriodicWorkerCall, period time.Duration) error {
    79  	timer := w.newTimer(0)
    80  	stop := w.tomb.Dying()
    81  	for {
    82  		select {
    83  		case <-stop:
    84  			return tomb.ErrDying
    85  		case <-timer.CountDown():
    86  			if err := call(stop); err != nil {
    87  				if err == ErrKilled {
    88  					return tomb.ErrDying
    89  				}
    90  				return err
    91  			}
    92  		}
    93  		timer.Reset(period)
    94  	}
    95  }
    96  
    97  // Kill implements Worker.Kill() and will close the channel given to the doWork
    98  // function.
    99  func (w *periodicWorker) Kill() {
   100  	w.tomb.Kill(nil)
   101  }
   102  
   103  // Wait implements Worker.Wait(), and will return the error returned by
   104  // the doWork function.
   105  func (w *periodicWorker) Wait() error {
   106  	return w.tomb.Wait()
   107  }