ratelimit

package
v0.1.61 Latest Latest
Warning

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

 Go to latest Published: Jul 26, 2026 License: MIT Imports: 16 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; live counters are 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 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) ReplaceConfigRules 

func (s *MongoDBStore) ReplaceConfigRules(ctx context.Context, rules []Rule) 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) (*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"`
	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"`
}

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) 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) ReplaceConfigRules added in v0.1.60 

func (s *SQLStore) ReplaceConfigRules(ctx context.Context, rules []Rule) 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) (*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) 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() error

ResetAll clears every live window counter.

func (*Service) ResetRule 

func (s *Service) ResetRule(scope RuleScope, subject string, periodSeconds int64) error

ResetRule clears the live window counters for one rule.

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) 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 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
	Close() error
}

Store persists rate limit rule definitions. Live counters are in-memory and never stored.

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)

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.