geo/server/internal/scheduler/schedule_dispatch_worker.go

package scheduler

import (
	"context"
	"encoding/json"
	"sync"
	"time"

	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"
	"go.uber.org/zap"

	sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
	"github.com/geo-platform/tenant-api/internal/shared/config"
	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
	sharedschedule "github.com/geo-platform/tenant-api/internal/shared/schedule"
	tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
)

const (
	defaultScheduleDispatchInterval = 15 * time.Second
	defaultScheduleDispatchTimeout  = 30 * time.Second
	defaultScheduleDispatchBatch    = 20
	scheduleCacheCleanupTimeout     = 30 * time.Second
)

type ScheduleDispatchWorker struct {
	pool              *pgxpool.Pool
	rabbitMQ          *rabbitmq.Client
	promptRuleService *tenantapp.PromptRuleGenerationService
	cache             sharedcache.Cache
	logger            *zap.Logger
	cfg               scheduleDispatchRuntimeConfig
	configProvider    config.Provider
}

type dueScheduleTask struct {
	ID                int64
	TenantID          int64
	WorkspaceID       int64
	OperatorID        *int64
	PromptRuleID      int64
	BrandID           *int64
	Name              string
	CronExpr          string
	AutoPublish       bool
	PublishAccountIDs []string
	CoverAssetURL     *string
	CoverImageAssetID *int64
	EnableWebSearch   bool
	GenerateCount     int
	StartAt           *time.Time
	EndAt             *time.Time
	ScheduledFor      time.Time
}

type scheduleTaskCacheUpdate struct {
	id       int64
	tenantID int64
}

type scheduleDispatchRuntimeConfig struct {
	Interval          time.Duration
	Timeout           time.Duration
	BatchSize         int
	DispatchWorkers   int
	QueueBackpressure int
}

func NewScheduleDispatchWorker(
	pool *pgxpool.Pool,
	rabbitMQClient *rabbitmq.Client,
	promptRuleService *tenantapp.PromptRuleGenerationService,
	cache sharedcache.Cache,
	logger *zap.Logger,
	cfg config.SchedulerConfig,
) *ScheduleDispatchWorker {
	return &ScheduleDispatchWorker{
		pool:              pool,
		rabbitMQ:          rabbitMQClient,
		promptRuleService: promptRuleService,
		cache:             cache,
		logger:            logger,
		cfg:               scheduleDispatchRuntimeFromConfig(cfg),
	}
}

func (w *ScheduleDispatchWorker) WithConfigProvider(provider config.Provider) *ScheduleDispatchWorker {
	w.configProvider = provider
	return w
}

func (w *ScheduleDispatchWorker) runtimeConfig() scheduleDispatchRuntimeConfig {
	if w != nil && w.configProvider != nil {
		if cfg := w.configProvider.Current(); cfg != nil {
			return scheduleDispatchRuntimeFromConfig(cfg.Scheduler)
		}
	}
	if w == nil {
		return scheduleDispatchRuntimeFromConfig(config.SchedulerConfig{})
	}
	return w.cfg
}

func (w *ScheduleDispatchWorker) Start(ctx context.Context) {
	go w.Run(ctx)
}

func (w *ScheduleDispatchWorker) Run(ctx context.Context) {
	if w == nil || w.pool == nil || w.promptRuleService == nil {
		return
	}
	w.run(ctx)
}

func (w *ScheduleDispatchWorker) run(ctx context.Context) {
	w.runOnce(context.Background())

	for {
		interval := w.runtimeConfig().Interval
		timer := time.NewTimer(interval)
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case <-timer.C:
			w.runOnce(context.Background())
		}
	}
}

func (w *ScheduleDispatchWorker) runOnce(parent context.Context) {
	runtimeCfg := w.runtimeConfig()
	ctx, cancel := w.stageContext(parent, runtimeCfg)
	count, err := w.hydrateMissingNextRuns(ctx, runtimeCfg)
	cancel()
	if err != nil {
		if w.logger != nil {
			w.logger.Warn("schedule dispatch hydrate missing next_run_at failed", zap.Error(err))
		}
	} else if count > 0 && w.logger != nil {
		w.logger.Info("schedule dispatch hydrated next_run_at", zap.Int("task_count", count))
	}

	ctx, cancel = w.stageContext(parent, runtimeCfg)
	paused, stats, err := w.shouldPauseDispatch(ctx, runtimeCfg)
	cancel()
	if err != nil {
		if w.logger != nil {
			w.logger.Warn("schedule dispatch inspect generation queue failed", zap.Error(err))
		}
	} else if paused {
		if w.logger != nil {
			runtimeCfg := w.runtimeConfig()
			w.logger.Info("schedule dispatch paused by generation queue backpressure",
				zap.Int("queue_depth", stats.Messages),
				zap.Int("queue_consumers", stats.Consumers),
				zap.Int("queue_limit", runtimeCfg.QueueBackpressure),
			)
		}
		return
	}

	ctx, cancel = w.stageContext(parent, runtimeCfg)
	tasks, err := w.claimDueSchedules(ctx, runtimeCfg)
	cancel()
	if err != nil {
		if w.logger != nil {
			w.logger.Warn("schedule dispatch claim due schedules failed", zap.Error(err))
		}
		return
	}

	w.dispatchBatch(parent, tasks)
}

func (w *ScheduleDispatchWorker) RunOnce(ctx context.Context, run JobRunContext) (map[string]any, error) {
	if w == nil || w.pool == nil || w.promptRuleService == nil {
		return map[string]any{"skipped": true, "reason": "worker_not_ready"}, nil
	}
	startedAt := time.Now()
	runtimeCfg := w.runtimeConfig()
	runtimeCfg.ApplyJobRun(run)

	stageCtx, cancel := w.stageContext(ctx, runtimeCfg)
	hydratedCount, err := w.hydrateMissingNextRuns(stageCtx, runtimeCfg)
	cancel()
	if err != nil {
		return map[string]any{
			"stage": "hydrate_missing_next_runs",
		}, err
	}

	stageCtx, cancel = w.stageContext(ctx, runtimeCfg)
	paused, stats, err := w.shouldPauseDispatch(stageCtx, runtimeCfg)
	cancel()
	if err != nil {
		return map[string]any{
			"stage":          "queue_backpressure",
			"hydrated_count": hydratedCount,
		}, err
	}
	if paused {
		return map[string]any{
			"hydrated_count":  hydratedCount,
			"paused":          true,
			"queue_depth":     stats.Messages,
			"queue_consumers": stats.Consumers,
			"duration_ms":     time.Since(startedAt).Milliseconds(),
			"batch_size":      runtimeCfg.BatchSize,
		}, nil
	}

	stageCtx, cancel = w.stageContext(ctx, runtimeCfg)
	tasks, err := w.claimDueSchedules(stageCtx, runtimeCfg)
	cancel()
	if err != nil {
		return map[string]any{
			"stage":          "claim_due_schedules",
			"hydrated_count": hydratedCount,
		}, err
	}
	w.dispatchBatch(ctx, tasks)
	return map[string]any{
		"hydrated_count": hydratedCount,
		"claimed_count":  len(tasks),
		"duration_ms":    time.Since(startedAt).Milliseconds(),
		"batch_size":     runtimeCfg.BatchSize,
	}, nil
}

func (w *ScheduleDispatchWorker) hydrateMissingNextRuns(ctx context.Context, runtimeCfg scheduleDispatchRuntimeConfig) (int, error) {
	tx, err := w.pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback(ctx)

	rows, err := tx.Query(ctx, `
		SELECT id, tenant_id, cron_expr, start_at, end_at, status
		FROM schedule_tasks
		WHERE deleted_at IS NULL
		  AND status = 'enabled'
		  AND next_run_at IS NULL
		ORDER BY created_at ASC, id ASC
		FOR UPDATE SKIP LOCKED
		LIMIT $1
	`, runtimeCfg.BatchSize)
	if err != nil {
		return 0, err
	}

	now := time.Now()
	count := 0
	updates := make([]struct {
		cacheUpdate scheduleTaskCacheUpdate
		nextRunAt   *time.Time
	}, 0, runtimeCfg.BatchSize)
	for rows.Next() {
		var (
			id       int64
			tenantID int64
			cronExpr string
			status   string
			startAt  pgtype.Timestamptz
			endAt    pgtype.Timestamptz
		)
		if err := rows.Scan(&id, &tenantID, &cronExpr, &startAt, &endAt, &status); err != nil {
			rows.Close()
			return 0, err
		}

		var nextRunAt *time.Time
		if status == "enabled" {
			nextRunAt, err = sharedschedule.ComputeNextRun(cronExpr, timePtrFromTimestamp(startAt), timePtrFromTimestamp(endAt), now)
			if err != nil && w.logger != nil {
				w.logger.Warn("schedule dispatch skipped invalid cron while hydrating next_run_at",
					zap.Int64("schedule_task_id", id),
					zap.String("cron_expr", cronExpr),
					zap.Error(err),
				)
			}
		}

		updates = append(updates, struct {
			cacheUpdate scheduleTaskCacheUpdate
			nextRunAt   *time.Time
		}{
			cacheUpdate: scheduleTaskCacheUpdate{
				id:       id,
				tenantID: tenantID,
			},
			nextRunAt: nextRunAt,
		})
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return 0, err
	}
	rows.Close()

	for _, update := range updates {
		if _, err := tx.Exec(ctx, `
			UPDATE schedule_tasks
			SET next_run_at = $1,
			    updated_at = NOW()
			WHERE id = $2
		`, update.nextRunAt, update.cacheUpdate.id); err != nil {
			return 0, err
		}
		count++
	}

	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	cacheUpdates := make([]scheduleTaskCacheUpdate, 0, len(updates))
	for _, update := range updates {
		cacheUpdates = append(cacheUpdates, update.cacheUpdate)
	}
	w.invalidateScheduleTaskCaches(cacheUpdates)
	return count, nil
}

func (w *ScheduleDispatchWorker) claimDueSchedules(ctx context.Context, runtimeCfg scheduleDispatchRuntimeConfig) ([]dueScheduleTask, error) {
	tx, err := w.pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback(ctx)

	rows, err := tx.Query(ctx, `
		SELECT id, tenant_id, workspace_id, operator_id, prompt_rule_id, brand_id, name, cron_expr,
		       auto_publish, publish_account_ids, cover_asset_url, cover_image_asset_id,
		       enable_web_search, generate_count, start_at, end_at, next_run_at
		FROM schedule_tasks
		WHERE deleted_at IS NULL
		  AND status = 'enabled'
		  AND next_run_at IS NOT NULL
		  AND next_run_at <= $1
		ORDER BY next_run_at ASC, id ASC
		FOR UPDATE SKIP LOCKED
		LIMIT $2
	`, time.Now(), runtimeCfg.BatchSize)
	if err != nil {
		return nil, err
	}

	now := time.Now()
	tasks := make([]dueScheduleTask, 0, runtimeCfg.BatchSize)
	updates := make([]struct {
		cacheUpdate scheduleTaskCacheUpdate
		nextRun     *time.Time
	}, 0, runtimeCfg.BatchSize)
	for rows.Next() {
		var (
			task                  dueScheduleTask
			workspaceID           pgtype.Int8
			operatorID            pgtype.Int8
			publishAccountIDsJSON []byte
			startAt               pgtype.Timestamptz
			endAt                 pgtype.Timestamptz
			nextRunAt             pgtype.Timestamptz
		)
		if err := rows.Scan(
			&task.ID,
			&task.TenantID,
			&workspaceID,
			&operatorID,
			&task.PromptRuleID,
			&task.BrandID,
			&task.Name,
			&task.CronExpr,
			&task.AutoPublish,
			&publishAccountIDsJSON,
			&task.CoverAssetURL,
			&task.CoverImageAssetID,
			&task.EnableWebSearch,
			&task.GenerateCount,
			&startAt,
			&endAt,
			&nextRunAt,
		); err != nil {
			rows.Close()
			return nil, err
		}

		task.StartAt = timePtrFromTimestamp(startAt)
		task.EndAt = timePtrFromTimestamp(endAt)
		if workspaceID.Valid {
			task.WorkspaceID = workspaceID.Int64
		}
		task.OperatorID = int64PtrFromInt8(operatorID)
		task.PublishAccountIDs = decodeDueSchedulePublishAccountIDs(publishAccountIDsJSON)
		if task.GenerateCount <= 0 {
			task.GenerateCount = 1
		}
		scheduledFor := timePtrFromTimestamp(nextRunAt)
		if scheduledFor == nil {
			continue
		}
		task.ScheduledFor = *scheduledFor

		nextRun, nextErr := sharedschedule.ComputeNextRun(task.CronExpr, task.StartAt, task.EndAt, now.Add(time.Nanosecond))
		if nextErr != nil && w.logger != nil {
			w.logger.Warn("schedule dispatch encountered invalid cron while claiming due schedule",
				zap.Int64("schedule_task_id", task.ID),
				zap.String("cron_expr", task.CronExpr),
				zap.Error(nextErr),
			)
			nextRun = nil
		}

		updates = append(updates, struct {
			cacheUpdate scheduleTaskCacheUpdate
			nextRun     *time.Time
		}{
			cacheUpdate: scheduleTaskCacheUpdate{
				id:       task.ID,
				tenantID: task.TenantID,
			},
			nextRun: nextRun,
		})

		if sharedschedule.IsRunAllowed(task.ScheduledFor, task.StartAt, task.EndAt) {
			tasks = append(tasks, task)
		}
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return nil, err
	}
	rows.Close()

	for _, update := range updates {
		if _, err := tx.Exec(ctx, `
			UPDATE schedule_tasks
			SET next_run_at = $1,
			    updated_at = NOW()
			WHERE id = $2
		`, update.nextRun, update.cacheUpdate.id); err != nil {
			return nil, err
		}
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, err
	}
	cacheUpdates := make([]scheduleTaskCacheUpdate, 0, len(updates))
	for _, update := range updates {
		cacheUpdates = append(cacheUpdates, update.cacheUpdate)
	}
	w.invalidateScheduleTaskCaches(cacheUpdates)
	return tasks, nil
}

func (w *ScheduleDispatchWorker) dispatch(parent context.Context, task dueScheduleTask) {
	generateCount := task.GenerateCount
	if generateCount <= 0 {
		generateCount = 1
	}

	for idx := 0; idx < generateCount; idx++ {
		ctx, cancel := w.stageContext(parent)
		resp, err := w.promptRuleService.EnqueueScheduledGeneration(ctx, tenantapp.SchedulePromptRuleGenerationInput{
			ScheduleTaskID:    task.ID,
			OperatorID:        task.OperatorID,
			TenantID:          task.TenantID,
			PromptRuleID:      task.PromptRuleID,
			BrandID:           task.BrandID,
			Name:              task.Name,
			WorkspaceID:       task.WorkspaceID,
			AutoPublish:       task.AutoPublish,
			PublishAccountIDs: task.PublishAccountIDs,
			CoverAssetURL:     task.CoverAssetURL,
			CoverImageAssetID: task.CoverImageAssetID,
			EnableWebSearch:   task.EnableWebSearch,
			GenerateCount:     generateCount,
			ScheduledFor:      task.ScheduledFor,
		})
		cancel()
		if err != nil {
			if w.logger != nil {
				w.logger.Warn("schedule dispatch enqueue generation failed",
					zap.Int64("schedule_task_id", task.ID),
					zap.Int64("tenant_id", task.TenantID),
					zap.Int64("prompt_rule_id", task.PromptRuleID),
					zap.Int("generate_index", idx+1),
					zap.Int("generate_count", generateCount),
					zap.Time("scheduled_for", task.ScheduledFor.UTC()),
					zap.Error(err),
				)
			}
			continue
		}

		if w.logger != nil {
			w.logger.Info("schedule dispatch enqueued generation",
				zap.Int64("schedule_task_id", task.ID),
				zap.Int64("tenant_id", task.TenantID),
				zap.Int64("prompt_rule_id", task.PromptRuleID),
				zap.Int64("article_id", resp.ArticleID),
				zap.Int64("generation_task_id", resp.TaskID),
				zap.Int("generate_index", idx+1),
				zap.Int("generate_count", generateCount),
				zap.Time("scheduled_for", task.ScheduledFor.UTC()),
			)
		}
	}
}

func (w *ScheduleDispatchWorker) dispatchBatch(parent context.Context, tasks []dueScheduleTask) {
	if len(tasks) == 0 {
		return
	}
	runtimeCfg := w.runtimeConfig()
	if runtimeCfg.DispatchWorkers <= 1 {
		for _, task := range tasks {
			if parent.Err() != nil {
				return
			}
			w.dispatch(parent, task)
		}
		return
	}

	sem := make(chan struct{}, runtimeCfg.DispatchWorkers)
	var wg sync.WaitGroup
	for _, task := range tasks {
		if parent.Err() != nil {
			break
		}
		task := task
		sem <- struct{}{}
		wg.Add(1)
		go func() {
			defer wg.Done()
			defer func() { <-sem }()
			w.dispatch(parent, task)
		}()
	}
	wg.Wait()
}

func (w *ScheduleDispatchWorker) shouldPauseDispatch(ctx context.Context, runtimeCfg scheduleDispatchRuntimeConfig) (bool, rabbitmq.QueueStats, error) {
	if w == nil || w.rabbitMQ == nil {
		return false, rabbitmq.QueueStats{}, nil
	}
	if runtimeCfg.QueueBackpressure <= 0 {
		return false, rabbitmq.QueueStats{}, nil
	}

	stats, err := w.rabbitMQ.InspectGenerationQueue(ctx)
	if err != nil {
		return false, rabbitmq.QueueStats{}, err
	}
	return stats.Messages >= runtimeCfg.QueueBackpressure, stats, nil
}

func scheduleDispatchRuntimeFromConfig(cfg config.SchedulerConfig) scheduleDispatchRuntimeConfig {
	return scheduleDispatchRuntimeConfig{
		Interval:          schedulerDispatchInterval(cfg),
		Timeout:           schedulerDispatchTimeout(cfg),
		BatchSize:         schedulerDispatchBatchSize(cfg),
		DispatchWorkers:   schedulerDispatchConcurrency(cfg),
		QueueBackpressure: cfg.GenerationQueueBackpressureLimit,
	}
}

func (cfg *scheduleDispatchRuntimeConfig) ApplyJobRun(run JobRunContext) {
	if cfg == nil {
		return
	}
	if run.Job != nil {
		if run.Job.BatchSize != nil && *run.Job.BatchSize > 0 {
			cfg.BatchSize = *run.Job.BatchSize
		}
		if run.Job.TimeoutSeconds > 0 {
			cfg.Timeout = time.Duration(run.Job.TimeoutSeconds) * time.Second
		}
	}
	if value := AsInt(run.Config, "dispatch_workers", 0); value > 0 {
		cfg.DispatchWorkers = value
	}
	if value := AsInt(run.Config, "queue_backpressure", 0); value > 0 {
		cfg.QueueBackpressure = value
	}
	if value := AsInt(run.Config, "batch_size", 0); value > 0 && (run.Job == nil || run.Job.BatchSize == nil) {
		cfg.BatchSize = value
	}
	if timeout := AsDuration(run.Config, "timeout", 0); timeout > 0 && (run.Job == nil || run.Job.TimeoutSeconds <= 0) {
		cfg.Timeout = timeout
	}
}

func schedulerDispatchInterval(cfg config.SchedulerConfig) time.Duration {
	if cfg.DispatchInterval > 0 {
		return cfg.DispatchInterval
	}
	return defaultScheduleDispatchInterval
}

func schedulerDispatchTimeout(cfg config.SchedulerConfig) time.Duration {
	if cfg.DispatchTimeout > 0 {
		return cfg.DispatchTimeout
	}
	return defaultScheduleDispatchTimeout
}

func schedulerDispatchBatchSize(cfg config.SchedulerConfig) int {
	if cfg.DispatchBatchSize > 0 {
		return cfg.DispatchBatchSize
	}
	return defaultScheduleDispatchBatch
}

func schedulerDispatchConcurrency(cfg config.SchedulerConfig) int {
	if cfg.DispatchConcurrency > 0 {
		return cfg.DispatchConcurrency
	}
	return 1
}

func (w *ScheduleDispatchWorker) stageContext(parent context.Context, runtimeCfg ...scheduleDispatchRuntimeConfig) (context.Context, context.CancelFunc) {
	if parent == nil {
		parent = context.Background()
	}
	cfg := w.runtimeConfig()
	if len(runtimeCfg) > 0 {
		cfg = runtimeCfg[0]
	}
	return context.WithTimeout(parent, cfg.Timeout)
}

func (w *ScheduleDispatchWorker) invalidateScheduleTaskCaches(updates []scheduleTaskCacheUpdate) {
	if w == nil || w.cache == nil || len(updates) == 0 {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), scheduleCacheCleanupTimeout)
	defer cancel()

	for _, update := range updates {
		taskID := update.id
		tenantapp.InvalidateScheduleTaskCaches(ctx, w.cache, update.tenantID, &taskID)
	}
}

func timePtrFromTimestamp(value pgtype.Timestamptz) *time.Time {
	if !value.Valid {
		return nil
	}
	result := value.Time.Round(0)
	return &result
}

func int64PtrFromInt8(value pgtype.Int8) *int64 {
	if !value.Valid {
		return nil
	}
	result := value.Int64
	return &result
}

func decodeDueSchedulePublishAccountIDs(raw []byte) []string {
	if len(raw) == 0 {
		return []string{}
	}
	var values []string
	if err := json.Unmarshal(raw, &values); err != nil {
		return []string{}
	}
	result := make([]string, 0, len(values))
	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		if value == "" {
			continue
		}
		if _, ok := seen[value]; ok {
			continue
		}
		seen[value] = struct{}{}
		result = append(result, value)
	}
	return result
}