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" "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() 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, ) promptRuleSvc := tenantapp.NewPromptRuleGenerationService( app.DB, app.LLM, app.RabbitMQ, knowledgeSvc, app.GenerationStreams, generationCfg, app.Config.LLM.MaxOutputTokens, ).WithCache(app.Cache) monitoringCallbackService := tenantapp.NewMonitoringCallbackService( app.DB, app.MonitoringDB, app.RabbitMQ, app.LLM, app.Logger, ) monitoringService := tenantapp.NewMonitoringService( app.DB, app.MonitoringDB, app.RabbitMQ, app.Config.MonitoringDispatch, app.Config.BrandLibrary, app.Logger, ).WithRedis(app.Redis) ctx, cancel := context.WithCancel(context.Background()) defer cancel() var workerWG sync.WaitGroup scheduleDispatchWorker := internalscheduler.NewScheduleDispatchWorker(app.DB, app.RabbitMQ, promptRuleSvc, app.Cache, app.Logger, app.Config.Scheduler) 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) 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) 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(token) app.Engine.GET("/api/internal/metrics/monitoring/daily-tasks", authMiddleware, func(c *gin.Context) { response.Success(c, tenantapp.MonitoringDailyTaskMetricsSnapshotValue()) }) 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(token string) gin.HandlerFunc { expected := strings.TrimSpace(token) return func(c *gin.Context) { 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] }