Compare commits
40 commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
781e0ee3d6 | ||
|
|
2ad4870bd0 | ||
|
|
ed5949ec3a | ||
| 61ccc226b2 | |||
|
|
b3a6279f35 | ||
| 0e8d02a528 | |||
|
|
1ec0ea4d28 | ||
| eea295f017 | |||
|
|
75f27a4906 | ||
| 27723ce8e9 | |||
|
|
ed1cdc11b2 | ||
| 52316d5377 | |||
|
|
36cc0a4750 | ||
| bd6c6e5136 | |||
|
|
c5e2ed8b7e | ||
| 62d8fd0f5b | |||
|
|
86dc04258a | ||
| a24fa5bad3 | |||
|
|
22ab4edc86 | ||
| d82736204e | |||
|
|
d1c90b937d | ||
| 0db35c4ce9 | |||
|
|
d4d9d7a798 | ||
|
|
d7655561b8 | ||
|
|
f5a83d774d | ||
|
|
b45262132b | ||
|
|
92fe978b1a | ||
|
|
d3a47eaecc | ||
| 5ad481d30d | |||
|
|
a003e532f7 | ||
|
|
e6a3ba9810 | ||
|
|
4bb1cb96d4 | ||
|
|
ee6c5aa69d | ||
|
|
b0b686bd3f | ||
|
|
25da438ca0 | ||
|
|
9cddfe46b7 | ||
|
|
9572425e89 | ||
|
|
ae2cb96d38 | ||
|
|
79448bf3f3 | ||
|
|
2eb0559ecb |
23 changed files with 3384 additions and 723 deletions
24
.core/build.yaml
Normal file
24
.core/build.yaml
Normal file
|
|
@ -0,0 +1,24 @@
|
||||||
|
version: 1
|
||||||
|
|
||||||
|
project:
|
||||||
|
name: go-ratelimit
|
||||||
|
description: Rate limiting
|
||||||
|
binary: ""
|
||||||
|
|
||||||
|
build:
|
||||||
|
cgo: false
|
||||||
|
flags:
|
||||||
|
- -trimpath
|
||||||
|
ldflags:
|
||||||
|
- -s
|
||||||
|
- -w
|
||||||
|
|
||||||
|
targets:
|
||||||
|
- os: linux
|
||||||
|
arch: amd64
|
||||||
|
- os: linux
|
||||||
|
arch: arm64
|
||||||
|
- os: darwin
|
||||||
|
arch: arm64
|
||||||
|
- os: windows
|
||||||
|
arch: amd64
|
||||||
20
.core/release.yaml
Normal file
20
.core/release.yaml
Normal file
|
|
@ -0,0 +1,20 @@
|
||||||
|
version: 1
|
||||||
|
|
||||||
|
project:
|
||||||
|
name: go-ratelimit
|
||||||
|
repository: core/go-ratelimit
|
||||||
|
|
||||||
|
publishers: []
|
||||||
|
|
||||||
|
changelog:
|
||||||
|
include:
|
||||||
|
- feat
|
||||||
|
- fix
|
||||||
|
- perf
|
||||||
|
- refactor
|
||||||
|
exclude:
|
||||||
|
- chore
|
||||||
|
- docs
|
||||||
|
- style
|
||||||
|
- test
|
||||||
|
- ci
|
||||||
54
.github/workflows/ci.yml
vendored
Normal file
54
.github/workflows/ci.yml
vendored
Normal file
|
|
@ -0,0 +1,54 @@
|
||||||
|
name: CI
|
||||||
|
|
||||||
|
on:
|
||||||
|
push:
|
||||||
|
branches: [main, dev]
|
||||||
|
pull_request:
|
||||||
|
branches: [main]
|
||||||
|
pull_request_review:
|
||||||
|
types: [submitted]
|
||||||
|
|
||||||
|
jobs:
|
||||||
|
test:
|
||||||
|
if: github.event_name != 'pull_request_review'
|
||||||
|
runs-on: ubuntu-latest
|
||||||
|
steps:
|
||||||
|
- uses: actions/checkout@v4
|
||||||
|
- uses: dAppCore/build/actions/build/core@dev
|
||||||
|
with:
|
||||||
|
go-version: "1.26"
|
||||||
|
run-vet: "true"
|
||||||
|
|
||||||
|
auto-fix:
|
||||||
|
if: >
|
||||||
|
github.event_name == 'pull_request_review' &&
|
||||||
|
github.event.review.user.login == 'coderabbitai' &&
|
||||||
|
github.event.review.state == 'changes_requested'
|
||||||
|
runs-on: ubuntu-latest
|
||||||
|
permissions:
|
||||||
|
contents: write
|
||||||
|
pull-requests: write
|
||||||
|
steps:
|
||||||
|
- uses: actions/checkout@v4
|
||||||
|
with:
|
||||||
|
ref: ${{ github.event.pull_request.head.ref }}
|
||||||
|
fetch-depth: 0
|
||||||
|
- uses: dAppCore/build/actions/fix@dev
|
||||||
|
with:
|
||||||
|
go-version: "1.26"
|
||||||
|
|
||||||
|
auto-merge:
|
||||||
|
if: >
|
||||||
|
github.event_name == 'pull_request_review' &&
|
||||||
|
github.event.review.user.login == 'coderabbitai' &&
|
||||||
|
github.event.review.state == 'approved'
|
||||||
|
runs-on: ubuntu-latest
|
||||||
|
permissions:
|
||||||
|
contents: write
|
||||||
|
pull-requests: write
|
||||||
|
steps:
|
||||||
|
- uses: actions/checkout@v4
|
||||||
|
- name: Merge PR
|
||||||
|
run: gh pr merge ${{ github.event.pull_request.number }} --merge --delete-branch
|
||||||
|
env:
|
||||||
|
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||||
2
.gitignore
vendored
Normal file
2
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,2 @@
|
||||||
|
.core/
|
||||||
|
.idea/
|
||||||
72
CLAUDE.md
72
CLAUDE.md
|
|
@ -1,28 +1,76 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
# CLAUDE.md
|
# CLAUDE.md
|
||||||
|
|
||||||
Token counting, model quotas, and sliding window rate limiter.
|
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
|
||||||
|
|
||||||
Module: `forge.lthn.ai/core/go-ratelimit`
|
## Overview
|
||||||
|
|
||||||
|
Provider-agnostic sliding window rate limiter for LLM API calls. Single Go package (no sub-packages) with two persistence backends: YAML (single-process, default) and SQLite (multi-process, WAL mode). Enforces RPM, TPM, and RPD quotas per model. Ships default profiles for Gemini, OpenAI, Anthropic, and Local providers.
|
||||||
|
|
||||||
|
Module: `dappco.re/go/core/go-ratelimit` — Go 1.26, no CGO required.
|
||||||
|
|
||||||
## Commands
|
## Commands
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
go test ./... # run all tests
|
go test ./... # run all tests
|
||||||
go test -race ./... # race detector (required before commit)
|
go test -race ./... # race detector (required before commit)
|
||||||
go test -v -run Name ./... # single test
|
go test -v -run TestCanSend ./... # single test
|
||||||
go vet ./... # vet check
|
go test -v -run "TestCanSend/RPM_at_exact_limit" ./... # single subtest
|
||||||
|
go test -bench=. -benchmem ./... # benchmarks
|
||||||
|
go vet ./... # vet check
|
||||||
|
golangci-lint run ./... # lint
|
||||||
```
|
```
|
||||||
|
|
||||||
|
Pre-commit gate: `go test -race ./...` and `go vet ./...` must both pass.
|
||||||
|
|
||||||
## Standards
|
## Standards
|
||||||
|
|
||||||
- UK English
|
- **UK English** everywhere: colour, organisation, serialise, initialise, behaviour
|
||||||
- `go test -race ./...` and `go vet ./...` must pass before commit
|
- **Conventional commits**: `type(scope): description` — scopes: `ratelimit`, `sqlite`, `persist`, `config`
|
||||||
- Conventional commits: `type(scope): description`
|
- **Co-Author line** on every commit: `Co-Authored-By: Virgil <virgil@lethean.io>`
|
||||||
- Co-Author: `Co-Authored-By: Virgil <virgil@lethean.io>`
|
- **Coverage** must not drop below 95%
|
||||||
- Coverage must not drop below 95%
|
- **Error format**: `core.E("ratelimit.FunctionName", "what", err)` via `dappco.re/go/core` — lowercase, no trailing punctuation
|
||||||
|
- **No `init()` functions**, no global mutable state
|
||||||
|
- **Mutex discipline**: lock at the top of public methods, never inside helpers. Helpers that need the lock document "Caller must hold the lock". `prune()` mutates state, so even "read-only" methods that call it take the write lock. Never call a public method from another public method while holding the lock.
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
All code lives in the root package. Key files:
|
||||||
|
|
||||||
|
- `ratelimit.go` — core types (`RateLimiter`, `ModelQuota`, `UsageStats`, `Config`, `Provider`), sliding window logic (`prune`, `CanSend`, `RecordUsage`), YAML persistence, `CountTokens` (Gemini-specific), iterators (`Models`, `Iter`)
|
||||||
|
- `sqlite.go` — `sqliteStore` internal type, schema creation, load/save for quotas and state
|
||||||
|
|
||||||
|
Constructor matrix: `New()` / `NewWithConfig()` for YAML, `NewWithSQLite()` / `NewWithSQLiteConfig()` for SQLite. Always `defer rl.Close()` with SQLite.
|
||||||
|
|
||||||
|
### Sliding window
|
||||||
|
|
||||||
|
1-minute window pruned on every `CanSend`/`Stats`/`RecordUsage` call. Daily counter is a rolling 24h window from first request, not a calendar boundary. Empty state entries are garbage-collected by `prune()` to prevent memory leaks.
|
||||||
|
|
||||||
|
## Test Organisation
|
||||||
|
|
||||||
|
White-box tests (`package ratelimit`), all assertions via `testify` (`require` for fatal, `assert` for non-fatal). Do not use `t.Error`/`t.Fatal` directly.
|
||||||
|
|
||||||
|
| File | Scope |
|
||||||
|
|------|-------|
|
||||||
|
| `ratelimit_test.go` | Core logic, provider profiles, concurrency, benchmarks |
|
||||||
|
| `sqlite_test.go` | SQLite backend, migration, concurrent persistence |
|
||||||
|
| `error_test.go` | Error paths for SQLite and YAML |
|
||||||
|
| `iter_test.go` | Iterators, `CountTokens` edge cases |
|
||||||
|
|
||||||
|
SQLite tests use `_Good`/`_Bad`/`_Ugly` suffixes (happy path / expected errors / edge cases). Core tests use plain descriptive names with table-driven subtests. Use `t.TempDir()` for all file paths.
|
||||||
|
|
||||||
|
## Dependencies
|
||||||
|
|
||||||
|
Four direct dependencies — do not add more without justification:
|
||||||
|
|
||||||
|
- `dappco.re/go/core` — file I/O helpers, structured errors, JSON helpers, path/environment utilities
|
||||||
|
- `gopkg.in/yaml.v3` — YAML backend
|
||||||
|
- `modernc.org/sqlite` — pure Go SQLite (no CGO)
|
||||||
|
- `github.com/stretchr/testify` — test-only
|
||||||
|
|
||||||
## Docs
|
## Docs
|
||||||
|
|
||||||
- `docs/architecture.md` — sliding window algorithm, provider quotas, YAML/SQLite backends
|
- `docs/architecture.md` — sliding window algorithm, provider quotas, YAML/SQLite backends, concurrency model
|
||||||
- `docs/development.md` — prerequisites, test patterns, coding standards
|
- `docs/development.md` — prerequisites, test patterns, coding standards
|
||||||
- `docs/history.md` — completed phases with commit hashes, known limitations
|
- `docs/history.md` — completed phases with commit hashes, known limitations
|
||||||
|
|
|
||||||
|
|
@ -1,15 +1,21 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
# Contributing
|
# Contributing
|
||||||
|
|
||||||
Thank you for your interest in contributing!
|
Thank you for your interest in contributing!
|
||||||
|
|
||||||
## Requirements
|
## Requirements
|
||||||
- **Go Version**: 1.26 or higher is required.
|
- **Go Version**: 1.26 or higher is required.
|
||||||
- **Tools**: `golangci-lint` and `task` (Taskfile.dev) are recommended.
|
- **Tools**: `golangci-lint` is recommended.
|
||||||
|
|
||||||
## Development Workflow
|
## Development Workflow
|
||||||
1. **Testing**: Ensure all tests pass before submitting changes.
|
1. **Testing**: Ensure all tests pass before submitting changes.
|
||||||
```bash
|
```bash
|
||||||
|
go build ./...
|
||||||
go test ./...
|
go test ./...
|
||||||
|
go test -race ./...
|
||||||
|
go test -cover ./...
|
||||||
|
go mod tidy
|
||||||
```
|
```
|
||||||
2. **Code Style**: All code must follow standard Go formatting.
|
2. **Code Style**: All code must follow standard Go formatting.
|
||||||
```bash
|
```bash
|
||||||
|
|
@ -22,14 +28,22 @@ Thank you for your interest in contributing!
|
||||||
```
|
```
|
||||||
|
|
||||||
## Commit Message Format
|
## Commit Message Format
|
||||||
We follow the [Conventional Commits](https://www.conventionalcommits.org/) specification:
|
We follow the [Conventional Commits](https://www.conventionalcommits.org/) specification using the repository format `type(scope): description`:
|
||||||
- `feat`: A new feature
|
- `feat`: A new feature
|
||||||
- `fix`: A bug fix
|
- `fix`: A bug fix
|
||||||
- `docs`: Documentation changes
|
- `docs`: Documentation changes
|
||||||
- `refactor`: A code change that neither fixes a bug nor adds a feature
|
- `refactor`: A code change that neither fixes a bug nor adds a feature
|
||||||
- `chore`: Changes to the build process or auxiliary tools and libraries
|
- `chore`: Changes to the build process or auxiliary tools and libraries
|
||||||
|
|
||||||
Example: `feat: add new endpoint for health check`
|
Common scopes: `ratelimit`, `sqlite`, `persist`, `config`
|
||||||
|
|
||||||
## License
|
Example:
|
||||||
|
|
||||||
|
```text
|
||||||
|
fix(ratelimit): align module metadata with dappco.re
|
||||||
|
|
||||||
|
Co-Authored-By: Virgil <virgil@lethean.io>
|
||||||
|
```
|
||||||
|
|
||||||
|
## Licence
|
||||||
By contributing to this project, you agree that your contributions will be licensed under the **European Union Public Licence (EUPL-1.2)**.
|
By contributing to this project, you agree that your contributions will be licensed under the **European Union Public Licence (EUPL-1.2)**.
|
||||||
|
|
|
||||||
44
README.md
44
README.md
|
|
@ -1,30 +1,50 @@
|
||||||
[](https://pkg.go.dev/forge.lthn.ai/core/go-ratelimit)
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
[](LICENSE.md)
|
|
||||||
|
[](https://pkg.go.dev/dappco.re/go/core/go-ratelimit)
|
||||||
|

|
||||||
[](go.mod)
|
[](go.mod)
|
||||||
|
|
||||||
# go-ratelimit
|
# go-ratelimit
|
||||||
|
|
||||||
Provider-agnostic sliding window rate limiter for LLM API calls. Enforces requests per minute (RPM), tokens per minute (TPM), and requests per day (RPD) quotas per model using an in-memory sliding window. Ships with default quota profiles for Gemini, OpenAI, Anthropic, and a local inference provider. State persists across process restarts via YAML (single-process) or SQLite (multi-process, WAL mode). Includes a Gemini-specific token counting helper and a YAML-to-SQLite migration path.
|
Provider-agnostic sliding window rate limiter for LLM API calls. Enforces requests per minute (RPM), tokens per minute (TPM), and requests per day (RPD) quotas per model using an in-memory sliding window. Ships with default quota profiles for Gemini, OpenAI, Anthropic, and a local inference provider. State persists across process restarts via YAML (single-process) or SQLite (multi-process, WAL mode). Includes a Gemini-specific token counting helper and a YAML-to-SQLite migration path.
|
||||||
|
|
||||||
**Module**: `forge.lthn.ai/core/go-ratelimit`
|
**Module**: `dappco.re/go/core/go-ratelimit`
|
||||||
**Licence**: EUPL-1.2
|
**Licence**: EUPL-1.2
|
||||||
**Language**: Go 1.25
|
**Language**: Go 1.26
|
||||||
|
|
||||||
## Quick Start
|
## Quick Start
|
||||||
|
|
||||||
```go
|
```go
|
||||||
import "forge.lthn.ai/core/go-ratelimit"
|
import "dappco.re/go/core/go-ratelimit"
|
||||||
|
|
||||||
// YAML backend (default, single-process)
|
// YAML backend (default, single-process)
|
||||||
rl, err := ratelimit.New()
|
rl, err := ratelimit.New()
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
// SQLite backend (multi-process)
|
// SQLite backend (multi-process)
|
||||||
rl, err := ratelimit.NewWithSQLite("~/.core/ratelimits.db")
|
rl, err = ratelimit.NewWithSQLite("/tmp/ratelimits.db")
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
defer rl.Close()
|
defer rl.Close()
|
||||||
|
|
||||||
ok, reason := rl.CanSend("gemini-2.0-flash", 1500)
|
if rl.CanSend("gemini-2.0-flash", 1500) {
|
||||||
if ok {
|
rl.RecordUsage("gemini-2.0-flash", 1000, 500)
|
||||||
rl.RecordUsage("gemini-2.0-flash", 1500)
|
}
|
||||||
|
|
||||||
|
if err := rl.Persist(); err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
For agent workflows, `Decide` returns a structured verdict with retry guidance:
|
||||||
|
|
||||||
|
```go
|
||||||
|
decision := rl.Decide("gemini-2.0-flash", 1500)
|
||||||
|
if !decision.Allowed {
|
||||||
|
log.Printf("throttled (%s); retry after %s", decision.Code, decision.RetryAfter)
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
@ -37,12 +57,14 @@ if ok {
|
||||||
## Build & Test
|
## Build & Test
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
go build ./...
|
||||||
go test ./...
|
go test ./...
|
||||||
go test -race ./...
|
go test -race ./...
|
||||||
go vet ./...
|
go vet ./...
|
||||||
go build ./...
|
go test -cover ./...
|
||||||
|
go mod tidy
|
||||||
```
|
```
|
||||||
|
|
||||||
## Licence
|
## Licence
|
||||||
|
|
||||||
European Union Public Licence 1.2 — see [LICENCE](LICENCE) for details.
|
European Union Public Licence 1.2.
|
||||||
|
|
|
||||||
40
docs/api-contract.md
Normal file
40
docs/api-contract.md
Normal file
|
|
@ -0,0 +1,40 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
|
# API Contract
|
||||||
|
|
||||||
|
Test coverage is marked `yes` when the symbol is exercised by the existing test suite in `ratelimit_test.go`, `sqlite_test.go`, `error_test.go`, or `iter_test.go`.
|
||||||
|
|
||||||
|
| Kind | Name | Signature | Description | Test coverage |
|
||||||
|
| --- | --- | --- | --- | --- |
|
||||||
|
| Type | `Provider` | `type Provider string` | Identifies an LLM provider used to select default quota profiles. | yes |
|
||||||
|
| Type | `ModelQuota` | `type ModelQuota struct { MaxRPM int; MaxTPM int; MaxRPD int }` | Declares per-model RPM, TPM, and RPD limits; `0` means unlimited. | yes |
|
||||||
|
| Type | `ProviderProfile` | `type ProviderProfile struct { Provider Provider; Models map[string]ModelQuota }` | Bundles a provider identifier with its default model quota map. | yes |
|
||||||
|
| Type | `Config` | `type Config struct { FilePath string; Backend string; Quotas map[string]ModelQuota; Providers []Provider }` | Configures limiter initialisation, persistence settings, explicit quotas, and provider defaults. | yes |
|
||||||
|
| Type | `TokenEntry` | `type TokenEntry struct { Time time.Time; Count int }` | Records a timestamped token-usage event. | yes |
|
||||||
|
| Type | `UsageStats` | `type UsageStats struct { Requests []time.Time; Tokens []TokenEntry; DayStart time.Time; DayCount int }` | Stores per-model sliding-window request and token history plus rolling daily usage state. | yes |
|
||||||
|
| Type | `RateLimiter` | `type RateLimiter struct { Quotas map[string]ModelQuota; State map[string]*UsageStats }` | Manages quotas, usage state, persistence, and concurrency across models. | yes |
|
||||||
|
| Type | `ModelStats` | `type ModelStats struct { RPM int; MaxRPM int; TPM int; MaxTPM int; RPD int; MaxRPD int; DayStart time.Time }` | Represents a snapshot of current usage and configured limits for a model. | yes |
|
||||||
|
| Type | `DecisionCode` | `type DecisionCode string` | Machine-readable allow/deny codes returned by `Decide` (e.g., `ok`, `rpm_exceeded`). | yes |
|
||||||
|
| Type | `Decision` | `type Decision struct { Allowed bool; Code DecisionCode; Reason string; RetryAfter time.Duration; Stats ModelStats }` | Structured decision result with a code, human-readable reason, optional retry guidance, and a stats snapshot. | yes |
|
||||||
|
| Function | `DefaultProfiles` | `func DefaultProfiles() map[Provider]ProviderProfile` | Returns the built-in quota profiles for the supported providers. | yes |
|
||||||
|
| Function | `New` | `func New() (*RateLimiter, error)` | Creates a new limiter with Gemini defaults for backward-compatible YAML-backed usage. | yes |
|
||||||
|
| Function | `NewWithConfig` | `func NewWithConfig(cfg Config) (*RateLimiter, error)` | Creates a YAML-backed limiter from explicit configuration, defaulting to Gemini when config is empty. | yes |
|
||||||
|
| Function | `NewWithSQLite` | `func NewWithSQLite(dbPath string) (*RateLimiter, error)` | Creates a SQLite-backed limiter with Gemini defaults and opens or creates the database. | yes |
|
||||||
|
| Function | `NewWithSQLiteConfig` | `func NewWithSQLiteConfig(dbPath string, cfg Config) (*RateLimiter, error)` | Creates a SQLite-backed limiter from explicit configuration and always uses SQLite persistence. | yes |
|
||||||
|
| Function | `MigrateYAMLToSQLite` | `func MigrateYAMLToSQLite(yamlPath, sqlitePath string) error` | Migrates quotas and usage state from a YAML state file into a SQLite database. | yes |
|
||||||
|
| Function | `CountTokens` | `func CountTokens(ctx context.Context, apiKey, model, text string) (int, error)` | Calls the Gemini `countTokens` API for a prompt and returns the reported token count. | yes |
|
||||||
|
| Method | `SetQuota` | `func (rl *RateLimiter) SetQuota(model string, quota ModelQuota)` | Sets or replaces a model quota at runtime. | yes |
|
||||||
|
| Method | `AddProvider` | `func (rl *RateLimiter) AddProvider(provider Provider)` | Loads a built-in provider profile and overwrites any matching model quotas. | yes |
|
||||||
|
| Method | `Load` | `func (rl *RateLimiter) Load() error` | Loads quotas and usage state from YAML or SQLite into memory. | yes |
|
||||||
|
| Method | `Persist` | `func (rl *RateLimiter) Persist() error` | Persists a snapshot of quotas and usage state to YAML or SQLite. | yes |
|
||||||
|
| Method | `BackgroundPrune` | `func (rl *RateLimiter) BackgroundPrune(interval time.Duration) func()` | Starts periodic pruning of expired usage state and returns a stop function. | yes |
|
||||||
|
| Method | `CanSend` | `func (rl *RateLimiter) CanSend(model string, estimatedTokens int) bool` | Reports whether a request with the estimated token count fits within current limits. | yes |
|
||||||
|
| Method | `Decide` | `func (rl *RateLimiter) Decide(model string, estimatedTokens int) Decision` | Returns structured allow/deny information including code, reason, retry guidance, and stats snapshot without recording usage. | yes |
|
||||||
|
| Method | `RecordUsage` | `func (rl *RateLimiter) RecordUsage(model string, promptTokens, outputTokens int)` | Records a successful request into the sliding-window and daily counters. | yes |
|
||||||
|
| Method | `WaitForCapacity` | `func (rl *RateLimiter) WaitForCapacity(ctx context.Context, model string, tokens int) error` | Blocks until `Decide` allows the request, sleeping according to `RetryAfter` hints or one-second polls. | yes |
|
||||||
|
| Method | `Reset` | `func (rl *RateLimiter) Reset(model string)` | Clears usage state for one model or for all models when `model` is empty. | yes |
|
||||||
|
| Method | `Models` | `func (rl *RateLimiter) Models() iter.Seq[string]` | Returns a sorted iterator of all model names known from quotas or state. | yes |
|
||||||
|
| Method | `Iter` | `func (rl *RateLimiter) Iter() iter.Seq2[string, ModelStats]` | Returns a sorted iterator of model names paired with current stats snapshots. | yes |
|
||||||
|
| Method | `Stats` | `func (rl *RateLimiter) Stats(model string) ModelStats` | Returns current stats for a single model after pruning expired usage. | yes |
|
||||||
|
| Method | `AllStats` | `func (rl *RateLimiter) AllStats() map[string]ModelStats` | Returns stats snapshots for every tracked model. | yes |
|
||||||
|
| Method | `Close` | `func (rl *RateLimiter) Close() error` | Closes SQLite resources for SQLite-backed limiters and is a no-op for YAML-backed limiters. | yes |
|
||||||
|
|
@ -1,81 +1,74 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
|
---
|
||||||
|
title: Architecture
|
||||||
|
description: Internals of go-ratelimit -- sliding window algorithm, provider quota system, persistence backends, and concurrency model.
|
||||||
|
---
|
||||||
|
|
||||||
# Architecture
|
# Architecture
|
||||||
|
|
||||||
go-ratelimit is a provider-agnostic rate limiter for LLM API calls. It enforces
|
go-ratelimit is a provider-agnostic rate limiter for LLM API calls. It enforces
|
||||||
three independent quota dimensions per model — requests per minute (RPM), tokens
|
three independent quota dimensions per model -- requests per minute (RPM), tokens
|
||||||
per minute (TPM), and requests per day (RPD) — using an in-memory sliding window
|
per minute (TPM), and requests per day (RPD) -- using an in-memory sliding window
|
||||||
that can be persisted across process restarts via YAML or SQLite.
|
that can be persisted across process restarts via YAML or SQLite.
|
||||||
|
|
||||||
Module path: `forge.lthn.ai/core/go-ratelimit`
|
Module path: `dappco.re/go/core/go-ratelimit`
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Sliding Window Algorithm
|
## Key Types
|
||||||
|
|
||||||
The limiter maintains per-model `UsageStats` structs in memory:
|
### RateLimiter
|
||||||
|
|
||||||
|
The central struct. Holds the quota definitions, current usage state, a mutex for
|
||||||
|
thread safety, and an optional SQLite backend reference.
|
||||||
|
|
||||||
|
```go
|
||||||
|
type RateLimiter struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
Quotas map[string]ModelQuota // per-model quota definitions
|
||||||
|
State map[string]*UsageStats // per-model sliding window state
|
||||||
|
filePath string // YAML file path (ignored when SQLite is active)
|
||||||
|
sqlite *sqliteStore // non-nil when using SQLite backend
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### ModelQuota
|
||||||
|
|
||||||
|
Defines the rate limits for a single model. A zero value in any field means that
|
||||||
|
dimension is unlimited.
|
||||||
|
|
||||||
|
```go
|
||||||
|
type ModelQuota struct {
|
||||||
|
MaxRPM int `yaml:"max_rpm"` // Requests per minute (0 = unlimited)
|
||||||
|
MaxTPM int `yaml:"max_tpm"` // Tokens per minute (0 = unlimited)
|
||||||
|
MaxRPD int `yaml:"max_rpd"` // Requests per day (0 = unlimited)
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### UsageStats
|
||||||
|
|
||||||
|
Tracks the sliding window state for a single model.
|
||||||
|
|
||||||
```go
|
```go
|
||||||
type UsageStats struct {
|
type UsageStats struct {
|
||||||
Requests []time.Time // timestamps of recent requests (1-minute window)
|
Requests []time.Time // timestamps of recent requests (1-minute window)
|
||||||
Tokens []TokenEntry // token counts with timestamps (1-minute window)
|
Tokens []TokenEntry // token counts with timestamps (1-minute window)
|
||||||
DayStart time.Time // when the current daily window started
|
DayStart time.Time // when the current 24-hour window started
|
||||||
DayCount int // total requests recorded since DayStart
|
DayCount int // total requests since DayStart
|
||||||
|
}
|
||||||
|
|
||||||
|
type TokenEntry struct {
|
||||||
|
Time time.Time
|
||||||
|
Count int // prompt + output tokens for this request
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
Every call to `CanSend()` or `Stats()` first calls `prune()`, which scans both
|
### Config
|
||||||
slices and discards entries older than `now - 1 minute`. Pruning is done
|
|
||||||
in-place to avoid allocation on the hot path:
|
Controls `RateLimiter` initialisation.
|
||||||
|
|
||||||
```go
|
```go
|
||||||
validReqs := 0
|
|
||||||
for _, t := range stats.Requests {
|
|
||||||
if t.After(window) {
|
|
||||||
stats.Requests[validReqs] = t
|
|
||||||
validReqs++
|
|
||||||
}
|
|
||||||
}
|
|
||||||
stats.Requests = stats.Requests[:validReqs]
|
|
||||||
```
|
|
||||||
|
|
||||||
The same loop runs for token entries. After pruning, `CanSend()` checks each
|
|
||||||
quota dimension in priority order: RPD first (cheapest check), then RPM, then
|
|
||||||
TPM. A zero value for any dimension means that dimension is unlimited. If all
|
|
||||||
three are zero the model is treated as fully unlimited and the check short-circuits
|
|
||||||
before touching any state.
|
|
||||||
|
|
||||||
### Daily Reset
|
|
||||||
|
|
||||||
The daily counter resets automatically inside `prune()`. When
|
|
||||||
`now - stats.DayStart >= 24h`, `DayCount` is set to zero and `DayStart` is set
|
|
||||||
to the current time. This means the daily window is a rolling 24-hour period
|
|
||||||
anchored to the first request of the day, not a calendar boundary.
|
|
||||||
|
|
||||||
### Concurrency
|
|
||||||
|
|
||||||
All reads and writes are protected by a single `sync.RWMutex`. Methods that
|
|
||||||
write state — `CanSend()`, `RecordUsage()`, `Reset()`, `Load()` — acquire a
|
|
||||||
full write lock. `Persist()`, `Stats()`, and `AllStats()` acquire a read lock
|
|
||||||
where possible. The `CanSend()` method acquires a write lock because it calls
|
|
||||||
`prune()`, which mutates the state slices.
|
|
||||||
|
|
||||||
`go test -race ./...` passes clean with 20 goroutines performing concurrent
|
|
||||||
`CanSend()`, `RecordUsage()`, and `Stats()` calls.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Provider and Quota Configuration
|
|
||||||
|
|
||||||
### Types
|
|
||||||
|
|
||||||
```go
|
|
||||||
type Provider string // "gemini", "openai", "anthropic", "local"
|
|
||||||
|
|
||||||
type ModelQuota struct {
|
|
||||||
MaxRPM int `yaml:"max_rpm"` // 0 = unlimited
|
|
||||||
MaxTPM int `yaml:"max_tpm"`
|
|
||||||
MaxRPD int `yaml:"max_rpd"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type Config struct {
|
type Config struct {
|
||||||
FilePath string // default: ~/.core/ratelimits.yaml
|
FilePath string // default: ~/.core/ratelimits.yaml
|
||||||
Backend string // "yaml" (default) or "sqlite"
|
Backend string // "yaml" (default) or "sqlite"
|
||||||
|
|
@ -84,32 +77,116 @@ type Config struct {
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
### Quota Resolution
|
### Provider
|
||||||
|
|
||||||
1. Provider profiles are loaded first (from `DefaultProfiles()`).
|
A string type identifying an LLM provider. Four constants are defined:
|
||||||
2. Explicit `Config.Quotas` are merged on top, overriding any matching model.
|
|
||||||
|
```go
|
||||||
|
type Provider string
|
||||||
|
|
||||||
|
const (
|
||||||
|
ProviderGemini Provider = "gemini"
|
||||||
|
ProviderOpenAI Provider = "openai"
|
||||||
|
ProviderAnthropic Provider = "anthropic"
|
||||||
|
ProviderLocal Provider = "local"
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Sliding Window Algorithm
|
||||||
|
|
||||||
|
Every call to `CanSend()` or `Stats()` first calls `prune()`, which removes
|
||||||
|
entries older than one minute from the `Requests` and `Tokens` slices. Pruning
|
||||||
|
is done in-place using `slices.DeleteFunc` to minimise allocations:
|
||||||
|
|
||||||
|
```go
|
||||||
|
window := now.Add(-1 * time.Minute)
|
||||||
|
|
||||||
|
stats.Requests = slices.DeleteFunc(stats.Requests, func(t time.Time) bool {
|
||||||
|
return t.Before(window)
|
||||||
|
})
|
||||||
|
stats.Tokens = slices.DeleteFunc(stats.Tokens, func(t TokenEntry) bool {
|
||||||
|
return t.Time.Before(window)
|
||||||
|
})
|
||||||
|
```
|
||||||
|
|
||||||
|
After pruning, `CanSend()` checks each quota dimension. If all three limits
|
||||||
|
(RPM, TPM, RPD) are zero, the model is treated as fully unlimited and the
|
||||||
|
check short-circuits before touching any state.
|
||||||
|
|
||||||
|
The check order is: RPD, then RPM, then TPM. RPD is checked first because it
|
||||||
|
is the cheapest comparison (a single integer). TPM is checked last because it
|
||||||
|
requires summing the token counts in the sliding window.
|
||||||
|
|
||||||
|
`Decide()` follows the same path as `CanSend()` but returns a structured
|
||||||
|
`Decision` containing a machine-readable code, reason, `RetryAfter` guidance,
|
||||||
|
and a `ModelStats` snapshot. It is agent-facing and does not record usage;
|
||||||
|
`WaitForCapacity()` consumes its `RetryAfter` hint to avoid unnecessary
|
||||||
|
one-second polling when limits are saturated.
|
||||||
|
|
||||||
|
### Daily Reset
|
||||||
|
|
||||||
|
The daily counter resets automatically inside `prune()`. When
|
||||||
|
`now - stats.DayStart >= 24h`, `DayCount` is set to zero and `DayStart` is
|
||||||
|
updated to the current time. The daily window is a rolling 24-hour period
|
||||||
|
anchored to the first request of the day, not a calendar boundary.
|
||||||
|
|
||||||
|
### Background Pruning
|
||||||
|
|
||||||
|
`BackgroundPrune(interval)` starts a goroutine that periodically prunes all
|
||||||
|
model states on a configurable interval. It returns a cancel function to stop
|
||||||
|
the pruner:
|
||||||
|
|
||||||
|
```go
|
||||||
|
stop := rl.BackgroundPrune(30 * time.Second)
|
||||||
|
defer stop()
|
||||||
|
```
|
||||||
|
|
||||||
|
This prevents memory growth in long-running processes where some models may
|
||||||
|
accumulate stale entries between calls to `CanSend()`.
|
||||||
|
|
||||||
|
### Memory Cleanup
|
||||||
|
|
||||||
|
When `prune()` empties both the `Requests` and `Tokens` slices for a model,
|
||||||
|
and `DayCount` is also zero, the entire `UsageStats` entry is deleted from
|
||||||
|
the `State` map. This prevents memory leaks from models that were used once
|
||||||
|
and never again.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Provider and Quota Configuration
|
||||||
|
|
||||||
|
### Quota Resolution Order
|
||||||
|
|
||||||
|
1. Provider profiles are loaded first from `DefaultProfiles()`.
|
||||||
|
2. Explicit `Config.Quotas` are merged on top using `maps.Copy`, overriding any
|
||||||
|
matching model.
|
||||||
3. If neither `Providers` nor `Quotas` are specified, Gemini defaults are used.
|
3. If neither `Providers` nor `Quotas` are specified, Gemini defaults are used.
|
||||||
|
|
||||||
`SetQuota()` and `AddProvider()` allow runtime modification; both are
|
`SetQuota()` and `AddProvider()` allow runtime modification. Both acquire the
|
||||||
mutex-protected. `AddProvider()` is additive — it does not remove existing
|
write lock. `AddProvider()` is additive -- it does not remove existing quotas for
|
||||||
quotas for models outside the new provider's profile.
|
models outside the new provider's profile.
|
||||||
|
|
||||||
### Default Quotas (as of February 2026)
|
### Default Quotas (as of February 2026)
|
||||||
|
|
||||||
| Provider | Model | MaxRPM | MaxTPM | MaxRPD |
|
| Provider | Model | MaxRPM | MaxTPM | MaxRPD |
|
||||||
|-----------|------------------------|-----------|-----------|-----------|
|
|-----------|------------------------|-----------|-------------|-----------|
|
||||||
| Gemini | gemini-3-pro-preview | 150 | 1,000,000 | 1,000 |
|
| Gemini | gemini-3-pro-preview | 150 | 1,000,000 | 1,000 |
|
||||||
| Gemini | gemini-3-flash-preview | 150 | 1,000,000 | 1,000 |
|
| Gemini | gemini-3-flash-preview | 150 | 1,000,000 | 1,000 |
|
||||||
| Gemini | gemini-2.5-pro | 150 | 1,000,000 | 1,000 |
|
| Gemini | gemini-2.5-pro | 150 | 1,000,000 | 1,000 |
|
||||||
| Gemini | gemini-2.0-flash | 150 | 1,000,000 | unlimited |
|
| Gemini | gemini-2.0-flash | 150 | 1,000,000 | unlimited |
|
||||||
| Gemini | gemini-2.0-flash-lite | unlimited | unlimited | unlimited |
|
| Gemini | gemini-2.0-flash-lite | unlimited | unlimited | unlimited |
|
||||||
| OpenAI | gpt-4o, gpt-4-turbo | 500 | 30,000 | unlimited |
|
| OpenAI | gpt-4o | 500 | 30,000 | unlimited |
|
||||||
| OpenAI | gpt-4o-mini, o1-mini | 500 | 200,000 | unlimited |
|
| OpenAI | gpt-4o-mini | 500 | 200,000 | unlimited |
|
||||||
| OpenAI | o1, o3-mini | 500 | varies | unlimited |
|
| OpenAI | gpt-4-turbo | 500 | 30,000 | unlimited |
|
||||||
| Anthropic | claude-opus-4 | 50 | 40,000 | unlimited |
|
| OpenAI | o1 | 500 | 30,000 | unlimited |
|
||||||
| Anthropic | claude-sonnet-4 | 50 | 40,000 | unlimited |
|
| OpenAI | o1-mini | 500 | 200,000 | unlimited |
|
||||||
| Anthropic | claude-haiku-3.5 | 50 | 50,000 | unlimited |
|
| OpenAI | o3-mini | 500 | 200,000 | unlimited |
|
||||||
| Local | (none by default) | user-defined |
|
| Anthropic | claude-opus-4 | 50 | 40,000 | unlimited |
|
||||||
|
| Anthropic | claude-sonnet-4 | 50 | 40,000 | unlimited |
|
||||||
|
| Anthropic | claude-haiku-3.5 | 50 | 50,000 | unlimited |
|
||||||
|
| Local | (none by default) | user-defined |
|
||||||
|
|
||||||
The Local provider exists for local inference backends (Ollama, MLX, llama.cpp)
|
The Local provider exists for local inference backends (Ollama, MLX, llama.cpp)
|
||||||
where the throttle limit is hardware rather than an API quota. No defaults are
|
where the throttle limit is hardware rather than an API quota. No defaults are
|
||||||
|
|
@ -117,10 +194,54 @@ provided; callers add per-model limits via `Config.Quotas` or `SetQuota()`.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## YAML Persistence (Legacy)
|
## Constructors
|
||||||
|
|
||||||
The default backend serialises the entire `RateLimiter` struct — both the
|
| Function | Backend | Default Provider |
|
||||||
`Quotas` map and the `State` map — to a YAML file at `~/.core/ratelimits.yaml`.
|
|----------|---------|------------------|
|
||||||
|
| `New()` | YAML | Gemini |
|
||||||
|
| `NewWithConfig(cfg)` | YAML | Configurable (Gemini if empty) |
|
||||||
|
| `NewWithSQLite(dbPath)` | SQLite | Gemini |
|
||||||
|
| `NewWithSQLiteConfig(dbPath, cfg)` | SQLite | Configurable (Gemini if empty) |
|
||||||
|
|
||||||
|
`Close()` releases the database connection for SQLite-backed limiters. It is a
|
||||||
|
no-op on YAML-backed limiters. Always call `Close()` (or `defer rl.Close()`)
|
||||||
|
when using the SQLite backend.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Data Flow
|
||||||
|
|
||||||
|
A typical request lifecycle:
|
||||||
|
|
||||||
|
```
|
||||||
|
1. CanSend(model, estimatedTokens)
|
||||||
|
|-- acquires write lock
|
||||||
|
|-- looks up ModelQuota for the model
|
||||||
|
|-- if unknown model or all-zero quota: returns true (allowed)
|
||||||
|
|-- calls prune(model) to discard stale entries
|
||||||
|
|-- checks RPD, RPM, TPM against the pruned state
|
||||||
|
'-- returns true/false
|
||||||
|
|
||||||
|
2. (caller makes the API call)
|
||||||
|
|
||||||
|
3. RecordUsage(model, promptTokens, outputTokens)
|
||||||
|
|-- acquires write lock
|
||||||
|
|-- calls prune(model)
|
||||||
|
|-- appends to Requests and Tokens slices
|
||||||
|
'-- increments DayCount
|
||||||
|
|
||||||
|
4. Persist()
|
||||||
|
|-- acquires write lock, clones state, releases lock
|
||||||
|
|-- YAML: marshals to file
|
||||||
|
'-- SQLite: saves quotas and state in transactions
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## YAML Persistence
|
||||||
|
|
||||||
|
The default backend serialises both the `Quotas` map and the `State` map to a
|
||||||
|
YAML file at `~/.core/ratelimits.yaml` (configurable via `Config.FilePath`).
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
quotas:
|
quotas:
|
||||||
|
|
@ -139,39 +260,34 @@ state:
|
||||||
day_count: 42
|
day_count: 42
|
||||||
```
|
```
|
||||||
|
|
||||||
`Persist()` creates parent directories with `os.MkdirAll` before writing.
|
`Persist()` creates parent directories with the `core.Fs` helper before writing.
|
||||||
`Load()` treats a missing file as an empty state (no error). Corrupt or
|
`Load()` treats a missing file as an empty state (no error). Corrupt or
|
||||||
unreadable files return an error.
|
unreadable files return an error.
|
||||||
|
|
||||||
**Limitations of YAML backend:**
|
**Limitations of the YAML backend:**
|
||||||
|
|
||||||
- Single-process only. Concurrent writes from multiple processes corrupt the
|
- Single-process only. Concurrent writes from multiple processes corrupt the
|
||||||
file because the write is not atomic at the OS level.
|
file because the write is not atomic at the OS level.
|
||||||
- The entire state is serialised on every `Persist()` call, which grows linearly
|
- The entire state is serialised on every `Persist()` call.
|
||||||
with the number of tracked models and entries.
|
- Timestamps are serialised as RFC 3339 strings.
|
||||||
- Timestamps are serialised as RFC3339 strings; sub-nanosecond precision is
|
|
||||||
preserved by Go's time marshaller but depends on the YAML library.
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## SQLite Backend
|
## SQLite Backend
|
||||||
|
|
||||||
The SQLite backend was added in Phase 2 to support multi-process scenarios and
|
The SQLite backend supports multi-process scenarios. It uses `modernc.org/sqlite`,
|
||||||
provide a more robust persistence layer. It uses `modernc.org/sqlite` — a pure
|
a pure Go port of SQLite that compiles without CGO.
|
||||||
Go port of SQLite that compiles without CGO.
|
|
||||||
|
|
||||||
### Connection Settings
|
### Connection Settings
|
||||||
|
|
||||||
```go
|
```go
|
||||||
db.SetMaxOpenConns(1) // single connection for PRAGMA consistency
|
db.SetMaxOpenConns(1) // single connection for PRAGMA consistency
|
||||||
db.Exec("PRAGMA journal_mode=WAL") // WAL mode for concurrent readers
|
db.Exec("PRAGMA journal_mode=WAL") // concurrent readers alongside a single writer
|
||||||
db.Exec("PRAGMA busy_timeout=5000") // 5-second busy timeout
|
db.Exec("PRAGMA busy_timeout=5000") // 5-second wait on lock contention
|
||||||
```
|
```
|
||||||
|
|
||||||
WAL mode allows one writer and multiple concurrent readers. The 5-second busy
|
WAL mode allows one writer and multiple concurrent readers. The 5-second busy
|
||||||
timeout prevents immediate failure when a second process is mid-commit. A single
|
timeout prevents immediate failure when a second process is mid-commit.
|
||||||
`sql.DB` connection is used because SQLite's WAL mode handles reader concurrency
|
|
||||||
at the file level; multiple Go connections to the same file through a single
|
|
||||||
process would not add throughput but would complicate locking.
|
|
||||||
|
|
||||||
### Schema
|
### Schema
|
||||||
|
|
||||||
|
|
@ -196,7 +312,7 @@ CREATE TABLE IF NOT EXISTS tokens (
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS daily (
|
CREATE TABLE IF NOT EXISTS daily (
|
||||||
model TEXT PRIMARY KEY,
|
model TEXT PRIMARY KEY,
|
||||||
day_start INTEGER NOT NULL, -- UnixNano
|
day_start INTEGER NOT NULL, -- UnixNano
|
||||||
day_count INTEGER NOT NULL DEFAULT 0
|
day_count INTEGER NOT NULL DEFAULT 0
|
||||||
);
|
);
|
||||||
|
|
||||||
|
|
@ -205,51 +321,22 @@ CREATE INDEX IF NOT EXISTS idx_tokens_model_ts ON tokens(model, ts);
|
||||||
```
|
```
|
||||||
|
|
||||||
Timestamps are stored as `INTEGER` UnixNano values. This preserves nanosecond
|
Timestamps are stored as `INTEGER` UnixNano values. This preserves nanosecond
|
||||||
precision without relying on SQLite's text date format, and allows efficient
|
precision and allows efficient range queries using the composite indices.
|
||||||
range queries using the composite indices.
|
|
||||||
|
|
||||||
### Save Strategy
|
### Save Strategy
|
||||||
|
|
||||||
`saveState()` uses a delete-then-insert pattern inside a single transaction.
|
- **Quotas**: full snapshot replace inside a single transaction. `saveQuotas()`
|
||||||
All three state tables are truncated and rewritten atomically:
|
clears the table and reinserts the current quota map.
|
||||||
|
- **State**: Delete-then-insert inside a single transaction. All three state
|
||||||
```go
|
tables (`requests`, `tokens`, `daily`) are truncated and rewritten atomically.
|
||||||
tx.Exec("DELETE FROM requests")
|
|
||||||
tx.Exec("DELETE FROM tokens")
|
|
||||||
tx.Exec("DELETE FROM daily")
|
|
||||||
// then INSERT for every model in state
|
|
||||||
tx.Commit()
|
|
||||||
```
|
|
||||||
|
|
||||||
`saveQuotas()` uses `INSERT ... ON CONFLICT(model) DO UPDATE` (upsert) so
|
|
||||||
existing quota rows are updated in place without deleting unrelated models.
|
|
||||||
|
|
||||||
### Constructors
|
|
||||||
|
|
||||||
```go
|
|
||||||
// YAML backend (default)
|
|
||||||
rl, err := ratelimit.New()
|
|
||||||
rl, err := ratelimit.NewWithConfig(cfg)
|
|
||||||
|
|
||||||
// SQLite backend
|
|
||||||
rl, err := ratelimit.NewWithSQLite(dbPath)
|
|
||||||
rl, err := ratelimit.NewWithSQLiteConfig(dbPath, cfg)
|
|
||||||
|
|
||||||
defer rl.Close() // releases the database connection
|
|
||||||
```
|
|
||||||
|
|
||||||
`Close()` is a no-op on YAML-backed limiters.
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Migration Path
|
## Migration Path
|
||||||
|
|
||||||
`MigrateYAMLToSQLite(yamlPath, sqlitePath string) error` reads an existing YAML
|
`MigrateYAMLToSQLite(yamlPath, sqlitePath)` reads an existing YAML state file
|
||||||
state file and writes all quotas and usage state to a new SQLite database. The
|
and writes all quotas and usage state to a new SQLite database. The function is
|
||||||
function is idempotent — running it again on the same YAML file overwrites the
|
idempotent -- running it again overwrites the SQLite database state.
|
||||||
SQLite database state.
|
|
||||||
|
|
||||||
Typical one-time migration:
|
|
||||||
|
|
||||||
```go
|
```go
|
||||||
err := ratelimit.MigrateYAMLToSQLite(
|
err := ratelimit.MigrateYAMLToSQLite(
|
||||||
|
|
@ -258,29 +345,59 @@ err := ratelimit.MigrateYAMLToSQLite(
|
||||||
)
|
)
|
||||||
```
|
```
|
||||||
|
|
||||||
After migration, switch the constructor:
|
After migration, switch the constructor from `New()` to `NewWithSQLite()`. The
|
||||||
|
YAML file can be kept as a backup; the two backends do not share state.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Iterators
|
||||||
|
|
||||||
|
Two Go 1.26+ iterators are provided for inspecting the limiter state:
|
||||||
|
|
||||||
|
- `Models() iter.Seq[string]` -- returns a sorted sequence of all model names
|
||||||
|
(from both `Quotas` and `State` maps, deduplicated).
|
||||||
|
- `Iter() iter.Seq2[string, ModelStats]` -- returns sorted model names paired
|
||||||
|
with their current `ModelStats` snapshot.
|
||||||
|
|
||||||
```go
|
```go
|
||||||
// Before
|
for model, stats := range rl.Iter() {
|
||||||
rl, _ := ratelimit.New()
|
fmt.Printf("%s: %d/%d RPM, %d/%d TPM\n",
|
||||||
|
model, stats.RPM, stats.MaxRPM, stats.TPM, stats.MaxTPM)
|
||||||
// After
|
}
|
||||||
rl, _ := ratelimit.NewWithSQLite(filepath.Join(home, ".core", "ratelimits.db"))
|
|
||||||
defer rl.Close()
|
|
||||||
```
|
```
|
||||||
|
|
||||||
The YAML file can be kept as a backup; the two backends do not share state.
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## CountTokens
|
## CountTokens
|
||||||
|
|
||||||
`CountTokens(apiKey, model, text string) (int, error)` calls the Google
|
`CountTokens(ctx, apiKey, model, text)` calls the Google Generative Language API
|
||||||
Generative Language API to obtain an exact token count for a prompt string. It
|
to obtain an exact token count for a prompt string. It is Gemini-specific and
|
||||||
is Gemini-specific and hardcodes the `generativelanguage.googleapis.com`
|
hardcodes the `generativelanguage.googleapis.com` endpoint.
|
||||||
endpoint. The URL is not configurable, which prevents unit testing of the
|
|
||||||
success path without network access.
|
|
||||||
|
|
||||||
For other providers, callers must supply `estimatedTokens` directly to
|
For other providers, callers must supply `estimatedTokens` directly to
|
||||||
`CanSend()` and `RecordUsage()`. Accurate token counts are typically available
|
`CanSend()`. Accurate token counts are typically available in API response
|
||||||
in API response metadata after a call completes.
|
metadata after a call completes.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Concurrency Model
|
||||||
|
|
||||||
|
All reads and writes are protected by a single `sync.RWMutex` on the
|
||||||
|
`RateLimiter` struct.
|
||||||
|
|
||||||
|
| Method | Lock type | Reason |
|
||||||
|
|--------|-----------|--------|
|
||||||
|
| `CanSend()` | Write | Calls `prune()`, which mutates state slices |
|
||||||
|
| `RecordUsage()` | Write | Appends to state slices |
|
||||||
|
| `Reset()` | Write | Deletes state entries |
|
||||||
|
| `Load()` | Write | Replaces in-memory state |
|
||||||
|
| `SetQuota()` | Write | Modifies quota map |
|
||||||
|
| `AddProvider()` | Write | Modifies quota map |
|
||||||
|
| `Persist()` | Write (brief) | Clones state, then releases lock before I/O |
|
||||||
|
| `Stats()` | Write | Calls `prune()` |
|
||||||
|
| `AllStats()` | Write | Prunes inline |
|
||||||
|
| `Models()` | Read | Reads keys only |
|
||||||
|
|
||||||
|
`Persist()` minimises lock contention by cloning the state under a write lock,
|
||||||
|
then performing I/O after releasing the lock. The test suite passes clean under
|
||||||
|
`go test -race ./...` with 20 goroutines performing concurrent operations.
|
||||||
|
|
|
||||||
75
docs/convention-drift-check-2026-03-23.md
Normal file
75
docs/convention-drift-check-2026-03-23.md
Normal file
|
|
@ -0,0 +1,75 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
|
# Convention Drift Check
|
||||||
|
|
||||||
|
Date: 2026-03-23
|
||||||
|
|
||||||
|
`CODEX.md` was not present anywhere under `/workspace`, so this check used
|
||||||
|
`CLAUDE.md`, `docs/development.md`, and the current tree.
|
||||||
|
|
||||||
|
## stdlib -> core.*
|
||||||
|
|
||||||
|
- `CLAUDE.md:65` still documents `forge.lthn.ai/core/go-io`, while the direct
|
||||||
|
dependency in `go.mod:6` is `dappco.re/go/core/io`.
|
||||||
|
- `CLAUDE.md:66` still documents `forge.lthn.ai/core/go-log`, while the direct
|
||||||
|
dependency in `go.mod:7` is `dappco.re/go/core/log`.
|
||||||
|
- `docs/development.md:175` still mandates `fmt.Errorf("ratelimit.Function: what: %w", err)`,
|
||||||
|
but the implementation and repo guidance use `coreerr.E(...)`
|
||||||
|
(`ratelimit.go:19`, `sqlite.go:8`, `CLAUDE.md:31`).
|
||||||
|
- `docs/development.md:195` omits the `core.*` direct dependencies entirely;
|
||||||
|
`go.mod:6-7` now declares both `dappco.re/go/core/io` and
|
||||||
|
`dappco.re/go/core/log`.
|
||||||
|
- `docs/index.md:102` likewise omits the `core.*` direct dependencies that are
|
||||||
|
present in `go.mod:6-7`.
|
||||||
|
|
||||||
|
## UK English
|
||||||
|
|
||||||
|
- `README.md:2` uses `License` in the badge alt text.
|
||||||
|
- `CONTRIBUTING.md:34` uses `License` as the section heading.
|
||||||
|
|
||||||
|
## Missing Tests
|
||||||
|
|
||||||
|
- `docs/development.md:150` says current coverage is `95.1%` and the floor is
|
||||||
|
`95%`, but `go test -coverprofile=/tmp/convention-cover.out ./...` currently
|
||||||
|
reports `94.4%`.
|
||||||
|
- `docs/history.md:28` still records coverage as `95.1%`; the current tree is
|
||||||
|
below that figure at `94.4%`.
|
||||||
|
- `docs/history.md:66` and `docs/development.md:157` say the
|
||||||
|
`os.UserHomeDir()` branch in `NewWithConfig()` is untestable, but
|
||||||
|
`error_test.go:648` now exercises that path.
|
||||||
|
- `ratelimit.go:202` (`Load`) is only `90.0%` covered.
|
||||||
|
- `ratelimit.go:242` (`Persist`) is only `90.0%` covered.
|
||||||
|
- `ratelimit.go:650` (`CountTokens`) is only `71.4%` covered.
|
||||||
|
- `sqlite.go:20` (`newSQLiteStore`) is only `64.3%` covered.
|
||||||
|
- `sqlite.go:110` (`loadQuotas`) is only `92.9%` covered.
|
||||||
|
- `sqlite.go:194` (`loadState`) is only `88.6%` covered.
|
||||||
|
- `ratelimit_test.go:746` starts a mock server for `CountTokens`, but the test
|
||||||
|
contains no assertion that exercises the success path.
|
||||||
|
- `iter_test.go:108` starts a second mock server for `CountTokens`, but again
|
||||||
|
does not exercise the mocked path.
|
||||||
|
- `error_test.go:42` defines `TestSQLiteInitErrors`, but the `WAL pragma failure`
|
||||||
|
subtest is still an empty placeholder.
|
||||||
|
|
||||||
|
## SPDX Headers
|
||||||
|
|
||||||
|
- `CLAUDE.md:1` is missing an SPDX header.
|
||||||
|
- `CONTRIBUTING.md:1` is missing an SPDX header.
|
||||||
|
- `README.md:1` is missing an SPDX header.
|
||||||
|
- `docs/architecture.md:1` is missing an SPDX header.
|
||||||
|
- `docs/development.md:1` is missing an SPDX header.
|
||||||
|
- `docs/history.md:1` is missing an SPDX header.
|
||||||
|
- `docs/index.md:1` is missing an SPDX header.
|
||||||
|
- `error_test.go:1` is missing an SPDX header.
|
||||||
|
- `go.mod:1` is missing an SPDX header.
|
||||||
|
- `iter_test.go:1` is missing an SPDX header.
|
||||||
|
- `ratelimit.go:1` is missing an SPDX header.
|
||||||
|
- `ratelimit_test.go:1` is missing an SPDX header.
|
||||||
|
- `sqlite.go:1` is missing an SPDX header.
|
||||||
|
- `sqlite_test.go:1` is missing an SPDX header.
|
||||||
|
|
||||||
|
## Verification
|
||||||
|
|
||||||
|
- `go test ./...`
|
||||||
|
- `go test -coverprofile=/tmp/convention-cover.out ./...`
|
||||||
|
- `go test -race ./...`
|
||||||
|
- `go vet ./...`
|
||||||
|
|
@ -1,9 +1,16 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
|
---
|
||||||
|
title: Development Guide
|
||||||
|
description: How to build, test, and contribute to go-ratelimit -- prerequisites, test patterns, coding standards, and commit conventions.
|
||||||
|
---
|
||||||
|
|
||||||
# Development Guide
|
# Development Guide
|
||||||
|
|
||||||
## Prerequisites
|
## Prerequisites
|
||||||
|
|
||||||
- Go 1.25 or later (the module declares `go 1.25.5`)
|
- **Go 1.26** or later (the module declares `go 1.26.0`)
|
||||||
- No CGO required — `modernc.org/sqlite` is a pure Go port
|
- No CGO required -- `modernc.org/sqlite` is a pure Go port
|
||||||
|
|
||||||
No C toolchain, no system SQLite library, no external build tools. A plain
|
No C toolchain, no system SQLite library, no external build tools. A plain
|
||||||
`go build ./...` is sufficient.
|
`go build ./...` is sufficient.
|
||||||
|
|
@ -13,6 +20,9 @@ No C toolchain, no system SQLite library, no external build tools. A plain
|
||||||
## Build and Test
|
## Build and Test
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
# Compile all packages
|
||||||
|
go build ./...
|
||||||
|
|
||||||
# Run all tests
|
# Run all tests
|
||||||
go test ./...
|
go test ./...
|
||||||
|
|
||||||
|
|
@ -34,40 +44,51 @@ go test -bench=BenchmarkCanSend -benchmem ./...
|
||||||
# Check for vet issues
|
# Check for vet issues
|
||||||
go vet ./...
|
go vet ./...
|
||||||
|
|
||||||
|
# Lint (requires golangci-lint)
|
||||||
|
golangci-lint run ./...
|
||||||
|
|
||||||
|
# Coverage check
|
||||||
|
go test -cover ./...
|
||||||
|
|
||||||
# Tidy dependencies
|
# Tidy dependencies
|
||||||
go mod tidy
|
go mod tidy
|
||||||
```
|
```
|
||||||
|
|
||||||
All three commands (`go test -race ./...`, `go vet ./...`, and `go mod tidy`)
|
Before a commit is pushed, `go build ./...`, `go test -race ./...`,
|
||||||
must produce no errors or warnings before a commit is pushed.
|
`go vet ./...`, `go test -cover ./...`, and `go mod tidy` must all pass
|
||||||
|
without errors, and coverage must remain at or above 95%.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Test Patterns
|
## Test Organisation
|
||||||
|
|
||||||
### File Organisation
|
### File Layout
|
||||||
|
|
||||||
- `ratelimit_test.go` — Phase 0 (core logic) and Phase 1 (provider profiles)
|
| File | Scope |
|
||||||
- `sqlite_test.go` — Phase 2 (SQLite backend)
|
|------|-------|
|
||||||
|
| `ratelimit_test.go` | Core sliding window logic, provider profiles, concurrency, benchmarks |
|
||||||
|
| `sqlite_test.go` | SQLite backend, migration, concurrent persistence |
|
||||||
|
| `error_test.go` | SQLite and YAML error paths |
|
||||||
|
| `iter_test.go` | `Models()` and `Iter()` iterators, `CountTokens` edge cases |
|
||||||
|
|
||||||
Both files are in `package ratelimit` (white-box tests) so they can access
|
All test files are in `package ratelimit` (white-box tests), giving access to
|
||||||
unexported fields and methods such as `prune()`, `filePath`, and `sqlite`.
|
unexported fields and methods such as `prune()`, `filePath`, and `sqlite`.
|
||||||
|
|
||||||
### Naming Convention
|
### Naming Convention
|
||||||
|
|
||||||
SQLite tests follow the `_Good`, `_Bad`, `_Ugly` suffix pattern:
|
SQLite tests follow the `_Good`, `_Bad`, `_Ugly` suffix pattern:
|
||||||
|
|
||||||
- `_Good` — happy path
|
- `_Good` -- happy path
|
||||||
- `_Bad` — expected error conditions (invalid paths, corrupt input)
|
- `_Bad` -- expected error conditions (invalid paths, corrupt input)
|
||||||
- `_Ugly` — panic-adjacent edge cases (corrupt DB files, truncated files)
|
- `_Ugly` -- panic-adjacent edge cases (corrupt database files, truncated files)
|
||||||
|
|
||||||
Core logic tests use plain descriptive names without suffixes, grouped by
|
Core logic tests use plain descriptive names without suffixes, grouped by method
|
||||||
method with table-driven subtests.
|
with table-driven subtests.
|
||||||
|
|
||||||
### Test Helpers
|
### Test Helper
|
||||||
|
|
||||||
`newTestLimiter(t *testing.T)` creates a `RateLimiter` with Gemini defaults and
|
`newTestLimiter(t)` creates a `RateLimiter` with Gemini defaults and redirects
|
||||||
redirects the YAML file path into `t.TempDir()`:
|
the YAML file path into `t.TempDir()`:
|
||||||
|
|
||||||
```go
|
```go
|
||||||
func newTestLimiter(t *testing.T) *RateLimiter {
|
func newTestLimiter(t *testing.T) *RateLimiter {
|
||||||
|
|
@ -86,45 +107,57 @@ after each test completes.
|
||||||
|
|
||||||
Tests use `github.com/stretchr/testify` exclusively:
|
Tests use `github.com/stretchr/testify` exclusively:
|
||||||
|
|
||||||
- `require.NoError(t, err)` — fail immediately on setup errors
|
- `require.NoError(t, err)` -- fail immediately on setup errors
|
||||||
- `assert.NoError(t, err)` — record failure but continue
|
- `assert.NoError(t, err)` -- record failure but continue
|
||||||
- `assert.Equal(t, expected, actual, "message")` — prefer over raw comparisons
|
- `assert.Equal(t, expected, actual, "message")` -- prefer over raw comparisons
|
||||||
- `assert.True / assert.False` — for boolean checks
|
- `assert.True` / `assert.False` -- for boolean checks
|
||||||
- `assert.Empty / assert.Len` — for slice length checks
|
- `assert.Empty` / `assert.Len` -- for slice length checks
|
||||||
- `assert.ErrorIs(t, err, context.DeadlineExceeded)` — for sentinel errors
|
- `assert.ErrorIs(t, err, target)` -- for sentinel errors
|
||||||
|
|
||||||
Do not use `t.Error`, `t.Fatal`, or `t.Log` directly.
|
Do not use `t.Error`, `t.Fatal`, or `t.Log` directly.
|
||||||
|
|
||||||
### Race Tests
|
### Concurrency Tests
|
||||||
|
|
||||||
Concurrency tests spin up goroutines and use `sync.WaitGroup`. They do not
|
Race tests spin up goroutines and use `sync.WaitGroup`. Some assert specific
|
||||||
assert anything beyond absence of data races (the race detector does the work):
|
outcomes (e.g., correct RPD count after concurrent recordings), while others
|
||||||
|
rely solely on the race detector to catch data races:
|
||||||
|
|
||||||
```go
|
```go
|
||||||
var wg sync.WaitGroup
|
var wg sync.WaitGroup
|
||||||
for i := range 20 {
|
for range 20 {
|
||||||
wg.Add(1)
|
wg.Go(func() {
|
||||||
go func() {
|
for range 50 {
|
||||||
defer wg.Done()
|
rl.CanSend(model, 10)
|
||||||
// concurrent operations
|
rl.RecordUsage(model, 5, 5)
|
||||||
}()
|
rl.Stats(model)
|
||||||
|
}
|
||||||
|
})
|
||||||
}
|
}
|
||||||
wg.Wait()
|
wg.Wait()
|
||||||
```
|
```
|
||||||
|
|
||||||
Run every concurrency test with `-race`. The CI baseline is `go test -race ./...`
|
Always run concurrency tests with `-race`.
|
||||||
clean.
|
|
||||||
|
### Benchmarks
|
||||||
|
|
||||||
|
The following benchmarks are included:
|
||||||
|
|
||||||
|
| Benchmark | What it measures |
|
||||||
|
|-----------|------------------|
|
||||||
|
| `BenchmarkCanSend` | CanSend with a 1,000-entry sliding window |
|
||||||
|
| `BenchmarkRecordUsage` | Recording usage on a single model |
|
||||||
|
| `BenchmarkCanSendConcurrent` | Parallel CanSend across goroutines |
|
||||||
|
| `BenchmarkCanSendWithPrune` | CanSend with 500 old + 500 new entries |
|
||||||
|
| `BenchmarkStats` | Stats retrieval with a 1,000-entry window |
|
||||||
|
| `BenchmarkAllStats` | AllStats across 5 models x 200 entries each |
|
||||||
|
| `BenchmarkPersist` | YAML persistence I/O |
|
||||||
|
| `BenchmarkSQLitePersist` | SQLite persistence I/O |
|
||||||
|
| `BenchmarkSQLiteLoad` | SQLite state loading |
|
||||||
|
|
||||||
### Coverage
|
### Coverage
|
||||||
|
|
||||||
Current coverage: 95.1%. The remaining 5% consists of three paths that cannot
|
Maintain at least 95% statement coverage. Verify it with `go test -cover ./...`
|
||||||
be covered in unit tests without modifying the production code:
|
and document any justified exception in the commit or PR that introduces it.
|
||||||
|
|
||||||
1. `CountTokens` success path — hardcoded Google API URL requires network access
|
|
||||||
2. `yaml.Marshal` error path in `Persist()` — cannot be triggered with valid Go structs
|
|
||||||
3. `os.UserHomeDir()` error path in `NewWithConfig()` — requires unsetting `$HOME`
|
|
||||||
|
|
||||||
Do not lower coverage below 95% without a documented reason.
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|
@ -137,26 +170,25 @@ Do not use American spellings in identifiers, comments, or documentation.
|
||||||
|
|
||||||
### Go Style
|
### Go Style
|
||||||
|
|
||||||
- All exported types, functions, and fields must have doc comments
|
- All exported types, functions, and fields must have doc comments.
|
||||||
- Error strings must be lowercase and not end with punctuation (Go convention)
|
- Error strings must be lowercase and not end with punctuation (Go convention).
|
||||||
- Contextual errors use `fmt.Errorf("package.Function: what: %w", err)` — the
|
- Contextual errors use `core.E("ratelimit.Function", "what", err)` so errors
|
||||||
prefix `ratelimit.` is included so errors identify their origin clearly
|
identify their origin clearly.
|
||||||
- No `init()` functions
|
- No `init()` functions.
|
||||||
- No global mutable state outside of `DefaultProfiles()` (which returns a fresh
|
- No global mutable state. `DefaultProfiles()` returns a fresh map on each call.
|
||||||
map on each call)
|
|
||||||
|
|
||||||
### Mutex Discipline
|
### Mutex Discipline
|
||||||
|
|
||||||
The `RateLimiter.mu` mutex is the only synchronisation primitive. Rules:
|
The `RateLimiter.mu` mutex is the only synchronisation primitive. Rules:
|
||||||
|
|
||||||
- Methods that call `prune()` always acquire the write lock (`mu.Lock()`),
|
- Methods that call `prune()` always acquire the write lock (`mu.Lock()`), even
|
||||||
even if they appear read-only, because `prune()` mutates slices
|
if they appear read-only, because `prune()` mutates state slices.
|
||||||
- `Persist()` acquires only the read lock (`mu.RLock()`) because it reads a
|
- `Persist()` acquires the write lock briefly to clone state, then releases it
|
||||||
snapshot of state
|
before performing I/O.
|
||||||
- Lock acquisition always happens at the top of the public method, never inside
|
- Lock acquisition always happens at the top of the public method, never inside
|
||||||
a helper — helpers document "Caller must hold the lock"
|
a helper. Helpers document "Caller must hold the lock".
|
||||||
- Never call a public method from inside another public method while holding
|
- Never call a public method from inside another public method while holding the
|
||||||
the lock (deadlock risk)
|
lock (deadlock risk).
|
||||||
|
|
||||||
### Dependencies
|
### Dependencies
|
||||||
|
|
||||||
|
|
@ -164,6 +196,7 @@ Direct dependencies are intentionally minimal:
|
||||||
|
|
||||||
| Dependency | Purpose |
|
| Dependency | Purpose |
|
||||||
|------------|---------|
|
|------------|---------|
|
||||||
|
| `dappco.re/go/core` | File I/O helpers, structured errors, JSON helpers, path/environment utilities |
|
||||||
| `gopkg.in/yaml.v3` | YAML serialisation for legacy backend |
|
| `gopkg.in/yaml.v3` | YAML serialisation for legacy backend |
|
||||||
| `modernc.org/sqlite` | Pure Go SQLite for persistent backend |
|
| `modernc.org/sqlite` | Pure Go SQLite for persistent backend |
|
||||||
| `github.com/stretchr/testify` | Test assertions (test-only) |
|
| `github.com/stretchr/testify` | Test assertions (test-only) |
|
||||||
|
|
@ -175,9 +208,8 @@ client libraries; the existing `CountTokens` function uses the standard library.
|
||||||
|
|
||||||
## Licence
|
## Licence
|
||||||
|
|
||||||
EUPL-1.2. Every new source file must carry the standard header if the project
|
EUPL-1.2. Confirm with the project lead before adding files under a different
|
||||||
adopts per-file headers in future. Confirm with the project lead before adding
|
licence.
|
||||||
files under a different licence.
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|
@ -205,3 +237,17 @@ Co-Authored-By: Virgil <virgil@lethean.io>
|
||||||
|
|
||||||
Commits must not be pushed unless `go test -race ./...` and `go vet ./...` both
|
Commits must not be pushed unless `go test -race ./...` and `go vet ./...` both
|
||||||
pass. `go mod tidy` must produce no changes.
|
pass. `go mod tidy` must produce no changes.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Linting
|
||||||
|
|
||||||
|
The project uses `golangci-lint` with the following enabled linters (see
|
||||||
|
`.golangci.yml`):
|
||||||
|
|
||||||
|
- `govet`, `errcheck`, `staticcheck`, `unused`, `gosimple`
|
||||||
|
- `ineffassign`, `typecheck`, `gocritic`, `gofmt`
|
||||||
|
|
||||||
|
Disabled linters: `exhaustive`, `wrapcheck`.
|
||||||
|
|
||||||
|
Run `golangci-lint run ./...` to check before committing.
|
||||||
|
|
|
||||||
|
|
@ -1,10 +1,13 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
# Project History
|
# Project History
|
||||||
|
|
||||||
## Origin
|
## Origin
|
||||||
|
|
||||||
go-ratelimit was extracted from the `pkg/ratelimit` package inside
|
go-ratelimit was extracted from the `pkg/ratelimit` package inside
|
||||||
`forge.lthn.ai/core/go` on 19 February 2026. The extraction gave the package
|
`forge.lthn.ai/core/go` on 19 February 2026. The package now lives at
|
||||||
its own module path, repository, and independent development cadence.
|
`dappco.re/go/core/go-ratelimit`, with its own repository and independent
|
||||||
|
development cadence.
|
||||||
|
|
||||||
Initial commit: `fa1a6fc` — `feat: extract go-ratelimit from core/go pkg/ratelimit`
|
Initial commit: `fa1a6fc` — `feat: extract go-ratelimit from core/go pkg/ratelimit`
|
||||||
|
|
||||||
|
|
@ -25,7 +28,7 @@ Commit: `3c63b10` — `feat(ratelimit): generalise beyond Gemini with provider p
|
||||||
|
|
||||||
Supplementary commit: `db958f2` — `test: expand race coverage and benchmarks`
|
Supplementary commit: `db958f2` — `test: expand race coverage and benchmarks`
|
||||||
|
|
||||||
Coverage increased from 77.1% to 95.1%. The test suite was rewritten using
|
Coverage increased from 77.1% to above the 95% floor. The test suite was rewritten using
|
||||||
testify with table-driven subtests throughout.
|
testify with table-driven subtests throughout.
|
||||||
|
|
||||||
### Tests added
|
### Tests added
|
||||||
|
|
@ -58,18 +61,6 @@ testify with table-driven subtests throughout.
|
||||||
- `BenchmarkAllStats` — 5 models x 200 entries
|
- `BenchmarkAllStats` — 5 models x 200 entries
|
||||||
- `BenchmarkPersist` — YAML I/O
|
- `BenchmarkPersist` — YAML I/O
|
||||||
|
|
||||||
### Remaining uncovered paths (5%)
|
|
||||||
|
|
||||||
These three paths are structurally impossible to cover in unit tests without
|
|
||||||
modifying production code:
|
|
||||||
|
|
||||||
1. `CountTokens` success path — the Google API URL is hardcoded; unit tests
|
|
||||||
cannot intercept the HTTP call without URL injection support
|
|
||||||
2. `yaml.Marshal` error path in `Persist()` — `yaml.Marshal` does not fail on
|
|
||||||
valid Go structs; the error branch exists for correctness only
|
|
||||||
3. `os.UserHomeDir()` error path in `NewWithConfig()` — triggered only when
|
|
||||||
`$HOME` is unset, which test infrastructure prevents
|
|
||||||
|
|
||||||
`go test -race ./...` passed clean. `go vet ./...` produced no warnings.
|
`go test -race ./...` passed clean. `go vet ./...` produced no warnings.
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
@ -139,7 +130,7 @@ established elsewhere in the ecosystem.
|
||||||
|
|
||||||
- `TestNewSQLiteStore_Good / _Bad` — creation and invalid path handling
|
- `TestNewSQLiteStore_Good / _Bad` — creation and invalid path handling
|
||||||
- `TestSQLiteQuotasRoundTrip_Good` — save/load round-trip
|
- `TestSQLiteQuotasRoundTrip_Good` — save/load round-trip
|
||||||
- `TestSQLiteQuotasUpsert_Good` — upsert replaces existing rows
|
- `TestSQLite_QuotasOverwrite_Good` — the latest quota snapshot replaces previous rows
|
||||||
- `TestSQLiteStateRoundTrip_Good` — multi-model state with nanosecond precision
|
- `TestSQLiteStateRoundTrip_Good` — multi-model state with nanosecond precision
|
||||||
- `TestSQLiteStateOverwrite_Good` — delete-then-insert atomicity
|
- `TestSQLiteStateOverwrite_Good` — delete-then-insert atomicity
|
||||||
- `TestSQLiteEmptyState_Good` — fresh database returns empty maps
|
- `TestSQLiteEmptyState_Good` — fresh database returns empty maps
|
||||||
|
|
@ -168,11 +159,10 @@ Not yet implemented. Intended downstream integrations:
|
||||||
|
|
||||||
## Known Limitations
|
## Known Limitations
|
||||||
|
|
||||||
**CountTokens URL is hardcoded.** The `CountTokens` helper calls
|
**CountTokens URL is hardcoded.** The exported `CountTokens` helper calls
|
||||||
`generativelanguage.googleapis.com` directly. There is no way to override the
|
`generativelanguage.googleapis.com` directly. Callers cannot redirect it to
|
||||||
base URL, which prevents testing the success path in unit tests and prevents
|
Gemini-compatible proxies or alternate endpoints without going through an
|
||||||
use with Gemini-compatible proxies. A future refactor would accept a base URL
|
internal helper or refactoring the API to accept a base URL or `http.Client`.
|
||||||
parameter or an `http.Client`.
|
|
||||||
|
|
||||||
**saveState is a full table replace.** On every `Persist()` call, the `requests`,
|
**saveState is a full table replace.** On every `Persist()` call, the `requests`,
|
||||||
`tokens`, and `daily` tables are truncated and rewritten. For a limiter tracking
|
`tokens`, and `daily` tables are truncated and rewritten. For a limiter tracking
|
||||||
|
|
@ -186,10 +176,10 @@ SQLite on `Persist()`. The database does not grow unboundedly between persist
|
||||||
cycles because `saveState` replaces all rows, but if `Persist()` is called
|
cycles because `saveState` replaces all rows, but if `Persist()` is called
|
||||||
frequently the WAL file can grow transiently.
|
frequently the WAL file can grow transiently.
|
||||||
|
|
||||||
**WaitForCapacity polling interval is fixed at 1 second.** This is appropriate
|
**WaitForCapacity now sleeps using `Decide`’s `RetryAfter` hint** (with a
|
||||||
for RPM-scale limits but is coarse for sub-second limits. If a caller needs
|
one-second fallback when no hint exists). This reduces busy looping on long
|
||||||
finer-grained waiting (e.g., smoothing requests within a minute), they must
|
windows but remains coarse for sub-second smoothing; callers that need
|
||||||
implement their own loop.
|
sub-second pacing should implement their own loop.
|
||||||
|
|
||||||
**No automatic persistence.** `Persist()` must be called explicitly. If a
|
**No automatic persistence.** `Persist()` must be called explicitly. If a
|
||||||
process exits without calling `Persist()`, any usage recorded since the last
|
process exits without calling `Persist()`, any usage recorded since the last
|
||||||
|
|
|
||||||
122
docs/index.md
Normal file
122
docs/index.md
Normal file
|
|
@ -0,0 +1,122 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
|
---
|
||||||
|
title: go-ratelimit
|
||||||
|
description: Provider-agnostic sliding window rate limiter for LLM API calls, with YAML and SQLite persistence backends.
|
||||||
|
---
|
||||||
|
|
||||||
|
# go-ratelimit
|
||||||
|
|
||||||
|
**Module**: `dappco.re/go/core/go-ratelimit`
|
||||||
|
**Licence**: EUPL-1.2
|
||||||
|
**Go version**: 1.26+
|
||||||
|
|
||||||
|
go-ratelimit enforces requests-per-minute (RPM), tokens-per-minute (TPM), and
|
||||||
|
requests-per-day (RPD) quotas on a per-model basis using an in-memory sliding
|
||||||
|
window. It ships with default quota profiles for Gemini, OpenAI, Anthropic, and
|
||||||
|
a local inference provider. State persists across process restarts via YAML
|
||||||
|
(single-process) or SQLite with WAL mode (multi-process). A YAML-to-SQLite
|
||||||
|
migration helper is included.
|
||||||
|
|
||||||
|
## Quick Start
|
||||||
|
|
||||||
|
```go
|
||||||
|
import "dappco.re/go/core/go-ratelimit"
|
||||||
|
|
||||||
|
// Create a limiter with Gemini defaults (YAML backend).
|
||||||
|
rl, err := ratelimit.New()
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check capacity before sending.
|
||||||
|
if rl.CanSend("gemini-2.0-flash", 1500) {
|
||||||
|
// Make the API call...
|
||||||
|
rl.RecordUsage("gemini-2.0-flash", 1000, 500) // promptTokens, outputTokens
|
||||||
|
}
|
||||||
|
|
||||||
|
// Persist state to disk for recovery across restarts.
|
||||||
|
if err := rl.Persist(); err != nil {
|
||||||
|
log.Printf("persist failed: %v", err)
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### Multi-provider configuration
|
||||||
|
|
||||||
|
```go
|
||||||
|
rl, err := ratelimit.NewWithConfig(ratelimit.Config{
|
||||||
|
Providers: []ratelimit.Provider{
|
||||||
|
ratelimit.ProviderGemini,
|
||||||
|
ratelimit.ProviderAnthropic,
|
||||||
|
},
|
||||||
|
Quotas: map[string]ratelimit.ModelQuota{
|
||||||
|
// Override a specific model's limits.
|
||||||
|
"gemini-3-pro-preview": {MaxRPM: 50, MaxTPM: 500000, MaxRPD: 200},
|
||||||
|
// Add a custom model not in any profile.
|
||||||
|
"llama-3.3-70b": {MaxRPM: 5, MaxTPM: 50000, MaxRPD: 0},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
```
|
||||||
|
|
||||||
|
### SQLite backend (multi-process safe)
|
||||||
|
|
||||||
|
```go
|
||||||
|
rl, err := ratelimit.NewWithSQLite("~/.core/ratelimits.db")
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
defer rl.Close()
|
||||||
|
|
||||||
|
// Load persisted state.
|
||||||
|
if err := rl.Load(); err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Use exactly as the YAML backend -- CanSend, RecordUsage, Persist, etc.
|
||||||
|
```
|
||||||
|
|
||||||
|
### Blocking until capacity is available
|
||||||
|
|
||||||
|
```go
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
if err := rl.WaitForCapacity(ctx, "claude-opus-4", 2000); err != nil {
|
||||||
|
log.Printf("timed out waiting for capacity: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// Capacity is available; proceed with the API call.
|
||||||
|
|
||||||
|
// WaitForCapacity uses Decide's RetryAfter hint to avoid tight polling.
|
||||||
|
```
|
||||||
|
|
||||||
|
## Package Layout
|
||||||
|
|
||||||
|
The module is a single package with no sub-packages.
|
||||||
|
|
||||||
|
| File | Purpose |
|
||||||
|
|------|---------|
|
||||||
|
| `ratelimit.go` | Core types (`RateLimiter`, `ModelQuota`, `Config`, `Provider`), sliding window logic, provider profiles, YAML persistence, `CountTokens` helper |
|
||||||
|
| `sqlite.go` | SQLite persistence backend (`sqliteStore`), schema creation, load/save operations |
|
||||||
|
| `ratelimit_test.go` | Tests for core logic, provider profiles, concurrency, and benchmarks |
|
||||||
|
| `sqlite_test.go` | Tests for SQLite backend, migration, and error recovery |
|
||||||
|
| `error_test.go` | Tests for SQLite and YAML error paths |
|
||||||
|
| `iter_test.go` | Tests for `Models()` and `Iter()` iterators, plus `CountTokens` edge cases |
|
||||||
|
|
||||||
|
## Dependencies
|
||||||
|
|
||||||
|
| Dependency | Purpose | Category |
|
||||||
|
|------------|---------|----------|
|
||||||
|
| `dappco.re/go/core` | File I/O helpers, structured errors, JSON helpers, path/environment utilities | Direct |
|
||||||
|
| `gopkg.in/yaml.v3` | YAML serialisation for the legacy persistence backend | Direct |
|
||||||
|
| `modernc.org/sqlite` | Pure Go SQLite driver (no CGO required) | Direct |
|
||||||
|
| `github.com/stretchr/testify` | Test assertions (`assert`, `require`) | Test only |
|
||||||
|
|
||||||
|
All indirect dependencies are pulled in by `modernc.org/sqlite`. No C toolchain
|
||||||
|
or system SQLite library is needed.
|
||||||
|
|
||||||
|
## Further Reading
|
||||||
|
|
||||||
|
- [Architecture](architecture.md) -- sliding window algorithm, provider quotas, YAML and SQLite backends, concurrency model
|
||||||
|
- [Development](development.md) -- build commands, test patterns, coding standards, commit conventions
|
||||||
|
- [History](history.md) -- completed phases with commit hashes, known limitations
|
||||||
46
docs/security-attack-vector-mapping.md
Normal file
46
docs/security-attack-vector-mapping.md
Normal file
|
|
@ -0,0 +1,46 @@
|
||||||
|
<!-- SPDX-License-Identifier: EUPL-1.2 -->
|
||||||
|
|
||||||
|
# Security Attack Vector Mapping
|
||||||
|
|
||||||
|
Scope: external inputs that cross into this package from callers, persisted storage, or the network. This is a mapping only; it does not propose or apply fixes.
|
||||||
|
|
||||||
|
Note: `CODEX.md` was not present anywhere under `/workspace` during this scan, so conventions were taken from `CLAUDE.md` and the existing repository layout.
|
||||||
|
|
||||||
|
## Caller-Controlled API Inputs
|
||||||
|
|
||||||
|
| Function | File:Line | Input Source | What It Flows Into | Current Validation | Potential Attack Vector |
|
||||||
|
| --- | --- | --- | --- | --- | --- |
|
||||||
|
| `NewWithConfig(cfg Config)` | `ratelimit.go:145` | Caller-controlled `cfg.FilePath`, `cfg.Providers`, `cfg.Quotas` | `cfg.FilePath` is stored in `rl.filePath` and later used by `Load()` / `Persist()`; `cfg.Providers` selects built-in profiles; `cfg.Quotas` is copied straight into `rl.Quotas` | Empty `FilePath` defaults to `~/.core/ratelimits.yaml`; unknown providers are ignored; `cfg.Backend` is not read here; quota values and model keys are copied verbatim | Untrusted `FilePath` can later steer YAML reads and writes to arbitrary local paths because the package uses unsandboxed `coreio.Local`; negative quota values act like "unlimited" because enforcement only checks `> 0`; very large maps or model strings can drive memory and disk growth |
|
||||||
|
| `(*RateLimiter).SetQuota(model string, quota ModelQuota)` | `ratelimit.go:183` | Caller-controlled `model` and `quota` | Direct write to `rl.Quotas[model]`; later consumed by `CanSend()`, `Stats()`, `Persist()`, and SQLite/YAML persistence | None beyond Go type checking | Negative quota fields disable enforcement for that dimension; extreme values can skew blocking behaviour; arbitrary model names expand the in-memory and persisted keyspace |
|
||||||
|
| `(*RateLimiter).AddProvider(provider Provider)` | `ratelimit.go:191` | Caller-controlled `provider` selector | Looks up `DefaultProfiles()[provider]` and overwrites matching entries in `rl.Quotas` via `maps.Copy` | Unknown providers are silently ignored; no authorisation or policy guard in this package | If a higher-level service exposes this call, an attacker can overwrite stricter runtime quotas for a provider's models with the shipped defaults and relax the intended rate-limit policy |
|
||||||
|
| `(*RateLimiter).BackgroundPrune(interval time.Duration)` | `ratelimit.go:328` | Caller-controlled `interval` | Passed to `time.NewTicker(interval)` and drives a background goroutine that repeatedly locks and prunes state | None | `interval <= 0` causes a panic; very small intervals can create CPU and lock-contention DoS; repeated calls without using the returned cancel function leak goroutines |
|
||||||
|
| `(*RateLimiter).CanSend(model string, estimatedTokens int)` | `ratelimit.go:350` | Caller-controlled `model` and `estimatedTokens` | `model` indexes `rl.Quotas` / `rl.State`; `estimatedTokens` is added to the current token total before the TPM comparison | Unknown models are allowed immediately; no non-negative or range checks on `estimatedTokens` | Passing an unconfigured model name bypasses throttling entirely; negative or overflowed token values can undercount the TPM check and permit oversend |
|
||||||
|
| `(*RateLimiter).RecordUsage(model string, promptTokens, outputTokens int)` | `ratelimit.go:396` | Caller-controlled `model`, `promptTokens`, `outputTokens` | Creates or updates `rl.State[model]`; stores `promptTokens + outputTokens` in the token window and increments `DayCount` | None | Arbitrary model names create unbounded state that will later persist to YAML/SQLite; negative or overflowed token totals poison accounting and can reduce future TPM totals below the real usage |
|
||||||
|
| `(*RateLimiter).WaitForCapacity(ctx context.Context, model string, tokens int)` | `ratelimit.go:429` | Caller-controlled `ctx`, `model`, `tokens` | Calls `Decide(model, tokens)` in a loop and sleeps for the returned `RetryAfter` (or 1s fallback) until allowed or `ctx.Done()` fires | No direct validation beyond negative-token guard; relies on downstream `Decide()` and caller-supplied context cancellation | Long `RetryAfter` values can delay rechecks; repeated calls with long-lived contexts can still accumulate goroutines and lock pressure |
|
||||||
|
| `(*RateLimiter).Reset(model string)` | `ratelimit.go:433` | Caller-controlled `model` | `model == ""` replaces the entire `rl.State` map; otherwise `delete(rl.State, model)` | Empty string is treated as a wildcard reset | If reachable by an untrusted actor, an empty string clears all rate-limit history and targeted resets erase throttling state for chosen models |
|
||||||
|
| `(*RateLimiter).Stats(model string)` | `ratelimit.go:484` | Caller-controlled `model` | Prunes `rl.State[model]`, reads `rl.Quotas[model]`, and returns a usage snapshot | None | If exposed through a service boundary, it discloses per-model quota ceilings and live usage counts that can help an attacker tune evasion or timing |
|
||||||
|
| `NewWithSQLite(dbPath string)` | `ratelimit.go:567` | Caller-controlled `dbPath` | Thin wrapper that forwards `dbPath` into `NewWithSQLiteConfig()` and then `newSQLiteStore()` | No additional validation in the wrapper | Untrusted `dbPath` can steer database creation/opening to unintended local filesystem locations, including companion `-wal` and `-shm` files |
|
||||||
|
| `NewWithSQLiteConfig(dbPath string, cfg Config)` | `ratelimit.go:576` | Caller-controlled `dbPath`, `cfg.Providers`, `cfg.Quotas` | `dbPath` goes straight to `newSQLiteStore()`; provider and quota inputs are copied into `rl.Quotas` exactly as in `NewWithConfig()` | `cfg.Backend` is ignored; unknown providers are ignored; no path, range, or size checks | Combines arbitrary database-path selection with the same quota poisoning risks as `NewWithConfig()` and `SetQuota()` |
|
||||||
|
|
||||||
|
## Filesystem And YAML Inputs
|
||||||
|
|
||||||
|
| Function | File:Line | Input Source | What It Flows Into | Current Validation | Potential Attack Vector |
|
||||||
|
| --- | --- | --- | --- | --- | --- |
|
||||||
|
| `(*RateLimiter).Load()` (YAML backend) | `ratelimit.go:210` | File bytes read from `rl.filePath` | `coreio.Local.Read(rl.filePath)` feeds `yaml.Unmarshal(..., rl)`, which overwrites `rl.Quotas` and `rl.State` | Missing file is ignored; YAML syntax and type mismatches return an error; no semantic checks on counts, limits, model names, or slice sizes | A malicious YAML file can inject negative quotas or counters to bypass enforcement, preload very large maps/slices for memory DoS, or replace the in-memory policy/state with attacker-chosen values |
|
||||||
|
| `MigrateYAMLToSQLite(yamlPath, sqlitePath string)` | `ratelimit.go:617` | Caller-controlled `yamlPath` and `sqlitePath`, plus YAML file bytes from `yamlPath` | Reads YAML with `coreio.Local.Read(yamlPath)`, unmarshals into a temporary `RateLimiter`, then opens `sqlitePath` and writes the imported quotas/state into SQLite | Read/open errors and YAML syntax/type errors are surfaced; no path restrictions and no semantic validation of imported quotas/state | Untrusted paths enable arbitrary local file reads and database creation/clobbering; attacker-controlled YAML can permanently seed the SQLite backend with quota-bypass values or oversized state |
|
||||||
|
|
||||||
|
## SQLite Inputs
|
||||||
|
|
||||||
|
| Function | File:Line | Input Source | What It Flows Into | Current Validation | Potential Attack Vector |
|
||||||
|
| --- | --- | --- | --- | --- | --- |
|
||||||
|
| `newSQLiteStore(dbPath string)` | `sqlite.go:20` | Caller-controlled database path passed from `NewWithSQLite*()` or `MigrateYAMLToSQLite()` | `sql.Open("sqlite", dbPath)` opens or creates the database, then applies PRAGMAs and creates tables and indexes | Only driver/open/PRAGMA/schema errors are checked; there is no allowlist, path normalisation, or sandboxing in this package | An attacker-chosen `dbPath` can redirect state into unintended files or directories and create matching SQLite sidecar files, which is useful for tampering, data placement, or storage-exhaustion attacks |
|
||||||
|
| `(*RateLimiter).Load()` (SQLite backend via `loadSQLite()`) | `ratelimit.go:223` | SQLite content reachable through the already-open `rl.sqlite` handle | `loadQuotas()` and `loadState()` results are copied into `rl.Quotas` and `rl.State`; loaded quotas override in-memory defaults | Only lower-level scan/query errors are checked; no semantic validation after load | A tampered SQLite database can override intended quotas/state, including negative limits and poisoned counters, before any enforcement call runs |
|
||||||
|
| `(*sqliteStore).loadQuotas()` | `sqlite.go:110` | Rows from the `quotas` table (`model`, `max_rpm`, `max_tpm`, `max_rpd`) | Scanned into `map[string]ModelQuota`, then copied into `rl.Quotas` by `loadSQLite()` | SQL scan and row iteration errors only; no range or length checks | Negative or extreme values disable or destabilise later quota enforcement; a large number of rows or large model strings can cause memory growth |
|
||||||
|
| `(*sqliteStore).loadState()` | `sqlite.go:194` | Rows from the `daily`, `requests`, and `tokens` tables | Scanned into `UsageStats` maps/slices and then copied into `rl.State` by `loadSQLite()` | SQL scan and row iteration errors only; no bounds, ordering, or semantic checks on timestamps and counts | Crafted counts or timestamps can poison later `CanSend()` / `Stats()` results; oversized tables can drive memory and CPU exhaustion during load |
|
||||||
|
|
||||||
|
## Network Inputs
|
||||||
|
|
||||||
|
| Function | File:Line | Input Source | What It Flows Into | Current Validation | Potential Attack Vector |
|
||||||
|
| --- | --- | --- | --- | --- | --- |
|
||||||
|
| `CountTokens(ctx, apiKey, model, text)` (request construction) | `ratelimit.go:650` | Caller-controlled `ctx`, `apiKey`, `model`, `text` | `model` is interpolated directly into the Google API URL path; `text` is marshalled into JSON request body; `apiKey` goes into `x-goog-api-key`; `ctx` governs request lifetime | JSON marshalling must succeed; `http.NewRequestWithContext()` rejects some malformed URLs; there is no path escaping, length check, or output-size cap on the request body | Untrusted prompt text is exfiltrated to a remote API; very large text can consume memory/bandwidth; unescaped model strings can alter the path or query on the fixed Google host; repeated calls burn external quota |
|
||||||
|
| `CountTokens(ctx, apiKey, model, text)` (response handling) | `ratelimit.go:681` | Remote HTTP status, response body, and JSON `totalTokens` value from `generativelanguage.googleapis.com` | Non-200 bodies are read fully with `io.ReadAll()` and embedded into the returned error; 200 responses are decoded into `result.TotalTokens` and returned to the caller | Checks for HTTP 200 and JSON decode errors only; no response-body size limit and no sanity check on `TotalTokens` | A very large error body can cause memory pressure and log/telemetry pollution; a negative or extreme `totalTokens` value would be returned unchanged and could poison downstream rate-limit accounting if the caller trusts it |
|
||||||
693
error_test.go
Normal file
693
error_test.go
Normal file
|
|
@ -0,0 +1,693 @@
|
||||||
|
// SPDX-License-Identifier: EUPL-1.2
|
||||||
|
|
||||||
|
package ratelimit
|
||||||
|
|
||||||
|
import (
|
||||||
|
"syscall"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestError_SQLiteErrorPaths_Bad(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "error.db")
|
||||||
|
rl, err := NewWithSQLite(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
// Close the underlying DB to trigger errors.
|
||||||
|
rl.sqlite.close()
|
||||||
|
|
||||||
|
t.Run("loadQuotas error", func(t *testing.T) {
|
||||||
|
_, err := rl.sqlite.loadQuotas()
|
||||||
|
assert.Error(t, err)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("saveQuotas error", func(t *testing.T) {
|
||||||
|
err := rl.sqlite.saveQuotas(map[string]ModelQuota{"test": {}})
|
||||||
|
assert.Error(t, err)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("saveState error", func(t *testing.T) {
|
||||||
|
err := rl.sqlite.saveState(map[string]*UsageStats{"test": {}})
|
||||||
|
assert.Error(t, err)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("loadState error", func(t *testing.T) {
|
||||||
|
_, err := rl.sqlite.loadState()
|
||||||
|
assert.Error(t, err)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteInitErrors_Bad(t *testing.T) {
|
||||||
|
t.Run("WAL pragma failure", func(t *testing.T) {
|
||||||
|
// This is hard to trigger without mocking sql.DB, but we can try an invalid connection string
|
||||||
|
// modernc.org/sqlite doesn't support all DSN options that might cause PRAGMA to fail but connection to succeed.
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_PersistYAML_Good(t *testing.T) {
|
||||||
|
t.Run("successful YAML persist and load", func(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
path := testPath(tmpDir, "ratelimits.yaml")
|
||||||
|
rl, _ := New()
|
||||||
|
rl.filePath = path
|
||||||
|
rl.Quotas["test"] = ModelQuota{MaxRPM: 1}
|
||||||
|
rl.RecordUsage("test", 1, 1)
|
||||||
|
|
||||||
|
require.NoError(t, rl.Persist())
|
||||||
|
|
||||||
|
rl2, _ := New()
|
||||||
|
rl2.filePath = path
|
||||||
|
require.NoError(t, rl2.Load())
|
||||||
|
assert.Equal(t, 1, rl2.Quotas["test"].MaxRPM)
|
||||||
|
assert.Equal(t, 1, rl2.State["test"].DayCount)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteLoadViaLimiter_Bad(t *testing.T) {
|
||||||
|
t.Run("Load returns error when SQLite DB is closed", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "load-err.db")
|
||||||
|
rl, err := NewWithSQLite(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
// Close the underlying DB to trigger errors on Load.
|
||||||
|
rl.sqlite.close()
|
||||||
|
|
||||||
|
err = rl.Load()
|
||||||
|
assert.Error(t, err, "Load should fail with closed DB")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Load returns error when loadState fails", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "load-state-err.db")
|
||||||
|
rl, err := NewWithSQLite(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
// Save some quotas so loadQuotas succeeds, then corrupt the state tables.
|
||||||
|
require.NoError(t, rl.Persist())
|
||||||
|
|
||||||
|
// Drop the daily table to cause loadState to fail.
|
||||||
|
_, execErr := rl.sqlite.db.Exec("DROP TABLE daily")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
err = rl.Load()
|
||||||
|
assert.Error(t, err, "Load should fail when loadState fails")
|
||||||
|
rl.Close()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLitePersistViaLimiter_Bad(t *testing.T) {
|
||||||
|
t.Run("Persist returns error when SQLite saveQuotas fails", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "persist-err.db")
|
||||||
|
rl, err := NewWithSQLite(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
// Drop the quotas table to cause saveQuotas to fail.
|
||||||
|
_, execErr := rl.sqlite.db.Exec("DROP TABLE quotas")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
err = rl.Persist()
|
||||||
|
assert.Error(t, err, "Persist should fail when saveQuotas fails")
|
||||||
|
assert.Contains(t, err.Error(), "ratelimit.Persist")
|
||||||
|
rl.Close()
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Persist returns error when SQLite saveState fails", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "persist-state-err.db")
|
||||||
|
rl, err := NewWithSQLite(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
rl.RecordUsage("test-model", 10, 10)
|
||||||
|
|
||||||
|
// Drop a state table to cause saveState to fail.
|
||||||
|
_, execErr := rl.sqlite.db.Exec("DROP TABLE requests")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
err = rl.Persist()
|
||||||
|
assert.Error(t, err, "Persist should fail when saveState fails")
|
||||||
|
assert.Contains(t, err.Error(), "ratelimit.Persist")
|
||||||
|
rl.Close()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_NewWithSQLite_Bad(t *testing.T) {
|
||||||
|
t.Run("NewWithSQLite with invalid path", func(t *testing.T) {
|
||||||
|
_, err := NewWithSQLite("/nonexistent/deep/nested/dir/test.db")
|
||||||
|
assert.Error(t, err, "should fail with invalid path")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("NewWithSQLiteConfig with invalid path", func(t *testing.T) {
|
||||||
|
_, err := NewWithSQLiteConfig("/nonexistent/deep/nested/dir/test.db", Config{
|
||||||
|
Providers: []Provider{ProviderGemini},
|
||||||
|
})
|
||||||
|
assert.Error(t, err, "should fail with invalid path")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteSaveState_Bad(t *testing.T) {
|
||||||
|
t.Run("saveState fails when tokens table is dropped", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "tokens-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
_, execErr := store.db.Exec("DROP TABLE tokens")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
Tokens: []TokenEntry{{Time: time.Now(), Count: 100}},
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
err = store.saveState(state)
|
||||||
|
assert.Error(t, err, "saveState should fail when tokens table is missing")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("saveState fails when daily table is dropped", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "daily-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
_, execErr := store.db.Exec("DROP TABLE daily")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
err = store.saveState(state)
|
||||||
|
assert.Error(t, err, "saveState should fail when daily table is missing")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("saveState fails on request insert with renamed column", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "req-insert-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Rename the ts column so INSERT INTO requests (model, ts) fails.
|
||||||
|
_, execErr := store.db.Exec("ALTER TABLE requests RENAME COLUMN ts TO timestamp")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
Requests: []time.Time{time.Now()},
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
err = store.saveState(state)
|
||||||
|
assert.Error(t, err, "saveState should fail on request insert with renamed column")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("saveState fails on token insert with renamed column", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "tok-insert-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Rename the count column so INSERT INTO tokens (model, ts, count) fails.
|
||||||
|
_, execErr := store.db.Exec("ALTER TABLE tokens RENAME COLUMN count TO amount")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
Tokens: []TokenEntry{{Time: time.Now(), Count: 100}},
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
err = store.saveState(state)
|
||||||
|
assert.Error(t, err, "saveState should fail on token insert with renamed column")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("saveState fails on daily insert with renamed column", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "day-insert-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Rename day_count column so INSERT INTO daily fails.
|
||||||
|
_, execErr := store.db.Exec("ALTER TABLE daily RENAME COLUMN day_count TO total")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
err = store.saveState(state)
|
||||||
|
assert.Error(t, err, "saveState should fail on daily insert with renamed column")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteLoadState_Bad(t *testing.T) {
|
||||||
|
t.Run("loadState fails when requests table is dropped", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "req-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
Requests: []time.Time{time.Now()},
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
require.NoError(t, store.saveState(state))
|
||||||
|
|
||||||
|
_, execErr := store.db.Exec("DROP TABLE requests")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
_, err = store.loadState()
|
||||||
|
assert.Error(t, err, "loadState should fail when requests table is missing")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("loadState fails when tokens table is dropped", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "tok-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
Tokens: []TokenEntry{{Time: time.Now(), Count: 100}},
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
require.NoError(t, store.saveState(state))
|
||||||
|
|
||||||
|
_, execErr := store.db.Exec("DROP TABLE tokens")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
_, err = store.loadState()
|
||||||
|
assert.Error(t, err, "loadState should fail when tokens table is missing")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("loadState fails when daily table is dropped", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "daily-load-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
require.NoError(t, store.saveState(state))
|
||||||
|
|
||||||
|
_, execErr := store.db.Exec("DROP TABLE daily")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
_, err = store.loadState()
|
||||||
|
assert.Error(t, err, "loadState should fail when daily table is missing")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteSaveQuotasExec_Bad(t *testing.T) {
|
||||||
|
t.Run("saveQuotas fails with renamed column at prepare", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "quota-exec-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Rename column so INSERT fails at prepare.
|
||||||
|
_, execErr := store.db.Exec("ALTER TABLE quotas RENAME COLUMN max_rpm TO rpm")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
err = store.saveQuotas(map[string]ModelQuota{
|
||||||
|
"test": {MaxRPM: 10, MaxTPM: 100, MaxRPD: 50},
|
||||||
|
})
|
||||||
|
assert.Error(t, err, "saveQuotas should fail with renamed column")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("saveQuotas fails at exec via trigger", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "quota-trigger.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Add trigger to abort INSERT.
|
||||||
|
_, execErr := store.db.Exec(`CREATE TRIGGER fail_quota BEFORE INSERT ON quotas
|
||||||
|
BEGIN SELECT RAISE(ABORT, 'forced quota insert failure'); END`)
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
err = store.saveQuotas(map[string]ModelQuota{
|
||||||
|
"test": {MaxRPM: 10, MaxTPM: 100, MaxRPD: 50},
|
||||||
|
})
|
||||||
|
assert.Error(t, err, "saveQuotas should fail when trigger fires")
|
||||||
|
assert.Contains(t, err.Error(), "exec test")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteSaveStateExec_Bad(t *testing.T) {
|
||||||
|
t.Run("request insert exec fails via trigger", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "trigger-req.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Add a trigger that aborts INSERT on requests.
|
||||||
|
_, execErr := store.db.Exec(`CREATE TRIGGER fail_req_insert BEFORE INSERT ON requests
|
||||||
|
BEGIN SELECT RAISE(ABORT, 'forced request insert failure'); END`)
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
Requests: []time.Time{time.Now()},
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
err = store.saveState(state)
|
||||||
|
assert.Error(t, err, "saveState should fail when request insert trigger fires")
|
||||||
|
assert.Contains(t, err.Error(), "insert request")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("token insert exec fails via trigger", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "trigger-tok.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Add a trigger that aborts INSERT on tokens.
|
||||||
|
_, execErr := store.db.Exec(`CREATE TRIGGER fail_tok_insert BEFORE INSERT ON tokens
|
||||||
|
BEGIN SELECT RAISE(ABORT, 'forced token insert failure'); END`)
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
Tokens: []TokenEntry{{Time: time.Now(), Count: 100}},
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
err = store.saveState(state)
|
||||||
|
assert.Error(t, err, "saveState should fail when token insert trigger fires")
|
||||||
|
assert.Contains(t, err.Error(), "insert token")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("daily insert exec fails via trigger", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "trigger-day.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Add a trigger that aborts INSERT on daily.
|
||||||
|
_, execErr := store.db.Exec(`CREATE TRIGGER fail_day_insert BEFORE INSERT ON daily
|
||||||
|
BEGIN SELECT RAISE(ABORT, 'forced daily insert failure'); END`)
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
state := map[string]*UsageStats{
|
||||||
|
"model": {
|
||||||
|
DayStart: time.Now(),
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
err = store.saveState(state)
|
||||||
|
assert.Error(t, err, "saveState should fail when daily insert trigger fires")
|
||||||
|
assert.Contains(t, err.Error(), "insert daily")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteLoadQuotasScan_Bad(t *testing.T) {
|
||||||
|
t.Run("loadQuotas fails with renamed column", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "quota-scan-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Save valid quotas first.
|
||||||
|
require.NoError(t, store.saveQuotas(map[string]ModelQuota{
|
||||||
|
"test": {MaxRPM: 10},
|
||||||
|
}))
|
||||||
|
|
||||||
|
// Rename column so SELECT ... FROM quotas fails.
|
||||||
|
_, execErr := store.db.Exec("ALTER TABLE quotas RENAME COLUMN max_rpm TO rpm")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
_, err = store.loadQuotas()
|
||||||
|
assert.Error(t, err, "loadQuotas should fail with renamed column")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_NewSQLiteStoreInReadOnlyDir_Bad(t *testing.T) {
|
||||||
|
if isRootUser() {
|
||||||
|
t.Skip("chmod restrictions do not apply to root")
|
||||||
|
}
|
||||||
|
|
||||||
|
t.Run("fails when parent directory is read-only", func(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
readonlyDir := testPath(tmpDir, "readonly")
|
||||||
|
ensureTestDir(t, readonlyDir)
|
||||||
|
setPathMode(t, readonlyDir, 0o555)
|
||||||
|
defer func() {
|
||||||
|
_ = syscall.Chmod(readonlyDir, 0o755)
|
||||||
|
}()
|
||||||
|
|
||||||
|
dbPath := testPath(readonlyDir, "test.db")
|
||||||
|
_, err := newSQLiteStore(dbPath)
|
||||||
|
assert.Error(t, err, "should fail when directory is read-only")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteCreateSchema_Bad(t *testing.T) {
|
||||||
|
t.Run("createSchema fails on closed DB", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "schema-err.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
db := store.db
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
err = createSchema(db)
|
||||||
|
assert.Error(t, err, "createSchema should fail on closed DB")
|
||||||
|
assert.Contains(t, err.Error(), "ratelimit.createSchema")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteLoadStateScan_Bad(t *testing.T) {
|
||||||
|
t.Run("scan daily fails with NULL values", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "scan-daily.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Recreate daily without NOT NULL constraints so we can insert NULLs.
|
||||||
|
_, _ = store.db.Exec("DROP TABLE daily")
|
||||||
|
_, execErr := store.db.Exec("CREATE TABLE daily (model TEXT PRIMARY KEY, day_start INTEGER, day_count INTEGER)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
// Insert a NULL day_start; scanning NULL into int64 returns an error.
|
||||||
|
_, execErr = store.db.Exec("INSERT INTO daily VALUES ('test', NULL, NULL)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
_, err = store.loadState()
|
||||||
|
if err != nil {
|
||||||
|
assert.Contains(t, err.Error(), "ratelimit.loadState")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("scan requests fails with NULL ts", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "scan-req.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Insert a valid daily entry so loadState gets past the daily phase.
|
||||||
|
_, execErr := store.db.Exec("INSERT INTO daily VALUES ('test', 0, 1)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
// Recreate requests without NOT NULL, insert NULL ts.
|
||||||
|
_, _ = store.db.Exec("DROP TABLE requests")
|
||||||
|
_, execErr = store.db.Exec("CREATE TABLE requests (model TEXT, ts INTEGER)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
_, execErr = store.db.Exec("CREATE INDEX IF NOT EXISTS idx_requests_model_ts ON requests(model, ts)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
_, execErr = store.db.Exec("INSERT INTO requests VALUES ('test', NULL)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
_, err = store.loadState()
|
||||||
|
if err != nil {
|
||||||
|
assert.Contains(t, err.Error(), "ratelimit.loadState")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("scan tokens fails with NULL values", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "scan-tok.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Insert valid daily entry.
|
||||||
|
_, execErr := store.db.Exec("INSERT INTO daily VALUES ('test', 0, 1)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
// Recreate tokens without NOT NULL, insert NULL values.
|
||||||
|
_, _ = store.db.Exec("DROP TABLE tokens")
|
||||||
|
_, execErr = store.db.Exec("CREATE TABLE tokens (model TEXT, ts INTEGER, count INTEGER)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
_, execErr = store.db.Exec("CREATE INDEX IF NOT EXISTS idx_tokens_model_ts ON tokens(model, ts)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
_, execErr = store.db.Exec("INSERT INTO tokens VALUES ('test', NULL, NULL)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
_, err = store.loadState()
|
||||||
|
if err != nil {
|
||||||
|
assert.Contains(t, err.Error(), "ratelimit.loadState")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_SQLiteLoadQuotasScanWithBadSchema_Bad(t *testing.T) {
|
||||||
|
t.Run("scan fails with NULL quota values", func(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "scan-quota.db")
|
||||||
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
// Recreate quotas without NOT NULL constraints.
|
||||||
|
_, _ = store.db.Exec("DROP TABLE quotas")
|
||||||
|
_, execErr := store.db.Exec("CREATE TABLE quotas (model TEXT PRIMARY KEY, max_rpm INTEGER, max_tpm INTEGER, max_rpd INTEGER)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
_, execErr = store.db.Exec("INSERT INTO quotas VALUES ('test', NULL, NULL, NULL)")
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
|
||||||
|
_, err = store.loadQuotas()
|
||||||
|
if err != nil {
|
||||||
|
assert.Contains(t, err.Error(), "ratelimit.loadQuotas")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_MigrateYAMLToSQLiteWithSaveErrors_Bad(t *testing.T) {
|
||||||
|
t.Run("saveQuotas failure during migration via trigger", func(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
yamlPath := testPath(tmpDir, "with-quotas.yaml")
|
||||||
|
sqlitePath := testPath(tmpDir, "migrate-quota-err.db")
|
||||||
|
|
||||||
|
// Write a YAML file with quotas.
|
||||||
|
yamlData := `quotas:
|
||||||
|
test-model:
|
||||||
|
max_rpm: 10
|
||||||
|
max_tpm: 100
|
||||||
|
max_rpd: 50
|
||||||
|
`
|
||||||
|
writeTestFile(t, yamlPath, yamlData)
|
||||||
|
|
||||||
|
// Pre-create DB with a trigger that aborts quota inserts.
|
||||||
|
store, err := newSQLiteStore(sqlitePath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
_, execErr := store.db.Exec(`CREATE TRIGGER fail_quota_migrate BEFORE INSERT ON quotas
|
||||||
|
BEGIN SELECT RAISE(ABORT, 'forced quota failure'); END`)
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
// Migration re-opens the DB; tables already exist, trigger persists.
|
||||||
|
err = MigrateYAMLToSQLite(yamlPath, sqlitePath)
|
||||||
|
assert.Error(t, err, "migration should fail when saveQuotas fails")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("saveState failure during migration via trigger", func(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
yamlPath := testPath(tmpDir, "with-state.yaml")
|
||||||
|
sqlitePath := testPath(tmpDir, "migrate-state-err.db")
|
||||||
|
|
||||||
|
// Write YAML with state.
|
||||||
|
yamlData := `state:
|
||||||
|
test-model:
|
||||||
|
requests:
|
||||||
|
- time: 2026-01-01T00:00:00Z
|
||||||
|
day_start: 2026-01-01T00:00:00Z
|
||||||
|
day_count: 1
|
||||||
|
`
|
||||||
|
writeTestFile(t, yamlPath, yamlData)
|
||||||
|
|
||||||
|
// Pre-create DB with a trigger that aborts daily inserts.
|
||||||
|
store, err := newSQLiteStore(sqlitePath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
_, execErr := store.db.Exec(`CREATE TRIGGER fail_daily_migrate BEFORE INSERT ON daily
|
||||||
|
BEGIN SELECT RAISE(ABORT, 'forced daily failure'); END`)
|
||||||
|
require.NoError(t, execErr)
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
err = MigrateYAMLToSQLite(yamlPath, sqlitePath)
|
||||||
|
assert.Error(t, err, "migration should fail when saveState fails")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_MigrateYAMLToSQLiteNilQuotasAndState_Good(t *testing.T) {
|
||||||
|
t.Run("YAML with empty quotas and state migrates cleanly", func(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
yamlPath := testPath(tmpDir, "empty.yaml")
|
||||||
|
writeTestFile(t, yamlPath, "{}")
|
||||||
|
|
||||||
|
sqlitePath := testPath(tmpDir, "empty.db")
|
||||||
|
require.NoError(t, MigrateYAMLToSQLite(yamlPath, sqlitePath))
|
||||||
|
|
||||||
|
store, err := newSQLiteStore(sqlitePath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
quotas, err := store.loadQuotas()
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Empty(t, quotas)
|
||||||
|
|
||||||
|
state, err := store.loadState()
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Empty(t, state)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_NewWithConfigHomeUnavailable_Bad(t *testing.T) {
|
||||||
|
// Clear all supported home env vars so defaultStatePath cannot resolve a home directory.
|
||||||
|
t.Setenv("CORE_HOME", "")
|
||||||
|
t.Setenv("HOME", "")
|
||||||
|
t.Setenv("home", "")
|
||||||
|
t.Setenv("USERPROFILE", "")
|
||||||
|
|
||||||
|
_, err := NewWithConfig(Config{})
|
||||||
|
assert.Error(t, err, "should fail when HOME is unset")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_PersistMarshal_Good(t *testing.T) {
|
||||||
|
// yaml.Marshal on a struct with map[string]ModelQuota and map[string]*UsageStats
|
||||||
|
// should not fail in practice. We test the error path by using a type that
|
||||||
|
// yaml.Marshal cannot handle: a channel.
|
||||||
|
// Since we cannot inject a channel into the typed struct, this path is
|
||||||
|
// unreachable in production. Instead, exercise the Persist YAML path
|
||||||
|
// with valid data to confirm coverage of the non-error path.
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
rl.RecordUsage("test", 1, 1)
|
||||||
|
assert.NoError(t, rl.Persist(), "valid persist should succeed")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestError_MigrateErrorsExtended_Bad(t *testing.T) {
|
||||||
|
t.Run("unmarshal failure", func(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
path := testPath(tmpDir, "bad.yaml")
|
||||||
|
writeTestFile(t, path, "invalid: yaml: [")
|
||||||
|
err := MigrateYAMLToSQLite(path, testPath(tmpDir, "out.db"))
|
||||||
|
assert.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "ratelimit.MigrateYAMLToSQLite: unmarshal")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("sqlite open failure", func(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
yamlPath := testPath(tmpDir, "ok.yaml")
|
||||||
|
writeTestFile(t, yamlPath, "quotas: {}")
|
||||||
|
// Use an invalid sqlite path (dir where file should be)
|
||||||
|
err := MigrateYAMLToSQLite(yamlPath, "/dev/null/not-a-db")
|
||||||
|
assert.Error(t, err)
|
||||||
|
})
|
||||||
|
}
|
||||||
17
go.mod
17
go.mod
|
|
@ -1,24 +1,25 @@
|
||||||
module forge.lthn.ai/core/go-ratelimit
|
// SPDX-License-Identifier: EUPL-1.2
|
||||||
|
|
||||||
|
module dappco.re/go/core/go-ratelimit
|
||||||
|
|
||||||
go 1.26.0
|
go 1.26.0
|
||||||
|
|
||||||
require (
|
require (
|
||||||
|
dappco.re/go/core v0.8.0-alpha.1
|
||||||
gopkg.in/yaml.v3 v3.0.1
|
gopkg.in/yaml.v3 v3.0.1
|
||||||
modernc.org/sqlite v1.46.1
|
modernc.org/sqlite v1.47.0
|
||||||
)
|
)
|
||||||
|
|
||||||
require (
|
require (
|
||||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||||
github.com/google/uuid v1.6.0 // indirect
|
github.com/google/uuid v1.6.0 // indirect
|
||||||
github.com/kr/pretty v0.3.1 // indirect
|
github.com/kr/text v0.2.0 // indirect
|
||||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||||
github.com/rogpeppe/go-internal v1.14.1 // indirect
|
golang.org/x/sys v0.42.0 // indirect
|
||||||
golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa // indirect
|
golang.org/x/tools v0.43.0 // indirect
|
||||||
golang.org/x/sys v0.41.0 // indirect
|
modernc.org/libc v1.70.0 // indirect
|
||||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
|
|
||||||
modernc.org/libc v1.68.0 // indirect
|
|
||||||
modernc.org/mathutil v1.7.1 // indirect
|
modernc.org/mathutil v1.7.1 // indirect
|
||||||
modernc.org/memory v1.11.0 // indirect
|
modernc.org/memory v1.11.0 // indirect
|
||||||
)
|
)
|
||||||
|
|
|
||||||
41
go.sum
41
go.sum
|
|
@ -1,3 +1,5 @@
|
||||||
|
dappco.re/go/core v0.8.0-alpha.1 h1:gj7+Scv+L63Z7wMxbJYHhaRFkHJo2u4MMPuUSv/Dhtk=
|
||||||
|
dappco.re/go/core v0.8.0-alpha.1/go.mod h1:f2/tBZ3+3IqDrg2F5F598llv0nmb/4gJVCFzM5geE4A=
|
||||||
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
|
||||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
|
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
|
||||||
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||||
|
|
@ -9,38 +11,31 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||||
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||||
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
|
|
||||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||||
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
|
|
||||||
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
|
|
||||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||||
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
|
|
||||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
|
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
|
||||||
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||||
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
|
|
||||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||||
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
||||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||||
golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa h1:Zt3DZoOFFYkKhDT3v7Lm9FDMEV06GpzjG2jrqW+QTE0=
|
golang.org/x/mod v0.34.0 h1:xIHgNUUnW6sYkcM5Jleh05DvLOtwc6RitGHbDk4akRI=
|
||||||
golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa/go.mod h1:K79w1Vqn7PoiZn+TkNpx3BUWUQksGO3JcVX6qIjytmA=
|
golang.org/x/mod v0.34.0/go.mod h1:ykgH52iCZe79kzLLMhyCUzhMci+nQj+0XkbXpNYtVjY=
|
||||||
golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8=
|
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||||
golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w=
|
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||||
golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
|
|
||||||
golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
|
|
||||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||||
golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k=
|
golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo=
|
||||||
golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||||
golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
|
golang.org/x/tools v0.43.0 h1:12BdW9CeB3Z+J/I/wj34VMl8X+fEXBxVR90JeMX5E7s=
|
||||||
golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=
|
golang.org/x/tools v0.43.0/go.mod h1:uHkMso649BX2cZK6+RpuIPXS3ho2hZo4FVwfoy1vIk0=
|
||||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||||
|
|
@ -48,18 +43,18 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||||
modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis=
|
modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis=
|
||||||
modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0=
|
modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0=
|
||||||
modernc.org/ccgo/v4 v4.30.2 h1:4yPaaq9dXYXZ2V8s1UgrC3KIj580l2N4ClrLwnbv2so=
|
modernc.org/ccgo/v4 v4.32.0 h1:hjG66bI/kqIPX1b2yT6fr/jt+QedtP2fqojG2VrFuVw=
|
||||||
modernc.org/ccgo/v4 v4.30.2/go.mod h1:yZMnhWEdW0qw3EtCndG1+ldRrVGS+bIwyWmAWzS0XEw=
|
modernc.org/ccgo/v4 v4.32.0/go.mod h1:6F08EBCx5uQc38kMGl+0Nm0oWczoo1c7cgpzEry7Uc0=
|
||||||
modernc.org/fileutil v1.3.40 h1:ZGMswMNc9JOCrcrakF1HrvmergNLAmxOPjizirpfqBA=
|
modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
|
||||||
modernc.org/fileutil v1.3.40/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc=
|
modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
|
||||||
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
|
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
|
||||||
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
|
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
|
||||||
modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo=
|
modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo=
|
||||||
modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
|
modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
|
||||||
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
|
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
|
||||||
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
|
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
|
||||||
modernc.org/libc v1.68.0 h1:PJ5ikFOV5pwpW+VqCK1hKJuEWsonkIJhhIXyuF/91pQ=
|
modernc.org/libc v1.70.0 h1:U58NawXqXbgpZ/dcdS9kMshu08aiA6b7gusEusqzNkw=
|
||||||
modernc.org/libc v1.68.0/go.mod h1:NnKCYeoYgsEqnY3PgvNgAeaJnso968ygU8Z0DxjoEc0=
|
modernc.org/libc v1.70.0/go.mod h1:OVmxFGP1CI/Z4L3E0Q3Mf1PDE0BucwMkcXjjLntvHJo=
|
||||||
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||||
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||||
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
|
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
|
||||||
|
|
@ -68,8 +63,8 @@ modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8=
|
||||||
modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
|
modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
|
||||||
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
|
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
|
||||||
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
|
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
|
||||||
modernc.org/sqlite v1.46.1 h1:eFJ2ShBLIEnUWlLy12raN0Z1plqmFX9Qe3rjQTKt6sU=
|
modernc.org/sqlite v1.47.0 h1:R1XyaNpoW4Et9yly+I2EeX7pBza/w+pmYee/0HJDyKk=
|
||||||
modernc.org/sqlite v1.46.1/go.mod h1:CzbrU2lSB1DKUusvwGz7rqEKIq+NUd8GWuBBZDs9/nA=
|
modernc.org/sqlite v1.47.0/go.mod h1:hWjRO6Tj/5Ik8ieqxQybiEOUXy0NJFNp2tpvVpKlvig=
|
||||||
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
||||||
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||||
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||||
|
|
|
||||||
140
iter_test.go
Normal file
140
iter_test.go
Normal file
|
|
@ -0,0 +1,140 @@
|
||||||
|
// SPDX-License-Identifier: EUPL-1.2
|
||||||
|
|
||||||
|
package ratelimit
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestIter_Iterators_Good(t *testing.T) {
|
||||||
|
rl, err := NewWithConfig(Config{
|
||||||
|
Quotas: map[string]ModelQuota{
|
||||||
|
"model-c": {MaxRPM: 10},
|
||||||
|
"model-a": {MaxRPM: 10},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
rl.RecordUsage("model-b", 1, 1)
|
||||||
|
|
||||||
|
t.Run("Models iterator is sorted", func(t *testing.T) {
|
||||||
|
var models []string
|
||||||
|
for m := range rl.Models() {
|
||||||
|
models = append(models, m)
|
||||||
|
}
|
||||||
|
// Should include Gemini defaults (from NewWithConfig's default) + custom models
|
||||||
|
// and be sorted.
|
||||||
|
assert.Contains(t, models, "model-a")
|
||||||
|
assert.Contains(t, models, "model-b")
|
||||||
|
assert.Contains(t, models, "model-c")
|
||||||
|
|
||||||
|
// Check sorting of our specific models
|
||||||
|
foundA, foundB, foundC := -1, -1, -1
|
||||||
|
for i, m := range models {
|
||||||
|
if m == "model-a" {
|
||||||
|
foundA = i
|
||||||
|
}
|
||||||
|
if m == "model-b" {
|
||||||
|
foundB = i
|
||||||
|
}
|
||||||
|
if m == "model-c" {
|
||||||
|
foundC = i
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert.True(t, foundA < foundB && foundB < foundC, "models should be sorted: a < b < c")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Iter iterator is sorted", func(t *testing.T) {
|
||||||
|
var models []string
|
||||||
|
for m, stats := range rl.Iter() {
|
||||||
|
models = append(models, m)
|
||||||
|
if m == "model-a" {
|
||||||
|
assert.Equal(t, 10, stats.MaxRPM)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert.Contains(t, models, "model-a")
|
||||||
|
assert.Contains(t, models, "model-b")
|
||||||
|
assert.Contains(t, models, "model-c")
|
||||||
|
|
||||||
|
// Check sorting
|
||||||
|
foundA, foundB, foundC := -1, -1, -1
|
||||||
|
for i, m := range models {
|
||||||
|
if m == "model-a" {
|
||||||
|
foundA = i
|
||||||
|
}
|
||||||
|
if m == "model-b" {
|
||||||
|
foundB = i
|
||||||
|
}
|
||||||
|
if m == "model-c" {
|
||||||
|
foundC = i
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert.True(t, foundA < foundB && foundB < foundC, "iter should be sorted: a < b < c")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIter_IterEarlyBreak_Good(t *testing.T) {
|
||||||
|
rl, err := NewWithConfig(Config{
|
||||||
|
Quotas: map[string]ModelQuota{
|
||||||
|
"model-a": {MaxRPM: 10},
|
||||||
|
"model-b": {MaxRPM: 20},
|
||||||
|
"model-c": {MaxRPM: 30},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
t.Run("Iter breaks early", func(t *testing.T) {
|
||||||
|
var count int
|
||||||
|
for range rl.Iter() {
|
||||||
|
count++
|
||||||
|
if count == 1 {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert.Equal(t, 1, count, "should stop after first iteration")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Models early break via manual iteration", func(t *testing.T) {
|
||||||
|
var count int
|
||||||
|
for range rl.Models() {
|
||||||
|
count++
|
||||||
|
if count == 2 {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert.Equal(t, 2, count, "should stop after two models")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIter_CountTokensFull_Ugly(t *testing.T) {
|
||||||
|
t.Run("empty model is rejected", func(t *testing.T) {
|
||||||
|
_, err := CountTokens(context.Background(), "key", "", "text")
|
||||||
|
assert.Error(t, err)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("API error non-200", func(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusBadRequest)
|
||||||
|
w.Write([]byte("bad request"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
_, err := countTokensWithClient(context.Background(), server.Client(), server.URL, "key", "model", "text")
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "status 400")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("context cancelled", func(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel()
|
||||||
|
_, err := countTokensWithClient(ctx, http.DefaultClient, "https://generativelanguage.googleapis.com", "key", "model", "text")
|
||||||
|
assert.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "do request")
|
||||||
|
})
|
||||||
|
}
|
||||||
852
ratelimit.go
852
ratelimit.go
File diff suppressed because it is too large
Load diff
|
|
@ -1,33 +1,113 @@
|
||||||
|
// SPDX-License-Identifier: EUPL-1.2
|
||||||
|
|
||||||
package ratelimit
|
package ratelimit
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"io"
|
||||||
"fmt"
|
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"os"
|
|
||||||
"path/filepath"
|
|
||||||
"sync"
|
"sync"
|
||||||
|
"syscall"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
core "dappco.re/go/core"
|
||||||
"github.com/stretchr/testify/assert"
|
"github.com/stretchr/testify/assert"
|
||||||
"github.com/stretchr/testify/require"
|
"github.com/stretchr/testify/require"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
func testPath(parts ...string) string {
|
||||||
|
return core.Path(parts...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func pathExists(path string) bool {
|
||||||
|
var fs core.Fs
|
||||||
|
return fs.Exists(path)
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeTestFile(tb testing.TB, path, content string) {
|
||||||
|
tb.Helper()
|
||||||
|
require.NoError(tb, writeLocalFile(path, content))
|
||||||
|
}
|
||||||
|
|
||||||
|
func ensureTestDir(tb testing.TB, path string) {
|
||||||
|
tb.Helper()
|
||||||
|
require.NoError(tb, ensureDir(path))
|
||||||
|
}
|
||||||
|
|
||||||
|
func setPathMode(tb testing.TB, path string, mode uint32) {
|
||||||
|
tb.Helper()
|
||||||
|
require.NoError(tb, syscall.Chmod(path, mode))
|
||||||
|
}
|
||||||
|
|
||||||
|
func overwriteTestFile(tb testing.TB, path, content string) {
|
||||||
|
tb.Helper()
|
||||||
|
|
||||||
|
var fs core.Fs
|
||||||
|
writer := fs.Create(path)
|
||||||
|
require.NoError(tb, resultError(writer))
|
||||||
|
require.NoError(tb, resultError(core.WriteAll(writer.Value, content)))
|
||||||
|
}
|
||||||
|
|
||||||
|
func isRootUser() bool {
|
||||||
|
return syscall.Geteuid() == 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func repeatString(part string, count int) string {
|
||||||
|
builder := core.NewBuilder()
|
||||||
|
for i := 0; i < count; i++ {
|
||||||
|
builder.WriteString(part)
|
||||||
|
}
|
||||||
|
return builder.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func substringCount(s, substr string) int {
|
||||||
|
if substr == "" {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return len(core.Split(s, substr)) - 1
|
||||||
|
}
|
||||||
|
|
||||||
|
func decodeJSONBody(tb testing.TB, r io.Reader, target any) {
|
||||||
|
tb.Helper()
|
||||||
|
|
||||||
|
data, err := io.ReadAll(r)
|
||||||
|
require.NoError(tb, err)
|
||||||
|
require.NoError(tb, resultError(core.JSONUnmarshal(data, target)))
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeJSONBody(tb testing.TB, w io.Writer, value any) {
|
||||||
|
tb.Helper()
|
||||||
|
|
||||||
|
_, err := io.WriteString(w, core.JSONMarshalString(value))
|
||||||
|
require.NoError(tb, err)
|
||||||
|
}
|
||||||
|
|
||||||
// newTestLimiter returns a RateLimiter with file path set to a temp directory.
|
// newTestLimiter returns a RateLimiter with file path set to a temp directory.
|
||||||
func newTestLimiter(t *testing.T) *RateLimiter {
|
func newTestLimiter(t *testing.T) *RateLimiter {
|
||||||
t.Helper()
|
t.Helper()
|
||||||
rl, err := New()
|
rl, err := New()
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
rl.filePath = filepath.Join(t.TempDir(), "ratelimits.yaml")
|
rl.filePath = testPath(t.TempDir(), "ratelimits.yaml")
|
||||||
return rl
|
return rl
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||||
|
|
||||||
|
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
|
return f(req)
|
||||||
|
}
|
||||||
|
|
||||||
|
type errReader struct{}
|
||||||
|
|
||||||
|
func (errReader) Read([]byte) (int, error) {
|
||||||
|
return 0, io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
|
||||||
// --- Phase 0: CanSend boundary conditions ---
|
// --- Phase 0: CanSend boundary conditions ---
|
||||||
|
|
||||||
func TestCanSend(t *testing.T) {
|
func TestRatelimit_CanSend_Good(t *testing.T) {
|
||||||
t.Run("fresh state allows send", func(t *testing.T) {
|
t.Run("fresh state allows send", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
model := "test-model"
|
model := "test-model"
|
||||||
|
|
@ -162,11 +242,144 @@ func TestCanSend(t *testing.T) {
|
||||||
rl.RecordUsage(model, 50, 50) // exactly 100 tokens
|
rl.RecordUsage(model, 50, 50) // exactly 100 tokens
|
||||||
assert.True(t, rl.CanSend(model, 0), "zero estimated tokens should fit even at limit")
|
assert.True(t, rl.CanSend(model, 0), "zero estimated tokens should fit even at limit")
|
||||||
})
|
})
|
||||||
|
|
||||||
|
t.Run("negative estimated tokens are rejected", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "negative-est"
|
||||||
|
rl.Quotas[model] = ModelQuota{MaxRPM: 100, MaxTPM: 100, MaxRPD: 100}
|
||||||
|
rl.RecordUsage(model, 50, 50)
|
||||||
|
|
||||||
|
assert.False(t, rl.CanSend(model, -1), "negative estimated tokens should not bypass TPM limits")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Phase 0: Decide surface area ---
|
||||||
|
|
||||||
|
func TestRatelimit_Decide_Good(t *testing.T) {
|
||||||
|
t.Run("unknown model remains allowed with unknown code", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
|
||||||
|
decision := rl.Decide("unknown-model", 50)
|
||||||
|
|
||||||
|
assert.True(t, decision.Allowed)
|
||||||
|
assert.Equal(t, DecisionUnknownModel, decision.Code)
|
||||||
|
assert.Zero(t, decision.RetryAfter)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("unlimited quota reports unlimited decision", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "unlimited"
|
||||||
|
rl.Quotas[model] = ModelQuota{}
|
||||||
|
|
||||||
|
decision := rl.Decide(model, 100)
|
||||||
|
|
||||||
|
assert.True(t, decision.Allowed)
|
||||||
|
assert.Equal(t, DecisionUnlimited, decision.Code)
|
||||||
|
assert.Equal(t, 0, decision.Stats.MaxRPM)
|
||||||
|
assert.Equal(t, 0, decision.Stats.MaxTPM)
|
||||||
|
assert.Equal(t, 0, decision.Stats.MaxRPD)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("rpd limit returns retry window", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "rpd-limit"
|
||||||
|
now := time.Now()
|
||||||
|
rl.Quotas[model] = ModelQuota{MaxRPM: 10, MaxTPM: 1000, MaxRPD: 2}
|
||||||
|
rl.State[model] = &UsageStats{DayStart: now.Add(-23 * time.Hour), DayCount: 2}
|
||||||
|
|
||||||
|
decision := rl.Decide(model, 10)
|
||||||
|
|
||||||
|
assert.False(t, decision.Allowed)
|
||||||
|
assert.Equal(t, DecisionRPDLimit, decision.Code)
|
||||||
|
assert.InDelta(t, time.Hour.Seconds(), decision.RetryAfter.Seconds(), 2)
|
||||||
|
assert.Equal(t, 2, decision.Stats.MaxRPD)
|
||||||
|
assert.Equal(t, 2, decision.Stats.RPD)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("rpm limit includes retry-after estimate", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "rpm-limit"
|
||||||
|
now := time.Now()
|
||||||
|
rl.Quotas[model] = ModelQuota{MaxRPM: 1, MaxTPM: 1000, MaxRPD: 5}
|
||||||
|
rl.State[model] = &UsageStats{
|
||||||
|
Requests: []time.Time{now.Add(-10 * time.Second)},
|
||||||
|
Tokens: []TokenEntry{{Time: now.Add(-10 * time.Second), Count: 10}},
|
||||||
|
DayStart: now,
|
||||||
|
DayCount: 1,
|
||||||
|
}
|
||||||
|
|
||||||
|
decision := rl.Decide(model, 5)
|
||||||
|
|
||||||
|
assert.False(t, decision.Allowed)
|
||||||
|
assert.Equal(t, DecisionRPMLimit, decision.Code)
|
||||||
|
assert.InDelta(t, 50, decision.RetryAfter.Seconds(), 1)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("tpm limit surfaces earliest expiry", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "tpm-limit"
|
||||||
|
now := time.Now()
|
||||||
|
rl.Quotas[model] = ModelQuota{MaxRPM: 10, MaxTPM: 100, MaxRPD: 10}
|
||||||
|
rl.State[model] = &UsageStats{
|
||||||
|
Requests: []time.Time{now.Add(-30 * time.Second)},
|
||||||
|
Tokens: []TokenEntry{
|
||||||
|
{Time: now.Add(-50 * time.Second), Count: 70},
|
||||||
|
{Time: now.Add(-10 * time.Second), Count: 20},
|
||||||
|
},
|
||||||
|
DayStart: now,
|
||||||
|
DayCount: 2,
|
||||||
|
}
|
||||||
|
|
||||||
|
decision := rl.Decide(model, 20)
|
||||||
|
|
||||||
|
assert.False(t, decision.Allowed)
|
||||||
|
assert.Equal(t, DecisionTPMLimit, decision.Code)
|
||||||
|
assert.InDelta(t, 10, decision.RetryAfter.Seconds(), 1)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("allowed decision carries stats snapshot", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "decide-allowed"
|
||||||
|
rl.Quotas[model] = ModelQuota{MaxRPM: 5, MaxTPM: 200, MaxRPD: 3}
|
||||||
|
now := time.Now()
|
||||||
|
rl.State[model] = &UsageStats{
|
||||||
|
Requests: []time.Time{now.Add(-5 * time.Second)},
|
||||||
|
Tokens: []TokenEntry{{Time: now.Add(-5 * time.Second), Count: 30}},
|
||||||
|
DayStart: now,
|
||||||
|
DayCount: 1,
|
||||||
|
}
|
||||||
|
|
||||||
|
decision := rl.Decide(model, 20)
|
||||||
|
|
||||||
|
assert.True(t, decision.Allowed)
|
||||||
|
assert.Equal(t, DecisionAllowed, decision.Code)
|
||||||
|
assert.Equal(t, 1, decision.Stats.RPM)
|
||||||
|
assert.Equal(t, 30, decision.Stats.TPM)
|
||||||
|
assert.Equal(t, 1, decision.Stats.RPD)
|
||||||
|
assert.Equal(t, 5, decision.Stats.MaxRPM)
|
||||||
|
assert.Equal(t, 200, decision.Stats.MaxTPM)
|
||||||
|
assert.Equal(t, 3, decision.Stats.MaxRPD)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("negative estimate returns invalid decision", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "neg"
|
||||||
|
rl.Quotas[model] = ModelQuota{MaxRPM: 5, MaxTPM: 50, MaxRPD: 5}
|
||||||
|
|
||||||
|
decision := rl.Decide(model, -5)
|
||||||
|
|
||||||
|
assert.False(t, decision.Allowed)
|
||||||
|
assert.Equal(t, DecisionInvalidTokens, decision.Code)
|
||||||
|
assert.Zero(t, decision.RetryAfter)
|
||||||
|
require.Contains(t, rl.State, model)
|
||||||
|
require.NotNil(t, rl.State[model])
|
||||||
|
assert.Equal(t, 0, rl.State[model].DayCount)
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- Phase 0: Sliding window / prune tests ---
|
// --- Phase 0: Sliding window / prune tests ---
|
||||||
|
|
||||||
func TestPrune(t *testing.T) {
|
func TestRatelimit_Prune_Good(t *testing.T) {
|
||||||
t.Run("removes old entries", func(t *testing.T) {
|
t.Run("removes old entries", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
model := "test-prune"
|
model := "test-prune"
|
||||||
|
|
@ -281,7 +494,7 @@ func TestPrune(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 0: RecordUsage ---
|
// --- Phase 0: RecordUsage ---
|
||||||
|
|
||||||
func TestRecordUsage(t *testing.T) {
|
func TestRatelimit_RecordUsage_Good(t *testing.T) {
|
||||||
t.Run("records into fresh state", func(t *testing.T) {
|
t.Run("records into fresh state", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
model := "record-fresh"
|
model := "record-fresh"
|
||||||
|
|
@ -335,11 +548,24 @@ func TestRecordUsage(t *testing.T) {
|
||||||
assert.Len(t, stats.Tokens, 2)
|
assert.Len(t, stats.Tokens, 2)
|
||||||
assert.Equal(t, 6, stats.DayCount)
|
assert.Equal(t, 6, stats.DayCount)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
t.Run("negative token inputs are clamped to zero", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "record-negative"
|
||||||
|
|
||||||
|
rl.RecordUsage(model, -100, 25)
|
||||||
|
|
||||||
|
stats := rl.State[model]
|
||||||
|
require.NotNil(t, stats)
|
||||||
|
assert.Len(t, stats.Requests, 1)
|
||||||
|
assert.Len(t, stats.Tokens, 1)
|
||||||
|
assert.Equal(t, 25, stats.Tokens[0].Count)
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- Phase 0: Reset ---
|
// --- Phase 0: Reset ---
|
||||||
|
|
||||||
func TestReset(t *testing.T) {
|
func TestRatelimit_Reset_Good(t *testing.T) {
|
||||||
t.Run("reset single model", func(t *testing.T) {
|
t.Run("reset single model", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.RecordUsage("model-a", 10, 10)
|
rl.RecordUsage("model-a", 10, 10)
|
||||||
|
|
@ -373,7 +599,7 @@ func TestReset(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 0: WaitForCapacity ---
|
// --- Phase 0: WaitForCapacity ---
|
||||||
|
|
||||||
func TestWaitForCapacity(t *testing.T) {
|
func TestRatelimit_WaitForCapacity_Good(t *testing.T) {
|
||||||
t.Run("context cancelled returns error", func(t *testing.T) {
|
t.Run("context cancelled returns error", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
model := "wait-cancel"
|
model := "wait-cancel"
|
||||||
|
|
@ -422,11 +648,63 @@ func TestWaitForCapacity(t *testing.T) {
|
||||||
err := rl.WaitForCapacity(ctx, model, 10)
|
err := rl.WaitForCapacity(ctx, model, 10)
|
||||||
assert.Error(t, err, "should return error for already-cancelled context")
|
assert.Error(t, err, "should return error for already-cancelled context")
|
||||||
})
|
})
|
||||||
|
|
||||||
|
t.Run("negative tokens return error", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
err := rl.WaitForCapacity(context.Background(), "wait-negative", -1)
|
||||||
|
assert.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "negative tokens")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRatelimit_NilUsageStats_Ugly(t *testing.T) {
|
||||||
|
t.Run("CanSend replaces nil state without panicking", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "nil-cansend"
|
||||||
|
rl.Quotas[model] = ModelQuota{MaxRPM: 10, MaxTPM: 100, MaxRPD: 10}
|
||||||
|
rl.State[model] = nil
|
||||||
|
|
||||||
|
assert.NotPanics(t, func() {
|
||||||
|
assert.True(t, rl.CanSend(model, 10))
|
||||||
|
})
|
||||||
|
assert.NotNil(t, rl.State[model])
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("RecordUsage replaces nil state without panicking", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "nil-record"
|
||||||
|
rl.State[model] = nil
|
||||||
|
|
||||||
|
assert.NotPanics(t, func() {
|
||||||
|
rl.RecordUsage(model, 10, 10)
|
||||||
|
})
|
||||||
|
require.NotNil(t, rl.State[model])
|
||||||
|
assert.Equal(t, 1, rl.State[model].DayCount)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("Stats and AllStats tolerate nil state entries", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
rl.Quotas["nil-stats"] = ModelQuota{MaxRPM: 1, MaxTPM: 2, MaxRPD: 3}
|
||||||
|
rl.State["nil-stats"] = nil
|
||||||
|
rl.State["nil-all-stats"] = nil
|
||||||
|
|
||||||
|
assert.NotPanics(t, func() {
|
||||||
|
stats := rl.Stats("nil-stats")
|
||||||
|
assert.Equal(t, 1, stats.MaxRPM)
|
||||||
|
assert.Equal(t, 0, stats.TPM)
|
||||||
|
})
|
||||||
|
|
||||||
|
assert.NotPanics(t, func() {
|
||||||
|
all := rl.AllStats()
|
||||||
|
assert.Contains(t, all, "nil-stats")
|
||||||
|
assert.Contains(t, all, "nil-all-stats")
|
||||||
|
})
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- Phase 0: Stats ---
|
// --- Phase 0: Stats ---
|
||||||
|
|
||||||
func TestStats(t *testing.T) {
|
func TestRatelimit_Stats_Good(t *testing.T) {
|
||||||
t.Run("returns stats for known model with usage", func(t *testing.T) {
|
t.Run("returns stats for known model with usage", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
model := "stats-test"
|
model := "stats-test"
|
||||||
|
|
@ -466,7 +744,7 @@ func TestStats(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 0: AllStats ---
|
// --- Phase 0: AllStats ---
|
||||||
|
|
||||||
func TestAllStats(t *testing.T) {
|
func TestRatelimit_AllStats_Good(t *testing.T) {
|
||||||
t.Run("includes all default quotas plus state-only models", func(t *testing.T) {
|
t.Run("includes all default quotas plus state-only models", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.RecordUsage("gemini-3-pro-preview", 1000, 500)
|
rl.RecordUsage("gemini-3-pro-preview", 1000, 500)
|
||||||
|
|
@ -524,10 +802,10 @@ func TestAllStats(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 0: Persist and Load ---
|
// --- Phase 0: Persist and Load ---
|
||||||
|
|
||||||
func TestPersistAndLoad(t *testing.T) {
|
func TestRatelimit_PersistAndLoad_Ugly(t *testing.T) {
|
||||||
t.Run("round-trip preserves state", func(t *testing.T) {
|
t.Run("round-trip preserves state", func(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
path := filepath.Join(tmpDir, "ratelimits.yaml")
|
path := testPath(tmpDir, "ratelimits.yaml")
|
||||||
|
|
||||||
rl1, err := New()
|
rl1, err := New()
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
@ -550,7 +828,7 @@ func TestPersistAndLoad(t *testing.T) {
|
||||||
|
|
||||||
t.Run("load from non-existent file is not an error", func(t *testing.T) {
|
t.Run("load from non-existent file is not an error", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.filePath = filepath.Join(t.TempDir(), "does-not-exist.yaml")
|
rl.filePath = testPath(t.TempDir(), "does-not-exist.yaml")
|
||||||
|
|
||||||
err := rl.Load()
|
err := rl.Load()
|
||||||
assert.NoError(t, err, "loading non-existent file should not error")
|
assert.NoError(t, err, "loading non-existent file should not error")
|
||||||
|
|
@ -558,8 +836,8 @@ func TestPersistAndLoad(t *testing.T) {
|
||||||
|
|
||||||
t.Run("load from corrupt YAML returns error", func(t *testing.T) {
|
t.Run("load from corrupt YAML returns error", func(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
path := filepath.Join(tmpDir, "corrupt.yaml")
|
path := testPath(tmpDir, "corrupt.yaml")
|
||||||
require.NoError(t, os.WriteFile(path, []byte("{{{{invalid yaml!!!!"), 0644))
|
writeTestFile(t, path, "{{{{invalid yaml!!!!")
|
||||||
|
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.filePath = path
|
rl.filePath = path
|
||||||
|
|
@ -569,13 +847,13 @@ func TestPersistAndLoad(t *testing.T) {
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("load from unreadable file returns error", func(t *testing.T) {
|
t.Run("load from unreadable file returns error", func(t *testing.T) {
|
||||||
if os.Getuid() == 0 {
|
if isRootUser() {
|
||||||
t.Skip("chmod 000 does not restrict root")
|
t.Skip("chmod 000 does not restrict root")
|
||||||
}
|
}
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
path := filepath.Join(tmpDir, "unreadable.yaml")
|
path := testPath(tmpDir, "unreadable.yaml")
|
||||||
require.NoError(t, os.WriteFile(path, []byte("quotas: {}"), 0644))
|
writeTestFile(t, path, "quotas: {}")
|
||||||
require.NoError(t, os.Chmod(path, 0000))
|
setPathMode(t, path, 0o000)
|
||||||
|
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.filePath = path
|
rl.filePath = path
|
||||||
|
|
@ -584,12 +862,12 @@ func TestPersistAndLoad(t *testing.T) {
|
||||||
assert.Error(t, err, "unreadable file should produce an error")
|
assert.Error(t, err, "unreadable file should produce an error")
|
||||||
|
|
||||||
// Clean up permissions for temp dir cleanup
|
// Clean up permissions for temp dir cleanup
|
||||||
_ = os.Chmod(path, 0644)
|
_ = syscall.Chmod(path, 0o644)
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("persist to nested non-existent directory creates it", func(t *testing.T) {
|
t.Run("persist to nested non-existent directory creates it", func(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
path := filepath.Join(tmpDir, "nested", "deep", "ratelimits.yaml")
|
path := testPath(tmpDir, "nested", "deep", "ratelimits.yaml")
|
||||||
|
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.filePath = path
|
rl.filePath = path
|
||||||
|
|
@ -598,32 +876,32 @@ func TestPersistAndLoad(t *testing.T) {
|
||||||
err := rl.Persist()
|
err := rl.Persist()
|
||||||
assert.NoError(t, err, "should create nested directories")
|
assert.NoError(t, err, "should create nested directories")
|
||||||
|
|
||||||
_, statErr := os.Stat(path)
|
assert.True(t, pathExists(path), "file should exist")
|
||||||
assert.NoError(t, statErr, "file should exist")
|
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("persist to unwritable directory returns error", func(t *testing.T) {
|
t.Run("persist to unwritable directory returns error", func(t *testing.T) {
|
||||||
if os.Getuid() == 0 {
|
if isRootUser() {
|
||||||
t.Skip("chmod 0555 does not restrict root")
|
t.Skip("chmod 0555 does not restrict root")
|
||||||
}
|
}
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
unwritable := filepath.Join(tmpDir, "readonly")
|
unwritable := testPath(tmpDir, "readonly")
|
||||||
require.NoError(t, os.MkdirAll(unwritable, 0555))
|
ensureTestDir(t, unwritable)
|
||||||
|
setPathMode(t, unwritable, 0o555)
|
||||||
|
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.filePath = filepath.Join(unwritable, "sub", "ratelimits.yaml")
|
rl.filePath = testPath(unwritable, "sub", "ratelimits.yaml")
|
||||||
|
|
||||||
err := rl.Persist()
|
err := rl.Persist()
|
||||||
assert.Error(t, err, "should fail when directory is unwritable")
|
assert.Error(t, err, "should fail when directory is unwritable")
|
||||||
|
|
||||||
// Clean up
|
// Clean up
|
||||||
_ = os.Chmod(unwritable, 0755)
|
_ = syscall.Chmod(unwritable, 0o755)
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- Phase 0: Default quotas ---
|
// --- Phase 0: Default quotas ---
|
||||||
|
|
||||||
func TestDefaultQuotas(t *testing.T) {
|
func TestRatelimit_DefaultQuotas_Good(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
|
|
||||||
tests := []struct {
|
tests := []struct {
|
||||||
|
|
@ -652,7 +930,7 @@ func TestDefaultQuotas(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 0: Concurrent access (race test) ---
|
// --- Phase 0: Concurrent access (race test) ---
|
||||||
|
|
||||||
func TestConcurrentAccess(t *testing.T) {
|
func TestRatelimit_ConcurrentAccess_Good(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
model := "concurrent-test"
|
model := "concurrent-test"
|
||||||
rl.Quotas[model] = ModelQuota{MaxRPM: 1000, MaxTPM: 10000000, MaxRPD: 10000}
|
rl.Quotas[model] = ModelQuota{MaxRPM: 1000, MaxTPM: 10000000, MaxRPD: 10000}
|
||||||
|
|
@ -678,7 +956,7 @@ func TestConcurrentAccess(t *testing.T) {
|
||||||
assert.Equal(t, expected, stats.RPD, "all recordings should be counted")
|
assert.Equal(t, expected, stats.RPD, "all recordings should be counted")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestConcurrentResetAndRecord(t *testing.T) {
|
func TestRatelimit_ConcurrentResetAndRecord_Ugly(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
model := "concurrent-reset"
|
model := "concurrent-reset"
|
||||||
rl.Quotas[model] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 100000}
|
rl.Quotas[model] = ModelQuota{MaxRPM: 10000, MaxTPM: 100000000, MaxRPD: 100000}
|
||||||
|
|
@ -716,38 +994,194 @@ func TestConcurrentResetAndRecord(t *testing.T) {
|
||||||
// No assertion needed -- if we get here without -race flagging, mutex is sound
|
// No assertion needed -- if we get here without -race flagging, mutex is sound
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestRatelimit_BackgroundPrune_Good(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
model := "prune-me"
|
||||||
|
rl.Quotas[model] = ModelQuota{MaxRPM: 100}
|
||||||
|
|
||||||
|
// Set state with old usage.
|
||||||
|
old := time.Now().Add(-2 * time.Minute)
|
||||||
|
rl.State[model] = &UsageStats{
|
||||||
|
Requests: []time.Time{old},
|
||||||
|
Tokens: []TokenEntry{{Time: old, Count: 100}},
|
||||||
|
}
|
||||||
|
|
||||||
|
stop := rl.BackgroundPrune(10 * time.Millisecond)
|
||||||
|
defer stop()
|
||||||
|
|
||||||
|
// Wait for pruner to run.
|
||||||
|
assert.Eventually(t, func() bool {
|
||||||
|
rl.mu.Lock()
|
||||||
|
defer rl.mu.Unlock()
|
||||||
|
_, exists := rl.State[model]
|
||||||
|
return !exists
|
||||||
|
}, 1*time.Second, 20*time.Millisecond, "old empty state should be pruned")
|
||||||
|
|
||||||
|
t.Run("non-positive interval is a safe no-op", func(t *testing.T) {
|
||||||
|
rl := newTestLimiter(t)
|
||||||
|
rl.State["still-here"] = &UsageStats{
|
||||||
|
Requests: []time.Time{time.Now().Add(-2 * time.Minute)},
|
||||||
|
}
|
||||||
|
|
||||||
|
assert.NotPanics(t, func() {
|
||||||
|
stop := rl.BackgroundPrune(0)
|
||||||
|
stop()
|
||||||
|
})
|
||||||
|
assert.Contains(t, rl.State, "still-here")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
// --- Phase 0: CountTokens (with mock HTTP server) ---
|
// --- Phase 0: CountTokens (with mock HTTP server) ---
|
||||||
|
|
||||||
func TestCountTokens(t *testing.T) {
|
func TestRatelimit_CountTokens_Ugly(t *testing.T) {
|
||||||
t.Run("successful token count", func(t *testing.T) {
|
t.Run("successful token count", func(t *testing.T) {
|
||||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
assert.Equal(t, http.MethodPost, r.Method)
|
assert.Equal(t, http.MethodPost, r.Method)
|
||||||
assert.Equal(t, "application/json", r.Header.Get("Content-Type"))
|
assert.Equal(t, "application/json", r.Header.Get("Content-Type"))
|
||||||
assert.Equal(t, "test-api-key", r.Header.Get("x-goog-api-key"))
|
assert.Equal(t, "test-api-key", r.Header.Get("x-goog-api-key"))
|
||||||
|
assert.Equal(t, "/v1beta/models/test-model:countTokens", r.URL.EscapedPath())
|
||||||
|
|
||||||
|
var body struct {
|
||||||
|
Contents []struct {
|
||||||
|
Parts []struct {
|
||||||
|
Text string `json:"text"`
|
||||||
|
} `json:"parts"`
|
||||||
|
} `json:"contents"`
|
||||||
|
}
|
||||||
|
decodeJSONBody(t, r.Body, &body)
|
||||||
|
require.Len(t, body.Contents, 1)
|
||||||
|
require.Len(t, body.Contents[0].Parts, 1)
|
||||||
|
assert.Equal(t, "hello", body.Contents[0].Parts[0].Text)
|
||||||
|
|
||||||
w.Header().Set("Content-Type", "application/json")
|
w.Header().Set("Content-Type", "application/json")
|
||||||
json.NewEncoder(w).Encode(map[string]int{"totalTokens": 42})
|
writeJSONBody(t, w, map[string]int{"totalTokens": 42})
|
||||||
}))
|
}))
|
||||||
defer server.Close()
|
defer server.Close()
|
||||||
|
|
||||||
// We need to override the URL. Since CountTokens hardcodes the Google API URL,
|
tokens, err := countTokensWithClient(context.Background(), server.Client(), server.URL, "test-api-key", "test-model", "hello")
|
||||||
// we test it via the exported function with a test server.
|
require.NoError(t, err)
|
||||||
// For proper unit testing, we would need to make the base URL configurable.
|
assert.Equal(t, 42, tokens)
|
||||||
// For now, test the error paths that don't require a real API.
|
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("API error returns error", func(t *testing.T) {
|
t.Run("model name is path escaped", func(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
assert.Equal(t, "/v1beta/models/folder%2Fmodel%3Fdebug=1:countTokens", r.URL.EscapedPath())
|
||||||
|
assert.Empty(t, r.URL.RawQuery)
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
writeJSONBody(t, w, map[string]int{"totalTokens": 7})
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
tokens, err := countTokensWithClient(context.Background(), server.Client(), server.URL, "test-api-key", "folder/model?debug=1", "hello")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 7, tokens)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("API error body is truncated", func(t *testing.T) {
|
||||||
|
largeBody := repeatString("x", countTokensErrorBodyLimit+256)
|
||||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
w.WriteHeader(http.StatusUnauthorized)
|
w.WriteHeader(http.StatusUnauthorized)
|
||||||
fmt.Fprint(w, `{"error": "invalid API key"}`)
|
_, err := io.WriteString(w, largeBody)
|
||||||
|
require.NoError(t, err)
|
||||||
}))
|
}))
|
||||||
defer server.Close()
|
defer server.Close()
|
||||||
|
|
||||||
// Can't test directly due to hardcoded URL, but we can verify error
|
_, err := countTokensWithClient(context.Background(), server.Client(), server.URL, "fake-key", "test-model", "hello")
|
||||||
// handling with an unreachable endpoint
|
require.Error(t, err)
|
||||||
_, err := CountTokens("fake-key", "test-model", "hello")
|
assert.Contains(t, err.Error(), "api error status 401")
|
||||||
assert.Error(t, err, "should fail with invalid API endpoint")
|
assert.True(t, substringCount(err.Error(), "x") < len(largeBody), "error body should be bounded")
|
||||||
|
assert.Contains(t, err.Error(), "...")
|
||||||
})
|
})
|
||||||
|
|
||||||
|
t.Run("empty model is rejected before request", func(t *testing.T) {
|
||||||
|
_, err := CountTokens(context.Background(), "fake-key", "", "hello")
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "build url")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("invalid base URL returns error", func(t *testing.T) {
|
||||||
|
_, err := countTokensWithClient(context.Background(), http.DefaultClient, "://bad-url", "fake-key", "test-model", "hello")
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "build url")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("base URL without host returns error", func(t *testing.T) {
|
||||||
|
_, err := countTokensWithClient(context.Background(), http.DefaultClient, "/relative", "fake-key", "test-model", "hello")
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "build url")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("invalid JSON response returns error", func(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, err := w.Write([]byte(`{"totalTokens":`))
|
||||||
|
require.NoError(t, err)
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
_, err := countTokensWithClient(context.Background(), server.Client(), server.URL, "fake-key", "test-model", "hello")
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "decode response")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("error body read failures are returned", func(t *testing.T) {
|
||||||
|
client := &http.Client{
|
||||||
|
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||||
|
return &http.Response{
|
||||||
|
StatusCode: http.StatusBadGateway,
|
||||||
|
Body: io.NopCloser(errReader{}),
|
||||||
|
Header: make(http.Header),
|
||||||
|
}, nil
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err := countTokensWithClient(context.Background(), client, "https://generativelanguage.googleapis.com", "fake-key", "test-model", "hello")
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "read error body")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("nil client falls back to http.DefaultClient", func(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
writeJSONBody(t, w, map[string]int{"totalTokens": 11})
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
originalClient := http.DefaultClient
|
||||||
|
http.DefaultClient = server.Client()
|
||||||
|
defer func() {
|
||||||
|
http.DefaultClient = originalClient
|
||||||
|
}()
|
||||||
|
|
||||||
|
tokens, err := countTokensWithClient(context.Background(), nil, server.URL, "fake-key", "test-model", "hello")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 11, tokens)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRatelimit_PersistSkipsNilState_Good(t *testing.T) {
|
||||||
|
path := testPath(t.TempDir(), "nil-state.yaml")
|
||||||
|
|
||||||
|
rl, err := New()
|
||||||
|
require.NoError(t, err)
|
||||||
|
rl.filePath = path
|
||||||
|
rl.State["nil-model"] = nil
|
||||||
|
|
||||||
|
require.NoError(t, rl.Persist())
|
||||||
|
|
||||||
|
rl2, err := New()
|
||||||
|
require.NoError(t, err)
|
||||||
|
rl2.filePath = path
|
||||||
|
require.NoError(t, rl2.Load())
|
||||||
|
assert.NotContains(t, rl2.State, "nil-model")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRatelimit_TokenTotals_Good(t *testing.T) {
|
||||||
|
maxInt := int(^uint(0) >> 1)
|
||||||
|
|
||||||
|
assert.Equal(t, 25, safeTokenSum(-100, 25))
|
||||||
|
assert.Equal(t, maxInt, safeTokenSum(maxInt, 1))
|
||||||
|
assert.Equal(t, 10, totalTokenCount([]TokenEntry{{Count: -5}, {Count: 10}}))
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- Phase 0: Benchmarks ---
|
// --- Phase 0: Benchmarks ---
|
||||||
|
|
@ -812,7 +1246,7 @@ func BenchmarkCanSendConcurrent(b *testing.B) {
|
||||||
|
|
||||||
// --- Phase 1: Provider profiles and NewWithConfig ---
|
// --- Phase 1: Provider profiles and NewWithConfig ---
|
||||||
|
|
||||||
func TestDefaultProfiles(t *testing.T) {
|
func TestRatelimit_DefaultProfiles_Good(t *testing.T) {
|
||||||
profiles := DefaultProfiles()
|
profiles := DefaultProfiles()
|
||||||
|
|
||||||
t.Run("contains all four providers", func(t *testing.T) {
|
t.Run("contains all four providers", func(t *testing.T) {
|
||||||
|
|
@ -853,10 +1287,10 @@ func TestDefaultProfiles(t *testing.T) {
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestNewWithConfig(t *testing.T) {
|
func TestRatelimit_NewWithConfig_Ugly(t *testing.T) {
|
||||||
t.Run("empty config defaults to Gemini", func(t *testing.T) {
|
t.Run("empty config defaults to Gemini", func(t *testing.T) {
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "test.yaml"),
|
FilePath: testPath(t.TempDir(), "test.yaml"),
|
||||||
})
|
})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
|
@ -866,7 +1300,7 @@ func TestNewWithConfig(t *testing.T) {
|
||||||
|
|
||||||
t.Run("single provider loads only its models", func(t *testing.T) {
|
t.Run("single provider loads only its models", func(t *testing.T) {
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "test.yaml"),
|
FilePath: testPath(t.TempDir(), "test.yaml"),
|
||||||
Providers: []Provider{ProviderOpenAI},
|
Providers: []Provider{ProviderOpenAI},
|
||||||
})
|
})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
@ -880,7 +1314,7 @@ func TestNewWithConfig(t *testing.T) {
|
||||||
|
|
||||||
t.Run("multiple providers merge models", func(t *testing.T) {
|
t.Run("multiple providers merge models", func(t *testing.T) {
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "test.yaml"),
|
FilePath: testPath(t.TempDir(), "test.yaml"),
|
||||||
Providers: []Provider{ProviderGemini, ProviderAnthropic},
|
Providers: []Provider{ProviderGemini, ProviderAnthropic},
|
||||||
})
|
})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
@ -896,7 +1330,7 @@ func TestNewWithConfig(t *testing.T) {
|
||||||
|
|
||||||
t.Run("explicit quotas override provider defaults", func(t *testing.T) {
|
t.Run("explicit quotas override provider defaults", func(t *testing.T) {
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "test.yaml"),
|
FilePath: testPath(t.TempDir(), "test.yaml"),
|
||||||
Providers: []Provider{ProviderGemini},
|
Providers: []Provider{ProviderGemini},
|
||||||
Quotas: map[string]ModelQuota{
|
Quotas: map[string]ModelQuota{
|
||||||
"gemini-3-pro-preview": {MaxRPM: 999, MaxTPM: 888, MaxRPD: 777},
|
"gemini-3-pro-preview": {MaxRPM: 999, MaxTPM: 888, MaxRPD: 777},
|
||||||
|
|
@ -912,7 +1346,7 @@ func TestNewWithConfig(t *testing.T) {
|
||||||
|
|
||||||
t.Run("explicit quotas without providers", func(t *testing.T) {
|
t.Run("explicit quotas without providers", func(t *testing.T) {
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "test.yaml"),
|
FilePath: testPath(t.TempDir(), "test.yaml"),
|
||||||
Quotas: map[string]ModelQuota{
|
Quotas: map[string]ModelQuota{
|
||||||
"my-custom-model": {MaxRPM: 10, MaxTPM: 1000, MaxRPD: 50},
|
"my-custom-model": {MaxRPM: 10, MaxTPM: 1000, MaxRPD: 50},
|
||||||
},
|
},
|
||||||
|
|
@ -925,7 +1359,7 @@ func TestNewWithConfig(t *testing.T) {
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("custom file path is respected", func(t *testing.T) {
|
t.Run("custom file path is respected", func(t *testing.T) {
|
||||||
customPath := filepath.Join(t.TempDir(), "custom", "limits.yaml")
|
customPath := testPath(t.TempDir(), "custom", "limits.yaml")
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: customPath,
|
FilePath: customPath,
|
||||||
Providers: []Provider{ProviderLocal},
|
Providers: []Provider{ProviderLocal},
|
||||||
|
|
@ -935,13 +1369,12 @@ func TestNewWithConfig(t *testing.T) {
|
||||||
rl.RecordUsage("test", 1, 1)
|
rl.RecordUsage("test", 1, 1)
|
||||||
require.NoError(t, rl.Persist())
|
require.NoError(t, rl.Persist())
|
||||||
|
|
||||||
_, statErr := os.Stat(customPath)
|
assert.True(t, pathExists(customPath), "file should be created at custom path")
|
||||||
assert.NoError(t, statErr, "file should be created at custom path")
|
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("unknown provider is silently skipped", func(t *testing.T) {
|
t.Run("unknown provider is silently skipped", func(t *testing.T) {
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "test.yaml"),
|
FilePath: testPath(t.TempDir(), "test.yaml"),
|
||||||
Providers: []Provider{"nonexistent-provider"},
|
Providers: []Provider{"nonexistent-provider"},
|
||||||
})
|
})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
@ -950,7 +1383,7 @@ func TestNewWithConfig(t *testing.T) {
|
||||||
|
|
||||||
t.Run("local provider with custom quotas", func(t *testing.T) {
|
t.Run("local provider with custom quotas", func(t *testing.T) {
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "test.yaml"),
|
FilePath: testPath(t.TempDir(), "test.yaml"),
|
||||||
Providers: []Provider{ProviderLocal},
|
Providers: []Provider{ProviderLocal},
|
||||||
Quotas: map[string]ModelQuota{
|
Quotas: map[string]ModelQuota{
|
||||||
"llama-3.3-70b": {MaxRPM: 5, MaxTPM: 50000, MaxRPD: 0},
|
"llama-3.3-70b": {MaxRPM: 5, MaxTPM: 50000, MaxRPD: 0},
|
||||||
|
|
@ -963,9 +1396,28 @@ func TestNewWithConfig(t *testing.T) {
|
||||||
assert.Equal(t, 5, q.MaxRPM)
|
assert.Equal(t, 5, q.MaxRPM)
|
||||||
assert.Equal(t, 50000, q.MaxTPM)
|
assert.Equal(t, 50000, q.MaxTPM)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
t.Run("invalid backend returns error", func(t *testing.T) {
|
||||||
|
_, err := NewWithConfig(Config{
|
||||||
|
Backend: "bogus",
|
||||||
|
})
|
||||||
|
require.Error(t, err)
|
||||||
|
assert.Contains(t, err.Error(), "unknown backend")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("default YAML path uses home directory", func(t *testing.T) {
|
||||||
|
home := t.TempDir()
|
||||||
|
t.Setenv("HOME", home)
|
||||||
|
t.Setenv("USERPROFILE", "")
|
||||||
|
t.Setenv("home", "")
|
||||||
|
|
||||||
|
rl, err := NewWithConfig(Config{})
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, testPath(home, defaultStateDirName, defaultYAMLStateFile), rl.filePath)
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestNewBackwardCompatibility(t *testing.T) {
|
func TestRatelimit_NewBackwardCompatibility_Good(t *testing.T) {
|
||||||
// New() should produce the exact same result as before Phase 1
|
// New() should produce the exact same result as before Phase 1
|
||||||
rl, err := New()
|
rl, err := New()
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
@ -988,7 +1440,7 @@ func TestNewBackwardCompatibility(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSetQuota(t *testing.T) {
|
func TestRatelimit_SetQuota_Good(t *testing.T) {
|
||||||
t.Run("adds new model quota", func(t *testing.T) {
|
t.Run("adds new model quota", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.SetQuota("custom-model", ModelQuota{MaxRPM: 42, MaxTPM: 9999, MaxRPD: 100})
|
rl.SetQuota("custom-model", ModelQuota{MaxRPM: 42, MaxTPM: 9999, MaxRPD: 100})
|
||||||
|
|
@ -1016,7 +1468,7 @@ func TestSetQuota(t *testing.T) {
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go func(n int) {
|
go func(n int) {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
model := fmt.Sprintf("model-%d", n)
|
model := core.Sprintf("model-%d", n)
|
||||||
rl.SetQuota(model, ModelQuota{MaxRPM: n, MaxTPM: n * 100, MaxRPD: n * 10})
|
rl.SetQuota(model, ModelQuota{MaxRPM: n, MaxTPM: n * 100, MaxRPD: n * 10})
|
||||||
}(i)
|
}(i)
|
||||||
}
|
}
|
||||||
|
|
@ -1026,7 +1478,7 @@ func TestSetQuota(t *testing.T) {
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestAddProvider(t *testing.T) {
|
func TestRatelimit_AddProvider_Good(t *testing.T) {
|
||||||
t.Run("adds OpenAI models to existing limiter", func(t *testing.T) {
|
t.Run("adds OpenAI models to existing limiter", func(t *testing.T) {
|
||||||
rl := newTestLimiter(t) // starts with Gemini defaults
|
rl := newTestLimiter(t) // starts with Gemini defaults
|
||||||
geminiCount := len(rl.Quotas)
|
geminiCount := len(rl.Quotas)
|
||||||
|
|
@ -1088,7 +1540,7 @@ func TestAddProvider(t *testing.T) {
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestProviderConstants(t *testing.T) {
|
func TestRatelimit_ProviderConstants_Good(t *testing.T) {
|
||||||
// Verify the string values are stable (they may be used in YAML configs)
|
// Verify the string values are stable (they may be used in YAML configs)
|
||||||
assert.Equal(t, Provider("gemini"), ProviderGemini)
|
assert.Equal(t, Provider("gemini"), ProviderGemini)
|
||||||
assert.Equal(t, Provider("openai"), ProviderOpenAI)
|
assert.Equal(t, Provider("openai"), ProviderOpenAI)
|
||||||
|
|
@ -1098,7 +1550,7 @@ func TestProviderConstants(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 0 addendum: Additional concurrent and multi-model race tests ---
|
// --- Phase 0 addendum: Additional concurrent and multi-model race tests ---
|
||||||
|
|
||||||
func TestConcurrentMultipleModels(t *testing.T) {
|
func TestRatelimit_ConcurrentMultipleModels_Good(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
models := []string{"model-a", "model-b", "model-c", "model-d", "model-e"}
|
models := []string{"model-a", "model-b", "model-c", "model-d", "model-e"}
|
||||||
for _, m := range models {
|
for _, m := range models {
|
||||||
|
|
@ -1128,9 +1580,9 @@ func TestConcurrentMultipleModels(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestConcurrentPersistAndLoad(t *testing.T) {
|
func TestRatelimit_ConcurrentPersistAndLoad_Ugly(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
path := filepath.Join(tmpDir, "concurrent.yaml")
|
path := testPath(tmpDir, "concurrent.yaml")
|
||||||
|
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
rl.filePath = path
|
rl.filePath = path
|
||||||
|
|
@ -1162,7 +1614,7 @@ func TestConcurrentPersistAndLoad(t *testing.T) {
|
||||||
// No panics or data races = pass
|
// No panics or data races = pass
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestConcurrentAllStatsAndRecordUsage(t *testing.T) {
|
func TestRatelimit_ConcurrentAllStatsAndRecordUsage_Good(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
models := []string{"stats-a", "stats-b", "stats-c"}
|
models := []string{"stats-a", "stats-b", "stats-c"}
|
||||||
for _, m := range models {
|
for _, m := range models {
|
||||||
|
|
@ -1193,7 +1645,7 @@ func TestConcurrentAllStatsAndRecordUsage(t *testing.T) {
|
||||||
wg.Wait()
|
wg.Wait()
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestConcurrentWaitForCapacityAndRecordUsage(t *testing.T) {
|
func TestRatelimit_ConcurrentWaitForCapacityAndRecordUsage_Good(t *testing.T) {
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
model := "race-wait"
|
model := "race-wait"
|
||||||
rl.Quotas[model] = ModelQuota{MaxRPM: 100, MaxTPM: 10000000, MaxRPD: 10000}
|
rl.Quotas[model] = ModelQuota{MaxRPM: 100, MaxTPM: 10000000, MaxRPD: 10000}
|
||||||
|
|
@ -1290,7 +1742,7 @@ func BenchmarkAllStats(b *testing.B) {
|
||||||
|
|
||||||
func BenchmarkPersist(b *testing.B) {
|
func BenchmarkPersist(b *testing.B) {
|
||||||
tmpDir := b.TempDir()
|
tmpDir := b.TempDir()
|
||||||
path := filepath.Join(tmpDir, "bench.yaml")
|
path := testPath(tmpDir, "bench.yaml")
|
||||||
|
|
||||||
rl, _ := New()
|
rl, _ := New()
|
||||||
rl.filePath = path
|
rl.filePath = path
|
||||||
|
|
@ -1311,10 +1763,10 @@ func BenchmarkPersist(b *testing.B) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestEndToEndMultiProvider(t *testing.T) {
|
func TestRatelimit_EndToEndMultiProvider_Good(t *testing.T) {
|
||||||
// Simulate a real-world scenario: limiter for both Gemini and Anthropic
|
// Simulate a real-world scenario: limiter for both Gemini and Anthropic
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "multi.yaml"),
|
FilePath: testPath(t.TempDir(), "multi.yaml"),
|
||||||
Providers: []Provider{ProviderGemini, ProviderAnthropic},
|
Providers: []Provider{ProviderGemini, ProviderAnthropic},
|
||||||
})
|
})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
|
||||||
154
specs/RFC.md
Normal file
154
specs/RFC.md
Normal file
|
|
@ -0,0 +1,154 @@
|
||||||
|
# ratelimit
|
||||||
|
**Import:** `dappco.re/go/core/go-ratelimit`
|
||||||
|
**Files:** 2
|
||||||
|
|
||||||
|
## Types
|
||||||
|
|
||||||
|
### `Provider`
|
||||||
|
`type Provider string`
|
||||||
|
|
||||||
|
`Provider` identifies an LLM provider used to select built-in quota profiles. The package defines four exported provider values: `ProviderGemini`, `ProviderOpenAI`, `ProviderAnthropic`, and `ProviderLocal`.
|
||||||
|
|
||||||
|
### `ModelQuota`
|
||||||
|
`type ModelQuota struct`
|
||||||
|
|
||||||
|
`ModelQuota` defines the rate limits for a single model. A value of `0` means the corresponding limit is unlimited.
|
||||||
|
|
||||||
|
- `MaxRPM int`: requests per minute.
|
||||||
|
- `MaxTPM int`: tokens per minute.
|
||||||
|
- `MaxRPD int`: requests per rolling 24-hour window.
|
||||||
|
|
||||||
|
### `ProviderProfile`
|
||||||
|
`type ProviderProfile struct`
|
||||||
|
|
||||||
|
`ProviderProfile` bundles a provider identifier with the default quota table for that provider.
|
||||||
|
|
||||||
|
- `Provider Provider`: the provider that owns the profile.
|
||||||
|
- `Models map[string]ModelQuota`: built-in quotas keyed by model name.
|
||||||
|
|
||||||
|
### `Config`
|
||||||
|
`type Config struct`
|
||||||
|
|
||||||
|
`Config` controls `RateLimiter` initialisation, backend selection, and default quotas.
|
||||||
|
|
||||||
|
- `FilePath string`: overrides the default persistence path. When empty, `NewWithConfig` resolves a default path under `~/.core`, using `ratelimits.yaml` for the YAML backend and `ratelimits.db` for the SQLite backend.
|
||||||
|
- `Backend string`: selects the persistence backend. `NewWithConfig` accepts `""` or `"yaml"` for YAML and `"sqlite"` for SQLite. `NewWithSQLiteConfig` ignores this field and always uses SQLite.
|
||||||
|
- `Quotas map[string]ModelQuota`: explicit per-model quotas. These are merged on top of any provider defaults loaded from `Providers`.
|
||||||
|
- `Providers []Provider`: provider profiles to load from `DefaultProfiles`. If both `Providers` and `Quotas` are empty, Gemini defaults are used.
|
||||||
|
|
||||||
|
### `TokenEntry`
|
||||||
|
`type TokenEntry struct`
|
||||||
|
|
||||||
|
`TokenEntry` records a single token-usage event.
|
||||||
|
|
||||||
|
- `Time time.Time`: when the token event was recorded.
|
||||||
|
- `Count int`: how many tokens were counted for that event.
|
||||||
|
|
||||||
|
### `UsageStats`
|
||||||
|
`type UsageStats struct`
|
||||||
|
|
||||||
|
`UsageStats` stores the in-memory usage history for one model.
|
||||||
|
|
||||||
|
- `Requests []time.Time`: request timestamps inside the sliding one-minute window.
|
||||||
|
- `Tokens []TokenEntry`: token usage entries inside the sliding one-minute window.
|
||||||
|
- `DayStart time.Time`: the start of the current rolling 24-hour window.
|
||||||
|
- `DayCount int`: the number of requests recorded in the current rolling 24-hour window.
|
||||||
|
|
||||||
|
### `RateLimiter`
|
||||||
|
`type RateLimiter struct`
|
||||||
|
|
||||||
|
`RateLimiter` is the package’s main concurrency-safe limiter. It stores quotas, tracks usage state per model, supports YAML or SQLite persistence, and prunes expired state as part of normal operations.
|
||||||
|
|
||||||
|
- `Quotas map[string]ModelQuota`: configured per-model limits. If a model has no quota entry, `CanSend` allows it.
|
||||||
|
- `State map[string]*UsageStats`: tracked usage windows keyed by model name.
|
||||||
|
|
||||||
|
### `ModelStats`
|
||||||
|
`type ModelStats struct`
|
||||||
|
|
||||||
|
`ModelStats` is the read-only snapshot returned by `Stats`, `AllStats`, and `Iter`.
|
||||||
|
|
||||||
|
- `RPM int`: current requests counted in the one-minute window.
|
||||||
|
- `MaxRPM int`: configured requests-per-minute limit.
|
||||||
|
- `TPM int`: current tokens counted in the one-minute window.
|
||||||
|
- `MaxTPM int`: configured tokens-per-minute limit.
|
||||||
|
- `RPD int`: current requests counted in the rolling 24-hour window.
|
||||||
|
- `MaxRPD int`: configured requests-per-day limit.
|
||||||
|
- `DayStart time.Time`: start of the current rolling 24-hour window. This is zero if the model has no recorded state.
|
||||||
|
|
||||||
|
### `DecisionCode`
|
||||||
|
`type DecisionCode string`
|
||||||
|
|
||||||
|
`DecisionCode` enumerates machine-readable allow/deny codes returned by `Decide`. Defined values: `ok`, `unknown_model`, `unlimited`, `invalid_tokens`, `rpd_exceeded`, `rpm_exceeded`, and `tpm_exceeded`.
|
||||||
|
|
||||||
|
### `Decision`
|
||||||
|
`type Decision struct`
|
||||||
|
|
||||||
|
`Decision` bundles the outcome from `Decide`, including whether the request is allowed, a `DecisionCode`, a human-readable `Reason`, an optional `RetryAfter` duration when throttled, and a `ModelStats` snapshot at the time of evaluation.
|
||||||
|
|
||||||
|
## Functions
|
||||||
|
|
||||||
|
### `DefaultProfiles() map[Provider]ProviderProfile`
|
||||||
|
Returns a fresh map of built-in quota profiles for the supported providers. The returned map currently contains Gemini, OpenAI, Anthropic, and Local profiles. Because a new map is built on each call, callers can modify the result without mutating shared package state.
|
||||||
|
|
||||||
|
### `New() (*RateLimiter, error)`
|
||||||
|
Creates a new YAML-backed `RateLimiter` with Gemini defaults. This is equivalent to calling `NewWithConfig(Config{Providers: []Provider{ProviderGemini}})`. It initialises in-memory state only; it does not automatically restore persisted data, so callers that want previous state must call `Load()`.
|
||||||
|
|
||||||
|
### `NewWithConfig(cfg Config) (*RateLimiter, error)`
|
||||||
|
Creates a `RateLimiter` from explicit configuration. If `cfg.Backend` is empty it uses the YAML backend for backward compatibility. If both `cfg.Providers` and `cfg.Quotas` are empty, Gemini defaults are loaded. When `cfg.FilePath` is empty, the constructor resolves a default path under `~/.core`; for the implicit SQLite path it also ensures the parent directory exists. Like `New`, it does not call `Load()` automatically.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) SetQuota(model string, quota ModelQuota)`
|
||||||
|
Adds or replaces the quota for `model` in memory. This change affects later `CanSend`, `Stats`, and related calls immediately, but it is not persisted until `Persist()` is called.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) AddProvider(provider Provider)`
|
||||||
|
Loads the built-in quota profile for `provider` and copies its model quotas into `rl.Quotas`. Any existing quota entries for matching model names are overwritten. Unknown provider values are ignored.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) Load() error`
|
||||||
|
Loads persisted state into the limiter. For the YAML backend, it reads the configured file and unmarshals the stored quotas and state; a missing file is treated as an empty state and returns `nil`. For the SQLite backend, it loads persisted quotas and usage state from the database. If the database has stored quotas, those quotas replace the in-memory configuration; if no stored quotas exist, the current in-memory quotas are retained. In both cases, the loaded usage state replaces the current in-memory state.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) Persist() error`
|
||||||
|
Writes the current quotas and usage state to the configured backend. The method clones the in-memory snapshot while holding the lock, then performs I/O after releasing it. YAML persistence serialises the quota and state maps into the state file. SQLite persistence writes a full snapshot transactionally so quotas and usage move together.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) BackgroundPrune(interval time.Duration) func()`
|
||||||
|
Starts a background goroutine that prunes expired entries from every tracked model on the supplied interval and returns a stop function. If `interval <= 0`, it returns a no-op stop function and does not start a goroutine.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) CanSend(model string, estimatedTokens int) bool`
|
||||||
|
Reports whether a request for `model` can be sent without violating the configured limits. Negative token estimates are rejected. Models with no configured quota are allowed. If all three limits for a known model are `0`, the model is treated as unlimited. Before evaluating the request, the limiter prunes entries older than one minute and resets the rolling daily counter when its 24-hour window has elapsed. The method then checks requests-per-day, requests-per-minute, and tokens-per-minute against the estimated token count.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) Decide(model string, estimatedTokens int) Decision`
|
||||||
|
Returns a structured allow/deny decision for the estimated request. The result includes a `DecisionCode`, a human-readable `Reason`, optional `RetryAfter` guidance when throttled, and a `ModelStats` snapshot. It prunes expired state, initialises empty state for configured models, but does not record usage.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) RecordUsage(model string, promptTokens, outputTokens int)`
|
||||||
|
Records a successful request for `model`. The limiter prunes stale entries first, creates state for the model if needed, appends the current timestamp to the request window, appends a token entry containing the combined prompt and output token count, and increments the rolling daily counter. Negative token values are ignored by the internal token summation logic rather than reducing the recorded total.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) WaitForCapacity(ctx context.Context, model string, tokens int) error`
|
||||||
|
Blocks until `Decide(model, tokens)` allows the request or `ctx` is cancelled. The method uses the `RetryAfter` hint from `Decide` to sleep between checks, falling back to one-second polling when no hint is available. If `tokens` is negative, it returns an error immediately.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) Reset(model string)`
|
||||||
|
Clears usage state without changing quotas. If `model` is empty, it drops all tracked state. Otherwise it removes state only for the named model.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) Models() iter.Seq[string]`
|
||||||
|
Returns a sorted iterator of all model names currently known to the limiter. The result is the union of model names present in `rl.Quotas` and `rl.State`, so it includes models that only have stored state as well as models that only have configured quotas.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) Iter() iter.Seq2[string, ModelStats]`
|
||||||
|
Returns a sorted iterator of model names paired with their current `ModelStats` snapshots. Internally it builds the snapshot via `AllStats()` and yields entries in lexical model-name order.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) Stats(model string) ModelStats`
|
||||||
|
Returns the current snapshot for a single model after pruning expired entries. The result includes both current usage and configured maxima. If the model has no configured quota, the maximum fields are zero. If the model has no recorded state, the usage counters are zero and `DayStart` is the zero time.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) AllStats() map[string]ModelStats`
|
||||||
|
Returns a snapshot for every tracked model. The returned map includes model names found in either `rl.Quotas` or `rl.State`. Each model is pruned before its snapshot is computed, so expired one-minute entries are removed and stale daily windows are reset as part of the call.
|
||||||
|
|
||||||
|
### `NewWithSQLite(dbPath string) (*RateLimiter, error)`
|
||||||
|
Creates a SQLite-backed `RateLimiter` with Gemini defaults and opens or creates the database at `dbPath`. Like the YAML constructors, it initialises in-memory configuration but does not automatically call `Load()`. Callers should `defer rl.Close()` when they are done with the limiter.
|
||||||
|
|
||||||
|
### `NewWithSQLiteConfig(dbPath string, cfg Config) (*RateLimiter, error)`
|
||||||
|
Creates a SQLite-backed `RateLimiter` using `cfg` for provider and quota configuration. The `Backend` field in `cfg` is ignored because this constructor always uses SQLite. The database is opened or created at `dbPath`, and callers should `defer rl.Close()` to release the connection. Existing persisted data is not loaded until `Load()` is called.
|
||||||
|
|
||||||
|
### `func (rl *RateLimiter) Close() error`
|
||||||
|
Releases resources held by the limiter. For YAML-backed limiters this is a no-op that returns `nil`. For SQLite-backed limiters it closes the underlying database connection.
|
||||||
|
|
||||||
|
### `MigrateYAMLToSQLite(yamlPath, sqlitePath string) error`
|
||||||
|
Reads a YAML state file into a temporary `RateLimiter` and writes its quotas and usage state into a SQLite database. The SQLite database is created if it does not exist. The migration writes a complete snapshot, so any existing SQLite snapshot tables are replaced by the imported data.
|
||||||
|
|
||||||
|
### `CountTokens(ctx context.Context, apiKey, model, text string) (int, error)`
|
||||||
|
Calls Google’s Gemini `countTokens` API for `model` and returns the `totalTokens` value from the response. The function uses `http.DefaultClient`, posts to the Generative Language API base URL, and sends the API key through the `x-goog-api-key` header. It validates that `model` is non-empty, truncates oversized response bodies when building error messages, and wraps transport, request-building, and decoding failures with package-scoped errors.
|
||||||
169
sqlite.go
169
sqlite.go
|
|
@ -1,10 +1,12 @@
|
||||||
|
// SPDX-License-Identifier: EUPL-1.2
|
||||||
|
|
||||||
package ratelimit
|
package ratelimit
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"database/sql"
|
"database/sql"
|
||||||
"fmt"
|
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
core "dappco.re/go/core"
|
||||||
_ "modernc.org/sqlite"
|
_ "modernc.org/sqlite"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -19,7 +21,7 @@ type sqliteStore struct {
|
||||||
func newSQLiteStore(dbPath string) (*sqliteStore, error) {
|
func newSQLiteStore(dbPath string) (*sqliteStore, error) {
|
||||||
db, err := sql.Open("sqlite", dbPath)
|
db, err := sql.Open("sqlite", dbPath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.newSQLiteStore: open: %w", err)
|
return nil, core.E("ratelimit.newSQLiteStore", "open", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Single connection for PRAGMA consistency.
|
// Single connection for PRAGMA consistency.
|
||||||
|
|
@ -27,11 +29,11 @@ func newSQLiteStore(dbPath string) (*sqliteStore, error) {
|
||||||
|
|
||||||
if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil {
|
if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil {
|
||||||
db.Close()
|
db.Close()
|
||||||
return nil, fmt.Errorf("ratelimit.newSQLiteStore: WAL: %w", err)
|
return nil, core.E("ratelimit.newSQLiteStore", "WAL", err)
|
||||||
}
|
}
|
||||||
if _, err := db.Exec("PRAGMA busy_timeout=5000"); err != nil {
|
if _, err := db.Exec("PRAGMA busy_timeout=5000"); err != nil {
|
||||||
db.Close()
|
db.Close()
|
||||||
return nil, fmt.Errorf("ratelimit.newSQLiteStore: busy_timeout: %w", err)
|
return nil, core.E("ratelimit.newSQLiteStore", "busy_timeout", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := createSchema(db); err != nil {
|
if err := createSchema(db); err != nil {
|
||||||
|
|
@ -71,45 +73,36 @@ func createSchema(db *sql.DB) error {
|
||||||
|
|
||||||
for _, stmt := range stmts {
|
for _, stmt := range stmts {
|
||||||
if _, err := db.Exec(stmt); err != nil {
|
if _, err := db.Exec(stmt); err != nil {
|
||||||
return fmt.Errorf("ratelimit.createSchema: %w", err)
|
return core.E("ratelimit.createSchema", "exec", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// saveQuotas upserts all quotas into the quotas table.
|
// saveQuotas writes a complete quota snapshot to the quotas table.
|
||||||
func (s *sqliteStore) saveQuotas(quotas map[string]ModelQuota) error {
|
func (s *sqliteStore) saveQuotas(quotas map[string]ModelQuota) error {
|
||||||
tx, err := s.db.Begin()
|
tx, err := s.db.Begin()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("ratelimit.saveQuotas: begin: %w", err)
|
return core.E("ratelimit.saveQuotas", "begin", err)
|
||||||
}
|
}
|
||||||
defer tx.Rollback()
|
defer tx.Rollback()
|
||||||
|
|
||||||
stmt, err := tx.Prepare(`INSERT INTO quotas (model, max_rpm, max_tpm, max_rpd)
|
if _, err := tx.Exec("DELETE FROM quotas"); err != nil {
|
||||||
VALUES (?, ?, ?, ?)
|
return core.E("ratelimit.saveQuotas", "clear", err)
|
||||||
ON CONFLICT(model) DO UPDATE SET
|
|
||||||
max_rpm = excluded.max_rpm,
|
|
||||||
max_tpm = excluded.max_tpm,
|
|
||||||
max_rpd = excluded.max_rpd`)
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("ratelimit.saveQuotas: prepare: %w", err)
|
|
||||||
}
|
|
||||||
defer stmt.Close()
|
|
||||||
|
|
||||||
for model, q := range quotas {
|
|
||||||
if _, err := stmt.Exec(model, q.MaxRPM, q.MaxTPM, q.MaxRPD); err != nil {
|
|
||||||
return fmt.Errorf("ratelimit.saveQuotas: exec %s: %w", model, err)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return tx.Commit()
|
if err := insertQuotas(tx, quotas); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return commitTx(tx, "ratelimit.saveQuotas")
|
||||||
}
|
}
|
||||||
|
|
||||||
// loadQuotas reads all rows from the quotas table.
|
// loadQuotas reads all rows from the quotas table.
|
||||||
func (s *sqliteStore) loadQuotas() (map[string]ModelQuota, error) {
|
func (s *sqliteStore) loadQuotas() (map[string]ModelQuota, error) {
|
||||||
rows, err := s.db.Query("SELECT model, max_rpm, max_tpm, max_rpd FROM quotas")
|
rows, err := s.db.Query("SELECT model, max_rpm, max_tpm, max_rpd FROM quotas")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadQuotas: query: %w", err)
|
return nil, core.E("ratelimit.loadQuotas", "query", err)
|
||||||
}
|
}
|
||||||
defer rows.Close()
|
defer rows.Close()
|
||||||
|
|
||||||
|
|
@ -118,71 +111,137 @@ func (s *sqliteStore) loadQuotas() (map[string]ModelQuota, error) {
|
||||||
var model string
|
var model string
|
||||||
var q ModelQuota
|
var q ModelQuota
|
||||||
if err := rows.Scan(&model, &q.MaxRPM, &q.MaxTPM, &q.MaxRPD); err != nil {
|
if err := rows.Scan(&model, &q.MaxRPM, &q.MaxTPM, &q.MaxRPD); err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadQuotas: scan: %w", err)
|
return nil, core.E("ratelimit.loadQuotas", "scan", err)
|
||||||
}
|
}
|
||||||
result[model] = q
|
result[model] = q
|
||||||
}
|
}
|
||||||
if err := rows.Err(); err != nil {
|
if err := rows.Err(); err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadQuotas: rows: %w", err)
|
return nil, core.E("ratelimit.loadQuotas", "rows", err)
|
||||||
}
|
}
|
||||||
return result, nil
|
return result, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// saveState writes all usage state to SQLite in a single transaction.
|
// saveSnapshot writes quotas and state as a single atomic snapshot.
|
||||||
// It deletes existing rows and inserts fresh data for each model.
|
func (s *sqliteStore) saveSnapshot(quotas map[string]ModelQuota, state map[string]*UsageStats) error {
|
||||||
func (s *sqliteStore) saveState(state map[string]*UsageStats) error {
|
|
||||||
tx, err := s.db.Begin()
|
tx, err := s.db.Begin()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("ratelimit.saveState: begin: %w", err)
|
return core.E("ratelimit.saveSnapshot", "begin", err)
|
||||||
}
|
}
|
||||||
defer tx.Rollback()
|
defer tx.Rollback()
|
||||||
|
|
||||||
// Clear existing state.
|
if err := clearSnapshotTables(tx, true); err != nil {
|
||||||
if _, err := tx.Exec("DELETE FROM requests"); err != nil {
|
return err
|
||||||
return fmt.Errorf("ratelimit.saveState: clear requests: %w", err)
|
|
||||||
}
|
|
||||||
if _, err := tx.Exec("DELETE FROM tokens"); err != nil {
|
|
||||||
return fmt.Errorf("ratelimit.saveState: clear tokens: %w", err)
|
|
||||||
}
|
|
||||||
if _, err := tx.Exec("DELETE FROM daily"); err != nil {
|
|
||||||
return fmt.Errorf("ratelimit.saveState: clear daily: %w", err)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if err := insertQuotas(tx, quotas); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := insertState(tx, state); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return commitTx(tx, "ratelimit.saveSnapshot")
|
||||||
|
}
|
||||||
|
|
||||||
|
// saveState writes all usage state to SQLite in a single transaction.
|
||||||
|
// It uses a truncate-and-insert approach for simplicity in this version,
|
||||||
|
// but ensures atomicity via a single transaction.
|
||||||
|
func (s *sqliteStore) saveState(state map[string]*UsageStats) error {
|
||||||
|
tx, err := s.db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
return core.E("ratelimit.saveState", "begin", err)
|
||||||
|
}
|
||||||
|
defer tx.Rollback()
|
||||||
|
|
||||||
|
if err := clearSnapshotTables(tx, false); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := insertState(tx, state); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
return commitTx(tx, "ratelimit.saveState")
|
||||||
|
}
|
||||||
|
|
||||||
|
func clearSnapshotTables(tx *sql.Tx, includeQuotas bool) error {
|
||||||
|
if includeQuotas {
|
||||||
|
if _, err := tx.Exec("DELETE FROM quotas"); err != nil {
|
||||||
|
return core.E("ratelimit.saveSnapshot", "clear quotas", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if _, err := tx.Exec("DELETE FROM requests"); err != nil {
|
||||||
|
return core.E("ratelimit.saveState", "clear requests", err)
|
||||||
|
}
|
||||||
|
if _, err := tx.Exec("DELETE FROM tokens"); err != nil {
|
||||||
|
return core.E("ratelimit.saveState", "clear tokens", err)
|
||||||
|
}
|
||||||
|
if _, err := tx.Exec("DELETE FROM daily"); err != nil {
|
||||||
|
return core.E("ratelimit.saveState", "clear daily", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func insertQuotas(tx *sql.Tx, quotas map[string]ModelQuota) error {
|
||||||
|
stmt, err := tx.Prepare("INSERT INTO quotas (model, max_rpm, max_tpm, max_rpd) VALUES (?, ?, ?, ?)")
|
||||||
|
if err != nil {
|
||||||
|
return core.E("ratelimit.saveQuotas", "prepare", err)
|
||||||
|
}
|
||||||
|
defer stmt.Close()
|
||||||
|
|
||||||
|
for model, q := range quotas {
|
||||||
|
if _, err := stmt.Exec(model, q.MaxRPM, q.MaxTPM, q.MaxRPD); err != nil {
|
||||||
|
return core.E("ratelimit.saveQuotas", core.Concat("exec ", model), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func insertState(tx *sql.Tx, state map[string]*UsageStats) error {
|
||||||
reqStmt, err := tx.Prepare("INSERT INTO requests (model, ts) VALUES (?, ?)")
|
reqStmt, err := tx.Prepare("INSERT INTO requests (model, ts) VALUES (?, ?)")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("ratelimit.saveState: prepare requests: %w", err)
|
return core.E("ratelimit.saveState", "prepare requests", err)
|
||||||
}
|
}
|
||||||
defer reqStmt.Close()
|
defer reqStmt.Close()
|
||||||
|
|
||||||
tokStmt, err := tx.Prepare("INSERT INTO tokens (model, ts, count) VALUES (?, ?, ?)")
|
tokStmt, err := tx.Prepare("INSERT INTO tokens (model, ts, count) VALUES (?, ?, ?)")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("ratelimit.saveState: prepare tokens: %w", err)
|
return core.E("ratelimit.saveState", "prepare tokens", err)
|
||||||
}
|
}
|
||||||
defer tokStmt.Close()
|
defer tokStmt.Close()
|
||||||
|
|
||||||
dayStmt, err := tx.Prepare("INSERT INTO daily (model, day_start, day_count) VALUES (?, ?, ?)")
|
dayStmt, err := tx.Prepare("INSERT INTO daily (model, day_start, day_count) VALUES (?, ?, ?)")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("ratelimit.saveState: prepare daily: %w", err)
|
return core.E("ratelimit.saveState", "prepare daily", err)
|
||||||
}
|
}
|
||||||
defer dayStmt.Close()
|
defer dayStmt.Close()
|
||||||
|
|
||||||
for model, stats := range state {
|
for model, stats := range state {
|
||||||
|
if stats == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
for _, t := range stats.Requests {
|
for _, t := range stats.Requests {
|
||||||
if _, err := reqStmt.Exec(model, t.UnixNano()); err != nil {
|
if _, err := reqStmt.Exec(model, t.UnixNano()); err != nil {
|
||||||
return fmt.Errorf("ratelimit.saveState: insert request %s: %w", model, err)
|
return core.E("ratelimit.saveState", core.Concat("insert request ", model), err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
for _, te := range stats.Tokens {
|
for _, te := range stats.Tokens {
|
||||||
if _, err := tokStmt.Exec(model, te.Time.UnixNano(), te.Count); err != nil {
|
if _, err := tokStmt.Exec(model, te.Time.UnixNano(), te.Count); err != nil {
|
||||||
return fmt.Errorf("ratelimit.saveState: insert token %s: %w", model, err)
|
return core.E("ratelimit.saveState", core.Concat("insert token ", model), err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if _, err := dayStmt.Exec(model, stats.DayStart.UnixNano(), stats.DayCount); err != nil {
|
if _, err := dayStmt.Exec(model, stats.DayStart.UnixNano(), stats.DayCount); err != nil {
|
||||||
return fmt.Errorf("ratelimit.saveState: insert daily %s: %w", model, err)
|
return core.E("ratelimit.saveState", core.Concat("insert daily ", model), err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
return tx.Commit()
|
func commitTx(tx *sql.Tx, scope string) error {
|
||||||
|
if err := tx.Commit(); err != nil {
|
||||||
|
return core.E(scope, "commit", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// loadState reconstructs the UsageStats map from SQLite tables.
|
// loadState reconstructs the UsageStats map from SQLite tables.
|
||||||
|
|
@ -192,7 +251,7 @@ func (s *sqliteStore) loadState() (map[string]*UsageStats, error) {
|
||||||
// Load daily counters first (these define which models have state).
|
// Load daily counters first (these define which models have state).
|
||||||
rows, err := s.db.Query("SELECT model, day_start, day_count FROM daily")
|
rows, err := s.db.Query("SELECT model, day_start, day_count FROM daily")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: query daily: %w", err)
|
return nil, core.E("ratelimit.loadState", "query daily", err)
|
||||||
}
|
}
|
||||||
defer rows.Close()
|
defer rows.Close()
|
||||||
|
|
||||||
|
|
@ -201,7 +260,7 @@ func (s *sqliteStore) loadState() (map[string]*UsageStats, error) {
|
||||||
var dayStartNano int64
|
var dayStartNano int64
|
||||||
var dayCount int
|
var dayCount int
|
||||||
if err := rows.Scan(&model, &dayStartNano, &dayCount); err != nil {
|
if err := rows.Scan(&model, &dayStartNano, &dayCount); err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: scan daily: %w", err)
|
return nil, core.E("ratelimit.loadState", "scan daily", err)
|
||||||
}
|
}
|
||||||
result[model] = &UsageStats{
|
result[model] = &UsageStats{
|
||||||
DayStart: time.Unix(0, dayStartNano),
|
DayStart: time.Unix(0, dayStartNano),
|
||||||
|
|
@ -209,13 +268,13 @@ func (s *sqliteStore) loadState() (map[string]*UsageStats, error) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if err := rows.Err(); err != nil {
|
if err := rows.Err(); err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: daily rows: %w", err)
|
return nil, core.E("ratelimit.loadState", "daily rows", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load requests.
|
// Load requests.
|
||||||
reqRows, err := s.db.Query("SELECT model, ts FROM requests ORDER BY ts")
|
reqRows, err := s.db.Query("SELECT model, ts FROM requests ORDER BY ts")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: query requests: %w", err)
|
return nil, core.E("ratelimit.loadState", "query requests", err)
|
||||||
}
|
}
|
||||||
defer reqRows.Close()
|
defer reqRows.Close()
|
||||||
|
|
||||||
|
|
@ -223,7 +282,7 @@ func (s *sqliteStore) loadState() (map[string]*UsageStats, error) {
|
||||||
var model string
|
var model string
|
||||||
var tsNano int64
|
var tsNano int64
|
||||||
if err := reqRows.Scan(&model, &tsNano); err != nil {
|
if err := reqRows.Scan(&model, &tsNano); err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: scan requests: %w", err)
|
return nil, core.E("ratelimit.loadState", "scan requests", err)
|
||||||
}
|
}
|
||||||
if _, ok := result[model]; !ok {
|
if _, ok := result[model]; !ok {
|
||||||
result[model] = &UsageStats{}
|
result[model] = &UsageStats{}
|
||||||
|
|
@ -231,13 +290,13 @@ func (s *sqliteStore) loadState() (map[string]*UsageStats, error) {
|
||||||
result[model].Requests = append(result[model].Requests, time.Unix(0, tsNano))
|
result[model].Requests = append(result[model].Requests, time.Unix(0, tsNano))
|
||||||
}
|
}
|
||||||
if err := reqRows.Err(); err != nil {
|
if err := reqRows.Err(); err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: request rows: %w", err)
|
return nil, core.E("ratelimit.loadState", "request rows", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Load tokens.
|
// Load tokens.
|
||||||
tokRows, err := s.db.Query("SELECT model, ts, count FROM tokens ORDER BY ts")
|
tokRows, err := s.db.Query("SELECT model, ts, count FROM tokens ORDER BY ts")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: query tokens: %w", err)
|
return nil, core.E("ratelimit.loadState", "query tokens", err)
|
||||||
}
|
}
|
||||||
defer tokRows.Close()
|
defer tokRows.Close()
|
||||||
|
|
||||||
|
|
@ -246,7 +305,7 @@ func (s *sqliteStore) loadState() (map[string]*UsageStats, error) {
|
||||||
var tsNano int64
|
var tsNano int64
|
||||||
var count int
|
var count int
|
||||||
if err := tokRows.Scan(&model, &tsNano, &count); err != nil {
|
if err := tokRows.Scan(&model, &tsNano, &count); err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: scan tokens: %w", err)
|
return nil, core.E("ratelimit.loadState", "scan tokens", err)
|
||||||
}
|
}
|
||||||
if _, ok := result[model]; !ok {
|
if _, ok := result[model]; !ok {
|
||||||
result[model] = &UsageStats{}
|
result[model] = &UsageStats{}
|
||||||
|
|
@ -257,7 +316,7 @@ func (s *sqliteStore) loadState() (map[string]*UsageStats, error) {
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
if err := tokRows.Err(); err != nil {
|
if err := tokRows.Err(); err != nil {
|
||||||
return nil, fmt.Errorf("ratelimit.loadState: token rows: %w", err)
|
return nil, core.E("ratelimit.loadState", "token rows", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
return result, nil
|
return result, nil
|
||||||
|
|
|
||||||
303
sqlite_test.go
303
sqlite_test.go
|
|
@ -1,8 +1,8 @@
|
||||||
|
// SPDX-License-Identifier: EUPL-1.2
|
||||||
|
|
||||||
package ratelimit
|
package ratelimit
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"os"
|
|
||||||
"path/filepath"
|
|
||||||
"sync"
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
@ -14,18 +14,17 @@ import (
|
||||||
|
|
||||||
// --- Phase 2: SQLite basic tests ---
|
// --- Phase 2: SQLite basic tests ---
|
||||||
|
|
||||||
func TestNewSQLiteStore_Good(t *testing.T) {
|
func TestSQLite_NewSQLiteStore_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "test.db")
|
dbPath := testPath(t.TempDir(), "test.db")
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer store.close()
|
defer store.close()
|
||||||
|
|
||||||
// Verify the database file was created.
|
// Verify the database file was created.
|
||||||
_, statErr := os.Stat(dbPath)
|
assert.True(t, pathExists(dbPath), "database file should exist")
|
||||||
assert.NoError(t, statErr, "database file should exist")
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestNewSQLiteStore_Bad(t *testing.T) {
|
func TestSQLite_NewSQLiteStore_Bad(t *testing.T) {
|
||||||
t.Run("invalid path returns error", func(t *testing.T) {
|
t.Run("invalid path returns error", func(t *testing.T) {
|
||||||
// Path inside a non-existent directory with no parent.
|
// Path inside a non-existent directory with no parent.
|
||||||
_, err := newSQLiteStore("/nonexistent/deep/nested/dir/test.db")
|
_, err := newSQLiteStore("/nonexistent/deep/nested/dir/test.db")
|
||||||
|
|
@ -33,8 +32,8 @@ func TestNewSQLiteStore_Bad(t *testing.T) {
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteQuotasRoundTrip_Good(t *testing.T) {
|
func TestSQLite_QuotasRoundTrip_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "quotas.db")
|
dbPath := testPath(t.TempDir(), "quotas.db")
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer store.close()
|
defer store.close()
|
||||||
|
|
@ -60,8 +59,8 @@ func TestSQLiteQuotasRoundTrip_Good(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteQuotasUpsert_Good(t *testing.T) {
|
func TestSQLite_QuotasOverwrite_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "upsert.db")
|
dbPath := testPath(t.TempDir(), "overwrite.db")
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer store.close()
|
defer store.close()
|
||||||
|
|
@ -71,7 +70,7 @@ func TestSQLiteQuotasUpsert_Good(t *testing.T) {
|
||||||
"model-a": {MaxRPM: 100, MaxTPM: 50000, MaxRPD: 1000},
|
"model-a": {MaxRPM: 100, MaxTPM: 50000, MaxRPD: 1000},
|
||||||
}))
|
}))
|
||||||
|
|
||||||
// Upsert with updated values.
|
// Save a second snapshot with updated values.
|
||||||
require.NoError(t, store.saveQuotas(map[string]ModelQuota{
|
require.NoError(t, store.saveQuotas(map[string]ModelQuota{
|
||||||
"model-a": {MaxRPM: 999, MaxTPM: 888, MaxRPD: 777},
|
"model-a": {MaxRPM: 999, MaxTPM: 888, MaxRPD: 777},
|
||||||
}))
|
}))
|
||||||
|
|
@ -85,8 +84,8 @@ func TestSQLiteQuotasUpsert_Good(t *testing.T) {
|
||||||
assert.Equal(t, 777, q.MaxRPD, "should have updated RPD")
|
assert.Equal(t, 777, q.MaxRPD, "should have updated RPD")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteStateRoundTrip_Good(t *testing.T) {
|
func TestSQLite_StateRoundTrip_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "state.db")
|
dbPath := testPath(t.TempDir(), "state.db")
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer store.close()
|
defer store.close()
|
||||||
|
|
@ -144,8 +143,8 @@ func TestSQLiteStateRoundTrip_Good(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteStateOverwrite_Good(t *testing.T) {
|
func TestSQLite_StateOverwrite_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "overwrite.db")
|
dbPath := testPath(t.TempDir(), "overwrite.db")
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer store.close()
|
defer store.close()
|
||||||
|
|
@ -182,8 +181,8 @@ func TestSQLiteStateOverwrite_Good(t *testing.T) {
|
||||||
assert.Len(t, b.Requests, 1)
|
assert.Len(t, b.Requests, 1)
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteEmptyState_Good(t *testing.T) {
|
func TestSQLite_EmptyState_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "empty.db")
|
dbPath := testPath(t.TempDir(), "empty.db")
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer store.close()
|
defer store.close()
|
||||||
|
|
@ -198,8 +197,8 @@ func TestSQLiteEmptyState_Good(t *testing.T) {
|
||||||
assert.Empty(t, state, "should return empty state from fresh DB")
|
assert.Empty(t, state, "should return empty state from fresh DB")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteClose_Good(t *testing.T) {
|
func TestSQLite_Close_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "close.db")
|
dbPath := testPath(t.TempDir(), "close.db")
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
|
@ -208,8 +207,8 @@ func TestSQLiteClose_Good(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 2: SQLite integration tests ---
|
// --- Phase 2: SQLite integration tests ---
|
||||||
|
|
||||||
func TestNewWithSQLite_Good(t *testing.T) {
|
func TestSQLite_NewWithSQLite_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "limiter.db")
|
dbPath := testPath(t.TempDir(), "limiter.db")
|
||||||
rl, err := NewWithSQLite(dbPath)
|
rl, err := NewWithSQLite(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer rl.Close()
|
defer rl.Close()
|
||||||
|
|
@ -222,8 +221,8 @@ func TestNewWithSQLite_Good(t *testing.T) {
|
||||||
assert.NotNil(t, rl.sqlite, "SQLite store should be initialised")
|
assert.NotNil(t, rl.sqlite, "SQLite store should be initialised")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestNewWithSQLiteConfig_Good(t *testing.T) {
|
func TestSQLite_NewWithSQLiteConfig_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "config.db")
|
dbPath := testPath(t.TempDir(), "config.db")
|
||||||
rl, err := NewWithSQLiteConfig(dbPath, Config{
|
rl, err := NewWithSQLiteConfig(dbPath, Config{
|
||||||
Providers: []Provider{ProviderAnthropic},
|
Providers: []Provider{ProviderAnthropic},
|
||||||
Quotas: map[string]ModelQuota{
|
Quotas: map[string]ModelQuota{
|
||||||
|
|
@ -243,8 +242,8 @@ func TestNewWithSQLiteConfig_Good(t *testing.T) {
|
||||||
assert.False(t, hasGemini, "should not have Gemini models")
|
assert.False(t, hasGemini, "should not have Gemini models")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLitePersistAndLoad_Good(t *testing.T) {
|
func TestSQLite_PersistAndLoad_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "persist.db")
|
dbPath := testPath(t.TempDir(), "persist.db")
|
||||||
rl, err := NewWithSQLite(dbPath)
|
rl, err := NewWithSQLite(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
|
@ -272,8 +271,8 @@ func TestSQLitePersistAndLoad_Good(t *testing.T) {
|
||||||
assert.Equal(t, 500, stats.MaxRPD)
|
assert.Equal(t, 500, stats.MaxRPD)
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLitePersistMultipleModels_Good(t *testing.T) {
|
func TestSQLite_PersistMultipleModels_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "multi.db")
|
dbPath := testPath(t.TempDir(), "multi.db")
|
||||||
rl, err := NewWithSQLiteConfig(dbPath, Config{
|
rl, err := NewWithSQLiteConfig(dbPath, Config{
|
||||||
Providers: []Provider{ProviderGemini, ProviderAnthropic},
|
Providers: []Provider{ProviderGemini, ProviderAnthropic},
|
||||||
})
|
})
|
||||||
|
|
@ -302,8 +301,8 @@ func TestSQLitePersistMultipleModels_Good(t *testing.T) {
|
||||||
assert.Equal(t, 400, claude.TPM)
|
assert.Equal(t, 400, claude.TPM)
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteRecordUsageThenPersistReload_Good(t *testing.T) {
|
func TestSQLite_RecordUsageThenPersistReload_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "record.db")
|
dbPath := testPath(t.TempDir(), "record.db")
|
||||||
rl, err := NewWithSQLite(dbPath)
|
rl, err := NewWithSQLite(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
|
@ -340,7 +339,7 @@ func TestSQLiteRecordUsageThenPersistReload_Good(t *testing.T) {
|
||||||
assert.Equal(t, 1000, stats2.TPM, "TPM should survive reload")
|
assert.Equal(t, 1000, stats2.TPM, "TPM should survive reload")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteClose_Good_NoOp(t *testing.T) {
|
func TestSQLite_CloseNoOp_Good(t *testing.T) {
|
||||||
// Close on YAML-backed limiter is a no-op.
|
// Close on YAML-backed limiter is a no-op.
|
||||||
rl := newTestLimiter(t)
|
rl := newTestLimiter(t)
|
||||||
assert.NoError(t, rl.Close(), "Close on YAML limiter should be no-op")
|
assert.NoError(t, rl.Close(), "Close on YAML limiter should be no-op")
|
||||||
|
|
@ -348,8 +347,8 @@ func TestSQLiteClose_Good_NoOp(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 2: Concurrent SQLite ---
|
// --- Phase 2: Concurrent SQLite ---
|
||||||
|
|
||||||
func TestSQLiteConcurrent_Good(t *testing.T) {
|
func TestSQLite_Concurrent_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "concurrent.db")
|
dbPath := testPath(t.TempDir(), "concurrent.db")
|
||||||
rl, err := NewWithSQLite(dbPath)
|
rl, err := NewWithSQLite(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer rl.Close()
|
defer rl.Close()
|
||||||
|
|
@ -398,10 +397,10 @@ func TestSQLiteConcurrent_Good(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 2: YAML backward compatibility ---
|
// --- Phase 2: YAML backward compatibility ---
|
||||||
|
|
||||||
func TestYAMLBackwardCompat_Good(t *testing.T) {
|
func TestSQLite_YAMLBackwardCompat_Good(t *testing.T) {
|
||||||
// Verify that the default YAML backend still works after SQLite additions.
|
// Verify that the default YAML backend still works after SQLite additions.
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
path := filepath.Join(tmpDir, "compat.yaml")
|
path := testPath(tmpDir, "compat.yaml")
|
||||||
|
|
||||||
rl1, err := New()
|
rl1, err := New()
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
@ -425,22 +424,59 @@ func TestYAMLBackwardCompat_Good(t *testing.T) {
|
||||||
assert.Equal(t, 200, stats.TPM)
|
assert.Equal(t, 200, stats.TPM)
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestConfigBackendDefault_Good(t *testing.T) {
|
func TestSQLite_ConfigBackendDefault_Good(t *testing.T) {
|
||||||
// Empty Backend string should default to YAML behaviour.
|
// Empty Backend string should default to YAML behaviour.
|
||||||
rl, err := NewWithConfig(Config{
|
rl, err := NewWithConfig(Config{
|
||||||
FilePath: filepath.Join(t.TempDir(), "default.yaml"),
|
FilePath: testPath(t.TempDir(), "default.yaml"),
|
||||||
})
|
})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
||||||
assert.Nil(t, rl.sqlite, "empty backend should use YAML (no sqlite)")
|
assert.Nil(t, rl.sqlite, "empty backend should use YAML (no sqlite)")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestSQLite_ConfigBackendSQLite_Good(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "config-backend.db")
|
||||||
|
rl, err := NewWithConfig(Config{
|
||||||
|
Backend: backendSQLite,
|
||||||
|
FilePath: dbPath,
|
||||||
|
Quotas: map[string]ModelQuota{
|
||||||
|
"backend-model": {MaxRPM: 10, MaxTPM: 1000, MaxRPD: 50},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer rl.Close()
|
||||||
|
|
||||||
|
require.NotNil(t, rl.sqlite)
|
||||||
|
rl.RecordUsage("backend-model", 10, 10)
|
||||||
|
require.NoError(t, rl.Persist())
|
||||||
|
|
||||||
|
assert.True(t, pathExists(dbPath), "sqlite backend should persist to the configured DB path")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSQLite_ConfigBackendSQLiteDefaultPath_Good(t *testing.T) {
|
||||||
|
home := t.TempDir()
|
||||||
|
t.Setenv("HOME", home)
|
||||||
|
t.Setenv("USERPROFILE", "")
|
||||||
|
t.Setenv("home", "")
|
||||||
|
|
||||||
|
rl, err := NewWithConfig(Config{
|
||||||
|
Backend: backendSQLite,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer rl.Close()
|
||||||
|
|
||||||
|
require.NotNil(t, rl.sqlite)
|
||||||
|
require.NoError(t, rl.Persist())
|
||||||
|
|
||||||
|
assert.True(t, pathExists(testPath(home, defaultStateDirName, defaultSQLiteStateFile)), "sqlite backend should use the default home DB path")
|
||||||
|
}
|
||||||
|
|
||||||
// --- Phase 2: MigrateYAMLToSQLite ---
|
// --- Phase 2: MigrateYAMLToSQLite ---
|
||||||
|
|
||||||
func TestMigrateYAMLToSQLite_Good(t *testing.T) {
|
func TestSQLite_MigrateYAMLToSQLite_Good(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
yamlPath := filepath.Join(tmpDir, "state.yaml")
|
yamlPath := testPath(tmpDir, "state.yaml")
|
||||||
sqlitePath := filepath.Join(tmpDir, "migrated.db")
|
sqlitePath := testPath(tmpDir, "migrated.db")
|
||||||
|
|
||||||
// Create a YAML-backed limiter with state.
|
// Create a YAML-backed limiter with state.
|
||||||
rl, err := New()
|
rl, err := New()
|
||||||
|
|
@ -476,26 +512,89 @@ func TestMigrateYAMLToSQLite_Good(t *testing.T) {
|
||||||
assert.Equal(t, 2, stats.RPD, "should have 2 daily requests")
|
assert.Equal(t, 2, stats.RPD, "should have 2 daily requests")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestMigrateYAMLToSQLite_Bad(t *testing.T) {
|
func TestSQLite_MigrateYAMLToSQLite_Bad(t *testing.T) {
|
||||||
t.Run("non-existent YAML file", func(t *testing.T) {
|
t.Run("non-existent YAML file", func(t *testing.T) {
|
||||||
err := MigrateYAMLToSQLite("/nonexistent/state.yaml", filepath.Join(t.TempDir(), "out.db"))
|
err := MigrateYAMLToSQLite("/nonexistent/state.yaml", testPath(t.TempDir(), "out.db"))
|
||||||
assert.Error(t, err, "should fail with non-existent YAML file")
|
assert.Error(t, err, "should fail with non-existent YAML file")
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("corrupt YAML file", func(t *testing.T) {
|
t.Run("corrupt YAML file", func(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
yamlPath := filepath.Join(tmpDir, "corrupt.yaml")
|
yamlPath := testPath(tmpDir, "corrupt.yaml")
|
||||||
require.NoError(t, os.WriteFile(yamlPath, []byte("{{{{not yaml!"), 0644))
|
writeTestFile(t, yamlPath, "{{{{not yaml!")
|
||||||
|
|
||||||
err := MigrateYAMLToSQLite(yamlPath, filepath.Join(tmpDir, "out.db"))
|
err := MigrateYAMLToSQLite(yamlPath, testPath(tmpDir, "out.db"))
|
||||||
assert.Error(t, err, "should fail with corrupt YAML")
|
assert.Error(t, err, "should fail with corrupt YAML")
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestMigrateYAMLToSQLitePreservesAllGeminiModels_Good(t *testing.T) {
|
func TestSQLite_MigrateYAMLToSQLiteAtomic_Good(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
yamlPath := filepath.Join(tmpDir, "full.yaml")
|
yamlPath := testPath(tmpDir, "atomic.yaml")
|
||||||
sqlitePath := filepath.Join(tmpDir, "full.db")
|
sqlitePath := testPath(tmpDir, "atomic.db")
|
||||||
|
now := time.Now().UTC()
|
||||||
|
|
||||||
|
store, err := newSQLiteStore(sqlitePath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
originalQuotas := map[string]ModelQuota{
|
||||||
|
"old-model": {MaxRPM: 1, MaxTPM: 2, MaxRPD: 3},
|
||||||
|
}
|
||||||
|
originalState := map[string]*UsageStats{
|
||||||
|
"old-model": {
|
||||||
|
Requests: []time.Time{now},
|
||||||
|
Tokens: []TokenEntry{{Time: now, Count: 9}},
|
||||||
|
DayStart: now,
|
||||||
|
DayCount: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
require.NoError(t, store.saveSnapshot(originalQuotas, originalState))
|
||||||
|
|
||||||
|
_, err = store.db.Exec(`CREATE TRIGGER fail_daily_migrate BEFORE INSERT ON daily
|
||||||
|
BEGIN SELECT RAISE(ABORT, 'forced daily failure'); END`)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.NoError(t, store.close())
|
||||||
|
|
||||||
|
migrated := &RateLimiter{
|
||||||
|
Quotas: map[string]ModelQuota{
|
||||||
|
"new-model": {MaxRPM: 10, MaxTPM: 20, MaxRPD: 30},
|
||||||
|
},
|
||||||
|
State: map[string]*UsageStats{
|
||||||
|
"new-model": {
|
||||||
|
Requests: []time.Time{now.Add(5 * time.Second)},
|
||||||
|
Tokens: []TokenEntry{{Time: now.Add(5 * time.Second), Count: 99}},
|
||||||
|
DayStart: now.Add(5 * time.Second),
|
||||||
|
DayCount: 2,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
data, err := yaml.Marshal(migrated)
|
||||||
|
require.NoError(t, err)
|
||||||
|
writeTestFile(t, yamlPath, string(data))
|
||||||
|
|
||||||
|
err = MigrateYAMLToSQLite(yamlPath, sqlitePath)
|
||||||
|
require.Error(t, err)
|
||||||
|
|
||||||
|
store, err = newSQLiteStore(sqlitePath)
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer store.close()
|
||||||
|
|
||||||
|
quotas, err := store.loadQuotas()
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, originalQuotas, quotas)
|
||||||
|
|
||||||
|
state, err := store.loadState()
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Contains(t, state, "old-model")
|
||||||
|
assert.Equal(t, originalState["old-model"].DayCount, state["old-model"].DayCount)
|
||||||
|
assert.Equal(t, originalState["old-model"].Tokens[0].Count, state["old-model"].Tokens[0].Count)
|
||||||
|
assert.NotContains(t, state, "new-model")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSQLite_MigrateYAMLToSQLitePreservesAllGeminiModels_Good(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
yamlPath := testPath(tmpDir, "full.yaml")
|
||||||
|
sqlitePath := testPath(tmpDir, "full.db")
|
||||||
|
|
||||||
// Create a full YAML state with all Gemini models.
|
// Create a full YAML state with all Gemini models.
|
||||||
rl, err := New()
|
rl, err := New()
|
||||||
|
|
@ -524,12 +623,12 @@ func TestMigrateYAMLToSQLitePreservesAllGeminiModels_Good(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 2: Corrupt DB recovery ---
|
// --- Phase 2: Corrupt DB recovery ---
|
||||||
|
|
||||||
func TestSQLiteCorruptDB_Ugly(t *testing.T) {
|
func TestSQLite_CorruptDB_Ugly(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
dbPath := filepath.Join(tmpDir, "corrupt.db")
|
dbPath := testPath(tmpDir, "corrupt.db")
|
||||||
|
|
||||||
// Write garbage to the DB file.
|
// Write garbage to the DB file.
|
||||||
require.NoError(t, os.WriteFile(dbPath, []byte("THIS IS NOT A SQLITE DATABASE"), 0644))
|
writeTestFile(t, dbPath, "THIS IS NOT A SQLITE DATABASE")
|
||||||
|
|
||||||
// Opening a corrupt DB may succeed (sqlite is lazy about validation),
|
// Opening a corrupt DB may succeed (sqlite is lazy about validation),
|
||||||
// but operations on it should fail gracefully.
|
// but operations on it should fail gracefully.
|
||||||
|
|
@ -546,9 +645,9 @@ func TestSQLiteCorruptDB_Ugly(t *testing.T) {
|
||||||
assert.Error(t, err, "loading from corrupt DB should return an error")
|
assert.Error(t, err, "loading from corrupt DB should return an error")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteTruncatedDB_Ugly(t *testing.T) {
|
func TestSQLite_TruncatedDB_Ugly(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
dbPath := filepath.Join(tmpDir, "truncated.db")
|
dbPath := testPath(tmpDir, "truncated.db")
|
||||||
|
|
||||||
// Create a valid DB first.
|
// Create a valid DB first.
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
|
|
@ -559,11 +658,7 @@ func TestSQLiteTruncatedDB_Ugly(t *testing.T) {
|
||||||
require.NoError(t, store.close())
|
require.NoError(t, store.close())
|
||||||
|
|
||||||
// Truncate the file to simulate corruption.
|
// Truncate the file to simulate corruption.
|
||||||
f, err := os.OpenFile(dbPath, os.O_WRONLY|os.O_TRUNC, 0644)
|
overwriteTestFile(t, dbPath, "TRUNC")
|
||||||
require.NoError(t, err)
|
|
||||||
_, err = f.Write([]byte("TRUNC"))
|
|
||||||
require.NoError(t, err)
|
|
||||||
require.NoError(t, f.Close())
|
|
||||||
|
|
||||||
// Opening should either fail or operations should fail.
|
// Opening should either fail or operations should fail.
|
||||||
store2, err := newSQLiteStore(dbPath)
|
store2, err := newSQLiteStore(dbPath)
|
||||||
|
|
@ -577,9 +672,9 @@ func TestSQLiteTruncatedDB_Ugly(t *testing.T) {
|
||||||
assert.Error(t, err, "loading from truncated DB should return an error")
|
assert.Error(t, err, "loading from truncated DB should return an error")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestSQLiteEmptyModelState_Good(t *testing.T) {
|
func TestSQLite_EmptyModelState_Good(t *testing.T) {
|
||||||
// State with no requests or tokens but with a daily counter.
|
// State with no requests or tokens but with a daily counter.
|
||||||
dbPath := filepath.Join(t.TempDir(), "empty-state.db")
|
dbPath := testPath(t.TempDir(), "empty-state.db")
|
||||||
store, err := newSQLiteStore(dbPath)
|
store, err := newSQLiteStore(dbPath)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
defer store.close()
|
defer store.close()
|
||||||
|
|
@ -606,8 +701,8 @@ func TestSQLiteEmptyModelState_Good(t *testing.T) {
|
||||||
|
|
||||||
// --- Phase 2: End-to-end with persist cycle ---
|
// --- Phase 2: End-to-end with persist cycle ---
|
||||||
|
|
||||||
func TestSQLiteEndToEnd_Good(t *testing.T) {
|
func TestSQLite_EndToEnd_Good(t *testing.T) {
|
||||||
dbPath := filepath.Join(t.TempDir(), "e2e.db")
|
dbPath := testPath(t.TempDir(), "e2e.db")
|
||||||
|
|
||||||
// Session 1: Create limiter, record usage, persist.
|
// Session 1: Create limiter, record usage, persist.
|
||||||
rl1, err := NewWithSQLiteConfig(dbPath, Config{
|
rl1, err := NewWithSQLiteConfig(dbPath, Config{
|
||||||
|
|
@ -650,10 +745,82 @@ func TestSQLiteEndToEnd_Good(t *testing.T) {
|
||||||
assert.Equal(t, 5, custom.MaxRPM)
|
assert.Equal(t, 5, custom.MaxRPM)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestSQLite_LoadReplacesPersistedSnapshot_Good(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "replace.db")
|
||||||
|
rl, err := NewWithSQLiteConfig(dbPath, Config{
|
||||||
|
Quotas: map[string]ModelQuota{
|
||||||
|
"model-a": {MaxRPM: 1, MaxTPM: 100, MaxRPD: 10},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
rl.RecordUsage("model-a", 10, 10)
|
||||||
|
require.NoError(t, rl.Persist())
|
||||||
|
|
||||||
|
delete(rl.Quotas, "model-a")
|
||||||
|
rl.Quotas["model-b"] = ModelQuota{MaxRPM: 2, MaxTPM: 200, MaxRPD: 20}
|
||||||
|
rl.Reset("")
|
||||||
|
rl.RecordUsage("model-b", 5, 5)
|
||||||
|
require.NoError(t, rl.Persist())
|
||||||
|
require.NoError(t, rl.Close())
|
||||||
|
|
||||||
|
rl2, err := NewWithSQLiteConfig(dbPath, Config{
|
||||||
|
Providers: []Provider{ProviderGemini},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer rl2.Close()
|
||||||
|
|
||||||
|
rl2.State["stale-memory"] = &UsageStats{DayStart: time.Now(), DayCount: 99}
|
||||||
|
require.NoError(t, rl2.Load())
|
||||||
|
|
||||||
|
assert.NotContains(t, rl2.Quotas, "gemini-3-pro-preview")
|
||||||
|
assert.NotContains(t, rl2.Quotas, "model-a")
|
||||||
|
assert.Contains(t, rl2.Quotas, "model-b")
|
||||||
|
assert.NotContains(t, rl2.State, "stale-memory")
|
||||||
|
assert.NotContains(t, rl2.State, "model-a")
|
||||||
|
assert.Equal(t, 1, rl2.Stats("model-b").RPD)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSQLite_PersistAtomic_Good(t *testing.T) {
|
||||||
|
dbPath := testPath(t.TempDir(), "persist-atomic.db")
|
||||||
|
rl, err := NewWithSQLiteConfig(dbPath, Config{
|
||||||
|
Quotas: map[string]ModelQuota{
|
||||||
|
"old-model": {MaxRPM: 1, MaxTPM: 100, MaxRPD: 10},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
rl.RecordUsage("old-model", 10, 10)
|
||||||
|
require.NoError(t, rl.Persist())
|
||||||
|
|
||||||
|
_, err = rl.sqlite.db.Exec(`CREATE TRIGGER fail_daily_persist BEFORE INSERT ON daily
|
||||||
|
BEGIN SELECT RAISE(ABORT, 'forced daily failure'); END`)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
delete(rl.Quotas, "old-model")
|
||||||
|
rl.Quotas["new-model"] = ModelQuota{MaxRPM: 2, MaxTPM: 200, MaxRPD: 20}
|
||||||
|
rl.Reset("")
|
||||||
|
rl.RecordUsage("new-model", 50, 50)
|
||||||
|
|
||||||
|
err = rl.Persist()
|
||||||
|
require.Error(t, err)
|
||||||
|
require.NoError(t, rl.Close())
|
||||||
|
|
||||||
|
rl2, err := NewWithSQLiteConfig(dbPath, Config{})
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer rl2.Close()
|
||||||
|
require.NoError(t, rl2.Load())
|
||||||
|
|
||||||
|
assert.Contains(t, rl2.Quotas, "old-model")
|
||||||
|
assert.NotContains(t, rl2.Quotas, "new-model")
|
||||||
|
assert.Equal(t, 1, rl2.Stats("old-model").RPD)
|
||||||
|
assert.Equal(t, 0, rl2.Stats("new-model").RPD)
|
||||||
|
}
|
||||||
|
|
||||||
// --- Phase 2: Benchmark ---
|
// --- Phase 2: Benchmark ---
|
||||||
|
|
||||||
func BenchmarkSQLitePersist(b *testing.B) {
|
func BenchmarkSQLitePersist(b *testing.B) {
|
||||||
dbPath := filepath.Join(b.TempDir(), "bench.db")
|
dbPath := testPath(b.TempDir(), "bench.db")
|
||||||
rl, err := NewWithSQLite(dbPath)
|
rl, err := NewWithSQLite(dbPath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
b.Fatal(err)
|
b.Fatal(err)
|
||||||
|
|
@ -678,7 +845,7 @@ func BenchmarkSQLitePersist(b *testing.B) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func BenchmarkSQLiteLoad(b *testing.B) {
|
func BenchmarkSQLiteLoad(b *testing.B) {
|
||||||
dbPath := filepath.Join(b.TempDir(), "bench-load.db")
|
dbPath := testPath(b.TempDir(), "bench-load.db")
|
||||||
rl, err := NewWithSQLite(dbPath)
|
rl, err := NewWithSQLite(dbPath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
b.Fatal(err)
|
b.Fatal(err)
|
||||||
|
|
@ -709,10 +876,10 @@ func BenchmarkSQLiteLoad(b *testing.B) {
|
||||||
|
|
||||||
// TestMigrateYAMLToSQLiteWithFullState tests migration of a realistic YAML
|
// TestMigrateYAMLToSQLiteWithFullState tests migration of a realistic YAML
|
||||||
// file that contains the full serialised RateLimiter struct.
|
// file that contains the full serialised RateLimiter struct.
|
||||||
func TestMigrateYAMLToSQLiteWithFullState_Good(t *testing.T) {
|
func TestSQLite_MigrateYAMLToSQLiteWithFullState_Good(t *testing.T) {
|
||||||
tmpDir := t.TempDir()
|
tmpDir := t.TempDir()
|
||||||
yamlPath := filepath.Join(tmpDir, "realistic.yaml")
|
yamlPath := testPath(tmpDir, "realistic.yaml")
|
||||||
sqlitePath := filepath.Join(tmpDir, "realistic.db")
|
sqlitePath := testPath(tmpDir, "realistic.db")
|
||||||
|
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
|
|
||||||
|
|
@ -745,7 +912,7 @@ func TestMigrateYAMLToSQLiteWithFullState_Good(t *testing.T) {
|
||||||
|
|
||||||
data, err := yaml.Marshal(rl)
|
data, err := yaml.Marshal(rl)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
require.NoError(t, os.WriteFile(yamlPath, data, 0644))
|
writeTestFile(t, yamlPath, string(data))
|
||||||
|
|
||||||
// Migrate.
|
// Migrate.
|
||||||
require.NoError(t, MigrateYAMLToSQLite(yamlPath, sqlitePath))
|
require.NoError(t, MigrateYAMLToSQLite(yamlPath, sqlitePath))
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue