github.com/bshelton229/agent@v3.5.4+incompatible/retry/retry.go (about)

     1  package retry
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"math/rand"
     7  	"time"
     8  )
     9  
    10  type Stats struct {
    11  	Attempt   int
    12  	Interval  time.Duration
    13  	Config    *Config
    14  	breakNext bool
    15  }
    16  
    17  type Config struct {
    18  	Maximum  int
    19  	Interval time.Duration
    20  	Forever  bool
    21  	Jitter   bool
    22  }
    23  
    24  // A human readable representation often useful for debugging.
    25  func (s *Stats) String() string {
    26  	str := fmt.Sprintf("Attempt %d/", s.Attempt)
    27  
    28  	if s.Config.Forever {
    29  		str = str + "∞"
    30  	} else {
    31  		str = str + fmt.Sprintf("%d", s.Config.Maximum)
    32  	}
    33  
    34  	if s.Config.Interval > 0 {
    35  		str = str + fmt.Sprintf(" Retrying in %s", s.Interval)
    36  	}
    37  
    38  	return str
    39  }
    40  
    41  // Allow a retry loop to break out of itself
    42  func (s *Stats) Break() {
    43  	s.breakNext = true
    44  }
    45  
    46  func Do(callback func(*Stats) error, config *Config) error {
    47  	var err error
    48  
    49  	// Setup a default config for the retry
    50  	if config == nil {
    51  		config = &Config{Forever: true, Interval: 1 * time.Second, Jitter: false}
    52  	}
    53  
    54  	// If the config isn't set to run forever, and the maximum is 0, set a
    55  	// default of 0
    56  	if config.Maximum == 0 && config.Forever == false {
    57  		config.Maximum = 10
    58  	}
    59  
    60  	// Don't allow a forever retry without an interval
    61  	if config.Forever && config.Interval == 0 {
    62  		return errors.New("You can't do a forever retry with no interval")
    63  	}
    64  
    65  	// The stats struct that is passed to every attempt of the callback
    66  	stats := &Stats{Attempt: 1, Config: config}
    67  
    68  	// Needed for jitter calcs
    69  	random := rand.New(rand.NewSource(time.Now().UnixNano()))
    70  
    71  	for {
    72  		// Preconfigure the interval that will be used (so that we have
    73  		// access to it in the callback)
    74  		stats.Interval = config.Interval
    75  		if config.Jitter {
    76  			stats.Interval = stats.Interval + (time.Duration(1000*random.Float32()) * time.Millisecond)
    77  		}
    78  
    79  		// Attempt the callback
    80  		err = callback(stats)
    81  		if err == nil {
    82  			return nil
    83  		}
    84  
    85  		// If the loop has callen stats.Break(), we should cancel out
    86  		// of the loop
    87  		if stats.breakNext {
    88  			return err
    89  		}
    90  
    91  		// Bump the attempt number
    92  		stats.Attempt = stats.Attempt + 1
    93  
    94  		// Try the callback again after the interval
    95  		time.Sleep(stats.Interval)
    96  
    97  		if !stats.Config.Forever {
    98  			// Should we give up?
    99  			if stats.Attempt > stats.Config.Maximum {
   100  				break
   101  			}
   102  		}
   103  	}
   104  
   105  	return err
   106  }