geo/server/internal/scheduler/schedule_dispatch_worker.go

package scheduler

import (
	"context"
	"sync"
	"time"

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

	"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
)

type ScheduleDispatchWorker struct {
	pool              *pgxpool.Pool
	rabbitMQ          *rabbitmq.Client
	promptRuleService *tenantapp.PromptRuleGenerationService
	logger            *zap.Logger
	interval          time.Duration
	timeout           time.Duration
	batchSize         int
	dispatchWorkers   int
	queueBackpressure int
}

type dueScheduleTask struct {
	ID              int64
	TenantID        int64
	OperatorID      *int64
	PromptRuleID    int64
	BrandID         *int64
	Name            string
	CronExpr        string
	TargetPlatform  *string
	EnableWebSearch bool
	GenerateCount   int
	StartAt         *time.Time
	EndAt           *time.Time
	ScheduledFor    time.Time
}

func NewScheduleDispatchWorker(
	pool *pgxpool.Pool,
	rabbitMQClient *rabbitmq.Client,
	promptRuleService *tenantapp.PromptRuleGenerationService,
	logger *zap.Logger,
	cfg config.SchedulerConfig,
) *ScheduleDispatchWorker {
	return &ScheduleDispatchWorker{
		pool:              pool,
		rabbitMQ:          rabbitMQClient,
		promptRuleService: promptRuleService,
		logger:            logger,
		interval:          schedulerDispatchInterval(cfg),
		timeout:           schedulerDispatchTimeout(cfg),
		batchSize:         schedulerDispatchBatchSize(cfg),
		dispatchWorkers:   schedulerDispatchConcurrency(cfg),
		queueBackpressure: cfg.GenerationQueueBackpressureLimit,
	}
}

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

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

	ticker := time.NewTicker(w.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			w.runOnce(ctx)
		}
	}
}

func (w *ScheduleDispatchWorker) runOnce(parent context.Context) {
	ctx, cancel := context.WithTimeout(parent, w.timeout)
	defer cancel()

	if count, err := w.hydrateMissingNextRuns(ctx); 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))
	}

	if paused, stats, err := w.shouldPauseDispatch(ctx); 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 {
			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", w.queueBackpressure),
			)
		}
		return
	}

	tasks, err := w.claimDueSchedules(ctx)
	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) hydrateMissingNextRuns(ctx context.Context) (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, 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
	`, w.batchSize)
	if err != nil {
		return 0, err
	}

	now := time.Now()
	count := 0
	updates := make([]struct {
		id        int64
		nextRunAt *time.Time
	}, 0, w.batchSize)
	for rows.Next() {
		var (
			id       int64
			cronExpr string
			status   string
			startAt  pgtype.Timestamptz
			endAt    pgtype.Timestamptz
		)
		if err := rows.Scan(&id, &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 {
			id        int64
			nextRunAt *time.Time
		}{
			id:        id,
			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.id); err != nil {
			return 0, err
		}
		count++
	}

	if err := tx.Commit(ctx); err != nil {
		return 0, err
	}
	return count, nil
}

func (w *ScheduleDispatchWorker) claimDueSchedules(ctx context.Context) ([]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, operator_id, prompt_rule_id, brand_id, name, cron_expr, target_platform,
		       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(), w.batchSize)
	if err != nil {
		return nil, err
	}

	now := time.Now()
	tasks := make([]dueScheduleTask, 0, w.batchSize)
	updates := make([]struct {
		id      int64
		nextRun *time.Time
	}, 0, w.batchSize)
	for rows.Next() {
		var (
			task       dueScheduleTask
			operatorID pgtype.Int8
			startAt    pgtype.Timestamptz
			endAt      pgtype.Timestamptz
			nextRunAt  pgtype.Timestamptz
		)
		if err := rows.Scan(
			&task.ID,
			&task.TenantID,
			&operatorID,
			&task.PromptRuleID,
			&task.BrandID,
			&task.Name,
			&task.CronExpr,
			&task.TargetPlatform,
			&task.EnableWebSearch,
			&task.GenerateCount,
			&startAt,
			&endAt,
			&nextRunAt,
		); err != nil {
			rows.Close()
			return nil, err
		}

		task.StartAt = timePtrFromTimestamp(startAt)
		task.EndAt = timePtrFromTimestamp(endAt)
		task.OperatorID = int64PtrFromInt8(operatorID)
		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 {
			id      int64
			nextRun *time.Time
		}{
			id:      task.ID,
			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.id); err != nil {
			return nil, err
		}
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, err
	}
	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 := context.WithTimeout(parent, w.timeout)
		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,
			TargetPlatform:  task.TargetPlatform,
			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
	}
	if w.dispatchWorkers <= 1 {
		for _, task := range tasks {
			if parent.Err() != nil {
				return
			}
			w.dispatch(parent, task)
		}
		return
	}

	sem := make(chan struct{}, w.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) (bool, rabbitmq.QueueStats, error) {
	if w == nil || w.rabbitMQ == nil || w.queueBackpressure <= 0 {
		return false, rabbitmq.QueueStats{}, nil
	}

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

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 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
}