89 lines
2.5 KiB
Go
89 lines
2.5 KiB
Go
|
// Package requestbudget carries a durable request reservation to provider transports.
|
||
|
|
package requestbudget
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"net/http"
|
||
|
|
)
|
||
|
|
|
||
|
|
var ErrStopped = errors.New("wiki recovery budget stopped")
|
||
|
|
|
||
|
|
type Guard interface {
|
||
|
|
Key() string
|
||
|
|
Check(context.Context) error
|
||
|
|
Reserve(context.Context, string, string) error
|
||
|
|
}
|
||
|
|
|
||
|
|
type guardKey struct{}
|
||
|
|
type phaseKey struct{}
|
||
|
|
type modelKey struct{}
|
||
|
|
|
||
|
|
func WithGuard(ctx context.Context, guard Guard) context.Context {
|
||
|
|
return context.WithValue(ctx, guardKey{}, guard)
|
||
|
|
}
|
||
|
|
|
||
|
|
func FromContext(ctx context.Context) Guard {
|
||
|
|
guard, _ := ctx.Value(guardKey{}).(Guard)
|
||
|
|
return guard
|
||
|
|
}
|
||
|
|
|
||
|
|
func WithPhase(ctx context.Context, phase string) context.Context {
|
||
|
|
return context.WithValue(ctx, phaseKey{}, phase)
|
||
|
|
}
|
||
|
|
|
||
|
|
func Phase(ctx context.Context) string {
|
||
|
|
phase, _ := ctx.Value(phaseKey{}).(string)
|
||
|
|
return phase
|
||
|
|
}
|
||
|
|
|
||
|
|
func WithModel(ctx context.Context, modelID string) context.Context {
|
||
|
|
return context.WithValue(ctx, modelKey{}, modelID)
|
||
|
|
}
|
||
|
|
|
||
|
|
func Check(ctx context.Context) error {
|
||
|
|
if guard := FromContext(ctx); guard != nil {
|
||
|
|
return guard.Check(ctx)
|
||
|
|
}
|
||
|
|
return ctx.Err()
|
||
|
|
}
|
||
|
|
|
||
|
|
// RequireSupported prevents a controlled request from reaching an adapter
|
||
|
|
// whose actual sends are not instrumented. Ordinary requests are unaffected.
|
||
|
|
func RequireSupported(ctx context.Context, supported bool) error {
|
||
|
|
if FromContext(ctx) != nil && !supported {
|
||
|
|
return fmt.Errorf("%w: unsupported provider transport", ErrStopped)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Transport reserves each RoundTrip, including redirects and SDK fallbacks.
|
||
|
|
// Failed or cancelled sends intentionally consume their reservation.
|
||
|
|
type Transport struct{ Base http.RoundTripper }
|
||
|
|
|
||
|
|
func (t Transport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||
|
|
if err := req.Context().Err(); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
if guard := FromContext(req.Context()); guard != nil {
|
||
|
|
modelID, _ := req.Context().Value(modelKey{}).(string)
|
||
|
|
if modelID == "" {
|
||
|
|
return nil, fmt.Errorf("%w: missing provider model identity", ErrStopped)
|
||
|
|
}
|
||
|
|
if err := guard.Reserve(req.Context(), modelID, Phase(req.Context())); err != nil {
|
||
|
|
return nil, fmt.Errorf("%w: %w", ErrStopped, err)
|
||
|
|
}
|
||
|
|
// A replayable POST with an idempotency header can otherwise be retried
|
||
|
|
// inside http.Transport without reserving again. Clone so Client.Do
|
||
|
|
// keeps the original GetBody for separately accounted 307/308 redirects.
|
||
|
|
req = req.Clone(req.Context())
|
||
|
|
req.GetBody = nil
|
||
|
|
}
|
||
|
|
base := t.Base
|
||
|
|
if base == nil {
|
||
|
|
base = http.DefaultTransport
|
||
|
|
}
|
||
|
|
return base.RoundTrip(req)
|
||
|
|
}
|