ModelRouter/internal/gateway/stats.go

package gateway

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"sync"
	"time"
)

// Req is one recorded gateway request (audit trail + per-key/per-model stats).
type Req struct {
	Time  int64  `json:"time"`  // unix milliseconds
	Key   string `json:"key"`   // gateway key id
	Type  string `json:"type"`  // chat | stream | image
	Model string `json:"model"` // effective model used upstream
	// Source provider/source name
	Source string `json:"source"`
	// Prompt prompt tokens
	Prompt int64 `json:"prompt_tokens"`
	// Compl completion tokens
	Compl int64 `json:"completion_tokens"`
	// LatMs total handling time ms
	LatMs int64 `json:"latency_ms"`
	OK    bool   `json:"ok"`
	// Status http status code
	Status int `json:"status"`
	// Err short error message
	Err string `json:"error,omitempty"`
}

// Stat aggregates counters for one dimension row.
type Stat struct {
	Reqs   int64 `json:"reqs"`
	OK     int64 `json:"ok"`
	Err    int64 `json:"err"`
	Tokens int64 `json:"tokens"`
	Prompt int64 `json:"prompt_tokens"`
	Compl  int64 `json:"completion_tokens"`
	LatSum int64 `json:"latency_sum_ms"`
	LatMax int64 `json:"latency_max_ms"`
}

type agrRow struct {
	Name string `json:"name"`
	Stat
}

// Stats collects per-key / per-model / per-source aggregates plus a bounded
// ring of raw request records, all guarded by one mutex.
type Stats struct {
	mu         sync.Mutex
	active     int64
	byKey      map[string]*Stat
	byModel    map[string]*Stat
	bySrc      map[string]*Stat
	byKeyModel map[string]map[string]*Stat
	byKeySrc   map[string]map[string]*Stat
	byStatus   map[int]*Stat // per http status code aggregates (incl. 402/400)
	recs       []Req
	maxRecs    int
	auditPath  string
	modelHour  map[string]map[int64]int64 // model -> unix-hour bucket -> tokens
}

const hourSec = 3600

// auditRotateBytes rotates the audit file once it grows past this size (the
// file is renamed to <path>.<unix>.old and a fresh one is started); pruning
// keeps at most auditKeepOld rotated files. Both are vars so tests can shrink
// the threshold.
var (
	auditRotateBytes int64 = 64 << 20
	auditKeepOld           = 10
)

func NewStats(maxRecords int) *Stats {
	if maxRecords <= 0 {
		maxRecords = 3000
	}
	return &Stats{
		byKey:      map[string]*Stat{},
		byModel:    map[string]*Stat{},
		bySrc:      map[string]*Stat{},
		byKeyModel: map[string]map[string]*Stat{},
		byKeySrc:   map[string]map[string]*Stat{},
		byStatus:   map[int]*Stat{},
		modelHour:  map[string]map[int64]int64{},
		maxRecs:    maxRecords,
	}
}

// Begin accounts an in-flight request; the returned func must be called once
// the request finished (defer ok).
func (s *Stats) Begin() func() {
	s.mu.Lock()
	s.active++
	s.mu.Unlock()
	return func() {
		s.mu.Lock()
		s.active--
		s.mu.Unlock()
	}
}

func inc(m map[string]*Stat, name string, r Req) {
	a := m[name]
	if a == nil {
		a = &Stat{}
		m[name] = a
	}
	incStatus(a, name, r)
}

func incStatus(a *Stat, name string, r Req) {
	a.Reqs++
	if r.OK {
		a.OK++
	} else {
		a.Err++
	}
	a.Tokens += r.Prompt + r.Compl
	a.Prompt += r.Prompt
	a.Compl += r.Compl
	a.LatSum += r.LatMs
	if r.LatMs > a.LatMax {
		a.LatMax = r.LatMs
	}
}

func (s *Stats) LoadAudit(path string) {
	f, err := os.Open(path)
	if err == nil {
		var recs []Req
		sc := bufio.NewScanner(f)
		for sc.Scan() {
			var r Req
			if json.Unmarshal(sc.Bytes(), &r) == nil {
				recs = append(recs, r)
			}
		}
		_ = f.Close()
		if len(recs) > s.maxRecs {
			recs = recs[len(recs)-s.maxRecs:]
		}
		for _, r := range recs {
			s.Record(r)
		}
	}
	s.mu.Lock()
	s.auditPath = path
	s.mu.Unlock()
}

// Record appends a finished request to the aggregates and ring buffer.
func (s *Stats) Record(r Req) {
	s.mu.Lock()
	defer s.mu.Unlock()
	inc(s.byKey, r.Key, r)
	if r.Model != "" {
		inc(s.byModel, r.Model, r)
	}
	inc(s.bySrc, r.Source, r)
	km := s.byKeyModel[r.Key]
	if km == nil {
		km = map[string]*Stat{}
		s.byKeyModel[r.Key] = km
	}
	if r.Model != "" {
		inc(km, r.Model, r)
	}
	ks := s.byKeySrc[r.Key]
	if ks == nil {
		ks = map[string]*Stat{}
		s.byKeySrc[r.Key] = ks
	}
	inc(ks, r.Source, r)
	if r.Status > 0 {
		name := strconv.Itoa(r.Status)
		a := s.byStatus[r.Status]
		if a == nil {
			a = &Stat{}
			s.byStatus[r.Status] = a
		}
		incStatus(a, name, r)
	}
	// window bucket for quota enforcement (per source-model pair, per unix hour)
	tok := r.Prompt + r.Compl
	if tok > 0 && r.Model != "" {
		key := r.Model
		if r.Source != "" {
			key = r.Source + "::" + r.Model
		}
		h := r.Time / hourSec
		hm := s.modelHour[key]
		if hm == nil {
			hm = map[int64]int64{}
			s.modelHour[key] = hm
		}
		hm[h] += tok
		if len(hm) > 24*40 {
			for k := range hm {
				if k < h-24*40 {
					delete(hm, k)
				}
			}
		}
	}
	s.recs = append(s.recs, r)
	if len(s.recs) > s.maxRecs {
		s.recs = s.recs[len(s.recs)-s.maxRecs:]
	}
	if s.auditPath != "" {
		s.rotateAuditLocked()
		appendAuditLine(s.auditPath, r)
	}
}

// rotateAuditLocked renames the audit file to <path>.<unix>.old once it
// exceeds auditRotateBytes and prunes old files beyond auditKeepOld, keeping
// the newest ones. Caller must hold s.mu.
func (s *Stats) rotateAuditLocked() {
	if s.auditPath == "" || auditRotateBytes <= 0 {
		return
	}
	if fi, err := os.Stat(s.auditPath); err == nil && fi.Size() < auditRotateBytes {
		return
	}
	ts := time.Now().Unix()
	if os.Rename(s.auditPath, fmt.Sprintf("%s.%d.old", s.auditPath, ts)) == nil {
		old, _ := filepath.Glob(s.auditPath + ".*.old")
		sort.Sort(sort.Reverse(sort.StringSlice(old)))
		for i := auditKeepOld; i < len(old); i++ {
			_ = os.Remove(old[i])
		}
	}
}

func appendAuditLine(path string, row interface{}) {
	f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer f.Close()
	if b, err := json.Marshal(row); err == nil {
		_, _ = f.Write(append(b, '\n'))
	}
}

// AppendAudit writes a generic event line (access log entry, login event,
// config change, …) to the same audit file without touching the aggregates.
func (s *Stats) AppendAudit(obj string, data map[string]interface{}) {
	row := map[string]interface{}{"obj": obj, "time": time.Now().UnixMilli()}
	for k, v := range data {
		row[k] = v
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.auditPath == "" {
		return
	}
	s.rotateAuditLocked()
	appendAuditLine(s.auditPath, row)
}

// ModelTokens returns the tokens consumed per model for one gateway key id
// (used for per-model token quota enforcement).
func (s *Stats) ModelTokens(key string) map[string]int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := map[string]int64{}
	for k, v := range s.byKeyModel[key] {
		out[k] = v.Tokens
	}
	return out
}

// KeyTokens returns the total tokens consumed by one gateway key id.
func (s *Stats) KeyTokens(key string) int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	a := s.byKey[key]
	if a == nil {
		return 0
	}
	return a.Tokens
}

// AutoPeriodSeconds maps a quota reset period to its window length in
// seconds. "" → 0 (never resets); "hour" → 1h; "week" → 7d; "month" → 30d;
// "nhour" → Hours (>=1) hours.
func AutoPeriodSeconds(period string, hours int64) int64 {
	switch period {
	case "hour":
		return hourSec
	case "week":
		return 7 * 24 * hourSec
	case "month":
		return 30 * 24 * hourSec
	case "nhour":
		if hours < 1 {
			hours = 1
		}
		return hours * hourSec
	}
	return 0
}

// WindowTokens returns the tokens billed for the model within the last `sec`
// seconds (0 = since forever).
// WindowTokens returns the tokens consumed for one model (optionally pinned
// to a single source) within the window; sec <= 0 means all time.
func (s *Stats) WindowTokens(model, source string, sec int64) int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	key := model
	if source != "" {
		key = source + "::" + model
	}
	now := time.Now().Unix()
	hm := s.modelHour[key]
	if len(hm) == 0 {
		return 0
	}
	var total int64
	if sec <= 0 {
		for _, v := range hm {
			total += v
		}
		return total
	}
	cut := now - sec
	for h, v := range hm {
		if h*hourSec >= cut {
			total += v
		}
	}
	return total
}

func rows(m map[string]*Stat) []StatsRow {
	out := make([]StatsRow, 0, len(m))
	for k, v := range m {
		out = append(out, StatsRow{Name: k, Stat: *v})
	}
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j].Reqs > out[j-1].Reqs; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	return out
}

// StatsRow is one aggregated row for the dashboard.
type StatsRow struct {
	Name string `json:"name"`
	Stat
}

// Records returns request records filtered by unix-millisecond time range and key.
func (s *Stats) Records(from, to int64, key string) []Req {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]Req, 0, len(s.recs))
	for _, r := range s.recs {
		if key != "" && r.Key != key {
			continue
		}
		if from > 0 && r.Time < from {
			continue
		}
		if to > 0 && r.Time > to {
			continue
		}
		out = append(out, r)
	}
	return out
}

// Snapshot returns the whole dashboard payload; when key != "" the records
// and aggregate views are restricted to that gateway key.
func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
	s.mu.Lock()
	defer s.mu.Unlock()
	if limit <= 0 {
		limit = s.maxRecs
	}
	start := 0
	if len(s.recs) > limit {
		start = len(s.recs) - limit
	}
	recs := s.recs[start:]
	if key != "" {
		filt := recs[:0]
		for _, r := range recs {
			if r.Key == key {
				filt = append(filt, r)
			}
		}
		recs = filt
	}
	var total Stat
	var byModel, byKey, bySrc map[string]*Stat
	if key == "" {
		byKey, byModel, bySrc = s.byKey, s.byModel, s.bySrc
	} else {
		byKey = map[string]*Stat{key: s.byKey[key]}
		byModel = s.byKeyModel[key]
		bySrc = s.byKeySrc[key]
	}
	for _, a := range byModel {
		total.Reqs += a.Reqs
		total.OK += a.OK
		total.Err += a.Err
		total.Tokens += a.Tokens
		total.Prompt += a.Prompt
		total.Compl += a.Compl
		total.LatSum += a.LatSum
		if a.LatMax > total.LatMax {
			total.LatMax = a.LatMax
		}
	}
	bs := make([]agrRow, 0, len(s.byStatus))
	for code := range s.byStatus {
		bs = append(bs, agrRow{Name: strconv.Itoa(code), Stat: *s.byStatus[code]})
	}
	sort.Slice(bs, func(i, j int) bool {
		ci, _ := strconv.Atoi(bs[i].Name)
		cj, _ := strconv.Atoi(bs[j].Name)
		return ci < cj
	})
	return map[string]interface{}{
		"active":    s.active,
		"total":     total,
		"by_key":    rows(byKey),
		"by_model":  rows(byModel),
		"by_source": rows(bySrc),
		"by_status": bs,
		"records":   append([]Req(nil), recs...),
	}
}