1package 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 { 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 { 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 { 60 return Stop 61 } 62 return next 63} 64