blob: fcff86c1b3de82ed74f86ea9576711dd0aa27d31 [file] [log] [blame]
khenaidoo5e4fca32021-05-12 16:02:23 -04001package backoff
2
3import (
4 "context"
5 "time"
6)
7
8// BackOffContext is a backoff policy that stops retrying after the context
9// is canceled.
10type BackOffContext interface { // nolint: golint
11 BackOff
12 Context() context.Context
13}
14
15type backOffContext struct {
16 BackOff
17 ctx context.Context
18}
19
20// WithContext returns a BackOffContext with context ctx
21//
22// ctx must not be nil
23func WithContext(b BackOff, ctx context.Context) BackOffContext { // nolint: golint
24 if ctx == nil {
25 panic("nil context")
26 }
27
28 if b, ok := b.(*backOffContext); ok {
29 return &backOffContext{
30 BackOff: b.BackOff,
31 ctx: ctx,
32 }
33 }
34
35 return &backOffContext{
36 BackOff: b,
37 ctx: ctx,
38 }
39}
40
41func getContext(b BackOff) context.Context {
42 if cb, ok := b.(BackOffContext); ok {
43 return cb.Context()
44 }
45 if tb, ok := b.(*backOffTries); ok {
46 return getContext(tb.delegate)
47 }
48 return context.Background()
49}
50
51func (b *backOffContext) Context() context.Context {
52 return b.ctx
53}
54
55func (b *backOffContext) NextBackOff() time.Duration {
56 select {
57 case <-b.ctx.Done():
58 return Stop
59 default:
60 }
61 next := b.BackOff.NextBackOff()
62 if deadline, ok := b.ctx.Deadline(); ok && deadline.Sub(time.Now()) < next { // nolint: gosimple
63 return Stop
64 }
65 return next
66}