package budget import ( "context" "errors" "fmt" "strings" "time" ) var ErrNotFound = errors.New("budget not found") // SpendWindow asks for the tracked spend of one budget subject over one time // window. type SpendWindow struct { Scope Scope Subject string CostBasis CostBasis Start time.Time End time.Time // ExcludeCacheHits keeps billed response-cache hits out of a billed // window. Provider windows never count cache hits. ExcludeCacheHits bool } // billed reports whether the window sums billed rather than provider cost. func (w SpendWindow) billed() bool { return w.CostBasis == CostBasisBilled } // billedCacheHits reports whether the window adds the billed cost of // response-cache hits. func (w SpendWindow) billedCacheHits() bool { return w.billed() && !w.ExcludeCacheHits } // Spend is the total tracked cost for one SpendWindow. HasUsage is false when // the window matched no priced usage at all, which keeps a budget with no // traffic from ever blocking. type Spend struct { Total float64 HasUsage bool } // spendBounds returns the union of every window's time range, which is the // range a batched spend query has to scan. func spendBounds(windows []SpendWindow) (time.Time, time.Time) { start, end := windows[0].Start, windows[0].End for _, window := range windows[1:] { if window.Start.Before(start) { start = window.Start } if window.End.After(end) { end = window.End } } return start, end } // Store persists budget definitions, reset settings, and spend lookups. type Store interface { ListBudgets(ctx context.Context) ([]Budget, error) UpsertBudgets(ctx context.Context, budgets []Budget) error DeleteBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64) error ReplaceConfigBudgets(ctx context.Context, budgets []Budget) error GetSettings(ctx context.Context) (Settings, error) SaveSettings(ctx context.Context, settings Settings) (Settings, error) ResetBudget(ctx context.Context, scope Scope, subject string, periodSeconds int64, at time.Time) error ResetAllBudgets(ctx context.Context, at time.Time) error // SumSpend totals spend for every window in one round trip, in the order // given: uncached provider cost, or billed cost for a billed window. Budget enforcement runs on every request, so the whole // matching set is asked for at once rather than one query per budget. SumSpend(ctx context.Context, windows []SpendWindow) ([]Spend, error) Close() error } // normalizeBudgetKey validates and canonicalizes the identity of one budget. func normalizeBudgetKey(scope Scope, subject string, periodSeconds int64) (Scope, string, error) { scope, err := NormalizeScope(string(scope)) if err != nil { return "", "", err } subject, err = NormalizeSubject(scope, subject) if err != nil { return "", "", err } if periodSeconds <= 0 { return "", "", fmt.Errorf("period_seconds must be greater than 0") } return scope, subject, nil } func normalizeBudgetsForUpsert(budgets []Budget) ([]Budget, error) { if len(budgets) == 0 { return nil, nil } normalized := make([]Budget, 0, len(budgets)) seen := make(map[string]int, len(budgets)) for _, budget := range budgets { item, err := NormalizeBudget(budget) if err != nil { return nil, err } key := budgetKey(item.Scope, item.Subject, item.PeriodSeconds) if existing, ok := seen[key]; ok { normalized[existing] = item continue } seen[key] = len(normalized) normalized = append(normalized, item) } return normalized, nil } func budgetKey(scope Scope, subject string, periodSeconds int64) string { return string(scope) + "\x00" + strings.TrimSpace(subject) + "\x00" + fmt.Sprint(periodSeconds) } func normalizeLoadedSettings(settings Settings) (Settings, error) { defaults := DefaultSettings() if settings.MonthlyResetDay == 0 { settings.MonthlyResetDay = defaults.MonthlyResetDay } if settings.UpdatedAt.IsZero() { settings.UpdatedAt = time.Now().UTC() } if err := ValidateSettings(settings); err != nil { return Settings{}, err } return settings, nil }