geo/server/cmd/scheduler/main.go

package main

import (
	"context"
	"crypto/subtle"
	"errors"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/geo-platform/tenant-api/internal/bootstrap"
	internalscheduler "github.com/geo-platform/tenant-api/internal/scheduler"
	"github.com/geo-platform/tenant-api/internal/shared/response"
	tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
	generateworker "github.com/geo-platform/tenant-api/internal/worker/generate"
	"github.com/gin-gonic/gin"
)

const schedulerWorkerShutdownGracePeriod = 60 * time.Second

func main() {
	configPath := "configs/config.yaml"
	if p := os.Getenv("CONFIG_PATH"); p != "" {
		configPath = p
	}

	app, err := bootstrap.New(configPath)
	if err != nil {
		log.Fatalf("bootstrap: %v", err)
	}
	defer app.Close()

	generationProvider := tenantapp.NewGenerationStreamDisabledConfigProvider(app.ConfigStore)
	generationCfg := app.Config().Generation
	generationCfg.StreamEnabled = false

	knowledgeSvc := tenantapp.NewKnowledgeService(
		app.DB,
		app.RetrievalProvider,
		app.VectorStore,
		app.ObjectStorage,
		app.LLM,
		app.Logger,
		app.Config().Retrieval,
		app.Config().LLM,
		0,
		0,
	).WithConfigProvider(app.ConfigStore)

	promptRuleSvc := tenantapp.NewPromptRuleGenerationService(
		app.DB,
		app.LLM,
		app.RabbitMQ,
		knowledgeSvc,
		app.GenerationStreams,
		generationCfg,
		app.Config().LLM.MaxOutputTokens,
	).WithCache(app.Cache).WithLogger(app.Logger).WithConfigProvider(generationProvider)

	monitoringCallbackService := tenantapp.NewMonitoringCallbackService(
		app.DB,
		app.MonitoringDB,
		app.RabbitMQ,
		app.LLM,
		app.Logger,
	)
	monitoringService := app.MonitoringService

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	app.StartConfigWatcher(ctx)

	var workerWG sync.WaitGroup
	scheduleDispatchWorker := internalscheduler.NewScheduleDispatchWorker(
		app.DB,
		app.RabbitMQ,
		promptRuleSvc,
		app.Cache,
		app.Logger,
		app.Config().Scheduler,
	).WithConfigProvider(app.ConfigStore)
	monitoringDailyTaskWorker := tenantapp.NewMonitoringDailyTaskWorker(monitoringService, monitoringCallbackService, app.Logger)
	monitoringResultRecoveryWorker := internalscheduler.NewMonitoringResultRecoveryWorker(app.MonitoringDB, monitoringCallbackService, app.Logger)
	monitoringLeaseRecoveryWorker := internalscheduler.NewMonitoringLeaseRecoveryWorker(app.MonitoringDB, app.Logger)
	monitoringReceivedInspectionWorker := internalscheduler.NewMonitoringReceivedInspectionWorker(app.MonitoringDB, app.Logger)
	knowledgeDeletedCleanupWorker := internalscheduler.NewKnowledgeDeletedCleanupWorker(knowledgeSvc)
	generationStateCheckWorker := generateworker.NewGenerationTaskStateCheckWorker(app.DB, app.Logger, app.Config().Generation).
		WithConfigProvider(app.ConfigStore).
		WithCache(app.Cache)

	startSchedulerWorker(ctx, &workerWG, "schedule_dispatch", app.Logger.Sugar(), scheduleDispatchWorker.Run)
	startSchedulerWorker(ctx, &workerWG, "monitoring_daily_task", app.Logger.Sugar(), monitoringDailyTaskWorker.Run)
	startSchedulerWorker(ctx, &workerWG, "monitoring_result_recovery", app.Logger.Sugar(), monitoringResultRecoveryWorker.Run)
	startSchedulerWorker(ctx, &workerWG, "monitoring_lease_recovery", app.Logger.Sugar(), monitoringLeaseRecoveryWorker.Run)
	startSchedulerWorker(ctx, &workerWG, "monitoring_received_inspection", app.Logger.Sugar(), monitoringReceivedInspectionWorker.Run)
	startSchedulerWorker(ctx, &workerWG, "knowledge_deleted_cleanup", app.Logger.Sugar(), knowledgeDeletedCleanupWorker.Run)
	startSchedulerWorker(ctx, &workerWG, "generation_state_check", app.Logger.Sugar(), generationStateCheckWorker.Run)

	var metricsServer *http.Server
	if app.Config().Scheduler.HTTPPort > 0 {
		metricsServer = startSchedulerMetricsServer(app)
	} else {
		app.Logger.Sugar().Infof("scheduler metrics http server disabled by http_port=%d", app.Config().Scheduler.HTTPPort)
	}

	app.Logger.Info("scheduler started")

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	app.Logger.Info("scheduler shutdown signal received")
	shutdownSchedulerMetricsServer(metricsServer, app.Logger.Sugar())
	cancel()
	waitSchedulerWorkers(&workerWG, schedulerWorkerShutdownGracePeriod, app.Logger.Sugar())
	app.Logger.Info("scheduler shutting down...")
}

func startSchedulerWorker(
	ctx context.Context,
	wg *sync.WaitGroup,
	name string,
	logger interface {
		Infof(template string, args ...interface{})
		Errorf(template string, args ...interface{})
	},
	run func(context.Context),
) {
	if wg == nil || run == nil {
		return
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		logger.Infof("scheduler worker %s started", name)
		run(ctx)
		logger.Infof("scheduler worker %s stopped", name)
	}()
}

func waitSchedulerWorkers(wg *sync.WaitGroup, timeout time.Duration, logger interface {
	Infof(template string, args ...interface{})
	Errorf(template string, args ...interface{})
}) bool {
	if wg == nil {
		return true
	}

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		logger.Infof("scheduler workers stopped")
		return true
	case <-time.After(timeout):
		logger.Errorf("scheduler worker shutdown timed out after %s", timeout)
		return false
	}
}

func startSchedulerMetricsServer(app *bootstrap.App) *http.Server {
	token := strings.TrimSpace(app.Config().Scheduler.InternalMetricsToken)
	if token == "" {
		app.Logger.Warn("scheduler metrics http server not started: scheduler.internal_metrics_token is required")
		return nil
	}

	authMiddleware := schedulerMetricsAuthMiddleware(app)
	app.Engine.GET("/api/internal/metrics/monitoring/daily-tasks", authMiddleware, func(c *gin.Context) {
		response.Success(c, tenantapp.MonitoringDailyTaskMetricsSnapshotValue())
	})
	app.Engine.GET("/api/internal/metrics/generation/tasks", authMiddleware, func(c *gin.Context) {
		response.Success(c, tenantapp.GenerationTaskMetricsSnapshotValue())
	})
	prometheusHandler := tenantapp.MonitoringSchedulerMetricsPrometheusHandler()
	app.Engine.GET("/metrics", authMiddleware, func(c *gin.Context) {
		prometheusHandler.ServeHTTP(c.Writer, c.Request)
	})

	addr := schedulerHTTPAddr(app.Config().Scheduler.HTTPHost, app.Config().Scheduler.HTTPPort)
	server := &http.Server{
		Addr:              addr,
		Handler:           app.Engine,
		ReadHeaderTimeout: 5 * time.Second,
	}

	go func() {
		app.Logger.Sugar().Infof("scheduler metrics http listening on %s", addr)
		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			app.Logger.Sugar().Errorf("scheduler metrics http server stopped: %v", err)
		}
	}()
	return server
}

func shutdownSchedulerMetricsServer(server *http.Server, logger interface {
	Infof(template string, args ...interface{})
	Errorf(template string, args ...interface{})
}) {
	if server == nil {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := server.Shutdown(ctx); err != nil {
		logger.Errorf("scheduler metrics http shutdown failed: %v", err)
		return
	}
	logger.Infof("scheduler metrics http server stopped")
}

func schedulerHTTPAddr(host string, port int) string {
	host = strings.TrimSpace(host)
	if host == "" {
		host = "127.0.0.1"
	}
	return net.JoinHostPort(host, strconv.Itoa(port))
}

func schedulerMetricsAuthMiddleware(app *bootstrap.App) gin.HandlerFunc {
	return schedulerMetricsAuthMiddlewareWithToken(func() string {
		if cfg := app.Config(); cfg != nil {
			return cfg.Scheduler.InternalMetricsToken
		}
		return ""
	})
}

func schedulerMetricsAuthMiddlewareWithToken(token func() string) gin.HandlerFunc {
	return func(c *gin.Context) {
		expected := ""
		if token != nil {
			expected = strings.TrimSpace(token())
		}
		if expected == "" || !schedulerMetricsTokenMatches(c, expected) {
			response.Error(c, response.ErrUnauthorized(40181, "internal_metrics_unauthorized", "valid internal metrics token is required"))
			c.Abort()
			return
		}
		c.Next()
	}
}

func schedulerMetricsTokenMatches(c *gin.Context, expected string) bool {
	if c == nil {
		return false
	}

	candidates := []string{
		c.GetHeader("X-Internal-Metrics-Token"),
		bearerToken(c.GetHeader("Authorization")),
	}
	for _, candidate := range candidates {
		candidate = strings.TrimSpace(candidate)
		if candidate == "" || len(candidate) != len(expected) {
			continue
		}
		if subtle.ConstantTimeCompare([]byte(candidate), []byte(expected)) == 1 {
			return true
		}
	}
	return false
}

func bearerToken(header string) string {
	header = strings.TrimSpace(header)
	if header == "" {
		return ""
	}
	parts := strings.Fields(header)
	if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
		return ""
	}
	return parts[1]
}