feat(infra): add scheduler and worker-generate standalone processes
Extract monitoring recovery/inspection workers and schedule dispatch into a dedicated scheduler process. Add worker-generate process for article generation and template-assist queue consumption. Introduce shared/schedule runtime for cron-based worker lifecycle management. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,467 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user