ratelimit

package
v0.2.2 Latest Latest
Warning

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

 Go to latest Published: Nov 6, 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. Every instance saves what it counted to the same store, so counts survive restart and reload, and reads what the others counted, so instances sharing a store enforce one limit together. Admission itself stays in memory: it checks this instance's live counters plus the others' as of the last load, never the store.

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, since int64) ([]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 window rows for it, so no replica's row, including one written after the delete under the old epoch, brings a reset window back. In-flight rows stay (see Store.ResetCounter). 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 loads skip 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, since int64) ([]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 window 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. In-flight rows stay (see Store.ResetCounter).

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, saves what changed since the last flush 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 loads the counters every writer saved, earlier runs of this one included, and from then on keeps writing this generation's and reading the others'. 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 publish its counters too. Start is idempotent, and a failed load leaves the generation idle rather than persisting from a blank slate.

The loaded rows become the limiter's remote counters, not its own: this generation's rows only ever hold what it counted itself (see WindowSnapshot), so a restart counts its predecessor's requests once, through their rows, rather than again through its own.

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 its changed counters and reads the other writers'. Zero disables the periodic loop; Start still loads and Close of an active generation still writes once, but nothing counted elsewhere after Start is seen.

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 saved at or after since, in
	// unix seconds by the writer's clock (WindowSnapshot.UpdatedAt); since 0
	// returns them all.
	LoadCounters(ctx context.Context, since int64) ([]WindowSnapshot, error)
	// SaveCounters upserts the given snapshots, keyed by their Writer as well
	// as their window, so writers never overwrite each other's rows, and
	// stamps them with the current time. It also collects window rows no
	// writer has saved for two of their own periods — by then every count
	// they hold has rolled out — and in-flight rows no writer has saved for
	// a day, long after peers stopped counting them.
	// 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 loads skip every row saved before,
	// and removes the window rows of every writer. In-flight rows stay: they
	// hold no counts to reset, only slots live requests still take, which a
	// reset leaves to them in memory too (see limiter.reset).
	// 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"`
	// InFlight is the writer's in-flight gauge on a concurrency rule
	// (PeriodConcurrent), whose window fields stay zero.
	InFlight int64 `bson:"in_flight"`
	// UpdatedAt is when the writer last saved the row, in unix seconds by
	// its own clock. Stores stamp it on save, whatever the snapshot holds,
	// and report it on load.
	UpdatedAt int64 `bson:"updated_at"`
	// OwnCounts marks a row written under the rule above. Releases up to
	// v0.2.1 wrote rows that also held the counts their writer restored at
	// start, the maximum of every row, so summing them would count each
	// restart's predecessor again. They never set this field, so their rows,
	// and every row written while one of them still runs in a rolling
	// upgrade, read false.
	OwnCounts bool `bson:"own_counts"`
}

WindowSnapshot is one writer's persisted counters for one key: its request/token sliding windows, or for a concurrency rule its in-flight gauge. 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.

A row holds only what its writer counted itself, never counts it loaded from other rows, so the cluster's count for a key is the sum of its rows (see peerView.sum). A restarted process writes under a new ID, and its earlier run's rows are summed like any other writer's. Rows written by releases before that rule (OwnCounts false) are merged the way those releases merged them, by maximum.

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.