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"
	"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)
	monitoringDailyTaskWorker := tenantapp.NewMonitoringDailyTaskWorker(monitoringService, monitoringCallbackService, app.Logger)
	monitoringResultRecoveryWorker := internalscheduler.NewMonitoringResultRecoveryWorker(app.MonitoringDB, monitoringCallbackService, app.Logger)
	monitoringLeaseRecoveryWorker := internalscheduler.NewMonitoringLeaseRecoveryWorker(app.MonitoringDB, app.Logger)
	publishLeaseRecoveryWorker := internalscheduler.NewPublishLeaseRecoveryWorker(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)
	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: "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)

	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 {
	batchSize := 0
	if jobCtx.Job != nil && jobCtx.Job.BatchSize != nil {
		batchSize = *jobCtx.Job.BatchSize
	}
	if batchSize <= 0 {
		batchSize = internalscheduler.AsInt(jobCtx.Config, "max_materialize_per_run", 0)
	}
	if batchSize <= 0 {
		batchSize = internalscheduler.AsInt(jobCtx.Config, "batch_size", 0)
	}
	return tenantapp.MonitoringDailyTaskRunOptions{
		MaxMaterializePerRun:   batchSize,
		MaxMaterializePerPlan:  internalscheduler.AsInt(jobCtx.Config, "max_materialize_per_plan", 0),
		MaxMaterializePerBrand: internalscheduler.AsInt(jobCtx.Config, "max_materialize_per_brand", 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]
}