blob: 7f7842131a4f1d19934b89ecb5c3d27ea93440ba [file] [log] [blame]
Takahiro Suzukid7bf8202020-12-17 20:21:59 +09001package 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 ensureContext(b BackOff) BackOffContext {
42 if cb, ok := b.(BackOffContext); ok {
43 return cb
44 }
45 return WithContext(b, context.Background())
46}
47
48func (b *backOffContext) Context() context.Context {
49 return b.ctx
50}
51
52func (b *backOffContext) NextBackOff() time.Duration {
53 select {
54 case <-b.ctx.Done():
55 return Stop
56 default:
57 }
58 next := b.BackOff.NextBackOff()
59 if deadline, ok := b.ctx.Deadline(); ok && deadline.Sub(time.Now()) < next { // nolint: gosimple
60 return Stop
61 }
62 return next
63}