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/geo-platform/tenant-api/internal/tenant/repository" 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) publishJobSvc := tenantapp.NewPublishJobServiceWithConfig( app.DB, app.RabbitMQ, app.Redis, app.Logger, app.ConfigStore, ).WithCache(app.Cache) enterpriseSiteSvc := tenantapp.NewEnterpriseSiteService(app.DB, app.ConfigStore). WithCache(app.Cache). WithObjectStorage(app.ObjectStorage) promptRuleSvc.WithPublishJobService(publishJobSvc) promptRuleSvc.WithEnterpriseSiteService(enterpriseSiteSvc) kolGenerationSvc := tenantapp.NewKolGenerationService( app.DB, app.KolSubscriptions, app.KolPrompts, repository.NewKolUsageRepository(app.DB), repository.NewArticleRepository(app.DB), repository.NewAuditRepository(app.DB), repository.NewQuotaRepository(app.DB), app.KolPromptAsset, app.RabbitMQ, ).WithCache(app.Cache) 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, ).WithKolGenerationService(kolGenerationSvc).WithConfigProvider(app.ConfigStore) desktopTaskService := tenantapp.NewDesktopTaskService(app.DB, app.MonitoringDB, app.RabbitMQ, app.Logger). WithCache(app.Cache). WithRedis(app.Redis) publishLeaseRecoveryWorker := internalscheduler.NewPublishLeaseRecoveryWorker(desktopTaskService, app.Logger) monitoringDailyTaskWorker := tenantapp.NewMonitoringDailyTaskWorker(monitoringService, monitoringCallbackService, app.Logger) monitoringResultRecoveryWorker := internalscheduler.NewMonitoringResultRecoveryWorker(app.MonitoringDB, monitoringCallbackService, app.Logger) monitoringLeaseRecoveryWorker := internalscheduler.NewMonitoringLeaseRecoveryWorker(app.MonitoringDB, desktopTaskService, app.Logger) monitoringReceivedInspectionWorker := internalscheduler.NewMonitoringReceivedInspectionWorker(app.MonitoringDB, app.Logger) generationStateCheckWorker := generateworker.NewGenerationTaskStateCheckWorker(app.DB, app.Logger, app.Config().Generation). WithConfigProvider(app.ConfigStore). WithCache(app.Cache) monitoringRetentionWorker := internalscheduler.NewMonitoringRetentionWorker(app.MonitoringDB, app.Logger) brandAssetCleanupWorker := internalscheduler.NewBrandAssetCleanupWorker(app.DB, app.MonitoringDB, app.Cache, app.Logger) brandAssetCleanupEventWorker := internalscheduler.NewBrandAssetCleanupEventWorker(brandAssetCleanupWorker, app.DB, app.Logger) brandAssetCleanupReconcileWorker := internalscheduler.NewBrandAssetCleanupReconcileWorker(app.DB, app.Logger) aiPointUsageCleanupWorker := internalscheduler.NewAIPointUsageCleanupWorker(app.DB, app.Logger) schedulerRunRetentionWorker := internalscheduler.NewSchedulerRunRetentionWorker(app.DB, app.Logger) controlPlane := internalscheduler.NewControlPlane(app.DB, app.Logger, []internalscheduler.ControlledJob{ { Key: "schedule_dispatch", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return scheduleDispatchWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "monitoring_daily_task", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { timeout := 45 * time.Second if jobCtx.Job != nil && jobCtx.Job.TimeoutSeconds > 0 { timeout = time.Duration(jobCtx.Job.TimeoutSeconds) * time.Second } taskCtx, cancel := context.WithTimeout(runCtx, timeout) defer cancel() return monitoringDailyTaskWorker.RunOnce(taskCtx, monitoringDailyRunOptions(jobCtx)) }, }, { Key: "monitoring_result_recovery", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return monitoringResultRecoveryWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "monitoring_lease_recovery", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return monitoringLeaseRecoveryWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "publish_lease_recovery", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return publishLeaseRecoveryWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "monitoring_received_inspection", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return monitoringReceivedInspectionWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "generation_state_check", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return generationStateCheckWorker.RunOnce(runCtx, generationStateCheckRunOptions(jobCtx)) }, }, { Key: "monitoring_retention_cleanup", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return monitoringRetentionWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "brand_asset_cleanup", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return brandAssetCleanupWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "brand_asset_cleanup_reconcile", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return brandAssetCleanupReconcileWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "ai_point_usage_cleanup", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return aiPointUsageCleanupWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "scheduler_run_retention", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { return schedulerRunRetentionWorker.RunOnce(runCtx, jobCtx) }, }, { Key: "media_supply_cache_sync", Run: func(runCtx context.Context, jobCtx internalscheduler.JobRunContext) (map[string]any, error) { modelID := internalscheduler.AsInt(jobCtx.Config, "model_id", 1) if jobCtx.DryRun { return map[string]any{ "model_id": modelID, "synced": 0, }, nil } return app.MediaSupplyService.RunMediaResourceSyncOnce(runCtx, modelID) }, }, }) startSchedulerWorker(ctx, &workerWG, "control_plane", app.Logger.Sugar(), controlPlane.Run) startSchedulerWorker(ctx, &workerWG, "brand_asset_cleanup_events", app.Logger.Sugar(), brandAssetCleanupEventWorker.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 monitoringDailyRunOptions(jobCtx internalscheduler.JobRunContext) tenantapp.MonitoringDailyTaskRunOptions { return tenantapp.MonitoringDailyTaskRunOptions{ MaxDesktopDispatchPerRun: internalscheduler.AsInt(jobCtx.Config, "max_desktop_dispatch_per_run", 0), MaxDesktopClientBacklog: internalscheduler.AsInt(jobCtx.Config, "max_desktop_client_backlog", 0), } } func generationStateCheckRunOptions(jobCtx internalscheduler.JobRunContext) generateworker.GenerationTaskStateCheckRunOptions { batchSize := 0 if jobCtx.Job != nil && jobCtx.Job.BatchSize != nil { batchSize = *jobCtx.Job.BatchSize } if batchSize <= 0 { batchSize = internalscheduler.AsInt(jobCtx.Config, "batch_size", 0) } timeout := time.Duration(0) if jobCtx.Job != nil && jobCtx.Job.TimeoutSeconds > 0 { timeout = time.Duration(jobCtx.Job.TimeoutSeconds) * time.Second } if timeout <= 0 { timeout = internalscheduler.AsDuration(jobCtx.Config, "timeout", 0) } return generateworker.GenerationTaskStateCheckRunOptions{ Timeout: timeout, BatchSize: batchSize, Lookback: internalscheduler.AsDuration(jobCtx.Config, "lookback", 0), } } 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] }