gcsbuf/retry.go

package buf

import (
	"context"
	"errors"
	"math/rand/v2"
	"net/http"
	"time"

	"google.golang.org/api/googleapi"
)

const maxRetries = 5

func isRateLimited(err error) bool {
	var apiErr *googleapi.Error
	return errors.As(err, &apiErr) && apiErr.Code == http.StatusTooManyRequests
}

// retry retries fn on rate-limit (429) errors with exponential backoff and jitter.
// OnOp is fired after each attempt, so callers see every rate-limit hit.
func retry[T any](ctx context.Context, onOp OnOp, op string, fn func() (T, error)) (T, error) {
	for attempt := range maxRetries + 1 {
		start := time.Now()
		v, err := fn()
		onOp.observe(op, time.Since(start), err)
		if err == nil || !isRateLimited(err) {
			return v, err
		}
		if attempt == maxRetries {
			return v, err
		}
		backoff := min(time.Duration(1<<attempt)*100*time.Millisecond, 5*time.Second)
		jitter := time.Duration(rand.Int64N(int64(backoff) / 2))
		select {
		case <-time.After(backoff + jitter):
		case <-ctx.Done():
			var zero T
			return zero, ctx.Err()
		}
	}
	panic("unreachable")
}

func retryVoid(ctx context.Context, onOp OnOp, op string, fn func() error) error {
	_, err := retry(ctx, onOp, op, func() (struct{}, error) {
		return struct{}{}, fn()
	})
	return err
}