mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-20 05:56:56 +00:00
feat: stdlib RESP client and redis-backed limit store with fail-open
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.
This commit is contained in:
@@ -0,0 +1,81 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user