mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-15 05:29:28 +00:00
lukaszraczylo
5eab51bdef
Add resp.go: a minimal stdlib-only RESP2 client (respClient) for a single Redis-compatible server. It lazily dials one mutex-guarded connection, applies 2s dial/read/write deadlines, sends AUTH (if a password is set) then SELECT on every new connection, and reconnects once on any I/O error before giving up. do and pipeline encode commands as "*N\r\n$len\r\n..." and decode +/-/:/$/* replies. Add redis_store.go: redisStore implements counterStore over respClient. incrBy pipelines INCRBY and EXPIRE in one round trip and returns INCRBY's value; get treats a missing key (RESP null bulk) as 0. Wire the limiter's fail-open/fail-closed policy in limits.go, which Task 5 stored but never read: a configured store's error now routes through storeIncrBy/storeGet. failOpen=true falls back to the in-process memoryStore transparently; failOpen=false makes checkAndCount return a storeDown violation and makes account drop the sample. Both paths log through an injectable, 30s-rate-limited logf. Wire newGateway (llmgateway.go) to build a redisStore-backed limiter when Config.Redis is set (password through resolveSecret, failOpen defaulting to true unless RedisConfig.FailOpen says otherwise) and a memory-only limiter otherwise, and to route pricing.go's unknown-model warnings through the gateway's own logf. Validate LimitsConfig at construction: a new LimitsConfig.validate rejects a negative int64 count limit or a negative/NaN/Inf cost limit, called from newAuthStore for every group and every user (inline and file-sourced) so a malformed limit is a constructor error instead of silently behaving as unlimited. Add RequestsPerDay, TokensPerMonth, and CostPerMonthUSD enforcement tests alongside the existing per-minute/day/cost-per-day cases, plus an erroringStore-backed test pair proving the fail-open (fallback counted) and fail-closed (storeDown violation) branches directly. Folded the validation work into this single commit rather than a second one: it interleaves with the fail-open changes across limits.go, limits_test.go, and llmgateway_test.go closely enough that splitting by hunk would be artificial.
82 lines
2.5 KiB
Go
82 lines
2.5 KiB
Go
package traefikllmgateway
|
|
|
|
import (
|
|
"fmt"
|
|
"strconv"
|
|
"time"
|
|
)
|
|
|
|
// redisStore is a counterStore backed by a Redis-compatible server over
|
|
// respClient, so multiple Traefik instances share the same limit counters
|
|
// instead of each enforcing limits against its own in-process state.
|
|
type redisStore struct {
|
|
client *respClient
|
|
}
|
|
|
|
// newRedisStore returns a redisStore using client for storage.
|
|
func newRedisStore(client *respClient) *redisStore {
|
|
return &redisStore{client: client}
|
|
}
|
|
|
|
// incrBy implements counterStore: INCRBY key n, then EXPIRE key
|
|
// ttlSeconds, sent as one pipeline so the two commands share a single round
|
|
// trip. It returns INCRBY's resulting counter value; EXPIRE's reply is not
|
|
// otherwise inspected — a failed EXPIRE right after a successful INCRBY on
|
|
// the same key would only leave that key without a fresh TTL, not corrupt
|
|
// the count. ttl is rounded up to whole seconds (Redis EXPIRE's unit),
|
|
// with a floor of 1s so a sub-second ttl never turns into EXPIRE 0 (an
|
|
// immediate delete).
|
|
func (s *redisStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) {
|
|
ttlSeconds := int64(ttl / time.Second)
|
|
if ttl%time.Second != 0 {
|
|
ttlSeconds++
|
|
}
|
|
if ttlSeconds < 1 {
|
|
ttlSeconds = 1
|
|
}
|
|
|
|
replies, err := s.client.pipeline([][]string{
|
|
{"INCRBY", key, strconv.FormatInt(n, 10)},
|
|
{"EXPIRE", key, strconv.FormatInt(ttlSeconds, 10)},
|
|
})
|
|
if err != nil {
|
|
return 0, fmt.Errorf("redisStore: incrBy %q: %w", key, err)
|
|
}
|
|
if len(replies) == 0 {
|
|
return 0, fmt.Errorf("redisStore: incrBy %q: empty pipeline reply", key)
|
|
}
|
|
if e, ok := replies[0].(error); ok {
|
|
return 0, fmt.Errorf("redisStore: incrBy %q: INCRBY failed: %w", key, e)
|
|
}
|
|
v, ok := replies[0].(int64)
|
|
if !ok {
|
|
return 0, fmt.Errorf("redisStore: incrBy %q: unexpected INCRBY reply type %T", key, replies[0])
|
|
}
|
|
return v, nil
|
|
}
|
|
|
|
// get implements counterStore: GET key, treating a missing key (RESP null
|
|
// bulk reply) as 0, matching memoryStore's behavior for an absent or
|
|
// expired counter.
|
|
func (s *redisStore) get(key string) (int64, error) {
|
|
reply, err := s.client.do("GET", key)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("redisStore: get %q: %w", key, err)
|
|
}
|
|
if reply == nil {
|
|
return 0, nil
|
|
}
|
|
if e, ok := reply.(error); ok {
|
|
return 0, fmt.Errorf("redisStore: get %q: %w", key, e)
|
|
}
|
|
b, ok := reply.([]byte)
|
|
if !ok {
|
|
return 0, fmt.Errorf("redisStore: get %q: unexpected reply type %T", key, reply)
|
|
}
|
|
v, err := strconv.ParseInt(string(b), 10, 64)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("redisStore: get %q: non-integer value %q: %w", key, b, err)
|
|
}
|
|
return v, nil
|
|
}
|