at-container-registry/pkg/appview/storage/retryafter.go

package storage

import (
	"context"
	"sync"
	"time"
)

// RetryAfterCarrier is a request-scoped, mutable container for a Retry-After
// hint emitted by storage handlers (e.g., when an upstream PDS returns 429).
// HTTP middleware injects an empty carrier into the request context; deep
// handlers populate it via SetRetryAfter when they convert a rate-limit error
// into a 429 response. The middleware then reads it back to set the
// Retry-After response header.
type RetryAfterCarrier struct {
	mu       sync.Mutex
	duration time.Duration
}

const RetryAfterContextKey contextKey = "atcr.retry-after"

// NewRetryAfterCarrier returns an empty carrier ready to be stored in context.
func NewRetryAfterCarrier() *RetryAfterCarrier {
	return &RetryAfterCarrier{}
}

// Set records a retry-after hint. Largest value wins (a later, longer
// throttle window in a multi-write request shouldn't be clobbered by a
// shorter one).
func (c *RetryAfterCarrier) Set(d time.Duration) {
	if c == nil || d <= 0 {
		return
	}
	c.mu.Lock()
	if d > c.duration {
		c.duration = d
	}
	c.mu.Unlock()
}

// Duration returns the recorded retry-after value, or 0 if none was set.
func (c *RetryAfterCarrier) Duration() time.Duration {
	if c == nil {
		return 0
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.duration
}

// SetRetryAfter is a convenience helper for handlers that have a context but
// not a direct carrier reference.
func SetRetryAfter(ctx context.Context, d time.Duration) {
	if c, ok := ctx.Value(RetryAfterContextKey).(*RetryAfterCarrier); ok {
		c.Set(d)
	}
}