wait.go 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757
  1. /*
  2. Copyright 2014 The Kubernetes Authors.
  3. Licensed under the Apache License, Version 2.0 (the "License");
  4. you may not use this file except in compliance with the License.
  5. You may obtain a copy of the License at
  6. http://www.apache.org/licenses/LICENSE-2.0
  7. Unless required by applicable law or agreed to in writing, software
  8. distributed under the License is distributed on an "AS IS" BASIS,
  9. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  10. See the License for the specific language governing permissions and
  11. limitations under the License.
  12. */
  13. package wait
  14. import (
  15. "context"
  16. "errors"
  17. "math"
  18. "math/rand"
  19. "sync"
  20. "time"
  21. "k8s.io/apimachinery/pkg/util/runtime"
  22. "k8s.io/utils/clock"
  23. )
  24. // For any test of the style:
  25. //
  26. // ...
  27. // <- time.After(timeout):
  28. // t.Errorf("Timed out")
  29. //
  30. // The value for timeout should effectively be "forever." Obviously we don't want our tests to truly lock up forever, but 30s
  31. // is long enough that it is effectively forever for the things that can slow down a run on a heavily contended machine
  32. // (GC, seeks, etc), but not so long as to make a developer ctrl-c a test run if they do happen to break that test.
  33. var ForeverTestTimeout = time.Second * 30
  34. // NeverStop may be passed to Until to make it never stop.
  35. var NeverStop <-chan struct{} = make(chan struct{})
  36. // Group allows to start a group of goroutines and wait for their completion.
  37. type Group struct {
  38. wg sync.WaitGroup
  39. }
  40. func (g *Group) Wait() {
  41. g.wg.Wait()
  42. }
  43. // StartWithChannel starts f in a new goroutine in the group.
  44. // stopCh is passed to f as an argument. f should stop when stopCh is available.
  45. func (g *Group) StartWithChannel(stopCh <-chan struct{}, f func(stopCh <-chan struct{})) {
  46. g.Start(func() {
  47. f(stopCh)
  48. })
  49. }
  50. // StartWithContext starts f in a new goroutine in the group.
  51. // ctx is passed to f as an argument. f should stop when ctx.Done() is available.
  52. func (g *Group) StartWithContext(ctx context.Context, f func(context.Context)) {
  53. g.Start(func() {
  54. f(ctx)
  55. })
  56. }
  57. // Start starts f in a new goroutine in the group.
  58. func (g *Group) Start(f func()) {
  59. g.wg.Add(1)
  60. go func() {
  61. defer g.wg.Done()
  62. f()
  63. }()
  64. }
  65. // Forever calls f every period for ever.
  66. //
  67. // Forever is syntactic sugar on top of Until.
  68. func Forever(f func(), period time.Duration) {
  69. Until(f, period, NeverStop)
  70. }
  71. // Until loops until stop channel is closed, running f every period.
  72. //
  73. // Until is syntactic sugar on top of JitterUntil with zero jitter factor and
  74. // with sliding = true (which means the timer for period starts after the f
  75. // completes).
  76. func Until(f func(), period time.Duration, stopCh <-chan struct{}) {
  77. JitterUntil(f, period, 0.0, true, stopCh)
  78. }
  79. // UntilWithContext loops until context is done, running f every period.
  80. //
  81. // UntilWithContext is syntactic sugar on top of JitterUntilWithContext
  82. // with zero jitter factor and with sliding = true (which means the timer
  83. // for period starts after the f completes).
  84. func UntilWithContext(ctx context.Context, f func(context.Context), period time.Duration) {
  85. JitterUntilWithContext(ctx, f, period, 0.0, true)
  86. }
  87. // NonSlidingUntil loops until stop channel is closed, running f every
  88. // period.
  89. //
  90. // NonSlidingUntil is syntactic sugar on top of JitterUntil with zero jitter
  91. // factor, with sliding = false (meaning the timer for period starts at the same
  92. // time as the function starts).
  93. func NonSlidingUntil(f func(), period time.Duration, stopCh <-chan struct{}) {
  94. JitterUntil(f, period, 0.0, false, stopCh)
  95. }
  96. // NonSlidingUntilWithContext loops until context is done, running f every
  97. // period.
  98. //
  99. // NonSlidingUntilWithContext is syntactic sugar on top of JitterUntilWithContext
  100. // with zero jitter factor, with sliding = false (meaning the timer for period
  101. // starts at the same time as the function starts).
  102. func NonSlidingUntilWithContext(ctx context.Context, f func(context.Context), period time.Duration) {
  103. JitterUntilWithContext(ctx, f, period, 0.0, false)
  104. }
  105. // JitterUntil loops until stop channel is closed, running f every period.
  106. //
  107. // If jitterFactor is positive, the period is jittered before every run of f.
  108. // If jitterFactor is not positive, the period is unchanged and not jittered.
  109. //
  110. // If sliding is true, the period is computed after f runs. If it is false then
  111. // period includes the runtime for f.
  112. //
  113. // Close stopCh to stop. f may not be invoked if stop channel is already
  114. // closed. Pass NeverStop to if you don't want it stop.
  115. func JitterUntil(f func(), period time.Duration, jitterFactor float64, sliding bool, stopCh <-chan struct{}) {
  116. BackoffUntil(f, NewJitteredBackoffManager(period, jitterFactor, &clock.RealClock{}), sliding, stopCh)
  117. }
  118. // BackoffUntil loops until stop channel is closed, run f every duration given by BackoffManager.
  119. //
  120. // If sliding is true, the period is computed after f runs. If it is false then
  121. // period includes the runtime for f.
  122. func BackoffUntil(f func(), backoff BackoffManager, sliding bool, stopCh <-chan struct{}) {
  123. var t clock.Timer
  124. for {
  125. select {
  126. case <-stopCh:
  127. return
  128. default:
  129. }
  130. if !sliding {
  131. t = backoff.Backoff()
  132. }
  133. func() {
  134. defer runtime.HandleCrash()
  135. f()
  136. }()
  137. if sliding {
  138. t = backoff.Backoff()
  139. }
  140. // NOTE: b/c there is no priority selection in golang
  141. // it is possible for this to race, meaning we could
  142. // trigger t.C and stopCh, and t.C select falls through.
  143. // In order to mitigate we re-check stopCh at the beginning
  144. // of every loop to prevent extra executions of f().
  145. select {
  146. case <-stopCh:
  147. if !t.Stop() {
  148. <-t.C()
  149. }
  150. return
  151. case <-t.C():
  152. }
  153. }
  154. }
  155. // JitterUntilWithContext loops until context is done, running f every period.
  156. //
  157. // If jitterFactor is positive, the period is jittered before every run of f.
  158. // If jitterFactor is not positive, the period is unchanged and not jittered.
  159. //
  160. // If sliding is true, the period is computed after f runs. If it is false then
  161. // period includes the runtime for f.
  162. //
  163. // Cancel context to stop. f may not be invoked if context is already expired.
  164. func JitterUntilWithContext(ctx context.Context, f func(context.Context), period time.Duration, jitterFactor float64, sliding bool) {
  165. JitterUntil(func() { f(ctx) }, period, jitterFactor, sliding, ctx.Done())
  166. }
  167. // Jitter returns a time.Duration between duration and duration + maxFactor *
  168. // duration.
  169. //
  170. // This allows clients to avoid converging on periodic behavior. If maxFactor
  171. // is 0.0, a suggested default value will be chosen.
  172. func Jitter(duration time.Duration, maxFactor float64) time.Duration {
  173. if maxFactor <= 0.0 {
  174. maxFactor = 1.0
  175. }
  176. wait := duration + time.Duration(rand.Float64()*maxFactor*float64(duration))
  177. return wait
  178. }
  179. // ErrWaitTimeout is returned when the condition exited without success.
  180. var ErrWaitTimeout = errors.New("timed out waiting for the condition")
  181. // ConditionFunc returns true if the condition is satisfied, or an error
  182. // if the loop should be aborted.
  183. type ConditionFunc func() (done bool, err error)
  184. // ConditionWithContextFunc returns true if the condition is satisfied, or an error
  185. // if the loop should be aborted.
  186. //
  187. // The caller passes along a context that can be used by the condition function.
  188. type ConditionWithContextFunc func(context.Context) (done bool, err error)
  189. // WithContext converts a ConditionFunc into a ConditionWithContextFunc
  190. func (cf ConditionFunc) WithContext() ConditionWithContextFunc {
  191. return func(context.Context) (done bool, err error) {
  192. return cf()
  193. }
  194. }
  195. // runConditionWithCrashProtection runs a ConditionFunc with crash protection
  196. func runConditionWithCrashProtection(condition ConditionFunc) (bool, error) {
  197. return runConditionWithCrashProtectionWithContext(context.TODO(), condition.WithContext())
  198. }
  199. // runConditionWithCrashProtectionWithContext runs a
  200. // ConditionWithContextFunc with crash protection.
  201. func runConditionWithCrashProtectionWithContext(ctx context.Context, condition ConditionWithContextFunc) (bool, error) {
  202. defer runtime.HandleCrash()
  203. return condition(ctx)
  204. }
  205. // Backoff holds parameters applied to a Backoff function.
  206. type Backoff struct {
  207. // The initial duration.
  208. Duration time.Duration
  209. // Duration is multiplied by factor each iteration, if factor is not zero
  210. // and the limits imposed by Steps and Cap have not been reached.
  211. // Should not be negative.
  212. // The jitter does not contribute to the updates to the duration parameter.
  213. Factor float64
  214. // The sleep at each iteration is the duration plus an additional
  215. // amount chosen uniformly at random from the interval between
  216. // zero and `jitter*duration`.
  217. Jitter float64
  218. // The remaining number of iterations in which the duration
  219. // parameter may change (but progress can be stopped earlier by
  220. // hitting the cap). If not positive, the duration is not
  221. // changed. Used for exponential backoff in combination with
  222. // Factor and Cap.
  223. Steps int
  224. // A limit on revised values of the duration parameter. If a
  225. // multiplication by the factor parameter would make the duration
  226. // exceed the cap then the duration is set to the cap and the
  227. // steps parameter is set to zero.
  228. Cap time.Duration
  229. }
  230. // Step (1) returns an amount of time to sleep determined by the
  231. // original Duration and Jitter and (2) mutates the provided Backoff
  232. // to update its Steps and Duration.
  233. func (b *Backoff) Step() time.Duration {
  234. if b.Steps < 1 {
  235. if b.Jitter > 0 {
  236. return Jitter(b.Duration, b.Jitter)
  237. }
  238. return b.Duration
  239. }
  240. b.Steps--
  241. duration := b.Duration
  242. // calculate the next step
  243. if b.Factor != 0 {
  244. b.Duration = time.Duration(float64(b.Duration) * b.Factor)
  245. if b.Cap > 0 && b.Duration > b.Cap {
  246. b.Duration = b.Cap
  247. b.Steps = 0
  248. }
  249. }
  250. if b.Jitter > 0 {
  251. duration = Jitter(duration, b.Jitter)
  252. }
  253. return duration
  254. }
  255. // ContextForChannel derives a child context from a parent channel.
  256. //
  257. // The derived context's Done channel is closed when the returned cancel function
  258. // is called or when the parent channel is closed, whichever happens first.
  259. //
  260. // Note the caller must *always* call the CancelFunc, otherwise resources may be leaked.
  261. func ContextForChannel(parentCh <-chan struct{}) (context.Context, context.CancelFunc) {
  262. ctx, cancel := context.WithCancel(context.Background())
  263. go func() {
  264. select {
  265. case <-parentCh:
  266. cancel()
  267. case <-ctx.Done():
  268. }
  269. }()
  270. return ctx, cancel
  271. }
  272. // BackoffManager manages backoff with a particular scheme based on its underlying implementation. It provides
  273. // an interface to return a timer for backoff, and caller shall backoff until Timer.C() drains. If the second Backoff()
  274. // is called before the timer from the first Backoff() call finishes, the first timer will NOT be drained and result in
  275. // undetermined behavior.
  276. // The BackoffManager is supposed to be called in a single-threaded environment.
  277. type BackoffManager interface {
  278. Backoff() clock.Timer
  279. }
  280. type exponentialBackoffManagerImpl struct {
  281. backoff *Backoff
  282. backoffTimer clock.Timer
  283. lastBackoffStart time.Time
  284. initialBackoff time.Duration
  285. backoffResetDuration time.Duration
  286. clock clock.Clock
  287. }
  288. // NewExponentialBackoffManager returns a manager for managing exponential backoff. Each backoff is jittered and
  289. // backoff will not exceed the given max. If the backoff is not called within resetDuration, the backoff is reset.
  290. // This backoff manager is used to reduce load during upstream unhealthiness.
  291. func NewExponentialBackoffManager(initBackoff, maxBackoff, resetDuration time.Duration, backoffFactor, jitter float64, c clock.Clock) BackoffManager {
  292. return &exponentialBackoffManagerImpl{
  293. backoff: &Backoff{
  294. Duration: initBackoff,
  295. Factor: backoffFactor,
  296. Jitter: jitter,
  297. // the current impl of wait.Backoff returns Backoff.Duration once steps are used up, which is not
  298. // what we ideally need here, we set it to max int and assume we will never use up the steps
  299. Steps: math.MaxInt32,
  300. Cap: maxBackoff,
  301. },
  302. backoffTimer: nil,
  303. initialBackoff: initBackoff,
  304. lastBackoffStart: c.Now(),
  305. backoffResetDuration: resetDuration,
  306. clock: c,
  307. }
  308. }
  309. func (b *exponentialBackoffManagerImpl) getNextBackoff() time.Duration {
  310. if b.clock.Now().Sub(b.lastBackoffStart) > b.backoffResetDuration {
  311. b.backoff.Steps = math.MaxInt32
  312. b.backoff.Duration = b.initialBackoff
  313. }
  314. b.lastBackoffStart = b.clock.Now()
  315. return b.backoff.Step()
  316. }
  317. // Backoff implements BackoffManager.Backoff, it returns a timer so caller can block on the timer for exponential backoff.
  318. // The returned timer must be drained before calling Backoff() the second time
  319. func (b *exponentialBackoffManagerImpl) Backoff() clock.Timer {
  320. if b.backoffTimer == nil {
  321. b.backoffTimer = b.clock.NewTimer(b.getNextBackoff())
  322. } else {
  323. b.backoffTimer.Reset(b.getNextBackoff())
  324. }
  325. return b.backoffTimer
  326. }
  327. type jitteredBackoffManagerImpl struct {
  328. clock clock.Clock
  329. duration time.Duration
  330. jitter float64
  331. backoffTimer clock.Timer
  332. }
  333. // NewJitteredBackoffManager returns a BackoffManager that backoffs with given duration plus given jitter. If the jitter
  334. // is negative, backoff will not be jittered.
  335. func NewJitteredBackoffManager(duration time.Duration, jitter float64, c clock.Clock) BackoffManager {
  336. return &jitteredBackoffManagerImpl{
  337. clock: c,
  338. duration: duration,
  339. jitter: jitter,
  340. backoffTimer: nil,
  341. }
  342. }
  343. func (j *jitteredBackoffManagerImpl) getNextBackoff() time.Duration {
  344. jitteredPeriod := j.duration
  345. if j.jitter > 0.0 {
  346. jitteredPeriod = Jitter(j.duration, j.jitter)
  347. }
  348. return jitteredPeriod
  349. }
  350. // Backoff implements BackoffManager.Backoff, it returns a timer so caller can block on the timer for jittered backoff.
  351. // The returned timer must be drained before calling Backoff() the second time
  352. func (j *jitteredBackoffManagerImpl) Backoff() clock.Timer {
  353. backoff := j.getNextBackoff()
  354. if j.backoffTimer == nil {
  355. j.backoffTimer = j.clock.NewTimer(backoff)
  356. } else {
  357. j.backoffTimer.Reset(backoff)
  358. }
  359. return j.backoffTimer
  360. }
  361. // ExponentialBackoff repeats a condition check with exponential backoff.
  362. //
  363. // It repeatedly checks the condition and then sleeps, using `backoff.Step()`
  364. // to determine the length of the sleep and adjust Duration and Steps.
  365. // Stops and returns as soon as:
  366. // 1. the condition check returns true or an error,
  367. // 2. `backoff.Steps` checks of the condition have been done, or
  368. // 3. a sleep truncated by the cap on duration has been completed.
  369. // In case (1) the returned error is what the condition function returned.
  370. // In all other cases, ErrWaitTimeout is returned.
  371. func ExponentialBackoff(backoff Backoff, condition ConditionFunc) error {
  372. for backoff.Steps > 0 {
  373. if ok, err := runConditionWithCrashProtection(condition); err != nil || ok {
  374. return err
  375. }
  376. if backoff.Steps == 1 {
  377. break
  378. }
  379. time.Sleep(backoff.Step())
  380. }
  381. return ErrWaitTimeout
  382. }
  383. // Poll tries a condition func until it returns true, an error, or the timeout
  384. // is reached.
  385. //
  386. // Poll always waits the interval before the run of 'condition'.
  387. // 'condition' will always be invoked at least once.
  388. //
  389. // Some intervals may be missed if the condition takes too long or the time
  390. // window is too short.
  391. //
  392. // If you want to Poll something forever, see PollInfinite.
  393. func Poll(interval, timeout time.Duration, condition ConditionFunc) error {
  394. return PollWithContext(context.Background(), interval, timeout, condition.WithContext())
  395. }
  396. // PollWithContext tries a condition func until it returns true, an error,
  397. // or when the context expires or the timeout is reached, whichever
  398. // happens first.
  399. //
  400. // PollWithContext always waits the interval before the run of 'condition'.
  401. // 'condition' will always be invoked at least once.
  402. //
  403. // Some intervals may be missed if the condition takes too long or the time
  404. // window is too short.
  405. //
  406. // If you want to Poll something forever, see PollInfinite.
  407. func PollWithContext(ctx context.Context, interval, timeout time.Duration, condition ConditionWithContextFunc) error {
  408. return poll(ctx, false, poller(interval, timeout), condition)
  409. }
  410. // PollUntil tries a condition func until it returns true, an error or stopCh is
  411. // closed.
  412. //
  413. // PollUntil always waits interval before the first run of 'condition'.
  414. // 'condition' will always be invoked at least once.
  415. func PollUntil(interval time.Duration, condition ConditionFunc, stopCh <-chan struct{}) error {
  416. ctx, cancel := ContextForChannel(stopCh)
  417. defer cancel()
  418. return PollUntilWithContext(ctx, interval, condition.WithContext())
  419. }
  420. // PollUntilWithContext tries a condition func until it returns true,
  421. // an error or the specified context is cancelled or expired.
  422. //
  423. // PollUntilWithContext always waits interval before the first run of 'condition'.
  424. // 'condition' will always be invoked at least once.
  425. func PollUntilWithContext(ctx context.Context, interval time.Duration, condition ConditionWithContextFunc) error {
  426. return poll(ctx, false, poller(interval, 0), condition)
  427. }
  428. // PollInfinite tries a condition func until it returns true or an error
  429. //
  430. // PollInfinite always waits the interval before the run of 'condition'.
  431. //
  432. // Some intervals may be missed if the condition takes too long or the time
  433. // window is too short.
  434. func PollInfinite(interval time.Duration, condition ConditionFunc) error {
  435. return PollInfiniteWithContext(context.Background(), interval, condition.WithContext())
  436. }
  437. // PollInfiniteWithContext tries a condition func until it returns true or an error
  438. //
  439. // PollInfiniteWithContext always waits the interval before the run of 'condition'.
  440. //
  441. // Some intervals may be missed if the condition takes too long or the time
  442. // window is too short.
  443. func PollInfiniteWithContext(ctx context.Context, interval time.Duration, condition ConditionWithContextFunc) error {
  444. return poll(ctx, false, poller(interval, 0), condition)
  445. }
  446. // PollImmediate tries a condition func until it returns true, an error, or the timeout
  447. // is reached.
  448. //
  449. // PollImmediate always checks 'condition' before waiting for the interval. 'condition'
  450. // will always be invoked at least once.
  451. //
  452. // Some intervals may be missed if the condition takes too long or the time
  453. // window is too short.
  454. //
  455. // If you want to immediately Poll something forever, see PollImmediateInfinite.
  456. func PollImmediate(interval, timeout time.Duration, condition ConditionFunc) error {
  457. return PollImmediateWithContext(context.Background(), interval, timeout, condition.WithContext())
  458. }
  459. // PollImmediateWithContext tries a condition func until it returns true, an error,
  460. // or the timeout is reached or the specified context expires, whichever happens first.
  461. //
  462. // PollImmediateWithContext always checks 'condition' before waiting for the interval.
  463. // 'condition' will always be invoked at least once.
  464. //
  465. // Some intervals may be missed if the condition takes too long or the time
  466. // window is too short.
  467. //
  468. // If you want to immediately Poll something forever, see PollImmediateInfinite.
  469. func PollImmediateWithContext(ctx context.Context, interval, timeout time.Duration, condition ConditionWithContextFunc) error {
  470. return poll(ctx, true, poller(interval, timeout), condition)
  471. }
  472. // PollImmediateUntil tries a condition func until it returns true, an error or stopCh is closed.
  473. //
  474. // PollImmediateUntil runs the 'condition' before waiting for the interval.
  475. // 'condition' will always be invoked at least once.
  476. func PollImmediateUntil(interval time.Duration, condition ConditionFunc, stopCh <-chan struct{}) error {
  477. ctx, cancel := ContextForChannel(stopCh)
  478. defer cancel()
  479. return PollImmediateUntilWithContext(ctx, interval, condition.WithContext())
  480. }
  481. // PollImmediateUntilWithContext tries a condition func until it returns true,
  482. // an error or the specified context is cancelled or expired.
  483. //
  484. // PollImmediateUntilWithContext runs the 'condition' before waiting for the interval.
  485. // 'condition' will always be invoked at least once.
  486. func PollImmediateUntilWithContext(ctx context.Context, interval time.Duration, condition ConditionWithContextFunc) error {
  487. return poll(ctx, true, poller(interval, 0), condition)
  488. }
  489. // PollImmediateInfinite tries a condition func until it returns true or an error
  490. //
  491. // PollImmediateInfinite runs the 'condition' before waiting for the interval.
  492. //
  493. // Some intervals may be missed if the condition takes too long or the time
  494. // window is too short.
  495. func PollImmediateInfinite(interval time.Duration, condition ConditionFunc) error {
  496. return PollImmediateInfiniteWithContext(context.Background(), interval, condition.WithContext())
  497. }
  498. // PollImmediateInfiniteWithContext tries a condition func until it returns true
  499. // or an error or the specified context gets cancelled or expired.
  500. //
  501. // PollImmediateInfiniteWithContext runs the 'condition' before waiting for the interval.
  502. //
  503. // Some intervals may be missed if the condition takes too long or the time
  504. // window is too short.
  505. func PollImmediateInfiniteWithContext(ctx context.Context, interval time.Duration, condition ConditionWithContextFunc) error {
  506. return poll(ctx, true, poller(interval, 0), condition)
  507. }
  508. // Internally used, each of the public 'Poll*' function defined in this
  509. // package should invoke this internal function with appropriate parameters.
  510. // ctx: the context specified by the caller, for infinite polling pass
  511. // a context that never gets cancelled or expired.
  512. // immediate: if true, the 'condition' will be invoked before waiting for the interval,
  513. // in this case 'condition' will always be invoked at least once.
  514. // wait: user specified WaitFunc function that controls at what interval the condition
  515. // function should be invoked periodically and whether it is bound by a timeout.
  516. // condition: user specified ConditionWithContextFunc function.
  517. func poll(ctx context.Context, immediate bool, wait WaitWithContextFunc, condition ConditionWithContextFunc) error {
  518. if immediate {
  519. done, err := runConditionWithCrashProtectionWithContext(ctx, condition)
  520. if err != nil {
  521. return err
  522. }
  523. if done {
  524. return nil
  525. }
  526. }
  527. select {
  528. case <-ctx.Done():
  529. // returning ctx.Err() will break backward compatibility
  530. return ErrWaitTimeout
  531. default:
  532. return WaitForWithContext(ctx, wait, condition)
  533. }
  534. }
  535. // WaitFunc creates a channel that receives an item every time a test
  536. // should be executed and is closed when the last test should be invoked.
  537. type WaitFunc func(done <-chan struct{}) <-chan struct{}
  538. // WithContext converts the WaitFunc to an equivalent WaitWithContextFunc
  539. func (w WaitFunc) WithContext() WaitWithContextFunc {
  540. return func(ctx context.Context) <-chan struct{} {
  541. return w(ctx.Done())
  542. }
  543. }
  544. // WaitWithContextFunc creates a channel that receives an item every time a test
  545. // should be executed and is closed when the last test should be invoked.
  546. //
  547. // When the specified context gets cancelled or expires the function
  548. // stops sending item and returns immediately.
  549. type WaitWithContextFunc func(ctx context.Context) <-chan struct{}
  550. // WaitFor continually checks 'fn' as driven by 'wait'.
  551. //
  552. // WaitFor gets a channel from 'wait()”, and then invokes 'fn' once for every value
  553. // placed on the channel and once more when the channel is closed. If the channel is closed
  554. // and 'fn' returns false without error, WaitFor returns ErrWaitTimeout.
  555. //
  556. // If 'fn' returns an error the loop ends and that error is returned. If
  557. // 'fn' returns true the loop ends and nil is returned.
  558. //
  559. // ErrWaitTimeout will be returned if the 'done' channel is closed without fn ever
  560. // returning true.
  561. //
  562. // When the done channel is closed, because the golang `select` statement is
  563. // "uniform pseudo-random", the `fn` might still run one or multiple time,
  564. // though eventually `WaitFor` will return.
  565. func WaitFor(wait WaitFunc, fn ConditionFunc, done <-chan struct{}) error {
  566. ctx, cancel := ContextForChannel(done)
  567. defer cancel()
  568. return WaitForWithContext(ctx, wait.WithContext(), fn.WithContext())
  569. }
  570. // WaitForWithContext continually checks 'fn' as driven by 'wait'.
  571. //
  572. // WaitForWithContext gets a channel from 'wait()”, and then invokes 'fn'
  573. // once for every value placed on the channel and once more when the
  574. // channel is closed. If the channel is closed and 'fn'
  575. // returns false without error, WaitForWithContext returns ErrWaitTimeout.
  576. //
  577. // If 'fn' returns an error the loop ends and that error is returned. If
  578. // 'fn' returns true the loop ends and nil is returned.
  579. //
  580. // context.Canceled will be returned if the ctx.Done() channel is closed
  581. // without fn ever returning true.
  582. //
  583. // When the ctx.Done() channel is closed, because the golang `select` statement is
  584. // "uniform pseudo-random", the `fn` might still run one or multiple times,
  585. // though eventually `WaitForWithContext` will return.
  586. func WaitForWithContext(ctx context.Context, wait WaitWithContextFunc, fn ConditionWithContextFunc) error {
  587. waitCtx, cancel := context.WithCancel(context.Background())
  588. defer cancel()
  589. c := wait(waitCtx)
  590. for {
  591. select {
  592. case _, open := <-c:
  593. ok, err := runConditionWithCrashProtectionWithContext(ctx, fn)
  594. if err != nil {
  595. return err
  596. }
  597. if ok {
  598. return nil
  599. }
  600. if !open {
  601. return ErrWaitTimeout
  602. }
  603. case <-ctx.Done():
  604. // returning ctx.Err() will break backward compatibility
  605. return ErrWaitTimeout
  606. }
  607. }
  608. }
  609. // poller returns a WaitFunc that will send to the channel every interval until
  610. // timeout has elapsed and then closes the channel.
  611. //
  612. // Over very short intervals you may receive no ticks before the channel is
  613. // closed. A timeout of 0 is interpreted as an infinity, and in such a case
  614. // it would be the caller's responsibility to close the done channel.
  615. // Failure to do so would result in a leaked goroutine.
  616. //
  617. // Output ticks are not buffered. If the channel is not ready to receive an
  618. // item, the tick is skipped.
  619. func poller(interval, timeout time.Duration) WaitWithContextFunc {
  620. return WaitWithContextFunc(func(ctx context.Context) <-chan struct{} {
  621. ch := make(chan struct{})
  622. go func() {
  623. defer close(ch)
  624. tick := time.NewTicker(interval)
  625. defer tick.Stop()
  626. var after <-chan time.Time
  627. if timeout != 0 {
  628. // time.After is more convenient, but it
  629. // potentially leaves timers around much longer
  630. // than necessary if we exit early.
  631. timer := time.NewTimer(timeout)
  632. after = timer.C
  633. defer timer.Stop()
  634. }
  635. for {
  636. select {
  637. case <-tick.C:
  638. // If the consumer isn't ready for this signal drop it and
  639. // check the other channels.
  640. select {
  641. case ch <- struct{}{}:
  642. default:
  643. }
  644. case <-after:
  645. return
  646. case <-ctx.Done():
  647. return
  648. }
  649. }
  650. }()
  651. return ch
  652. })
  653. }
  654. // ExponentialBackoffWithContext works with a request context and a Backoff. It ensures that the retry wait never
  655. // exceeds the deadline specified by the request context.
  656. func ExponentialBackoffWithContext(ctx context.Context, backoff Backoff, condition ConditionFunc) error {
  657. for backoff.Steps > 0 {
  658. select {
  659. case <-ctx.Done():
  660. return ctx.Err()
  661. default:
  662. }
  663. if ok, err := runConditionWithCrashProtection(condition); err != nil || ok {
  664. return err
  665. }
  666. if backoff.Steps == 1 {
  667. break
  668. }
  669. waitBeforeRetry := backoff.Step()
  670. select {
  671. case <-ctx.Done():
  672. return ctx.Err()
  673. case <-time.After(waitBeforeRetry):
  674. }
  675. }
  676. return ErrWaitTimeout
  677. }