context.go raw

   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  }
  63