labmon/main.go

package main

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"net/http/pprof"
	"os"
	"os/signal"
	"time"

	"git.t-juice.club/torjus/labmon/config"
	"git.t-juice.club/torjus/labmon/otel"
	"git.t-juice.club/torjus/labmon/stepmon"
	"git.t-juice.club/torjus/labmon/tlsconmon"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const Version = "0.1.2"

func LoadConfig() (*config.Config, error) {
	path := "labmon.toml"
	if len(os.Args) > 1 {
		path = os.Args[1]
	}
	config, err := config.FromFile(path)
	if err != nil {
		return nil, err
	}

	return config, nil
}

func main() {
	// Load cfg
	cfg, err := LoadConfig()
	if err != nil {
		fmt.Printf("Error loading config: %v\n", err)
		os.Exit(1)
	}

	// Setup logger
	logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
		Level: slog.LevelDebug,
	}))

	// Setup graceful shutdown
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
	defer cancel()

	shutdownDone := make(chan struct{}, 1)

	// Setup otel
	otelShutdown := func(ctx context.Context) error { return nil }
	if cfg.Tracing {
		var err error
		otelShutdown, err = otel.SetupOTEL(ctx, cfg.TracingEndpoint)
		if err != nil {
			fmt.Printf("Error setting up OpenTelemetry: %v\n", err)
			os.Exit(1)
		}
	}

	// Setup stepmons
	var stepmons []*stepmon.StepMonitor
	for _, s := range cfg.StepMonitors {
		if s.Enabled {
			sm := stepmon.NewStepMonitor(s.BaseURL, s.RootID)
			sm.SetLogger(logger)
			stepmons = append(stepmons, sm)
		}
	}

	// Start stepmons
	for _, sm := range stepmons {
		go func(sm *stepmon.StepMonitor) {
			sm.Start()
		}(sm)
	}

	// Setup tlsconmons
	var tlsconmons []*tlsconmon.TLSConnectionMonitor
	for _, t := range cfg.TLSConnectionMonitors {
		if t.Enabled {
			duration, err := time.ParseDuration(t.Duration)
			if err != nil {
				logger.Error("Failed to parse duration", "duration", t.Duration, "error", err)
				os.Exit(1)
			}
			tm, err := tlsconmon.NewTLSConnectionMonitor(t.Address, t.Verify, t.ExtraCAPaths, duration)
			if err != nil {
				logger.Error("Failed to create TLSConnectionMonitor", "address", t.Address, "error", err)
				os.Exit(1)
			}
			tm.SetLogger(logger)

			tlsconmons = append(tlsconmons, tm)
		}
	}

	// Start tlsconmons
	for _, tm := range tlsconmons {
		go func(tm *tlsconmon.TLSConnectionMonitor) {
			tm.Start(ctx)
		}(tm)
	}

	// Start http server
	srv := &http.Server{}
	srv.Addr = cfg.ListenAddr

	mux := http.NewServeMux()
	srv.Handler = mux

	mux.Handle("/metrics", promhttp.Handler())

	if cfg.Profiling {
		logger.Info("Profiling enabled, exposing /debug/pprof")
		mux.HandleFunc("/debug/pprof/", pprof.Index)
		mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
	}

	// Start http server
	go func() {
		logger.Info("Starting HTTP server", "addr", cfg.ListenAddr)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logger.Error("HTTP server error", "error", err)
		}
	}()

	// Wait for shutdown signal
	go func() {
		<-ctx.Done()
		logger.Debug("Shutdown signal received")
		// Shutdown metrics server
		srv.Shutdown(context.Background())
		logger.Debug("HTTP server shutdown complete")

		// Shutdown stepmons
		for _, sm := range stepmons {
			sm.Shutdown()
			logger.Debug("StepMonitor shutdown complete", "root_id", sm.RootID)
		}
		shutdownDone <- struct{}{}

		// Shutdown tlsconmons
		for _, tm := range tlsconmons {
			tm.Shutdown()
			logger.Debug("TLSConnectionMonitor shutdown complete", "address", tm.Address)
		}

		if err := otelShutdown(context.Background()); err != nil {
			logger.Warn("Error shutting down OpenTelemetry", "error", err)
		}
	}()

	<-shutdownDone
	logger.Info("Shutdown complete")
}