backoffhandler.go 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. package retry
  2. import (
  3. "context"
  4. "math/rand"
  5. "time"
  6. )
  7. const (
  8. DefaultBaseTime time.Duration = time.Second
  9. )
  10. // Redeclare time functions so they can be overridden in tests.
  11. type Clock struct {
  12. Now func() time.Time
  13. After func(d time.Duration) <-chan time.Time
  14. }
  15. // BackoffHandler manages exponential backoff and limits the maximum number of retries.
  16. // The base time period is 1 second, doubling with each retry.
  17. // After initial success, a grace period can be set to reset the backoff timer if
  18. // a connection is maintained successfully for a long enough period. The base grace period
  19. // is 2 seconds, doubling with each retry.
  20. type BackoffHandler struct {
  21. // MaxRetries sets the maximum number of retries to perform. The default value
  22. // of 0 disables retry completely.
  23. maxRetries uint
  24. // RetryForever caps the exponential backoff period according to MaxRetries
  25. // but allows you to retry indefinitely.
  26. retryForever bool
  27. // BaseTime sets the initial backoff period.
  28. baseTime time.Duration
  29. retries uint
  30. resetDeadline time.Time
  31. Clock Clock
  32. }
  33. func NewBackoff(maxRetries uint, baseTime time.Duration, retryForever bool) BackoffHandler {
  34. return BackoffHandler{
  35. maxRetries: maxRetries,
  36. baseTime: baseTime,
  37. retryForever: retryForever,
  38. Clock: Clock{Now: time.Now, After: time.After},
  39. }
  40. }
  41. func (b BackoffHandler) GetMaxBackoffDuration(ctx context.Context) (time.Duration, bool) {
  42. // Follows the same logic as Backoff, but without mutating the receiver.
  43. // This select has to happen first to reflect the actual behaviour of the Backoff function.
  44. select {
  45. case <-ctx.Done():
  46. return time.Duration(0), false
  47. default:
  48. }
  49. if !b.resetDeadline.IsZero() && b.Clock.Now().After(b.resetDeadline) {
  50. // b.retries would be set to 0 at this point
  51. return time.Second, true
  52. }
  53. if b.retries >= b.maxRetries && !b.retryForever {
  54. return time.Duration(0), false
  55. }
  56. maxTimeToWait := b.GetBaseTime() * 1 << (b.retries + 1)
  57. return maxTimeToWait, true
  58. }
  59. // BackoffTimer returns a channel that sends the current time when the exponential backoff timeout expires.
  60. // Returns nil if the maximum number of retries have been used.
  61. func (b *BackoffHandler) BackoffTimer() <-chan time.Time {
  62. if !b.resetDeadline.IsZero() && b.Clock.Now().After(b.resetDeadline) {
  63. b.retries = 0
  64. b.resetDeadline = time.Time{}
  65. }
  66. if b.retries >= b.maxRetries {
  67. if !b.retryForever {
  68. return nil
  69. }
  70. } else {
  71. b.retries++
  72. }
  73. maxTimeToWait := time.Duration(b.GetBaseTime() * 1 << (b.retries))
  74. timeToWait := time.Duration(rand.Int63n(maxTimeToWait.Nanoseconds()))
  75. return b.Clock.After(timeToWait)
  76. }
  77. // Backoff is used to wait according to exponential backoff. Returns false if the
  78. // maximum number of retries have been used or if the underlying context has been cancelled.
  79. func (b *BackoffHandler) Backoff(ctx context.Context) bool {
  80. c := b.BackoffTimer()
  81. if c == nil {
  82. return false
  83. }
  84. select {
  85. case <-c:
  86. return true
  87. case <-ctx.Done():
  88. return false
  89. }
  90. }
  91. // Sets a grace period within which the the backoff timer is maintained. After the grace
  92. // period expires, the number of retries & backoff duration is reset.
  93. func (b *BackoffHandler) SetGracePeriod() time.Duration {
  94. maxTimeToWait := b.GetBaseTime() * 2 << (b.retries + 1)
  95. timeToWait := time.Duration(rand.Int63n(maxTimeToWait.Nanoseconds()))
  96. b.resetDeadline = b.Clock.Now().Add(timeToWait)
  97. return timeToWait
  98. }
  99. func (b BackoffHandler) GetBaseTime() time.Duration {
  100. if b.baseTime == 0 {
  101. return DefaultBaseTime
  102. }
  103. return b.baseTime
  104. }
  105. // Retries returns the number of retries consumed so far.
  106. func (b *BackoffHandler) Retries() int {
  107. return int(b.retries)
  108. }
  109. func (b *BackoffHandler) ReachedMaxRetries() bool {
  110. return b.retries == b.maxRetries
  111. }
  112. func (b *BackoffHandler) ResetNow() {
  113. b.resetDeadline = b.Clock.Now()
  114. b.retries = 0
  115. }