ratelimit

package
v0.2.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Oct 30, 2026 License: MIT Imports: 22 Imported by: 0

Documentation

Overview

Package ratelimit enforces request, token, and concurrency limits for the AI gateway. Rules are scoped to a consumer user-path subtree, a provider, or a model. Rule definitions are persisted. Request and token windows are snapshotted to the same store so they survive restart and reload; live admission stays in memory and per instance.

Index

Constants

View Source 
const (
	PeriodMinuteSeconds int64 = 60
	PeriodHourSeconds   int64 = 3600
	PeriodDaySeconds    int64 = 86400
	// PeriodConcurrent marks a window-less rule: MaxRequests caps in-flight
	// requests instead of requests per period.
	PeriodConcurrent int64 = 0
)
View Source 
const (
	// SourceConfig marks rules seeded from static configuration.
	SourceConfig = "config"
	// SourceManual marks rules created or changed through admin APIs.
	SourceManual = "manual"
)

Variables

View Source 
var ErrNotFound = errors.New("rate limit rule not found")
View Source 
var ErrQuotaTemplatesUnavailable = errors.New("per-child quota templates are not enabled")

ErrQuotaTemplatesUnavailable indicates that a per-child template was found while the deployment does not have the quota-templates capability.

View Source 
var ErrUnavailable = errors.New("rate limit service is unavailable")

ErrUnavailable indicates a rate limit service was used without an initialized store.

Functions

func NewUsageTap 

func NewUsageTap(inner usage.LoggerInterface, service *Service) usage.LoggerInterface

NewUsageTap wraps the logger, or returns it unchanged when there is no rate limit service to feed.

func NormalizeSubject 

func NormalizeSubject(scope RuleScope, subject string) (string, error)

NormalizeSubject canonicalizes a rule subject for its scope.

func NormalizeUserPath 

func NormalizeUserPath(raw string) (string, error)

func PeriodLabel 

func PeriodLabel(seconds int64) string

func PeriodSecondsFromName 

func PeriodSecondsFromName(period string) (int64, bool)

PeriodSecondsFromName resolves a named period. The bool reports whether the name is recognized.

Types

type ExceededError 

type ExceededError struct {
	Rule       Rule
	Scope      LimitScope
	Observed   int64
	Limit      int64
	RetryAfter time.Duration
}

ExceededError indicates a rate limit rejected the request.

func (*ExceededError) Error 

func (e *ExceededError) Error() string

type HeaderSnapshot 

type HeaderSnapshot struct {
	HasRequests       bool
	RequestLimit      int64
	RequestRemaining  int64
	RequestResetAfter time.Duration
	HasTokens         bool
	TokenLimit        int64
	TokenRemaining    int64
	TokenResetAfter   time.Duration
}

HeaderSnapshot carries the most-constrained matching limits for OpenAI-style x-ratelimit-* response headers.

type LimitScope 

type LimitScope string

LimitScope names which limit dimension a check or breach refers to. 

const (
	ScopeRequests    LimitScope = "requests"
	ScopeTokens      LimitScope = "tokens"
	ScopeConcurrency LimitScope = "concurrency"
)

type MongoDBStore 

type MongoDBStore struct {
	// contains filtered or unexported fields
}

func NewMongoDBStore 

func NewMongoDBStore(ctx context.Context, database *mongo.Database) (*MongoDBStore, error)

func (*MongoDBStore) Close 

func (s *MongoDBStore) Close() error

func (*MongoDBStore) DeleteRule 

func (s *MongoDBStore) DeleteRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error

func (*MongoDBStore) ListRules 

func (s *MongoDBStore) ListRules(ctx context.Context) ([]Rule, error)

func (*MongoDBStore) LoadCounters added in v0.1.77 

func (s *MongoDBStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error)

func (*MongoDBStore) ReplaceConfigRules 

func (s *MongoDBStore) ReplaceConfigRules(ctx context.Context, rules []Rule) error

func (*MongoDBStore) ResetAllCounters added in v0.2.1 

func (s *MongoDBStore) ResetAllCounters(ctx context.Context) error

ResetAllCounters gives every rule the same new epoch: rows are keyed by rule as well, so sharing it is safe.

func (*MongoDBStore) ResetCounter added in v0.2.1 

func (s *MongoDBStore) ResetCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error

ResetCounter draws a new counter epoch for the rule and removes every writer's rows for it, so no replica's row, including one written after the delete under the old epoch, brings a reset window back. The legacy collection is cleared too: an older replica in a rolling upgrade may still write there, and the next start would import it. It goes before the window collection, so a concurrent import either conflicts with this delete and retries, or commits before the window collection is cleared (see importLegacyCounters).

func (*MongoDBStore) SaveCounters added in v0.1.77 

func (s *MongoDBStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error

func (*MongoDBStore) UpsertRules 

func (s *MongoDBStore) UpsertRules(ctx context.Context, rules []Rule) error

type Reservation 

type Reservation struct {
	// contains filtered or unexported fields
}

Reservation represents one admitted request. Release must be called when the request finishes so concurrency slots return; it is idempotent.

func (*Reservation) Headers 

func (r *Reservation) Headers() HeaderSnapshot

Headers returns the x-ratelimit-* snapshot captured at admission.

func (*Reservation) Release 

func (r *Reservation) Release()

type Result 

type Result struct {
	Service *Service
	Store   Store
	// contains filtered or unexported fields
}

func New 

func New(ctx context.Context, cfg *config.Config, shared storage.Storage, quotaTemplatesEnabled bool) (*Result, error)

func (*Result) Close 

func (r *Result) Close() error

type Rule 

type Rule struct {
	Scope   RuleScope `json:"scope" bson:"scope"`
	Subject string    `json:"subject" bson:"subject"`
	// SubjectDisplay keeps the spelling the rule was written with for
	// provider and model scopes, whose Subject is case-folded to stay the
	// match key. It is empty when the two agree, and rules stored before it
	// existed simply fall back to Subject.
	SubjectDisplay string `json:"subject_display,omitempty" bson:"subject_display,omitempty"`
	PerChild       bool   `json:"per_child" bson:"per_child"`
	// EffectiveSubject identifies the runtime child partition for a per-child
	// rule. It is never persisted or exposed as part of the rule definition.
	EffectiveSubject string    `json:"-" bson:"-"`
	PeriodSeconds    int64     `json:"period_seconds" bson:"period_seconds"`
	MaxRequests      *int64    `json:"max_requests,omitempty" bson:"max_requests,omitempty"`
	MaxTokens        *int64    `json:"max_tokens,omitempty" bson:"max_tokens,omitempty"`
	Source           string    `json:"source,omitempty" bson:"source,omitempty"`
	CreatedAt        time.Time `json:"created_at" bson:"created_at"`
	UpdatedAt        time.Time `json:"updated_at" bson:"updated_at"`
	// CounterEpoch names the counters this rule definition owns. The store
	// draws a new one when the rule is created, switches between shared and
	// per-child, or has its counters reset, and keeps it on any other edit. Saved window rows carry
	// the epoch they were counted under, and restore skips rows of any other
	// epoch, so counts a change cleared never come back, whether or not
	// deleting their rows succeeded. Rules stored before it existed have the
	// empty epoch, as do the rows saved for them.
	CounterEpoch string `json:"-" bson:"counter_epoch,omitempty"`
}

Rule stores the limits for one scope, subject, and period. A period of PeriodConcurrent caps in-flight requests via MaxRequests.

func NormalizeRule 

func NormalizeRule(r Rule) (Rule, error)

func (Rule) DisplaySubject added in v0.1.92 

func (r Rule) DisplaySubject() string

DisplaySubject spells the subject the way the rule was written. Provider and model subjects are matched case-insensitively and stored folded, so the folded key must never be the one shown to clients: it can name a provider that appears in no configuration.

func (Rule) SubjectLabel 

func (r Rule) SubjectLabel() string

SubjectLabel names the rule subject for error messages and logs.

type RuleScope 

type RuleScope string

RuleScope names what a rule limits: a consumer user-path subtree, a provider instance, or a model. 

const (
	// ScopeUserPath limits a consumer subtree; the subject is a user path and
	// covers all its descendants.
	ScopeUserPath RuleScope = "user_path"
	// ScopeProvider limits everything routed to one configured provider
	// instance; the subject is the provider name (e.g. "openai").
	ScopeProvider RuleScope = "provider"
	// ScopeModel limits one model. The subject is a provider-qualified model
	// ("openai/gpt-4o") or a bare model id ("gpt-4o", matching any provider).
	ScopeModel RuleScope = "model"
)

func NormalizeScope 

func NormalizeScope(raw string) (RuleScope, error)

NormalizeScope canonicalizes a rule scope name. An empty scope means user_path, keeping pre-scope rule definitions and requests valid.

type SQLStore added in v0.1.60 

type SQLStore struct {
	// contains filtered or unexported fields
}

SQLStore stores rate limit rules in a SQL database.

func NewSQLStore added in v0.1.60 

func NewSQLStore(ctx context.Context, db sqlx.DB) (*SQLStore, error)

NewSQLStore creates the rate_limits table and index if needed, after migrating any pre-scope table shape.

func (*SQLStore) Close added in v0.1.60 

func (s *SQLStore) Close() error

func (*SQLStore) DeleteRule added in v0.1.60 

func (s *SQLStore) DeleteRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error

func (*SQLStore) ListRules added in v0.1.60 

func (s *SQLStore) ListRules(ctx context.Context) ([]Rule, error)

func (*SQLStore) LoadCounters added in v0.1.77 

func (s *SQLStore) LoadCounters(ctx context.Context) ([]WindowSnapshot, error)

func (*SQLStore) ReplaceConfigRules added in v0.1.60 

func (s *SQLStore) ReplaceConfigRules(ctx context.Context, rules []Rule) error

func (*SQLStore) ResetAllCounters added in v0.2.1 

func (s *SQLStore) ResetAllCounters(ctx context.Context) error

ResetAllCounters gives every rule the same new epoch: rows are keyed by rule as well, so sharing it is safe.

func (*SQLStore) ResetCounter added in v0.2.1 

func (s *SQLStore) ResetCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error

ResetCounter draws a new counter epoch for the rule and removes every writer's rows for it in one transaction, so no replica's row, including one written after the delete under the old epoch, brings a reset window back.

func (*SQLStore) SaveCounters added in v0.1.77 

func (s *SQLStore) SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error

func (*SQLStore) UpsertRules added in v0.1.60 

func (s *SQLStore) UpsertRules(ctx context.Context, rules []Rule) error

type Service 

type Service struct {
	// contains filtered or unexported fields
}

func NewService 

func NewService(ctx context.Context, store Store, options ...ServiceOption) (*Service, error)

func (*Service) Acquire 

func (s *Service) Acquire(subjects Subjects, now time.Time) (*Reservation, error)

Acquire admits or rejects a request for the given subjects. On rejection the returned error unwraps to *ExceededError. The returned reservation is always non-nil on success.

func (*Service) Close added in v0.1.77 

func (s *Service) Close()

Close stops the flush loop, writes a final snapshot if this generation was started, and stops the in-memory expiry cleanup worker.

func (*Service) DeleteRule 

func (s *Service) DeleteRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error

func (*Service) HasTokenRules 

func (s *Service) HasTokenRules() bool

HasTokenRules reports whether any rule limits tokens. Used at startup to warn when token limits cannot be enforced because usage tracking is off.

func (*Service) RecordTokens 

func (s *Service) RecordTokens(subjects Subjects, tokens int64, now time.Time)

RecordTokens adds consumed tokens to every token window matching the subjects. Called after responses complete, from the usage tap.

func (*Service) Refresh 

func (s *Service) Refresh(ctx context.Context) error

func (*Service) ReplaceConfigRules 

func (s *Service) ReplaceConfigRules(ctx context.Context, rules []Rule) error

func (*Service) ResetAll 

func (s *Service) ResetAll(ctx context.Context) error

ResetAll clears every window counter on every instance, as ResetRule does for one rule.

func (*Service) ResetRule 

func (s *Service) ResetRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error

ResetRule clears one rule's window counters on every instance: the store gives the rule a new counter epoch, and the reload that the change notifier triggers clears each instance's counters for it (see Refresh). The local counters are cleared first, so the reset holds here even when the store write fails.

func (*Service) RouteAvailable 

func (s *Service) RouteAvailable(providerName, model string) bool

RouteAvailable reports whether the provider/model route currently has rate-limit capacity. Load balancing and failover use it to skip saturated targets; user-path rules are intentionally ignored because switching targets cannot relieve a consumer limit.

func (*Service) Rules 

func (s *Service) Rules() []Rule

func (*Service) SetChangeNotifier added in v0.1.109 

func (s *Service) SetChangeNotifier(fn func(context.Context))

SetChangeNotifier installs fn, called after every successful write of the rules so other gateway instances reload them. Call it before serving; writes made before it was set (startup seeding) are not announced. Counter resets are rule writes too: they give the rule a new counter epoch, and the reload clears every instance's counters for it.

func (*Service) Start added in v0.1.77 

func (s *Service) Start(ctx context.Context)

Start restores persisted windows and, from then on, keeps writing them. It is separate from NewService because a reload builds the next generation while the current one still serves: only the generation that gets to serve may touch the snapshot, or a replacement that is built and then discarded would flush its empty windows over the live ones. Start is idempotent, and a failed load leaves the generation idle rather than persisting from a blank slate.

func (*Service) Statuses 

func (s *Service) Statuses(now time.Time) []Status

Statuses reports live counter state for every rule.

func (*Service) StatusesForUserPath 

func (s *Service) StatusesForUserPath(userPath string, now time.Time) []Status

StatusesForUserPath reports live counter state for every user-path-scoped rule covering userPath (subtree match). Provider- and model-scoped rules are excluded: they describe shared infrastructure capacity, not the consumer.

func (*Service) UpsertRules 

func (s *Service) UpsertRules(ctx context.Context, rules []Rule) error

type ServiceOption added in v0.1.77 

type ServiceOption func(*Service)

ServiceOption configures a rate limit service.

func WithFlushInterval added in v0.1.77 

func WithFlushInterval(interval time.Duration) ServiceOption

WithFlushInterval sets how often an active generation writes window snapshots. Zero disables the periodic loop; Start still loads and Close of an active generation still writes once.

func WithQuotaTemplates added in v0.1.77 

func WithQuotaTemplates(enabled bool) ServiceOption

WithQuotaTemplates controls whether per-child user-path rules are accepted.

type Status 

type Status struct {
	Rule              Rule
	WindowStart       time.Time
	WindowEnd         time.Time
	RequestsUsed      int64
	RequestsRemaining *int64
	TokensUsed        int64
	TokensRemaining   *int64
	InFlight          int64
}

Status reports the live counter state for one rule.

type Store 

type Store interface {
	ListRules(ctx context.Context) ([]Rule, error)
	UpsertRules(ctx context.Context, rules []Rule) error
	DeleteRule(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error
	ReplaceConfigRules(ctx context.Context, rules []Rule) error
	// LoadCounters returns every writer's rows; restore merges them per key.
	LoadCounters(ctx context.Context) ([]WindowSnapshot, error)
	// SaveCounters upserts the given snapshots, keyed by their Writer as well
	// as their window, so writers never overwrite each other's rows. It also
	// collects rows no writer has refreshed for two of their own periods — the
	// same staleness bound a load applies, so it only ever drops rows that
	// could no longer be restored.
	// Nothing is deleted to make room for a write, so a crash mid-save costs at
	// most the flush interval rather than a whole window.
	SaveCounters(ctx context.Context, snapshots []WindowSnapshot) error
	// ResetCounter gives the rule a new counter epoch, so every replica's
	// next Refresh clears its counts and restore skips every row saved
	// before, and removes the rows of every writer. ResetAllCounters does
	// the same for every rule.
	ResetCounter(ctx context.Context, scope RuleScope, subject string, periodSeconds int64) error
	ResetAllCounters(ctx context.Context) error
	Close() error
}

Store persists rate limit rule definitions and optional window snapshots.

type Subjects 

type Subjects struct {
	UserPath string
	Provider string
	Model    string
}

Subjects identifies the dimensions one request can be limited by. UserPath is always known at ingress; Provider and Model are set once the route is resolved (provider name and provider-qualified model).

type UsageTap 

type UsageTap struct {
	// contains filtered or unexported fields
}

UsageTap decorates a usage logger so every recorded entry also feeds token rate limit windows. All modalities (non-streaming, SSE streams, realtime, audio) funnel through Logger.Write, so one tap covers them all.

func (*UsageTap) Close 

func (t *UsageTap) Close() error

func (*UsageTap) Config 

func (t *UsageTap) Config() usage.Config

func (*UsageTap) Write 

func (t *UsageTap) Write(entry *usage.UsageEntry)

type WindowSnapshot added in v0.1.77 

type WindowSnapshot struct {
	Scope         string `bson:"scope"`
	Subject       string `bson:"subject"`
	Partition     string `bson:"partition"`
	PeriodSeconds int64  `bson:"period_seconds"`
	Writer        string `bson:"writer"`
	// Epoch is the CounterEpoch of the rule the row was counted under.
	Epoch               string `bson:"epoch"`
	RequestsWindowStart int64  `bson:"requests_window_start"`
	RequestsCurrent     int64  `bson:"requests_current"`
	RequestsPrevious    int64  `bson:"requests_previous"`
	TokensWindowStart   int64  `bson:"tokens_window_start"`
	TokensCurrent       int64  `bson:"tokens_current"`
	TokensPrevious      int64  `bson:"tokens_previous"`
}

WindowSnapshot is one persisted request/token sliding window. Partition is empty for a shared rule and the child path for a per-child template (Rule.EffectiveSubject). Writer is the random ID of the process that wrote the row, so replicas sharing a store keep separate rows instead of overwriting each other's; rows imported from before writer IDs existed carry an empty writer.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.