query-orchestration/internal/serviceconfig/observability/prometheus/prometheus.go

package prometheus

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// MetricType represents the type of Prometheus metric
type MetricType string

const (
	Counter   MetricType = "counter"
	Gauge     MetricType = "gauge"
	Histogram MetricType = "histogram"
	Summary   MetricType = "summary"
)

// MetricDefinition defines the structure for creating a new metric
type MetricDefinition struct {
	Type    MetricType
	Name    string
	Help    string
	Labels  []string
	Buckets []float64

	// Objectives is needed for Prometheus Summary metrics to define quantile calculations
	Objectives map[float64]float64
}

// Metrics represents the wrapper around Prometheus metrics
type Metrics struct {
	metrics map[string]interface{}
	mu      sync.RWMutex
	server  *http.Server
}

// MetricsConfig represents the configuration for the metrics wrapper
type MetricsConfig struct {
	Definitions []MetricDefinition
	Namespace   string
	Subsystem   string
	Port        int
	ServeHttp   bool
}

// New creates a new Metrics instance
func New(ctx context.Context, config MetricsConfig) (*Metrics, error) {
	m := &Metrics{
		metrics: make(map[string]interface{}),
	}

	for _, def := range config.Definitions {
		var metric interface{}
		var err error

		switch def.Type {
		case Counter:
			metric = prometheus.NewCounterVec(
				prometheus.CounterOpts{
					Namespace: config.Namespace,
					Subsystem: config.Subsystem,
					Name:      def.Name,
					Help:      def.Help,
				},
				def.Labels,
			)
		case Gauge:
			metric = prometheus.NewGaugeVec(
				prometheus.GaugeOpts{
					Namespace: config.Namespace,
					Subsystem: config.Subsystem,
					Name:      def.Name,
					Help:      def.Help,
				},
				def.Labels,
			)
		case Histogram:
			opts := prometheus.HistogramOpts{
				Namespace: config.Namespace,
				Subsystem: config.Subsystem,
				Name:      def.Name,
				Help:      def.Help,
			}
			if len(def.Buckets) > 0 {
				opts.Buckets = def.Buckets
			}
			metric = prometheus.NewHistogramVec(opts, def.Labels)
		case Summary:
			opts := prometheus.SummaryOpts{
				Namespace: config.Namespace,
				Subsystem: config.Subsystem,
				Name:      def.Name,
				Help:      def.Help,
			}
			if def.Objectives != nil {
				opts.Objectives = def.Objectives
			}
			metric = prometheus.NewSummaryVec(opts, def.Labels)
		default:
			return nil, fmt.Errorf("unsupported metric type: %s", def.Type)
		}

		collector, ok := metric.(prometheus.Collector)
		if !ok {
			return nil, fmt.Errorf("metric %s is not a valid prometheus collector", def.Name)
		}
		if err = prometheus.Register(collector); err != nil {
			return nil, fmt.Errorf("failed to register metric %s: %w", def.Name, err)
		}

		m.metrics[def.Name] = metric
	}

	if config.ServeHttp {
		mux := http.NewServeMux()
		mux.Handle("/metrics", promhttp.Handler())
		mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
		})

		m.server = &http.Server{
			Addr:              fmt.Sprintf(":%d", config.Port),
			Handler:           mux,
			ReadHeaderTimeout: time.Minute,
		}

		go func() {
			if err := m.server.ListenAndServe(); err != http.ErrServerClosed {
				fmt.Printf("Metrics server error: %v\n", err)
			}
		}()

		// Listen for context cancellation
		go func() {
			<-ctx.Done()
			if err := m.server.Shutdown(context.Background()); err != nil {
				fmt.Printf("Error shutting down metrics server: %v\n", err)
			}
		}()
	}

	return m, nil
}

// Shutdown gracefully shuts down the metrics server
func (m *Metrics) Shutdown(ctx context.Context) error {
	if m.server != nil {
		return m.server.Shutdown(ctx)
	}
	return nil
}

// RecordMetric records a value for a given metric
func (m *Metrics) RecordMetric(name string, value float64, labels map[string]string) error {
	m.mu.RLock()
	metric, exists := m.metrics[name]
	m.mu.RUnlock()

	if !exists {
		return fmt.Errorf("metric %s not found", name)
	}

	labelValues := make([]string, 0, len(labels))
	for _, v := range labels {
		labelValues = append(labelValues, v)
	}

	switch m := metric.(type) {
	case *prometheus.CounterVec:
		m.WithLabelValues(labelValues...).Add(value)
	case *prometheus.GaugeVec:
		m.WithLabelValues(labelValues...).Set(value)
	case *prometheus.HistogramVec:
		m.WithLabelValues(labelValues...).Observe(value)
	case *prometheus.SummaryVec:
		m.WithLabelValues(labelValues...).Observe(value)
	default:
		return fmt.Errorf("unknown metric type for %s", name)
	}

	return nil
}