feat: export request statistics to InfluxDB 3
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
Add an optional exporter that writes one point per proxied request to the InfluxDB 3 write endpoint. Points are batched in memory and flushed by size or interval. The exporter is off until GATEWAY_INFLUX_URL is set.
This commit is contained in:
@@ -25,6 +25,14 @@ GATEWAY_DB_URL=postgres://lmr:lmr@localhost:5432/lmr?sslmode=disable
|
||||
# GATEWAY_DB_MIN_CONNS=2
|
||||
# GATEWAY_CATALOG_SYNC_INTERVAL=30m
|
||||
|
||||
# InfluxDB 3 statistics exporter. Off until GATEWAY_INFLUX_URL is set.
|
||||
# GATEWAY_INFLUX_URL=http://influxdb:8181
|
||||
# GATEWAY_INFLUX_TOKEN=
|
||||
# GATEWAY_INFLUX_DATABASE=lmr
|
||||
# GATEWAY_INFLUX_MEASUREMENT=lmr_request
|
||||
# GATEWAY_INFLUX_FLUSH_INTERVAL=10s
|
||||
# GATEWAY_INFLUX_BATCH_SIZE=500
|
||||
|
||||
# Host port to expose (maps to container :8080)
|
||||
# PORT=8080
|
||||
|
||||
|
||||
24
README.md
24
README.md
@@ -117,10 +117,34 @@ is loaded automatically; in production, set real environment variables.
|
||||
| `GATEWAY_OIDC_ADMIN_GROUP` | — | **yes** | group required for admin access |
|
||||
| `GATEWAY_SECURITY_ENCRYPTION_KEY` | — | **yes** | 64 hex chars (32 bytes) for AES-256-GCM |
|
||||
| `GATEWAY_CATALOG_SYNC_INTERVAL` | `30m` | no | how often provider model lists are refreshed |
|
||||
| `GATEWAY_INFLUX_URL` | — | no | InfluxDB 3 base URL, e.g. `http://influxdb:8181`. Empty = exporter off |
|
||||
| `GATEWAY_INFLUX_TOKEN` | — | no | InfluxDB 3 token, sent as a Bearer token |
|
||||
| `GATEWAY_INFLUX_DATABASE` | — | with URL | target database |
|
||||
| `GATEWAY_INFLUX_MEASUREMENT` | `lmr_request` | no | measurement name for request points |
|
||||
| `GATEWAY_INFLUX_FLUSH_INTERVAL` | `10s` | no | max time a point waits in the buffer |
|
||||
| `GATEWAY_INFLUX_BATCH_SIZE` | `500` | no | points per write request |
|
||||
|
||||
Config is validated on startup; missing required values or a wrong-length
|
||||
encryption key cause the server to exit immediately.
|
||||
|
||||
### InfluxDB 3 export
|
||||
|
||||
Set `GATEWAY_INFLUX_URL` and `GATEWAY_INFLUX_DATABASE` to export one point per
|
||||
proxied request. Points are batched in memory and written with the
|
||||
`/api/v3/write_lp` endpoint. A full buffer drops points and logs a warning, so
|
||||
the exporter never delays a response. Each point has these tags and fields:
|
||||
|
||||
| Kind | Name | Notes |
|
||||
|---|---|---|
|
||||
| tag | `provider`, `model`, `key_prefix` | group by provider, model address, or key |
|
||||
| tag | `user_id`, `provider_id`, `provider_model_id` | UUIDs, set when known |
|
||||
| tag | `status_code`, `streamed` | HTTP status and `true`/`false` |
|
||||
| field | `requests`, `errors` | always `1` and `0` or `1`, so `sum()` gives counts |
|
||||
| field | `prompt_tokens`, `completion_tokens`, `cached_tokens`, `total_tokens` | integers |
|
||||
| field | `cost` | dollars, float |
|
||||
| field | `latency_ms`, `ttft_ms` | integers, `ttft_ms` only for streamed chat |
|
||||
| field | `request_id`, `finish_reason`, `error` | strings, set when known |
|
||||
|
||||
Generate an encryption key:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -34,3 +34,12 @@ GATEWAY_SECURITY_ENCRYPTION_KEY=
|
||||
# --- Catalog ---
|
||||
# How often to sync model lists from providers. Default: 30m
|
||||
GATEWAY_CATALOG_SYNC_INTERVAL=30m
|
||||
|
||||
# --- InfluxDB 3 statistics exporter (optional) ---
|
||||
# Leave GATEWAY_INFLUX_URL empty to disable the exporter.
|
||||
GATEWAY_INFLUX_URL=
|
||||
GATEWAY_INFLUX_TOKEN=
|
||||
GATEWAY_INFLUX_DATABASE=lmr
|
||||
GATEWAY_INFLUX_MEASUREMENT=lmr_request
|
||||
GATEWAY_INFLUX_FLUSH_INTERVAL=10s
|
||||
GATEWAY_INFLUX_BATCH_SIZE=500
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/config"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/crypto"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/db"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/influx"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/oidcauth"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/providermgr"
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/ratelimit"
|
||||
@@ -147,8 +148,29 @@ func run() error {
|
||||
}
|
||||
}()
|
||||
|
||||
// The exporter sends one point per request to InfluxDB 3 when configured.
|
||||
var stats api.RequestStats
|
||||
var exporter *influx.Exporter
|
||||
if cfg.Influx.Enabled() {
|
||||
exporter, err = influx.New(influx.Config{
|
||||
URL: cfg.Influx.URL,
|
||||
Token: cfg.Influx.Token,
|
||||
Database: cfg.Influx.Database,
|
||||
Measurement: cfg.Influx.Measurement,
|
||||
FlushInterval: cfg.Influx.FlushInterval,
|
||||
BatchSize: cfg.Influx.BatchSize,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("init influx exporter: %w", err)
|
||||
}
|
||||
exporter.Start()
|
||||
stats = exporter
|
||||
slog.Info("influx exporter enabled", "url", cfg.Influx.URL, "database", cfg.Influx.Database)
|
||||
}
|
||||
|
||||
deps := api.Deps{
|
||||
Store: st,
|
||||
Stats: stats,
|
||||
Cache: cache,
|
||||
Manager: mgr,
|
||||
Cipher: cipher,
|
||||
@@ -193,6 +215,11 @@ func run() error {
|
||||
if err := srv.Shutdown(shutdownCtx); err != nil {
|
||||
return fmt.Errorf("graceful shutdown: %w", err)
|
||||
}
|
||||
if exporter != nil {
|
||||
if err := exporter.Close(shutdownCtx); err != nil {
|
||||
slog.Warn("influx exporter did not flush in time", "err", err)
|
||||
}
|
||||
}
|
||||
slog.Info("server stopped cleanly")
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -53,11 +53,19 @@ type RequestLogSink interface {
|
||||
TouchVirtualKey(ctx context.Context, id uuid.UUID) error
|
||||
}
|
||||
|
||||
// RequestStats receives a copy of every request log row for export to a
|
||||
// metrics backend. Record must not block. *influx.Exporter satisfies it.
|
||||
type RequestStats interface {
|
||||
Record(e store.RequestLog)
|
||||
}
|
||||
|
||||
// Deps are the dependencies the API handlers need.
|
||||
type Deps struct {
|
||||
Store *store.Store
|
||||
// Logs receives request log rows. Defaults to Store when nil.
|
||||
Logs RequestLogSink
|
||||
Logs RequestLogSink
|
||||
// Stats receives request log rows for the statistics exporter. May be nil.
|
||||
Stats RequestStats
|
||||
Cache *catalog.Cache
|
||||
Manager *providermgr.Manager
|
||||
Cipher *crypto.Cipher
|
||||
|
||||
@@ -210,6 +210,10 @@ func (a *API) logRequest(r *http.Request, key *store.VirtualKey, entry *catalog.
|
||||
}
|
||||
}
|
||||
|
||||
if a.Stats != nil {
|
||||
a.Stats.Record(rec)
|
||||
}
|
||||
|
||||
// The log write must outlive the request, so it gets its own context.
|
||||
go func() { //nolint:gosec // G118: intentional, see above
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
|
||||
42
backend/internal/api/stats_test.go
Normal file
42
backend/internal/api/stats_test.go
Normal file
@@ -0,0 +1,42 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
||||
)
|
||||
|
||||
type fakeStats struct {
|
||||
mu sync.Mutex
|
||||
rows []store.RequestLog
|
||||
}
|
||||
|
||||
func (f *fakeStats) Record(e store.RequestLog) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.rows = append(f.rows, e)
|
||||
}
|
||||
|
||||
func TestLogRequestForwardsRowToStats(t *testing.T) {
|
||||
env := newTestEnv(t, chatReply, nil)
|
||||
stats := &fakeStats{}
|
||||
env.api.Stats = stats
|
||||
rec := call(env.api.chatCompletions, allowAllKey(), http.MethodPost,
|
||||
`{"model":"or/openai/gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rec.Code, rec.Body)
|
||||
}
|
||||
logged := env.logs.wait(t)
|
||||
|
||||
stats.mu.Lock()
|
||||
defer stats.mu.Unlock()
|
||||
if len(stats.rows) != 1 {
|
||||
t.Fatalf("stats rows = %d, want 1", len(stats.rows))
|
||||
}
|
||||
got := stats.rows[0]
|
||||
if got.ModelName != logged.ModelName || got.TotalTokens != logged.TotalTokens || got.Cost != logged.Cost || got.KeyPrefix != "sk-test" {
|
||||
t.Errorf("stats row = %+v, log row = %+v", got, logged)
|
||||
}
|
||||
}
|
||||
@@ -15,8 +15,22 @@ type Config struct {
|
||||
OIDC OIDCConfig
|
||||
Security SecurityConfig
|
||||
Catalog CatalogConfig
|
||||
Influx InfluxConfig
|
||||
}
|
||||
|
||||
// InfluxConfig enables the InfluxDB 3 statistics exporter when URL is set.
|
||||
type InfluxConfig struct {
|
||||
URL string
|
||||
Token string
|
||||
Database string
|
||||
Measurement string
|
||||
FlushInterval time.Duration
|
||||
BatchSize int
|
||||
}
|
||||
|
||||
// Enabled reports whether the exporter is configured.
|
||||
func (c InfluxConfig) Enabled() bool { return c.URL != "" }
|
||||
|
||||
type CatalogConfig struct {
|
||||
SyncInterval time.Duration // how often to refresh the model catalog
|
||||
}
|
||||
@@ -62,6 +76,9 @@ func Load() (*Config, error) {
|
||||
v.SetDefault("db.min_conns", 2)
|
||||
v.SetDefault("db.max_conn_lifetime", 30)
|
||||
v.SetDefault("catalog.sync_interval", "30m")
|
||||
v.SetDefault("influx.measurement", "lmr_request")
|
||||
v.SetDefault("influx.flush_interval", "10s")
|
||||
v.SetDefault("influx.batch_size", 500)
|
||||
|
||||
cfg := &Config{
|
||||
Server: ServerConfig{
|
||||
@@ -87,6 +104,14 @@ func Load() (*Config, error) {
|
||||
Catalog: CatalogConfig{
|
||||
SyncInterval: v.GetDuration("catalog.sync_interval"),
|
||||
},
|
||||
Influx: InfluxConfig{
|
||||
URL: v.GetString("influx.url"),
|
||||
Token: v.GetString("influx.token"),
|
||||
Database: v.GetString("influx.database"),
|
||||
Measurement: v.GetString("influx.measurement"),
|
||||
FlushInterval: v.GetDuration("influx.flush_interval"),
|
||||
BatchSize: v.GetInt("influx.batch_size"),
|
||||
},
|
||||
}
|
||||
|
||||
return cfg, cfg.validate()
|
||||
@@ -123,6 +148,9 @@ func (c *Config) validate() error {
|
||||
if len(missing) > 0 {
|
||||
return fmt.Errorf("missing required OIDC config: %s", strings.Join(missing, ", "))
|
||||
}
|
||||
if c.Influx.Enabled() && c.Influx.Database == "" {
|
||||
return fmt.Errorf("influx.database is required when influx.url is set")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
281
backend/internal/influx/influx.go
Normal file
281
backend/internal/influx/influx.go
Normal file
@@ -0,0 +1,281 @@
|
||||
// Package influx exports one point per proxied request to InfluxDB 3 over the
|
||||
// /api/v3/write_lp HTTP endpoint. Points are batched in memory and flushed by
|
||||
// size or by interval. The exporter never blocks a request: when the buffer is
|
||||
// full the point is dropped and counted.
|
||||
package influx
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
||||
)
|
||||
|
||||
// Config controls the exporter. URL, Database and Measurement are required.
|
||||
type Config struct {
|
||||
URL string
|
||||
Token string
|
||||
Database string
|
||||
Measurement string
|
||||
FlushInterval time.Duration
|
||||
BatchSize int
|
||||
// BufferSize is the number of points held between flushes. Zero = 10 × BatchSize.
|
||||
BufferSize int
|
||||
// Timeout bounds one write request. Zero = 10s.
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
const (
|
||||
defaultFlushInterval = 10 * time.Second
|
||||
defaultBatchSize = 500
|
||||
defaultTimeout = 10 * time.Second
|
||||
maxErrorLen = 512
|
||||
writeAttempts = 3
|
||||
)
|
||||
|
||||
// Exporter batches points and writes them to InfluxDB 3.
|
||||
type Exporter struct {
|
||||
cfg Config
|
||||
writeURL string
|
||||
client *http.Client
|
||||
queue chan Point
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
stopOnce sync.Once
|
||||
dropped atomic.Int64
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
// New validates cfg and returns a stopped exporter. Call Start to run it.
|
||||
func New(cfg Config) (*Exporter, error) {
|
||||
if cfg.URL == "" {
|
||||
return nil, errors.New("influx: url is required")
|
||||
}
|
||||
if cfg.Database == "" {
|
||||
return nil, errors.New("influx: database is required")
|
||||
}
|
||||
if cfg.Measurement == "" {
|
||||
return nil, errors.New("influx: measurement is required")
|
||||
}
|
||||
base, err := url.Parse(cfg.URL)
|
||||
if err != nil || (base.Scheme != "http" && base.Scheme != "https") || base.Host == "" {
|
||||
return nil, fmt.Errorf("influx: url %q must be an http(s) URL", cfg.URL)
|
||||
}
|
||||
if cfg.FlushInterval <= 0 {
|
||||
cfg.FlushInterval = defaultFlushInterval
|
||||
}
|
||||
if cfg.BatchSize <= 0 {
|
||||
cfg.BatchSize = defaultBatchSize
|
||||
}
|
||||
if cfg.BufferSize <= 0 {
|
||||
cfg.BufferSize = 10 * cfg.BatchSize
|
||||
}
|
||||
if cfg.Timeout <= 0 {
|
||||
cfg.Timeout = defaultTimeout
|
||||
}
|
||||
|
||||
q := url.Values{}
|
||||
q.Set("db", cfg.Database)
|
||||
q.Set("precision", "nanosecond")
|
||||
base.Path = strings.TrimRight(base.Path, "/") + "/api/v3/write_lp"
|
||||
base.RawQuery = q.Encode()
|
||||
|
||||
return &Exporter{
|
||||
cfg: cfg,
|
||||
writeURL: base.String(),
|
||||
client: &http.Client{Timeout: cfg.Timeout},
|
||||
queue: make(chan Point, cfg.BufferSize),
|
||||
stop: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
now: time.Now,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Start runs the flush loop in the background.
|
||||
func (e *Exporter) Start() {
|
||||
go e.loop()
|
||||
}
|
||||
|
||||
// Close stops the loop, writes the points still in the buffer, and waits for
|
||||
// that write or for ctx to end.
|
||||
func (e *Exporter) Close(ctx context.Context) error {
|
||||
e.stopOnce.Do(func() { close(e.stop) })
|
||||
select {
|
||||
case <-e.done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// Record converts a request log row to a point and queues it. It never blocks.
|
||||
func (e *Exporter) Record(rec store.RequestLog) {
|
||||
e.Enqueue(e.requestPoint(rec))
|
||||
}
|
||||
|
||||
// Enqueue queues an arbitrary point. It never blocks: a full buffer drops the point.
|
||||
func (e *Exporter) Enqueue(p Point) {
|
||||
select {
|
||||
case e.queue <- p:
|
||||
default:
|
||||
e.dropped.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
// requestPoint maps a log row to the request measurement. Bounded values are
|
||||
// tags, so dashboards can group by them. Per-request values are fields.
|
||||
func (e *Exporter) requestPoint(rec store.RequestLog) Point {
|
||||
tags := map[string]string{
|
||||
"provider": rec.ProviderName,
|
||||
"model": rec.ModelName,
|
||||
"key_prefix": rec.KeyPrefix,
|
||||
"status_code": strconv.Itoa(rec.StatusCode),
|
||||
"streamed": strconv.FormatBool(rec.Streamed),
|
||||
}
|
||||
if rec.UserID != nil {
|
||||
tags["user_id"] = rec.UserID.String()
|
||||
}
|
||||
if rec.ProviderID != nil {
|
||||
tags["provider_id"] = rec.ProviderID.String()
|
||||
}
|
||||
if rec.ProviderModelID != nil {
|
||||
tags["provider_model_id"] = rec.ProviderModelID.String()
|
||||
}
|
||||
|
||||
isErr := 0
|
||||
if rec.StatusCode >= 400 || rec.Error != "" {
|
||||
isErr = 1
|
||||
}
|
||||
fields := map[string]any{
|
||||
"requests": 1,
|
||||
"prompt_tokens": rec.PromptTokens,
|
||||
"completion_tokens": rec.CompletionTokens,
|
||||
"cached_tokens": rec.CachedTokens,
|
||||
"total_tokens": rec.TotalTokens,
|
||||
"cost": rec.Cost,
|
||||
"latency_ms": rec.LatencyMS,
|
||||
"errors": isErr,
|
||||
}
|
||||
if rec.TTFTMs != nil {
|
||||
fields["ttft_ms"] = *rec.TTFTMs
|
||||
}
|
||||
if rec.RequestID != "" {
|
||||
fields["request_id"] = rec.RequestID
|
||||
}
|
||||
if rec.FinishReason != "" {
|
||||
fields["finish_reason"] = rec.FinishReason
|
||||
}
|
||||
if rec.Error != "" {
|
||||
msg := rec.Error
|
||||
if len(msg) > maxErrorLen {
|
||||
msg = msg[:maxErrorLen]
|
||||
}
|
||||
fields["error"] = msg
|
||||
}
|
||||
return Point{Measurement: e.cfg.Measurement, Tags: tags, Fields: fields, Time: e.now()}
|
||||
}
|
||||
|
||||
func (e *Exporter) loop() {
|
||||
defer close(e.done)
|
||||
batch := make([]Point, 0, e.cfg.BatchSize)
|
||||
tick := time.NewTicker(e.cfg.FlushInterval)
|
||||
defer tick.Stop()
|
||||
|
||||
flush := func() {
|
||||
if n := e.dropped.Swap(0); n > 0 {
|
||||
slog.Warn("influx exporter dropped points: buffer full", "points", n)
|
||||
}
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
e.write(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case p := <-e.queue:
|
||||
batch = append(batch, p)
|
||||
if len(batch) >= e.cfg.BatchSize {
|
||||
flush()
|
||||
}
|
||||
case <-tick.C:
|
||||
flush()
|
||||
case <-e.stop:
|
||||
for {
|
||||
select {
|
||||
case p := <-e.queue:
|
||||
batch = append(batch, p)
|
||||
if len(batch) >= e.cfg.BatchSize {
|
||||
flush()
|
||||
}
|
||||
default:
|
||||
flush()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// write sends one batch. Network errors and 5xx responses are retried with a
|
||||
// short backoff. 4xx responses are not, because the same body would fail again.
|
||||
func (e *Exporter) write(batch []Point) {
|
||||
var b strings.Builder
|
||||
for _, p := range batch {
|
||||
p.encode(&b)
|
||||
}
|
||||
body := b.String()
|
||||
if body == "" {
|
||||
return
|
||||
}
|
||||
|
||||
var err error
|
||||
for attempt := 1; attempt <= writeAttempts; attempt++ {
|
||||
var retry bool
|
||||
retry, err = e.post(body)
|
||||
if err == nil {
|
||||
return
|
||||
}
|
||||
if !retry {
|
||||
break
|
||||
}
|
||||
time.Sleep(time.Duration(attempt) * 500 * time.Millisecond)
|
||||
}
|
||||
slog.Warn("influx write failed, batch dropped", "points", len(batch), "err", err)
|
||||
}
|
||||
|
||||
// post sends one write request. The bool reports whether a retry makes sense.
|
||||
func (e *Exporter) post(body string) (bool, error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), e.cfg.Timeout)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.writeURL, strings.NewReader(body))
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "text/plain; charset=utf-8")
|
||||
if e.cfg.Token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+e.cfg.Token)
|
||||
}
|
||||
resp, err := e.client.Do(req)
|
||||
if err != nil {
|
||||
return true, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
return false, nil
|
||||
}
|
||||
msg, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
|
||||
err = fmt.Errorf("influx returned %d: %s", resp.StatusCode, strings.TrimSpace(string(msg)))
|
||||
return resp.StatusCode >= 500 || resp.StatusCode == http.StatusTooManyRequests, err
|
||||
}
|
||||
258
backend/internal/influx/influx_test.go
Normal file
258
backend/internal/influx/influx_test.go
Normal file
@@ -0,0 +1,258 @@
|
||||
package influx
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
||||
)
|
||||
|
||||
func TestEncodeEscapesAndSorts(t *testing.T) {
|
||||
p := Point{
|
||||
Measurement: "my measure,ment",
|
||||
Tags: map[string]string{"z": "last", "a": "has space", "empty": ""},
|
||||
Fields: map[string]any{
|
||||
"s": `quote " back \ newline` + "\nend",
|
||||
"i": int(7),
|
||||
"i64": int64(-3),
|
||||
"f": 1.5,
|
||||
"b": true,
|
||||
"nan": 0.0 / zero(),
|
||||
"bad": struct{}{},
|
||||
},
|
||||
Time: time.Unix(1, 2),
|
||||
}
|
||||
var b strings.Builder
|
||||
if !p.encode(&b) {
|
||||
t.Fatal("encode returned false")
|
||||
}
|
||||
want := `my\ measure\,ment,a=has\ space,z=last b=true,f=1.5,i=7i,i64=-3i,s="quote \" back \\ newline end" 1000000002` + "\n"
|
||||
if b.String() != want {
|
||||
t.Fatalf("got\n%s\nwant\n%s", b.String(), want)
|
||||
}
|
||||
}
|
||||
|
||||
func zero() float64 { return 0 }
|
||||
|
||||
func TestEncodeNoFields(t *testing.T) {
|
||||
var b strings.Builder
|
||||
if (Point{Measurement: "m", Fields: map[string]any{"x": struct{}{}}}).encode(&b) {
|
||||
t.Fatal("expected false for a point without usable fields")
|
||||
}
|
||||
if b.Len() != 0 {
|
||||
t.Fatalf("expected no output, got %q", b.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewRejectsBadConfig(t *testing.T) {
|
||||
cases := []Config{
|
||||
{Database: "d", Measurement: "m"},
|
||||
{URL: "http://x", Measurement: "m"},
|
||||
{URL: "http://x", Database: "d"},
|
||||
{URL: "not a url", Database: "d", Measurement: "m"},
|
||||
{URL: "ftp://x", Database: "d", Measurement: "m"},
|
||||
}
|
||||
for i, c := range cases {
|
||||
if _, err := New(c); err == nil {
|
||||
t.Errorf("case %d: expected error for %+v", i, c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewWriteURL(t *testing.T) {
|
||||
e, err := New(Config{URL: "http://influx:8181/prefix/", Database: "lmr", Measurement: "m"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "http://influx:8181/prefix/api/v3/write_lp?db=lmr&precision=nanosecond"
|
||||
if e.writeURL != want {
|
||||
t.Fatalf("writeURL = %q, want %q", e.writeURL, want)
|
||||
}
|
||||
}
|
||||
|
||||
// captureServer records every write body and the auth header it saw.
|
||||
type captureServer struct {
|
||||
*httptest.Server
|
||||
mu sync.Mutex
|
||||
bodies []string
|
||||
auth string
|
||||
status atomic.Int32
|
||||
got chan struct{}
|
||||
}
|
||||
|
||||
func newCaptureServer(t *testing.T) *captureServer {
|
||||
t.Helper()
|
||||
cs := &captureServer{got: make(chan struct{}, 64)}
|
||||
cs.status.Store(http.StatusNoContent)
|
||||
cs.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/v3/write_lp" {
|
||||
t.Errorf("unexpected path %s", r.URL.Path)
|
||||
}
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
cs.mu.Lock()
|
||||
cs.bodies = append(cs.bodies, string(b))
|
||||
cs.auth = r.Header.Get("Authorization")
|
||||
cs.mu.Unlock()
|
||||
w.WriteHeader(int(cs.status.Load()))
|
||||
cs.got <- struct{}{}
|
||||
}))
|
||||
t.Cleanup(cs.Close)
|
||||
return cs
|
||||
}
|
||||
|
||||
func (cs *captureServer) wait(t *testing.T) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-cs.got:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("no write arrived")
|
||||
}
|
||||
}
|
||||
|
||||
func newTestExporter(t *testing.T, srv *captureServer, cfg Config) *Exporter {
|
||||
t.Helper()
|
||||
cfg.URL = srv.URL
|
||||
cfg.Database = "lmr"
|
||||
cfg.Measurement = "lmr_request"
|
||||
cfg.Token = "secret"
|
||||
e, err := New(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
e.now = func() time.Time { return time.Unix(100, 0) }
|
||||
return e
|
||||
}
|
||||
|
||||
func TestRecordWritesRequestPoint(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
|
||||
e.Start()
|
||||
|
||||
uid := uuid.MustParse("11111111-1111-1111-1111-111111111111")
|
||||
ttft := 42
|
||||
e.Record(store.RequestLog{
|
||||
UserID: &uid,
|
||||
KeyPrefix: "lmr_abc",
|
||||
ModelName: "openai/gpt-4o",
|
||||
ProviderName: "openai",
|
||||
StatusCode: 200,
|
||||
Streamed: true,
|
||||
PromptTokens: 10,
|
||||
CompletionTokens: 5,
|
||||
CachedTokens: 2,
|
||||
TotalTokens: 15,
|
||||
Cost: 0.00125,
|
||||
LatencyMS: 321,
|
||||
TTFTMs: &ttft,
|
||||
RequestID: "req-1",
|
||||
})
|
||||
srv.wait(t)
|
||||
if err := e.Close(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
srv.mu.Lock()
|
||||
defer srv.mu.Unlock()
|
||||
if srv.auth != "Bearer secret" {
|
||||
t.Fatalf("auth = %q", srv.auth)
|
||||
}
|
||||
want := `lmr_request,key_prefix=lmr_abc,model=openai/gpt-4o,provider=openai,status_code=200,streamed=true,user_id=11111111-1111-1111-1111-111111111111 ` +
|
||||
`cached_tokens=2i,completion_tokens=5i,cost=0.00125,errors=0i,latency_ms=321i,prompt_tokens=10i,request_id="req-1",requests=1i,total_tokens=15i,ttft_ms=42i 100000000000` + "\n"
|
||||
if len(srv.bodies) != 1 || srv.bodies[0] != want {
|
||||
t.Fatalf("bodies = %q\nwant %q", srv.bodies, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordMarksErrors(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
|
||||
long := strings.Repeat("x", maxErrorLen+10)
|
||||
p := e.requestPoint(store.RequestLog{StatusCode: 502, Error: long})
|
||||
if p.Fields["errors"] != 1 {
|
||||
t.Fatalf("errors = %v", p.Fields["errors"])
|
||||
}
|
||||
if got := p.Fields["error"].(string); len(got) != maxErrorLen {
|
||||
t.Fatalf("error length = %d", len(got))
|
||||
}
|
||||
p = e.requestPoint(store.RequestLog{StatusCode: 200, Error: "stream interrupted"})
|
||||
if p.Fields["errors"] != 1 {
|
||||
t.Fatal("a 200 with an error message must count as an error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchAndFlushOnClose(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 2})
|
||||
e.Start()
|
||||
for range 3 {
|
||||
e.Record(store.RequestLog{StatusCode: 200})
|
||||
}
|
||||
srv.wait(t) // first batch of 2
|
||||
if err := e.Close(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
srv.mu.Lock()
|
||||
defer srv.mu.Unlock()
|
||||
if len(srv.bodies) != 2 {
|
||||
t.Fatalf("expected 2 writes, got %d", len(srv.bodies))
|
||||
}
|
||||
if n := strings.Count(srv.bodies[0], "\n"); n != 2 {
|
||||
t.Fatalf("first batch has %d lines, want 2", n)
|
||||
}
|
||||
if n := strings.Count(srv.bodies[1], "\n"); n != 1 {
|
||||
t.Fatalf("second batch has %d lines, want 1", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFlushOnInterval(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: 20 * time.Millisecond, BatchSize: 100})
|
||||
e.Start()
|
||||
defer e.Close(context.Background()) //nolint:errcheck
|
||||
e.Record(store.RequestLog{StatusCode: 200})
|
||||
srv.wait(t)
|
||||
}
|
||||
|
||||
func TestEnqueueDropsWhenFull(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1, BufferSize: 1})
|
||||
// Not started: nothing drains the queue.
|
||||
e.Enqueue(Point{Measurement: "m", Fields: map[string]any{"v": 1}})
|
||||
e.Enqueue(Point{Measurement: "m", Fields: map[string]any{"v": 1}})
|
||||
if e.dropped.Load() != 1 {
|
||||
t.Fatalf("dropped = %d, want 1", e.dropped.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteDoesNotRetryClientErrors(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
srv.status.Store(http.StatusBadRequest)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
|
||||
e.write([]Point{{Measurement: "m", Fields: map[string]any{"v": 1}, Time: time.Unix(1, 0)}})
|
||||
srv.mu.Lock()
|
||||
defer srv.mu.Unlock()
|
||||
if len(srv.bodies) != 1 {
|
||||
t.Fatalf("expected 1 attempt, got %d", len(srv.bodies))
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteRetriesServerErrors(t *testing.T) {
|
||||
srv := newCaptureServer(t)
|
||||
srv.status.Store(http.StatusInternalServerError)
|
||||
e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
|
||||
e.write([]Point{{Measurement: "m", Fields: map[string]any{"v": 1}, Time: time.Unix(1, 0)}})
|
||||
srv.mu.Lock()
|
||||
defer srv.mu.Unlock()
|
||||
if len(srv.bodies) != writeAttempts {
|
||||
t.Fatalf("expected %d attempts, got %d", writeAttempts, len(srv.bodies))
|
||||
}
|
||||
}
|
||||
91
backend/internal/influx/lineprotocol.go
Normal file
91
backend/internal/influx/lineprotocol.go
Normal file
@@ -0,0 +1,91 @@
|
||||
package influx
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Point is one line-protocol row.
|
||||
type Point struct {
|
||||
Measurement string
|
||||
Tags map[string]string
|
||||
Fields map[string]any
|
||||
Time time.Time
|
||||
}
|
||||
|
||||
// encode appends the point as one line-protocol line. It returns false when
|
||||
// the point has no usable field, because InfluxDB rejects such a line.
|
||||
func (p Point) encode(b *strings.Builder) bool {
|
||||
fields := make([]string, 0, len(p.Fields))
|
||||
for _, k := range sortedKeys(p.Fields) {
|
||||
v, ok := formatField(p.Fields[k])
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fields = append(fields, escapeKey(k)+"="+v)
|
||||
}
|
||||
if len(fields) == 0 {
|
||||
return false
|
||||
}
|
||||
|
||||
b.WriteString(escapeMeasurement(p.Measurement))
|
||||
for _, k := range sortedKeys(p.Tags) {
|
||||
v := p.Tags[k]
|
||||
if v == "" {
|
||||
continue
|
||||
}
|
||||
b.WriteByte(',')
|
||||
b.WriteString(escapeKey(k))
|
||||
b.WriteByte('=')
|
||||
b.WriteString(escapeKey(v))
|
||||
}
|
||||
b.WriteByte(' ')
|
||||
b.WriteString(strings.Join(fields, ","))
|
||||
b.WriteByte(' ')
|
||||
b.WriteString(strconv.FormatInt(p.Time.UnixNano(), 10))
|
||||
b.WriteByte('\n')
|
||||
return true
|
||||
}
|
||||
|
||||
// formatField renders a field value. Unsupported types and non-finite floats
|
||||
// are dropped so one bad value does not reject the whole batch.
|
||||
func formatField(v any) (string, bool) {
|
||||
switch x := v.(type) {
|
||||
case int:
|
||||
return strconv.FormatInt(int64(x), 10) + "i", true
|
||||
case int64:
|
||||
return strconv.FormatInt(x, 10) + "i", true
|
||||
case float64:
|
||||
if math.IsNaN(x) || math.IsInf(x, 0) {
|
||||
return "", false
|
||||
}
|
||||
return strconv.FormatFloat(x, 'f', -1, 64), true
|
||||
case bool:
|
||||
return strconv.FormatBool(x), true
|
||||
case string:
|
||||
return `"` + escapeString(x) + `"`, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func sortedKeys[V any](m map[string]V) []string {
|
||||
keys := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
|
||||
var (
|
||||
measurementEscaper = strings.NewReplacer(",", `\,`, " ", `\ `, "\n", " ")
|
||||
keyEscaper = strings.NewReplacer(",", `\,`, "=", `\=`, " ", `\ `, "\n", " ")
|
||||
stringEscaper = strings.NewReplacer(`\`, `\\`, `"`, `\"`, "\n", " ", "\r", " ")
|
||||
)
|
||||
|
||||
func escapeMeasurement(s string) string { return measurementEscaper.Replace(s) }
|
||||
func escapeKey(s string) string { return keyEscaper.Replace(s) }
|
||||
func escapeString(s string) string { return stringEscaper.Replace(s) }
|
||||
@@ -47,6 +47,7 @@ internal/providers provider abstraction and the per-type packages
|
||||
internal/catalog model library: sync, cache, address resolution
|
||||
internal/ratelimit per-model fixed-window RPM and TPM counters
|
||||
internal/billing cost calculation
|
||||
internal/influx InfluxDB 3 exporter: one point per request, batched
|
||||
internal/api chi routes: /v1 proxy, /api/admin, /auth
|
||||
internal/oidcauth OIDC discovery, ID token verification, group check
|
||||
internal/web go:embed of the built dashboard, served as a SPA
|
||||
|
||||
Reference in New Issue
Block a user