AIRouter/internal/observability/metrics.go

package observability

import (
	"fmt"
	"net/http"
	"sync"
	"sync/atomic"
)

// Metrics holds all Prometheus-compatible metrics for the gateway.
type Metrics struct {
	mu sync.RWMutex

	// Counters
	requestsTotal           map[string]int64 // by status
	tasksTotal              map[string]int64 // by state
	tokensInputTotal        int64
	tokensOutputTotal       int64
	cancellationsTotal      int64
	queueTimeoutsTotal      int64
	firstTokenTimeoutsTotal int64
	inferenceTimeoutsTotal  int64
	degradedRequestsTotal   int64

	// Gauges
	queueLength       int64
	runningTasks      int64
	activeSessions    int64
	backpressureLevel int64

	// Histograms (simplified as buckets)
	gatewayLatencyBuckets    map[string]int64
	firstTokenLatencyBuckets map[string]int64
}

// NewMetrics creates a new Metrics instance.
func NewMetrics() *Metrics {
	return &Metrics{
		requestsTotal:            make(map[string]int64),
		tasksTotal:               make(map[string]int64),
		gatewayLatencyBuckets:    make(map[string]int64),
		firstTokenLatencyBuckets: make(map[string]int64),
	}
}

// IncRequest increments the request counter by status.
func (m *Metrics) IncRequest(status string) {
	key := fmt.Sprintf("status=%s", status)
	m.mu.Lock()
	m.requestsTotal[key]++
	m.mu.Unlock()
}

// IncTask increments the task counter by final state.
func (m *Metrics) IncTask(state string) {
	key := fmt.Sprintf("state=%s", state)
	m.mu.Lock()
	m.tasksTotal[key]++
	m.mu.Unlock()
}

// AddTokens adds to the token counters.
func (m *Metrics) AddTokens(input, output int) {
	atomic.AddInt64(&m.tokensInputTotal, int64(input))
	atomic.AddInt64(&m.tokensOutputTotal, int64(output))
}

// IncCancellation increments the cancellation counter.
func (m *Metrics) IncCancellation() {
	atomic.AddInt64(&m.cancellationsTotal, 1)
}

// IncQueueTimeout increments the queue timeout counter.
func (m *Metrics) IncQueueTimeout() {
	atomic.AddInt64(&m.queueTimeoutsTotal, 1)
}

// IncFirstTokenTimeout increments the first token timeout counter.
func (m *Metrics) IncFirstTokenTimeout() {
	atomic.AddInt64(&m.firstTokenTimeoutsTotal, 1)
}

// IncInferenceTimeout increments the inference timeout counter.
func (m *Metrics) IncInferenceTimeout() {
	atomic.AddInt64(&m.inferenceTimeoutsTotal, 1)
}

// IncDegraded increments the degraded request counter.
func (m *Metrics) IncDegraded() {
	atomic.AddInt64(&m.degradedRequestsTotal, 1)
}

// SetQueueLength sets the current queue length gauge.
func (m *Metrics) SetQueueLength(n int) {
	atomic.StoreInt64(&m.queueLength, int64(n))
}

// SetRunningTasks sets the running tasks gauge.
func (m *Metrics) SetRunningTasks(n int) {
	atomic.StoreInt64(&m.runningTasks, int64(n))
}

// SetActiveSessions sets the active sessions gauge.
func (m *Metrics) SetActiveSessions(n int) {
	atomic.StoreInt64(&m.activeSessions, int64(n))
}

// SetBackpressureLevel sets the backpressure level gauge.
func (m *Metrics) SetBackpressureLevel(level int) {
	atomic.StoreInt64(&m.backpressureLevel, int64(level))
}

// ObserveGatewayLatency records gateway latency in a histogram bucket.
func (m *Metrics) ObserveGatewayLatency(ms int64) {
	bucket := latencyBucket(ms)
	m.mu.Lock()
	m.gatewayLatencyBuckets[bucket]++
	m.mu.Unlock()
}

// ObserveFirstTokenLatency records first token latency in a histogram bucket.
func (m *Metrics) ObserveFirstTokenLatency(ms int64) {
	bucket := latencyBucket(ms)
	m.mu.Lock()
	m.firstTokenLatencyBuckets[bucket]++
	m.mu.Unlock()
}

func latencyBucket(ms int64) string {
	buckets := []int64{5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000}
	for _, b := range buckets {
		if ms <= b {
			return fmt.Sprintf("le_%d", b)
		}
	}
	return "le_inf"
}

// Handler returns an http.HandlerFunc that writes Prometheus-format metrics.
func (m *Metrics) Handler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")

		// Counters
		m.mu.RLock()
		for key, val := range m.requestsTotal {
			fmt.Fprintf(w, "edgeai_requests_total{%s} %d\n", key, val)
		}
		for key, val := range m.tasksTotal {
			fmt.Fprintf(w, "edgeai_tasks_total{%s} %d\n", key, val)
		}
		m.mu.RUnlock()

		fmt.Fprintf(w, "edgeai_tokens_input_total %d\n", atomic.LoadInt64(&m.tokensInputTotal))
		fmt.Fprintf(w, "edgeai_tokens_output_total %d\n", atomic.LoadInt64(&m.tokensOutputTotal))
		fmt.Fprintf(w, "edgeai_cancellations_total %d\n", atomic.LoadInt64(&m.cancellationsTotal))
		fmt.Fprintf(w, "edgeai_queue_timeouts_total %d\n", atomic.LoadInt64(&m.queueTimeoutsTotal))
		fmt.Fprintf(w, "edgeai_first_token_timeouts_total %d\n", atomic.LoadInt64(&m.firstTokenTimeoutsTotal))
		fmt.Fprintf(w, "edgeai_inference_timeouts_total %d\n", atomic.LoadInt64(&m.inferenceTimeoutsTotal))
		fmt.Fprintf(w, "edgeai_degraded_requests_total %d\n", atomic.LoadInt64(&m.degradedRequestsTotal))

		// Gauges
		fmt.Fprintf(w, "edgeai_queue_length %d\n", atomic.LoadInt64(&m.queueLength))
		fmt.Fprintf(w, "edgeai_running_tasks %d\n", atomic.LoadInt64(&m.runningTasks))
		fmt.Fprintf(w, "edgeai_active_sessions %d\n", atomic.LoadInt64(&m.activeSessions))
		fmt.Fprintf(w, "edgeai_backpressure_level %d\n", atomic.LoadInt64(&m.backpressureLevel))

		// Histograms
		m.mu.RLock()
		for bucket, count := range m.gatewayLatencyBuckets {
			fmt.Fprintf(w, "edgeai_gateway_latency_bucket{%s} %d\n", bucket, count)
		}
		for bucket, count := range m.firstTokenLatencyBuckets {
			fmt.Fprintf(w, "edgeai_first_token_latency_bucket{%s} %d\n", bucket, count)
		}
		m.mu.RUnlock()
	}
}