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

package observability

import (
	"context"
	"log/slog"

	"queryorchestration/internal/serviceconfig/build"
	"queryorchestration/internal/serviceconfig/observability/prometheus"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
)

type ConfigProvider interface {
	IsOtelEnabled() bool
	SetOtel(context.Context) func() error
	GetMetrics() *prometheus.Metrics
	SetMetrics(metrics *prometheus.Metrics)
}

type ObsConfig struct {
	EnableOtel bool `env:"ENABLE_OTEL" envDefault:"false"`
	metrics    *prometheus.Metrics
}

func (o *ObsConfig) GetMetrics() *prometheus.Metrics {
	return o.metrics
}

func (o *ObsConfig) SetMetrics(metrics *prometheus.Metrics) {
	o.metrics = metrics
}

func (o *ObsConfig) IsOtelEnabled() bool {
	return o.EnableOtel
}

func (o *ObsConfig) SetOtel(ctx context.Context) func() error {
	if !o.EnableOtel {
		slog.Warn("OpenTelemetry is disabled. Set ENABLE_OTEL to true to enable.")
		return func() error { return nil }
	}

	exporter, err := otlptracegrpc.New(ctx)
	if err != nil {
		return func() error { return nil }
	}

	tp := sdktrace.NewTracerProvider(
		sdktrace.WithBatcher(exporter),
	)

	otel.SetTracerProvider(tp)

	return func() error {
		err := tp.Shutdown(ctx)
		if err != nil {
			return err
		}
		return nil
	}
}

const (
	SubSystem = "RunnerService"
)

func InitializeMetrics(ctx context.Context, cfg ConfigProvider) error {
	standardDefinitions := prometheus.NewStandardMetricDefinitions()
	config := prometheus.MetricsConfig{
		Port:        8080,
		Namespace:   build.SystemName,
		Subsystem:   SubSystem,
		ServeHttp:   true,
		Definitions: standardDefinitions,
	}

	metrics, err := prometheus.New(ctx, config)
	if err != nil {
		slog.Error("Failed to create metrics object", "error", err)
		return err
	}
	slog.Info("Metrics server started on port 8080")
	cfg.SetMetrics(metrics)

	err = cfg.GetMetrics().RecordLifecycleEvent(1.0, "startup", "main")
	if err != nil {
		slog.Error("Failed to create metrics at startup", "error", err)
		return err
	}

	return nil
}