449 lines
14 KiB
Go
449 lines
14 KiB
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)
|
|
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 {
|
|
maxDesktopDispatch := 0
|
|
if jobCtx.Job != nil && jobCtx.Job.BatchSize != nil {
|
|
maxDesktopDispatch = *jobCtx.Job.BatchSize
|
|
}
|
|
if maxDesktopDispatch <= 0 {
|
|
maxDesktopDispatch = internalscheduler.AsInt(jobCtx.Config, "max_desktop_dispatch_per_run", 0)
|
|
}
|
|
if maxDesktopDispatch <= 0 {
|
|
maxDesktopDispatch = internalscheduler.AsInt(jobCtx.Config, "batch_size", 0)
|
|
}
|
|
return tenantapp.MonitoringDailyTaskRunOptions{
|
|
MaxDesktopDispatchPerRun: maxDesktopDispatch,
|
|
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]
|
|
}
|