localmodelroutex/backend/internal/influx/influx.go

// 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. Every
// field is written on every point, because InfluxDB 3 creates a column on its
// first write and a query that names a missing column fails.
func (e *Exporter) requestPoint(rec store.RequestLog) Point {
	keyPrefix := rec.KeyPrefix
	if keyPrefix == "" {
		keyPrefix = "(no key)"
	}
	tags := map[string]string{
		"provider":    rec.ProviderName,
		"model":       rec.ModelName,
		"key_prefix":  keyPrefix,
		"status_code": strconv.Itoa(rec.StatusCode),
		"streamed":    strconv.FormatBool(rec.Streamed),
		// fallback is true when a backup model served the request; then
		// served_model names it and cost is the backup's cost.
		"fallback": strconv.FormatBool(rec.ServedModel != ""),
	}
	if rec.ServedModel != "" {
		tags["served_model"] = rec.ServedModel
	}
	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
	}
	// ttft_ms is 0 when the response was not streamed.
	ttft := 0
	if rec.TTFTMs != nil {
		ttft = *rec.TTFTMs
	}
	msg := rec.Error
	if len(msg) > maxErrorLen {
		msg = msg[:maxErrorLen]
	}
	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,
		"ttft_ms":           ttft,
		"errors":            isErr,
		"request_id":        rec.RequestID,
		"finish_reason":     rec.FinishReason,
		"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
}