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
	breakerState      int64 // 0=closed, 1=open, 2=half_open

	// 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))
}

// SetBreakerState 设置熔断器状态 gauge（0=closed, 1=open, 2=half_open）。
func (m *Metrics) SetBreakerState(state int) {
	atomic.StoreInt64(&m.breakerState, int64(state))
}

// 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
		fmt.Fprintf(w, "# HELP edgeai_requests_total Total requests by status\n")
		fmt.Fprintf(w, "# TYPE edgeai_requests_total counter\n")
		m.mu.RLock()
		for key, val := range m.requestsTotal {
			fmt.Fprintf(w, "edgeai_requests_total{%s} %d\n", key, val)
		}
		fmt.Fprintf(w, "# HELP edgeai_tasks_total Total tasks by final state\n")
		fmt.Fprintf(w, "# TYPE edgeai_tasks_total counter\n")
		for key, val := range m.tasksTotal {
			fmt.Fprintf(w, "edgeai_tasks_total{%s} %d\n", key, val)
		}
		m.mu.RUnlock()

		fmt.Fprintf(w, "# HELP edgeai_tokens_input_total Total input tokens processed\n")
		fmt.Fprintf(w, "# TYPE edgeai_tokens_input_total counter\n")
		fmt.Fprintf(w, "edgeai_tokens_input_total %d\n", atomic.LoadInt64(&m.tokensInputTotal))

		fmt.Fprintf(w, "# HELP edgeai_tokens_output_total Total output tokens generated\n")
		fmt.Fprintf(w, "# TYPE edgeai_tokens_output_total counter\n")
		fmt.Fprintf(w, "edgeai_tokens_output_total %d\n", atomic.LoadInt64(&m.tokensOutputTotal))

		fmt.Fprintf(w, "# HELP edgeai_cancellations_total Total cancelled requests\n")
		fmt.Fprintf(w, "# TYPE edgeai_cancellations_total counter\n")
		fmt.Fprintf(w, "edgeai_cancellations_total %d\n", atomic.LoadInt64(&m.cancellationsTotal))

		fmt.Fprintf(w, "# HELP edgeai_queue_timeouts_total Total queue timeouts\n")
		fmt.Fprintf(w, "# TYPE edgeai_queue_timeouts_total counter\n")
		fmt.Fprintf(w, "edgeai_queue_timeouts_total %d\n", atomic.LoadInt64(&m.queueTimeoutsTotal))

		fmt.Fprintf(w, "# HELP edgeai_first_token_timeouts_total Total first token timeouts\n")
		fmt.Fprintf(w, "# TYPE edgeai_first_token_timeouts_total counter\n")
		fmt.Fprintf(w, "edgeai_first_token_timeouts_total %d\n", atomic.LoadInt64(&m.firstTokenTimeoutsTotal))

		fmt.Fprintf(w, "# HELP edgeai_inference_timeouts_total Total inference timeouts\n")
		fmt.Fprintf(w, "# TYPE edgeai_inference_timeouts_total counter\n")
		fmt.Fprintf(w, "edgeai_inference_timeouts_total %d\n", atomic.LoadInt64(&m.inferenceTimeoutsTotal))

		fmt.Fprintf(w, "# HELP edgeai_degraded_requests_total Total requests served with degradation\n")
		fmt.Fprintf(w, "# TYPE edgeai_degraded_requests_total counter\n")
		fmt.Fprintf(w, "edgeai_degraded_requests_total %d\n", atomic.LoadInt64(&m.degradedRequestsTotal))

		// Gauges
		fmt.Fprintf(w, "# HELP edgeai_queue_length Current queue length\n")
		fmt.Fprintf(w, "# TYPE edgeai_queue_length gauge\n")
		fmt.Fprintf(w, "edgeai_queue_length %d\n", atomic.LoadInt64(&m.queueLength))

		fmt.Fprintf(w, "# HELP edgeai_running_tasks Current running tasks\n")
		fmt.Fprintf(w, "# TYPE edgeai_running_tasks gauge\n")
		fmt.Fprintf(w, "edgeai_running_tasks %d\n", atomic.LoadInt64(&m.runningTasks))

		fmt.Fprintf(w, "# HELP edgeai_active_sessions Current active sessions\n")
		fmt.Fprintf(w, "# TYPE edgeai_active_sessions gauge\n")
		fmt.Fprintf(w, "edgeai_active_sessions %d\n", atomic.LoadInt64(&m.activeSessions))

		fmt.Fprintf(w, "# HELP edgeai_backpressure_level Current backpressure level (0-3)\n")
		fmt.Fprintf(w, "# TYPE edgeai_backpressure_level gauge\n")
		fmt.Fprintf(w, "edgeai_backpressure_level %d\n", atomic.LoadInt64(&m.backpressureLevel))

		fmt.Fprintf(w, "# HELP edgeai_circuit_breaker_state Circuit breaker state (0=closed, 1=open, 2=half_open)\n")
		fmt.Fprintf(w, "# TYPE edgeai_circuit_breaker_state gauge\n")
		fmt.Fprintf(w, "edgeai_circuit_breaker_state %d\n", atomic.LoadInt64(&m.breakerState))

		// Histograms
		fmt.Fprintf(w, "# HELP edgeai_gateway_latency_bucket Gateway latency distribution in milliseconds\n")
		fmt.Fprintf(w, "# TYPE edgeai_gateway_latency_bucket histogram\n")
		m.mu.RLock()
		for bucket, count := range m.gatewayLatencyBuckets {
			fmt.Fprintf(w, "edgeai_gateway_latency_bucket{%s} %d\n", bucket, count)
		}
		fmt.Fprintf(w, "# HELP edgeai_first_token_latency_bucket First token latency distribution in milliseconds\n")
		fmt.Fprintf(w, "# TYPE edgeai_first_token_latency_bucket histogram\n")
		for bucket, count := range m.firstTokenLatencyBuckets {
			fmt.Fprintf(w, "edgeai_first_token_latency_bucket{%s} %d\n", bucket, count)
		}
		m.mu.RUnlock()
	}
}