Go: Retries with Truncated Exponential Backoff and Jitter

truncated_exponential_backoff.go

package main

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

type backoff struct {
	initial    time.Duration
	max        time.Duration
	multiplier float64
	rand       *rand.Rand
}

// Creates a truncated exponential backoff with jitter. `source` can be nil.
func NewBackoff(initial time.Duration, max time.Duration, multiplier float64, source rand.Source) (*backoff, error) {
	if initial <= 0 {
		return nil, errors.New("initial must be positive")
	}
	if max <= 0 {
		return nil, errors.New("max must be positive")
	}
	if multiplier < 1.0 {
		return nil, errors.New("multiplier must be 1.0 or greater")
	}
	if source == nil {
		source = rand.NewPCG(rand.Uint64(), rand.Uint64())
	}
	return &backoff{initial, max, multiplier, rand.New(source)}, nil
}

func (b *backoff) Delay(i int) time.Duration {
	// The implementaion details can be found at https://pkg.go.dev/github.com/googleapis/gax-go/v2#Backoff
	factor := math.Pow(b.multiplier, float64(i))
	calculatedDelay := time.Duration(factor * float64(b.initial.Nanoseconds()))
	effectiveDelay := min(calculatedDelay, b.max)
	return time.Duration(b.rand.Int64N(effectiveDelay.Nanoseconds()) + 1)
}

func (b *backoff) Wait(ctx context.Context, i int) error {
	select {
	case <-time.Tick(b.Delay(i)):
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func BackoffExample(ctx context.Context) (any, error) {
	bo, err := NewBackoff(time.Second, 30*time.Second, 2.0, nil)
	if err != nil {
		return nil, err
	}
	for i := 0; ; i++ {
		v, err := DoSomething()
		if errors.Is(err, ErrRetryable) {
			slog.Error("DoSomething failed; retrying after delay...", "err", err, "delay", bo.Delay(i))
			if err := bo.Wait(ctx, i); err != nil {
				return nil, err
			}
			continue
		}
		if err != nil {
			return nil, err
		}
		return v, nil
	}
}