query-orchestration/cmd/metricsExample_test/main.go

// This main.go serves as an example of how to use the custom Prometheus metrics package.
// It demonstrates:
// 1. Setting up metrics with proper context handling
// 2. Configuring different types of metrics (Counter, Gauge, Histogram, Summary)
// 3. Recording metrics safely in a concurrent environment
// 4. Implementing graceful shutdown
// 5. Proper error handling
package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"queryorchestration/internal/serviceconfig/observability/prometheus"
	"syscall"
	"time"
)

func main() {
	// Create a context that we can cancel
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Define metrics configuration
	// This example shows all supported metric types and their configurations
	config := prometheus.MetricsConfig{
		Port:      8080,
		Namespace: "myapp",
		Subsystem: "api",
		// When ServeHttp is true, the package will start its own HTTP server
		// If you want to use your own HTTP server (e.g., with Echo middleware),
		// set this to false and configure your server to handle the /metrics endpoint
		ServeHttp: true,
		Definitions: []prometheus.MetricDefinition{
			{
				// Counter: Only increases, never decreases
				Type:   prometheus.Counter,
				Name:   "api_requests_total",
				Help:   "Total number of API requests",
				Labels: []string{"method", "endpoint", "status"},
			},
			{
				// Gauge: Can increase and decrease
				Type:   prometheus.Gauge,
				Name:   "active_users",
				Help:   "Number of currently active users",
				Labels: []string{"user_type"},
			},
			{
				// Histogram: Tracks distribution of values
				Type:    prometheus.Histogram,
				Name:    "request_duration_seconds",
				Help:    "Request duration distribution",
				Labels:  []string{"endpoint"},
				Buckets: []float64{0.1, 0.5, 1, 2, 5}, // Define your bucket boundaries
			},
			{
				// Summary: Calculates quantiles over a sliding time window
				Type:   prometheus.Summary,
				Name:   "response_size_bytes",
				Help:   "Response size distribution",
				Labels: []string{"endpoint"},
				Objectives: map[float64]float64{
					0.5:  0.05,  // 50th percentile with 5% error
					0.9:  0.01,  // 90th percentile with 1% error
					0.99: 0.001, // 99th percentile with 0.1% error
				},
			},
		},
	}

	// Create new metrics instance with context
	metrics, err := prometheus.New(ctx, config)
	if err != nil {
		log.Fatalf("Failed to create metrics: %v", err)
	}

	// Start a goroutine to simulate metric recording
	go simulateMetrics(ctx, metrics)

	// Set up graceful shutdown
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	fmt.Println("Running... Press Ctrl+C to exit")

	// Wait for interrupt signal
	<-sigChan
	fmt.Println("\nShutting down...")

	// Cancel context to stop metric simulation
	cancel()

	// Create a timeout context for shutdown
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer shutdownCancel()

	// Gracefully shutdown the metrics server
	if err := metrics.Shutdown(shutdownCtx); err != nil {
		log.Printf("Error shutting down metrics server: %v", err)
	}
}

// simulateMetrics demonstrates how to record different types of metrics
// in a concurrent environment. It takes a context for cancellation and
// the metrics instance.
func simulateMetrics(ctx context.Context, metrics *prometheus.Metrics) {
	// Sample data for simulation
	endpoints := []string{"/api/users", "/api/products", "/api/orders"}
	methods := []string{"GET", "POST", "PUT", "DELETE"}
	statuses := []string{"200", "400", "500"}
	userTypes := []string{"free", "premium", "enterprise"}

	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			// Simulate request counter (Counter type)
			endpoint := endpoints[rand.Intn(len(endpoints))]
			method := methods[rand.Intn(len(methods))]
			status := statuses[rand.Intn(len(statuses))]

			err := metrics.RecordMetric("api_requests_total", 1, map[string]string{
				"method":   method,
				"endpoint": endpoint,
				"status":   status,
			})
			if err != nil {
				log.Printf("Error recording request metric: %v", err)
			}

			// Simulate active users (Gauge type)
			userType := userTypes[rand.Intn(len(userTypes))]
			activeUsers := rand.Float64() * 100
			err = metrics.RecordMetric("active_users", activeUsers, map[string]string{
				"user_type": userType,
			})
			if err != nil {
				log.Printf("Error recording active users metric: %v", err)
			}

			// Simulate request duration (Histogram type)
			duration := rand.Float64() * 3 // Random duration between 0 and 3 seconds
			err = metrics.RecordMetric("request_duration_seconds", duration, map[string]string{
				"endpoint": endpoint,
			})
			if err != nil {
				log.Printf("Error recording duration metric: %v", err)
			}

			// Simulate response size (Summary type)
			responseSize := rand.Float64() * 5000 // Random size between 0 and 5000 bytes
			err = metrics.RecordMetric("response_size_bytes", responseSize, map[string]string{
				"endpoint": endpoint,
			})
			if err != nil {
				log.Printf("Error recording response size metric: %v", err)
			}
		}
	}
}