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