backoff.go

// Package backoff implements backoff logic for loop controls.
package backoff

import (
	"context"
	"iter"
	"math"
	"math/rand/v2"
	"time"
)

// StepFunc is the signature of the function that
// generates step values for a [Generator].
type StepFunc func(base time.Duration, count int) time.Duration

// Generator generates a sequence of duration values using a step function.
//
// A Generator is not safe for concurrent use by multiple goroutines.
type Generator struct {
	base    time.Duration
	ceiling time.Duration
	iter    int
	fn      StepFunc
}

// Base returns the base duration used by the generator.
func (g *Generator) Base() time.Duration {
	return g.base
}

// Ceiling returns the maximum duration that can be yielded.
func (g *Generator) Ceiling() time.Duration {
	return g.ceiling
}

// Delays returns an infinite iterator that yields the 1-based attempt
// number and the delay duration. Unlike [Generator.Attempts], this iterator
// does not sleep—the caller is responsible for sleeping the returned duration.
//
// The iterator stops when the context is cancelled.
func (g *Generator) Delays(ctx context.Context) iter.Seq2[int, time.Duration] {
	return func(yield func(int, time.Duration) bool) {
		for attempt := 1; ; attempt++ {
			select {
			case <-ctx.Done():
				return
			default:
				if !yield(attempt, g.next()) {
					return
				}
			}
		}
	}
}

// Attempts returns an iterator that yields attempt numbers (1-based),
// sleeping the backoff duration after each iteration before yielding
// the next attempt. The sleep respects context cancellation.
//
// The iterator stops when the context is cancelled or maxAttempts is reached.
func (g *Generator) Attempts(ctx context.Context, maxAttempts int) iter.Seq[int] {
	return func(yield func(int) bool) {
		t := time.NewTimer(0)
		t.Stop()
		defer t.Stop()

		for attempt := 1; attempt <= maxAttempts; attempt++ {
			select {
			case <-ctx.Done():
				return
			default:
			}

			if !yield(attempt) {
				return
			}

			if attempt < maxAttempts {
				t.Reset(g.next())
				select {
				case <-ctx.Done():
					return
				case <-t.C:
				}
			}
		}
	}
}

// next generates and returns the next duration.
func (g *Generator) next() time.Duration {
	val := g.fn(g.base, g.iter)
	if g.ceiling > 0 {
		val = min(val, g.ceiling)
	}
	g.iter++
	return val
}

// Reset resets the generator to its initial state.
func (g *Generator) Reset() {
	g.iter = 0
}

// New returns a new generator that uses the given step function.
//
// Base is the first value that will be returned, and the one from which
// the next values will be derived. Ceiling is the maximum value that can
// be generated. A ceiling of 0 disables the ceiling check.
func New(base, ceiling time.Duration, fn StepFunc) *Generator {
	return &Generator{
		base:    base,
		ceiling: ceiling,
		fn:      fn,
		iter:    0,
	}
}

// Delays returns an infinite iterator that yields the 1-based attempt
// number and the delay duration. The caller is responsible for sleeping
// the returned duration. The iterator stops when the context is cancelled.
func Delays(ctx context.Context, base, ceiling time.Duration, fn StepFunc) iter.Seq2[int, time.Duration] {
	return New(base, ceiling, fn).Delays(ctx)
}

// Attempts returns an iterator that yields attempt numbers (1-based),
// sleeping the backoff duration after each iteration before yielding
// the next attempt. The sleep respects context cancellation.
// The iterator stops when the context is cancelled or maxAttempts is reached.
func Attempts(ctx context.Context, base, ceiling time.Duration, fn StepFunc, maxAttempts int) iter.Seq[int] {
	return New(base, ceiling, fn).Attempts(ctx, maxAttempts)
}

// Fixed returns a StepFunc that always returns the base duration.
// It is equivalent to Linear(0).
func Fixed() StepFunc {
	return Linear(0)
}

var _ StepFunc = Fixed()

// Incremental returns a step function that scales the base by the given multiple
// each retry.
func Incremental(multiplier int) StepFunc {
	return func(base time.Duration, count int) time.Duration {
		if count == 0 {
			return base
		}
		factor := saturatingMulInt64(int64(multiplier), int64(count))
		return saturatingMulDuration(base, factor)
	}
}

var _ StepFunc = Incremental(0)

// Linear returns a step function that adds d each retry.
func Linear(d time.Duration) StepFunc {
	return func(base time.Duration, count int) time.Duration {
		if count == 0 {
			return base
		}
		increment := saturatingMulDurationInt(d, count)
		return saturatingAddDuration(base, increment)
	}
}

var _ StepFunc = Linear(0)

// Exponential returns an exponential step function.
func Exponential(factor int) StepFunc {
	if factor <= 0 {
		panic("backoff: exponential factor must be greater than 0")
	}

	return func(base time.Duration, count int) time.Duration {
		if count == 0 {
			return base
		}
		return saturatingExp(base, factor, count)
	}
}

var _ StepFunc = Exponential(2)

// Jitter wraps a StepFunc and returns a new StepFunc that applies
// "full-range" jitter in the range [d - factor*d, d].
//
// For example, with d=1s and factor=0.5, the result is uniformly
// distributed in [500ms, 1s]. Factor should be in [0, 1].
func Jitter(fn StepFunc, factor float64) StepFunc {
	if factor < 0 || factor > 1 {
		panic("backoff: jitter factor must be in [0, 1]")
	}

	return func(base time.Duration, count int) time.Duration {
		d := fn(base, count)
		j := time.Duration(float64(d) * rand.Float64() * factor) //nolint:gosec // jitter doesn't need crypto rand
		return d - j
	}
}

var _ StepFunc = Jitter(Linear(0), 1)

func saturatingMulDuration(d time.Duration, m int64) time.Duration {
	if d == 0 || m == 0 {
		return 0
	}

	d64 := int64(d)
	ad := absInt64(d64)
	am := absInt64(m)

	if am != 0 && ad > uint64(math.MaxInt64)/am {
		if (d64 < 0) == (m < 0) {
			return time.Duration(math.MaxInt64)
		}
		return time.Duration(math.MinInt64)
	}

	return time.Duration(d64 * m)
}

func saturatingMulDurationInt(d time.Duration, m int) time.Duration {
	return saturatingMulDuration(d, int64(m))
}

func saturatingMulInt64(a, b int64) int64 {
	if a == 0 || b == 0 {
		return 0
	}

	aa := absInt64(a)
	bb := absInt64(b)

	if bb != 0 && aa > uint64(math.MaxInt64)/bb {
		if (a < 0) == (b < 0) {
			return math.MaxInt64
		}
		return math.MinInt64
	}

	return a * b
}

func saturatingAddDuration(a, b time.Duration) time.Duration {
	if b > 0 && a > time.Duration(math.MaxInt64)-b {
		return time.Duration(math.MaxInt64)
	}
	if b < 0 && a < time.Duration(math.MinInt64)-b {
		return time.Duration(math.MinInt64)
	}
	return a + b
}

func saturatingExp(base time.Duration, factor, count int) time.Duration {
	result := base
	for i := 0; i < count; i++ {
		result = saturatingMulDurationInt(result, factor)
		if result == time.Duration(math.MaxInt64) || result == time.Duration(math.MinInt64) {
			return result
		}
	}
	return result
}

func absInt64(v int64) uint64 {
	if v >= 0 {
		return uint64(v)
	}
	if v == math.MinInt64 {
		return uint64(math.MaxInt64) + 1
	}
	return uint64(-v)
}