cloudflared-mirror/origin/backoffhandler.go

package origin

import (
	"time"

	"golang.org/x/net/context"
)

// Redeclare time functions so they can be overridden in tests.
var (
	timeNow   = time.Now
	timeAfter = time.After
)

// BackoffHandler manages exponential backoff and limits the maximum number of retries.
// The base time period is 1 second, doubling with each retry.
// After initial success, a grace period can be set to reset the backoff timer if
// a connection is maintained successfully for a long enough period. The base grace period
// is 2 seconds, doubling with each retry.
type BackoffHandler struct {
	// MaxRetries sets the maximum number of retries to perform. The default value
	// of 0 disables retry completely.
	MaxRetries uint
	// RetryForever caps the exponential backoff period according to MaxRetries
	// but allows you to retry indefinitely.
	RetryForever bool
	// BaseTime sets the initial backoff period.
	BaseTime time.Duration

	retries       uint
	resetDeadline time.Time
}

func (b BackoffHandler) GetBackoffDuration(ctx context.Context) (time.Duration, bool) {
	// Follows the same logic as Backoff, but without mutating the receiver.
	// This select has to happen first to reflect the actual behaviour of the Backoff function.
	select {
	case <-ctx.Done():
		return time.Duration(0), false
	default:
	}
	if !b.resetDeadline.IsZero() && timeNow().After(b.resetDeadline) {
		// b.retries would be set to 0 at this point
		return time.Second, true
	}
	if b.retries >= b.MaxRetries && !b.RetryForever {
		return time.Duration(0), false
	}
	return time.Duration(b.GetBaseTime() * 1 << b.retries), true
}

// BackoffTimer returns a channel that sends the current time when the exponential backoff timeout expires.
// Returns nil if the maximum number of retries have been used.
func (b *BackoffHandler) BackoffTimer() <-chan time.Time {
	if !b.resetDeadline.IsZero() && timeNow().After(b.resetDeadline) {
		b.retries = 0
		b.resetDeadline = time.Time{}
	}
	if b.retries >= b.MaxRetries {
		if !b.RetryForever {
			return nil
		}
	} else {
		b.retries++
	}
	return timeAfter(time.Duration(b.GetBaseTime() * 1 << (b.retries - 1)))
}

// Backoff is used to wait according to exponential backoff. Returns false if the
// maximum number of retries have been used or if the underlying context has been cancelled.
func (b *BackoffHandler) Backoff(ctx context.Context) bool {
	c := b.BackoffTimer()
	if c == nil {
		return false
	}
	select {
	case <-c:
		return true
	case <-ctx.Done():
		return false
	}
}

// Sets a grace period within which the the backoff timer is maintained. After the grace
// period expires, the number of retries & backoff duration is reset.
func (b *BackoffHandler) SetGracePeriod() {
	b.resetDeadline = timeNow().Add(time.Duration(b.GetBaseTime() * 2 << b.retries))
}

func (b BackoffHandler) GetBaseTime() time.Duration {
	if b.BaseTime == 0 {
		return time.Second
	}
	return b.BaseTime
}