exporter.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log/slog"
	"net/http"
	"net/http/pprof"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
	"github.com/prometheus/client_golang/prometheus/collectors/version"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	cversion "github.com/prometheus/common/version"

	cloudcost_exporter "github.com/grafana/cloudcost-exporter"
	"github.com/grafana/cloudcost-exporter/cmd/exporter/config"
	"github.com/grafana/cloudcost-exporter/cmd/exporter/web"
	"github.com/grafana/cloudcost-exporter/pkg/aws"
	"github.com/grafana/cloudcost-exporter/pkg/azure"
	"github.com/grafana/cloudcost-exporter/pkg/google"
	"github.com/grafana/cloudcost-exporter/pkg/logger"
	"github.com/grafana/cloudcost-exporter/pkg/provider"
)

func main() {
	var cfg config.Config
	providerFlags(flag.CommandLine, &cfg)
	operationalFlags(&cfg)
	flag.Parse()

	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer cancel()

	logs := setupLogger(cfg.LoggerOpts.Level, cfg.LoggerOpts.Output, cfg.LoggerOpts.Type)
	logs.LogAttrs(ctx, slog.LevelInfo, "Starting cloudcost-exporter",
		slog.String("version", cversion.Info()),
		slog.String("build_context", cversion.BuildContext()),
	)
	cfg.Logger = logs

	if cfg.Providers.GCP.BucketProjectsDeprecated {
		logs.LogAttrs(ctx, slog.LevelWarn, "'gcp.bucket-projects' is deprecated and will be removed in a future version. Use '--gcp.projects' instead.")
	}

	csp, err := selectProvider(ctx, &cfg)
	if err != nil {
		logs.LogAttrs(ctx, slog.LevelError, "Error selecting provider",
			slog.String("message", err.Error()),
			slog.String("provider", cfg.Provider),
		)
		os.Exit(1)
	}

	err = runServer(ctx, &cfg, csp, logs)
	if err != nil {
		logs.LogAttrs(ctx, slog.LevelError, "Error running server", slog.String("message", err.Error()))
		os.Exit(1)
	}
}

// providerFlags is a helper method that is responsible for setting up the flags that are used to configure the provider.
// TODO: This should probably be moved over to the config package.
func providerFlags(fs *flag.FlagSet, cfg *config.Config) {
	flag.StringVar(&cfg.Provider, "provider", "aws", "aws, gcp, or azure")
	fs.StringVar(&cfg.Providers.AWS.Profile, "aws.profile", "", "AWS Profile to authenticate with.")
	fs.Var(&cfg.Providers.GCP.Projects, "gcp.projects", "GCP project(s).")
	fs.Var(config.NewDeprecatedStringSliceFlag(&cfg.Providers.GCP.Projects, &cfg.Providers.GCP.BucketProjectsDeprecated), "gcp.bucket-projects", "GCP project(s). (deprecated: use --gcp.projects instead)")
	fs.Var(&cfg.Providers.AWS.Services, "aws.services", "AWS service(s).")
	fs.Var(&cfg.Providers.AWS.ExcludeRegions, "aws.exclude-regions", "AWS region(s) to exclude from cost collection.")
	fs.Var(&cfg.Providers.Azure.Services, "azure.services", "Azure service(s): AKS, blob (comma-separated and/or repeat flag; case-insensitive).")
	fs.Var(&cfg.Providers.GCP.Services, "gcp.services", "GCP service(s).")
	flag.StringVar(&cfg.Providers.AWS.Region, "aws.region", "", "AWS region")
	flag.StringVar(&cfg.Providers.AWS.RoleARN, "aws.roleARN", "", "Optional AWS role ARN to assume for cross-account access.")
	// TODO - PUT PROJECT-ID UNDER GCP
	flag.StringVar(&cfg.ProjectID, "project-id", "", "Project ID to target.")
	flag.StringVar(&cfg.Providers.Azure.SubscriptionId, "azure.subscription-id", "", "Azure subscription ID to pull data from.")
	flag.IntVar(&cfg.Providers.GCP.DefaultGCSDiscount, "gcp.default-discount", 19, "GCP default discount")
}

// operationalFlags is a helper method that is responsible for setting up the flags that are used to configure the operational aspects of the application.
// TODO: This should probably be moved over to the config package.
func operationalFlags(cfg *config.Config) {
	flag.DurationVar(&cfg.Collector.ScrapeInterval, "scrape-interval", 1*time.Hour, "Scrape interval")
	flag.DurationVar(&cfg.Collector.Timeout, "collector-interval", 1*time.Minute, "Context timeout for collectors")
	flag.DurationVar(&cfg.Server.Timeout, "server-timeout", 30*time.Second, "Server timeout")
	flag.StringVar(&cfg.Server.Address, "server.address", ":8080", "Default address for the server to listen on.")
	flag.StringVar(&cfg.Server.Path, "server.path", "/metrics", "Default path for the server to listen on.")
	flag.StringVar(&cfg.LoggerOpts.Level, "log.level", "info", "Log level: debug, info, warn, error")
	flag.StringVar(&cfg.LoggerOpts.Output, "log.output", "stdout", "Log output stream: stdout, stderr, file")
	flag.StringVar(&cfg.LoggerOpts.Type, "log.type", "text", "Log type: json, text")
}

// setupLogger is a helper method that is responsible for creating a structured logger that is used throughout the application.
// It sets the log level, output, and type of log.
func setupLogger(level string, output string, logtype string) *slog.Logger {
	handler := logger.NewLevelHandler(logger.GetLogLevel(level), logger.HandlerForOutput(logtype, logger.WriterForOutput(output)))
	return slog.New(handler)
}

// runServer is a helper method that is responsible for starting the metrics server and handling shutdown signals.
func runServer(ctx context.Context, cfg *config.Config, csp provider.Provider, log *slog.Logger) error {
	mux := http.NewServeMux()

	mux.HandleFunc("/debug/pprof/", pprof.Index)
	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)

	mux.HandleFunc("/", web.HomePageHandler(cfg.Server.Path)) // landing page

	registryHandler, err := createPromRegistryHandler(csp, regionFromConfig(cfg)) // prom metrics handler
	if err != nil {
		return err
	}
	mux.Handle(cfg.Server.Path, registryHandler) // prom metrics handler (/metrics)

	server := &http.Server{Addr: cfg.Server.Address, Handler: mux}
	errChan := make(chan error)

	go func() {
		log.LogAttrs(ctx, slog.LevelInfo, "Starting server",
			slog.String("address", cfg.Server.Address),
			slog.String("path", cfg.Server.Path))
		errChan <- server.ListenAndServe()
	}()

	select {
	case <-ctx.Done():
		log.LogAttrs(ctx, slog.LevelInfo, "Shutting down server")
		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.Server.Timeout)
		defer shutdownCancel()

		err := server.Shutdown(shutdownCtx)
		if err != nil {
			return fmt.Errorf("error shutting down server: %w", err)
		}
	case err := <-errChan:
		if !errors.Is(err, http.ErrServerClosed) {
			return fmt.Errorf("error running server: %w", err)
		}
	}

	return nil
}

func regionFromConfig(cfg *config.Config) string {
	switch cfg.Provider {
	case "aws":
		return cfg.Providers.AWS.Region
	case "gcp":
		return cfg.Providers.GCP.Region
	default:
		// TODO: add region support for Azure (currently has no region in config)
		return ""
	}
}

func createPromRegistryHandler(csp provider.Provider, region string) (http.Handler, error) {
	var subsystem = "metrics_handler"
	requestDuration := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:                           prometheus.BuildFQName(cloudcost_exporter.ExporterName, subsystem, "request_duration_seconds"),
			Help:                           "Duration of HTTP requests in seconds for the metrics endpoint",
			NativeHistogramBucketFactor:    1.1,
			NativeHistogramMaxBucketNumber: 10,
			Buckets:                        []float64{1, 2, 5, 7, 10, 20, 40, 80},
			ConstLabels:                    prometheus.Labels{"region": region},
		},
		[]string{"method"},
	)

	requestCounter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name:        prometheus.BuildFQName(cloudcost_exporter.ExporterName, subsystem, "requests_total"),
			Help:        "Total number of HTTP requests for the metrics endpoint",
			ConstLabels: prometheus.Labels{"region": region},
		},
		[]string{"code", "method"},
	)
	registry := prometheus.NewRegistry()
	registry.MustRegister(
		collectors.NewBuildInfoCollector(),
		collectors.NewGoCollector(),
		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
		version.NewCollector(cloudcost_exporter.ExporterName),
		csp,
		requestCounter,
		requestDuration,
	)
	err := csp.RegisterCollectors(registry)
	if err != nil {
		return nil, err
	}

	handler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{
		EnableOpenMetrics: true,
	})

	return promhttp.InstrumentHandlerDuration(
		requestDuration,
		promhttp.InstrumentHandlerCounter(
			requestCounter,
			handler,
		),
	), nil
}

func selectProvider(ctx context.Context, cfg *config.Config) (provider.Provider, error) {
	return selectProviderWith(ctx, cfg,
		func(ctx context.Context, cfg *aws.Config) (provider.Provider, error) { return aws.New(ctx, cfg) },
		func(ctx context.Context, cfg *azure.Config) (provider.Provider, error) { return azure.New(ctx, cfg) },
		func(ctx context.Context, cfg *google.Config) (provider.Provider, error) { return google.New(ctx, cfg) },
	)
}

type newProviderFunc[T any] func(context.Context, T) (provider.Provider, error)

// expandAzureServices flattens repeated -azure.services values and comma-separated tokens (parity with AWS/GCP flags).
func expandAzureServices(flags config.StringSliceFlag) []string {
	var out []string
	for _, s := range flags {
		for _, part := range strings.Split(s, ",") {
			if t := strings.TrimSpace(part); t != "" {
				out = append(out, t)
			}
		}
	}
	return out
}

func selectProviderWith(
	ctx context.Context,
	cfg *config.Config,
	newAWS newProviderFunc[*aws.Config],
	newAzure newProviderFunc[*azure.Config],
	newGCP newProviderFunc[*google.Config],
) (provider.Provider, error) {
	// Set collector timeout with 1 minute default
	collectorTimeout := cfg.Collector.Timeout
	if collectorTimeout == 0 {
		collectorTimeout = 1 * time.Minute
	}

	switch cfg.Provider {
	case "azure":
		return newAzure(ctx, &azure.Config{
			Logger:           cfg.Logger,
			SubscriptionId:   cfg.Providers.Azure.SubscriptionId,
			ScrapeInterval:   cfg.Collector.ScrapeInterval,
			Services:         expandAzureServices(cfg.Providers.Azure.Services),
			CollectorTimeout: collectorTimeout,
		})
	case "aws":
		return newAWS(ctx, &aws.Config{
			Logger:           cfg.Logger,
			Region:           cfg.Providers.AWS.Region,
			Profile:          cfg.Providers.AWS.Profile,
			RoleARN:          cfg.Providers.AWS.RoleARN,
			ScrapeInterval:   cfg.Collector.ScrapeInterval,
			Services:         strings.Split(cfg.Providers.AWS.Services.String(), ","),
			ExcludeRegions:   strings.Split(cfg.Providers.AWS.ExcludeRegions.String(), ","),
			CollectorTimeout: collectorTimeout,
		})

	case "gcp":
		return newGCP(ctx, &google.Config{
			Logger:           cfg.Logger,
			ProjectId:        cfg.ProjectID,
			Region:           cfg.Providers.GCP.Region,
			Projects:         cfg.Providers.GCP.Projects.String(),
			DefaultDiscount:  cfg.Providers.GCP.DefaultGCSDiscount,
			ScrapeInterval:   cfg.Collector.ScrapeInterval,
			Services:         strings.Split(cfg.Providers.GCP.Services.String(), ","),
			CollectorTimeout: collectorTimeout,
		})

	default:
		return nil, fmt.Errorf("unknown provider")
	}
}