github.com/cenkalti/backoff/v4@v4.2.1/context.go (about)

     1  package backoff
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  )
     7  
     8  // BackOffContext is a backoff policy that stops retrying after the context
     9  // is canceled.
    10  type BackOffContext interface { // nolint: golint
    11  	BackOff
    12  	Context() context.Context
    13  }
    14  
    15  type 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
    23  func 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  
    41  func 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  
    51  func (b *backOffContext) Context() context.Context {
    52  	return b.ctx
    53  }
    54  
    55  func (b *backOffContext) NextBackOff() time.Duration {
    56  	select {
    57  	case <-b.ctx.Done():
    58  		return Stop
    59  	default:
    60  		return b.BackOff.NextBackOff()
    61  	}
    62  }