From 3f5b5c1e010d1aa5252d245f95789212a6d29f7e Mon Sep 17 00:00:00 2001 From: liangxu Date: Wed, 15 Apr 2026 14:19:21 +0800 Subject: [PATCH] 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) --- server/Makefile | 8 +- server/cmd/scheduler/main.go | 77 +++ server/cmd/worker-generate/main.go | 95 ++++ .../knowledge_deleted_cleanup_worker.go | 43 ++ .../monitoring_lease_recovery_worker.go | 88 ++++ .../monitoring_received_inspection_worker.go | 114 +++++ .../monitoring_result_recovery_worker.go | 205 ++++++++ .../scheduler/schedule_dispatch_worker.go | 467 ++++++++++++++++++ server/internal/shared/schedule/runtime.go | 68 +++ .../generate/article_generation_worker.go | 261 ++++++++++ .../worker/generate/template_assist_worker.go | 163 ++++++ 11 files changed, 1588 insertions(+), 1 deletion(-) create mode 100644 server/cmd/scheduler/main.go create mode 100644 server/cmd/worker-generate/main.go create mode 100644 server/internal/scheduler/knowledge_deleted_cleanup_worker.go create mode 100644 server/internal/scheduler/monitoring_lease_recovery_worker.go create mode 100644 server/internal/scheduler/monitoring_received_inspection_worker.go create mode 100644 server/internal/scheduler/monitoring_result_recovery_worker.go create mode 100644 server/internal/scheduler/schedule_dispatch_worker.go create mode 100644 server/internal/shared/schedule/runtime.go create mode 100644 server/internal/worker/generate/article_generation_worker.go create mode 100644 server/internal/worker/generate/template_assist_worker.go diff --git a/server/Makefile b/server/Makefile index 7ed1cba..fe9e043 100644 --- a/server/Makefile +++ b/server/Makefile @@ -1,4 +1,4 @@ -.PHONY: dev-init dev-api migrate-up migrate-down migrate-monitoring-up migrate-monitoring-down migrate-create sqlc-generate tenant-scope-guard lint test tidy seed +.PHONY: dev-init dev-api dev-worker-generate dev-scheduler migrate-up migrate-down migrate-monitoring-up migrate-monitoring-down migrate-create sqlc-generate tenant-scope-guard lint test tidy seed DB_URL ?= postgres://geo:geo_dev@localhost:5432/geo?sslmode=disable MONITORING_DB_URL ?= postgres://geo:geo_dev@localhost:5433/geo_monitoring?sslmode=disable @@ -15,6 +15,12 @@ dev-init: ## Start infra, run migrations, seed data dev-api: ## Run tenant-api server go run ./cmd/tenant-api +dev-worker-generate: ## Run article generation worker + go run ./cmd/worker-generate + +dev-scheduler: ## Run scheduler workers + go run ./cmd/scheduler + migrate-up: ## Run all pending migrations $(MIGRATE_CMD) -path migrations -database "$(DB_URL)" up diff --git a/server/cmd/scheduler/main.go b/server/cmd/scheduler/main.go new file mode 100644 index 0000000..974dd2c --- /dev/null +++ b/server/cmd/scheduler/main.go @@ -0,0 +1,77 @@ +package main + +import ( + "context" + "log" + "os" + "os/signal" + "syscall" + + "github.com/geo-platform/tenant-api/internal/bootstrap" + internalscheduler "github.com/geo-platform/tenant-api/internal/scheduler" + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" +) + +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() + + 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, + ) + + promptRuleSvc := tenantapp.NewPromptRuleGenerationService( + app.DB, + app.LLM, + app.RabbitMQ, + knowledgeSvc, + app.GenerationStreams, + generationCfg, + app.Config.LLM.MaxOutputTokens, + ) + + monitoringCallbackService := tenantapp.NewMonitoringCallbackService( + app.DB, + app.MonitoringDB, + app.RabbitMQ, + app.LLM, + app.Logger, + ) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + internalscheduler.NewScheduleDispatchWorker(app.DB, app.RabbitMQ, promptRuleSvc, app.Logger, app.Config.Scheduler).Start(ctx) + internalscheduler.NewMonitoringResultRecoveryWorker(app.MonitoringDB, monitoringCallbackService, app.Logger).Start(ctx) + internalscheduler.NewMonitoringLeaseRecoveryWorker(app.MonitoringDB, app.Logger).Start(ctx) + internalscheduler.NewMonitoringReceivedInspectionWorker(app.MonitoringDB, app.Logger).Start(ctx) + internalscheduler.NewKnowledgeDeletedCleanupWorker(knowledgeSvc).Start(ctx) + + app.Logger.Info("scheduler started") + + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + <-quit + cancel() + app.Logger.Info("scheduler shutting down...") +} diff --git a/server/cmd/worker-generate/main.go b/server/cmd/worker-generate/main.go new file mode 100644 index 0000000..7b97649 --- /dev/null +++ b/server/cmd/worker-generate/main.go @@ -0,0 +1,95 @@ +package main + +import ( + "context" + "log" + "os" + "os/signal" + "syscall" + + "github.com/geo-platform/tenant-api/internal/bootstrap" + 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" +) + +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() + + 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, + ) + + templateSvc := tenantapp.NewTemplateService( + app.DB, + repository.NewCachedTemplateRepository( + repository.NewTemplateRepository(app.DB), + app.Cache, + ), + app.LLM, + nil, + knowledgeSvc, + app.GenerationStreams, + generationCfg, + app.Config.LLM.MaxOutputTokens, + ) + + promptRuleSvc := tenantapp.NewPromptRuleGenerationService( + app.DB, + app.LLM, + nil, + knowledgeSvc, + app.GenerationStreams, + generationCfg, + app.Config.LLM.MaxOutputTokens, + ) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + generateworker.NewArticleGenerationWorker( + app.DB, + app.RabbitMQ, + templateSvc, + promptRuleSvc, + app.Logger, + app.Config.Generation, + ).Start(ctx) + + generateworker.NewTemplateAssistWorker( + app.RabbitMQ, + templateSvc, + app.Logger, + app.Config.Generation.WorkerConcurrency, + app.Config.Generation.ArticleTimeout, + ).Start(ctx) + + app.Logger.Info("worker-generate started") + + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + <-quit + cancel() + app.Logger.Info("worker-generate shutting down...") +} diff --git a/server/internal/scheduler/knowledge_deleted_cleanup_worker.go b/server/internal/scheduler/knowledge_deleted_cleanup_worker.go new file mode 100644 index 0000000..60ef9c0 --- /dev/null +++ b/server/internal/scheduler/knowledge_deleted_cleanup_worker.go @@ -0,0 +1,43 @@ +package scheduler + +import ( + "context" + "time" + + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" +) + +type KnowledgeDeletedCleanupWorker struct { + service *tenantapp.KnowledgeService + interval time.Duration +} + +func NewKnowledgeDeletedCleanupWorker(service *tenantapp.KnowledgeService) *KnowledgeDeletedCleanupWorker { + return &KnowledgeDeletedCleanupWorker{ + service: service, + interval: 15 * time.Second, + } +} + +func (w *KnowledgeDeletedCleanupWorker) Start(ctx context.Context) { + if w == nil || w.service == nil { + return + } + go w.run(ctx) +} + +func (w *KnowledgeDeletedCleanupWorker) run(ctx context.Context) { + w.service.RunDeletedCleanupSweep() + + ticker := time.NewTicker(w.interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + w.service.RunDeletedCleanupSweep() + } + } +} diff --git a/server/internal/scheduler/monitoring_lease_recovery_worker.go b/server/internal/scheduler/monitoring_lease_recovery_worker.go new file mode 100644 index 0000000..65f38ec --- /dev/null +++ b/server/internal/scheduler/monitoring_lease_recovery_worker.go @@ -0,0 +1,88 @@ +package scheduler + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "go.uber.org/zap" + + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" +) + +const ( + defaultMonitoringLeaseRecoveryInterval = 30 * time.Minute + defaultMonitoringLeaseRecoveryTimeout = 30 * time.Second +) + +type MonitoringLeaseRecoveryWorker struct { + monitoringPool *pgxpool.Pool + logger *zap.Logger + interval time.Duration + timeout time.Duration +} + +func NewMonitoringLeaseRecoveryWorker(monitoringPool *pgxpool.Pool, logger *zap.Logger) *MonitoringLeaseRecoveryWorker { + return &MonitoringLeaseRecoveryWorker{ + monitoringPool: monitoringPool, + logger: logger, + interval: defaultMonitoringLeaseRecoveryInterval, + timeout: defaultMonitoringLeaseRecoveryTimeout, + } +} + +func (w *MonitoringLeaseRecoveryWorker) Start(ctx context.Context) { + if w == nil || w.monitoringPool == nil { + return + } + go w.run(ctx) +} + +func (w *MonitoringLeaseRecoveryWorker) 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 *MonitoringLeaseRecoveryWorker) runOnce(parent context.Context) { + ctx, cancel := context.WithTimeout(parent, w.timeout) + defer cancel() + + tx, err := w.monitoringPool.Begin(ctx) + if err != nil { + if w.logger != nil { + w.logger.Warn("monitoring lease recovery begin failed", zap.Error(err)) + } + return + } + defer tx.Rollback(ctx) + + expiredCount, err := tenantapp.ExpireMonitoringLeasedTasks(ctx, tx, nil, time.Now().UTC()) + if err != nil { + if w.logger != nil { + w.logger.Warn("monitoring lease recovery failed", tenantapp.MonitoringWorkerErrorFields(err)...) + } + return + } + + if err := tx.Commit(ctx); err != nil { + if w.logger != nil { + w.logger.Warn("monitoring lease recovery commit failed", zap.Error(err)) + } + return + } + + if expiredCount > 0 && w.logger != nil { + w.logger.Info("monitoring lease recovery completed", zap.Int64("expired_task_count", expiredCount)) + } +} diff --git a/server/internal/scheduler/monitoring_received_inspection_worker.go b/server/internal/scheduler/monitoring_received_inspection_worker.go new file mode 100644 index 0000000..994695c --- /dev/null +++ b/server/internal/scheduler/monitoring_received_inspection_worker.go @@ -0,0 +1,114 @@ +package scheduler + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "go.uber.org/zap" + + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" +) + +const ( + defaultMonitoringReceivedInspectionInterval = 5 * time.Minute + defaultMonitoringReceivedInspectionTimeout = 30 * time.Second + defaultMonitoringReceivedAlertThreshold = 15 * time.Minute + defaultMonitoringReceivedInspectLimit = 100 +) + +type MonitoringReceivedInspectionWorker struct { + monitoringPool *pgxpool.Pool + logger *zap.Logger + interval time.Duration + timeout time.Duration + threshold time.Duration + limit int +} + +func NewMonitoringReceivedInspectionWorker(monitoringPool *pgxpool.Pool, logger *zap.Logger) *MonitoringReceivedInspectionWorker { + return &MonitoringReceivedInspectionWorker{ + monitoringPool: monitoringPool, + logger: logger, + interval: defaultMonitoringReceivedInspectionInterval, + timeout: defaultMonitoringReceivedInspectionTimeout, + threshold: defaultMonitoringReceivedAlertThreshold, + limit: defaultMonitoringReceivedInspectLimit, + } +} + +func (w *MonitoringReceivedInspectionWorker) Start(ctx context.Context) { + if w == nil || w.monitoringPool == nil { + return + } + go w.run(ctx) +} + +func (w *MonitoringReceivedInspectionWorker) 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 *MonitoringReceivedInspectionWorker) runOnce(parent context.Context) { + ctx, cancel := context.WithTimeout(parent, w.timeout) + defer cancel() + + tx, err := w.monitoringPool.Begin(ctx) + if err != nil { + if w.logger != nil { + w.logger.Warn("monitoring received watchdog begin failed", zap.Error(err)) + } + return + } + defer tx.Rollback(ctx) + + now := time.Now().UTC() + items, err := tenantapp.MarkMonitoringStaleReceivedTasks(ctx, tx, now, w.threshold, w.limit) + if err != nil { + if w.logger != nil { + w.logger.Warn("monitoring received watchdog failed", tenantapp.MonitoringWorkerErrorFields(err)...) + } + return + } + + if err := tx.Commit(ctx); err != nil { + if w.logger != nil { + w.logger.Warn("monitoring received watchdog commit failed", zap.Error(err)) + } + return + } + + if len(items) == 0 || w.logger == nil { + return + } + + w.logger.Warn("monitoring received watchdog found overdue tasks", + zap.Int("task_count", len(items)), + zap.Duration("threshold", w.threshold), + ) + + for _, item := range items { + w.logger.Warn("monitoring task stuck in received state", + zap.Int64("task_id", item.ID), + zap.Int64("tenant_id", item.TenantID), + zap.Int64("brand_id", item.BrandID), + zap.Int64("question_id", item.QuestionID), + zap.String("collector_type", item.CollectorType), + zap.String("ai_platform_id", item.AIPlatformID), + zap.String("business_date", item.BusinessDate.Format("2006-01-02")), + zap.Time("callback_received_at", item.CallbackReceivedAt.UTC()), + zap.Duration("received_lag", now.Sub(item.CallbackReceivedAt.UTC())), + ) + } +} diff --git a/server/internal/scheduler/monitoring_result_recovery_worker.go b/server/internal/scheduler/monitoring_result_recovery_worker.go new file mode 100644 index 0000000..5215126 --- /dev/null +++ b/server/internal/scheduler/monitoring_result_recovery_worker.go @@ -0,0 +1,205 @@ +package scheduler + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "go.uber.org/zap" + + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" +) + +const ( + defaultMonitoringResultRecoveryInterval = 1 * time.Minute + defaultMonitoringResultRecoveryTimeout = 30 * time.Second + defaultMonitoringResultRecoveryLimit = 100 +) + +type MonitoringResultRecoveryWorker struct { + monitoringPool *pgxpool.Pool + service *tenantapp.MonitoringCallbackService + logger *zap.Logger + interval time.Duration + timeout time.Duration + limit int +} + +type monitoringRecoverableTaskResult struct { + ID int64 + CallbackReceivedAt time.Time + RequestPayloadJSON []byte +} + +func NewMonitoringResultRecoveryWorker( + monitoringPool *pgxpool.Pool, + service *tenantapp.MonitoringCallbackService, + logger *zap.Logger, +) *MonitoringResultRecoveryWorker { + return &MonitoringResultRecoveryWorker{ + monitoringPool: monitoringPool, + service: service, + logger: logger, + interval: defaultMonitoringResultRecoveryInterval, + timeout: defaultMonitoringResultRecoveryTimeout, + limit: defaultMonitoringResultRecoveryLimit, + } +} + +func (w *MonitoringResultRecoveryWorker) Start(ctx context.Context) { + if w == nil || w.monitoringPool == nil || w.service == nil { + return + } + go w.run(ctx) +} + +func (w *MonitoringResultRecoveryWorker) 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 *MonitoringResultRecoveryWorker) runOnce(parent context.Context) { + ctx, cancel := context.WithTimeout(parent, w.timeout) + defer cancel() + + tx, err := w.monitoringPool.Begin(ctx) + if err != nil { + if w.logger != nil { + w.logger.Warn("monitoring result recovery begin failed", zap.Error(err)) + } + return + } + defer tx.Rollback(ctx) + + items, err := loadMonitoringRecoverableTaskResults(ctx, tx, w.limit) + if err != nil { + if w.logger != nil { + w.logger.Warn("monitoring result recovery failed", tenantapp.MonitoringWorkerErrorFields(err)...) + } + return + } + + if err := tx.Commit(ctx); err != nil { + if w.logger != nil { + w.logger.Warn("monitoring result recovery commit failed", zap.Error(err)) + } + return + } + + if len(items) == 0 { + return + } + + republishedCount := 0 + for _, item := range items { + if w.tryRepublish(parent, item) { + republishedCount++ + } + } + + if republishedCount > 0 && w.logger != nil { + w.logger.Info("monitoring result recovery republished tasks", zap.Int("task_count", republishedCount)) + } +} + +func (w *MonitoringResultRecoveryWorker) tryRepublish(parent context.Context, item monitoringRecoverableTaskResult) bool { + if err := tenantapp.ValidateMonitoringTaskResultEnvelope(item.RequestPayloadJSON); err != nil { + w.service.PatchTaskResultQueueMeta(item.ID, item.CallbackReceivedAt, tenantapp.MonitoringTaskResultQueueMetaPatch{ + QueueStatus: "invalid_payload", + QueueClaimedAt: nil, + QueueEnqueuedAt: nil, + LastQueuePublishError: stringPtr(err.Error()), + }) + if w.logger != nil { + w.logger.Warn("monitoring result recovery skipped invalid payload", zap.Int64("task_id", item.ID), zap.Error(err)) + } + return false + } + + ctx, cancel := context.WithTimeout(parent, 5*time.Second) + defer cancel() + + if err := w.service.PublishTaskResultEnvelope(ctx, item.RequestPayloadJSON); err != nil { + w.service.PatchTaskResultQueueMeta(item.ID, item.CallbackReceivedAt, tenantapp.MonitoringTaskResultQueueMetaPatch{ + QueueStatus: "publish_failed", + QueueClaimedAt: nil, + QueueEnqueuedAt: nil, + LastQueuePublishError: stringPtr(err.Error()), + }) + if w.logger != nil { + w.logger.Warn("monitoring result recovery republish failed", zap.Int64("task_id", item.ID), zap.Error(err)) + } + return false + } + + enqueuedAt := time.Now().UTC().Round(0) + w.service.PatchTaskResultQueueMeta(item.ID, item.CallbackReceivedAt, tenantapp.MonitoringTaskResultQueueMetaPatch{ + QueueStatus: "queued", + QueueClaimedAt: nil, + QueueEnqueuedAt: &enqueuedAt, + LastQueuePublishError: nil, + }) + + return true +} + +func loadMonitoringRecoverableTaskResults(ctx context.Context, tx pgx.Tx, limit int) ([]monitoringRecoverableTaskResult, error) { + if limit <= 0 { + limit = defaultMonitoringResultRecoveryLimit + } + + rows, err := tx.Query(ctx, ` + SELECT + id, + callback_received_at, + request_payload_json + FROM monitoring_collect_tasks + WHERE collector_type = $1 + AND status = 'received' + AND callback_received_at IS NOT NULL + AND request_payload_json IS NOT NULL + AND jsonb_typeof(request_payload_json) = 'object' + AND COALESCE(NULLIF(request_payload_json ->> 'queue_status', ''), 'pending') IN ('pending', 'publish_failed') + ORDER BY callback_received_at ASC, id ASC + FOR UPDATE SKIP LOCKED + LIMIT $2 + `, "plugin", limit) + if err != nil { + return nil, tenantapp.InternalMonitoringError(50041, "result_recovery_query_failed", "failed to load recoverable monitoring results", err) + } + defer rows.Close() + + result := make([]monitoringRecoverableTaskResult, 0, limit) + for rows.Next() { + var item monitoringRecoverableTaskResult + if scanErr := rows.Scan(&item.ID, &item.CallbackReceivedAt, &item.RequestPayloadJSON); scanErr != nil { + return nil, tenantapp.InternalMonitoringError(50041, "result_recovery_scan_failed", "failed to parse recoverable monitoring results", scanErr) + } + result = append(result, item) + } + if err := rows.Err(); err != nil { + return nil, tenantapp.InternalMonitoringError(50041, "result_recovery_scan_failed", "failed to iterate recoverable monitoring results", err) + } + + return result, nil +} + +func stringPtr(value string) *string { + if value == "" { + return nil + } + result := value + return &result +} diff --git a/server/internal/scheduler/schedule_dispatch_worker.go b/server/internal/scheduler/schedule_dispatch_worker.go new file mode 100644 index 0000000..f647b85 --- /dev/null +++ b/server/internal/scheduler/schedule_dispatch_worker.go @@ -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 +} diff --git a/server/internal/shared/schedule/runtime.go b/server/internal/shared/schedule/runtime.go new file mode 100644 index 0000000..2c8f16d --- /dev/null +++ b/server/internal/shared/schedule/runtime.go @@ -0,0 +1,68 @@ +package schedule + +import ( + "fmt" + "strings" + "time" + + cron "github.com/robfig/cron/v3" +) + +var cronParser = cron.NewParser( + cron.SecondOptional | + cron.Minute | + cron.Hour | + cron.Dom | + cron.Month | + cron.Dow | + cron.Descriptor, +) + +func Parse(cronExpr string) (cron.Schedule, error) { + trimmed := strings.TrimSpace(cronExpr) + if trimmed == "" { + return nil, fmt.Errorf("cron expression is required") + } + + schedule, err := cronParser.Parse(trimmed) + if err != nil { + return nil, fmt.Errorf("parse cron expression %q: %w", trimmed, err) + } + return schedule, nil +} + +func ComputeNextRun(cronExpr string, startAt, endAt *time.Time, now time.Time) (*time.Time, error) { + schedule, err := Parse(cronExpr) + if err != nil { + return nil, err + } + + notBefore := now.In(time.Local) + if startAt != nil { + start := startAt.In(time.Local) + if start.After(notBefore) { + notBefore = start + } + } + + next := schedule.Next(notBefore.Add(-time.Nanosecond)).Round(0) + if next.IsZero() { + return nil, nil + } + if endAt != nil && next.After(endAt.In(time.Local)) { + return nil, nil + } + + return &next, nil +} + +func IsRunAllowed(runAt time.Time, startAt, endAt *time.Time) bool { + runAt = runAt.In(time.Local) + if startAt != nil && runAt.Before(startAt.In(time.Local)) { + return false + } + if endAt != nil && runAt.After(endAt.In(time.Local)) { + return false + } + return true +} diff --git a/server/internal/worker/generate/article_generation_worker.go b/server/internal/worker/generate/article_generation_worker.go new file mode 100644 index 0000000..fc6a558 --- /dev/null +++ b/server/internal/worker/generate/article_generation_worker.go @@ -0,0 +1,261 @@ +package generate + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "os" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + amqp "github.com/rabbitmq/amqp091-go" + "go.uber.org/zap" + + "github.com/geo-platform/tenant-api/internal/shared/config" + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" +) + +const defaultArticleGenerationWorkerRetryInterval = 5 * time.Second + +var errArticleGenerationConsumerClosed = errors.New("article generation consumer channel closed") + +type ArticleGenerationWorker struct { + pool *pgxpool.Pool + rabbitMQ *rabbitmq.Client + templateService *tenantapp.TemplateService + promptRuleService *tenantapp.PromptRuleGenerationService + logger *zap.Logger + retryInterval time.Duration + processTimeout time.Duration + consumerPrefix string + workerConcurrency int +} + +func NewArticleGenerationWorker( + pool *pgxpool.Pool, + rabbitMQClient *rabbitmq.Client, + templateService *tenantapp.TemplateService, + promptRuleService *tenantapp.PromptRuleGenerationService, + logger *zap.Logger, + cfg config.GenerationConfig, +) *ArticleGenerationWorker { + workerConcurrency := cfg.WorkerConcurrency + if workerConcurrency <= 0 { + workerConcurrency = 1 + } + + processTimeout := cfg.ArticleTimeout + 2*time.Minute + if processTimeout <= 2*time.Minute { + processTimeout = 10 * time.Minute + } + + return &ArticleGenerationWorker{ + pool: pool, + rabbitMQ: rabbitMQClient, + templateService: templateService, + promptRuleService: promptRuleService, + logger: logger, + retryInterval: defaultArticleGenerationWorkerRetryInterval, + processTimeout: processTimeout, + consumerPrefix: fmt.Sprintf("worker-generate-%d", os.Getpid()), + workerConcurrency: workerConcurrency, + } +} + +func (w *ArticleGenerationWorker) Start(ctx context.Context) { + if w == nil || w.pool == nil || w.rabbitMQ == nil { + return + } + + for i := 0; i < w.workerConcurrency; i++ { + consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1) + go w.run(ctx, consumerName) + } +} + +func (w *ArticleGenerationWorker) run(ctx context.Context, consumerName string) { + for { + if err := w.consumeOnce(ctx, consumerName); err != nil && ctx.Err() == nil && w.logger != nil { + if errors.Is(err, errArticleGenerationConsumerClosed) { + w.logger.Info("article generation consumer channel closed, retrying", + zap.String("consumer", consumerName), + ) + } else { + w.logger.Warn("article generation worker stopped unexpectedly", + zap.Error(err), + zap.String("consumer", consumerName), + ) + } + } + + select { + case <-ctx.Done(): + return + case <-time.After(w.retryInterval): + } + } +} + +func (w *ArticleGenerationWorker) consumeOnce(ctx context.Context, consumerName string) error { + deliveries, ch, err := w.rabbitMQ.ConsumeGenerationTask(consumerName) + if err != nil { + return err + } + defer ch.Close() + + for { + select { + case <-ctx.Done(): + return nil + case delivery, ok := <-deliveries: + if !ok { + return errArticleGenerationConsumerClosed + } + w.handleDelivery(ctx, delivery) + } + } +} + +func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, delivery amqp.Delivery) { + envelope, err := tenantapp.DecodeArticleGenerationTaskEnvelope(delivery.Body) + if err != nil { + w.rejectDelivery(delivery, false, err, 0) + return + } + + ctx, cancel := context.WithTimeout(parent, w.processTimeout) + defer cancel() + + task, terminal, err := w.loadTask(ctx, envelope.TaskID) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil { + w.logger.Warn("article generation ack failed for missing task", + zap.Error(ackErr), + zap.Int64("task_id", envelope.TaskID), + ) + } + return + } + if task.ID > 0 { + _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, err) + if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil { + w.logger.Warn("article generation ack failed after task load failure", + zap.Error(ackErr), + zap.Int64("task_id", envelope.TaskID), + ) + } + return + } + requeue := !delivery.Redelivered + w.rejectDelivery(delivery, requeue, err, envelope.TaskID) + return + } + if terminal { + if err := delivery.Ack(false); err != nil && w.logger != nil { + w.logger.Warn("article generation ack failed for terminal task", + zap.Error(err), + zap.Int64("task_id", envelope.TaskID), + ) + } + return + } + + switch { + case task.ArticleID <= 0 || task.QuotaReservationID <= 0: + _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("generation task references are incomplete")) + case task.TaskType == "template": + w.templateService.ProcessQueuedTask(ctx, task) + case task.TaskType == "custom_generation": + w.promptRuleService.ProcessQueuedTask(ctx, task) + default: + _ = tenantapp.HandleClaimedGenerationTaskFailure(ctx, w.pool, w.templateService, w.promptRuleService, task, fmt.Errorf("unsupported generation task type: %s", task.TaskType)) + } + + if err := delivery.Ack(false); err != nil && w.logger != nil { + w.logger.Warn("article generation ack failed", + zap.Error(err), + zap.Int64("task_id", task.ID), + zap.Int64("article_id", task.ArticleID), + ) + } +} + +func (w *ArticleGenerationWorker) loadTask(ctx context.Context, taskID int64) (tenantapp.ClaimedGenerationTask, bool, error) { + var ( + task tenantapp.ClaimedGenerationTask + status string + operatorID sql.NullInt64 + articleID sql.NullInt64 + reservationID sql.NullInt64 + inputParamsJSON []byte + ) + + err := w.pool.QueryRow(ctx, ` + SELECT id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, status + FROM generation_tasks + WHERE id = $1 + `, taskID).Scan( + &task.ID, + &task.TenantID, + &operatorID, + &articleID, + &reservationID, + &task.TaskType, + &inputParamsJSON, + &status, + ) + if err != nil { + return tenantapp.ClaimedGenerationTask{}, false, err + } + + if status == "completed" || status == "failed" { + return tenantapp.ClaimedGenerationTask{}, true, nil + } + + task.OperatorID = operatorID.Int64 + task.ArticleID = articleID.Int64 + task.QuotaReservationID = reservationID.Int64 + + params, err := parseJSONMap(inputParamsJSON) + if err != nil { + task.InputParams = map[string]interface{}{} + return task, false, fmt.Errorf("decode generation input params: %w", err) + } + task.InputParams = params + + return task, false, nil +} + +func (w *ArticleGenerationWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, taskID int64) { + if nackErr := delivery.Nack(false, requeue); nackErr != nil && w.logger != nil { + w.logger.Warn("article generation nack failed", + zap.Error(nackErr), + zap.Int64("task_id", taskID), + ) + } + if w.logger != nil { + w.logger.Warn("article generation processing failed", + zap.Error(err), + zap.Bool("requeue", requeue), + zap.Bool("redelivered", delivery.Redelivered), + zap.Int64("task_id", taskID), + ) + } +} + +func parseJSONMap(raw []byte) (map[string]interface{}, error) { + if len(raw) == 0 { + return map[string]interface{}{}, nil + } + + payload := make(map[string]interface{}) + if err := json.Unmarshal(raw, &payload); err != nil { + return nil, err + } + return payload, nil +} diff --git a/server/internal/worker/generate/template_assist_worker.go b/server/internal/worker/generate/template_assist_worker.go new file mode 100644 index 0000000..27e219d --- /dev/null +++ b/server/internal/worker/generate/template_assist_worker.go @@ -0,0 +1,163 @@ +package generate + +import ( + "context" + "errors" + "fmt" + "os" + "time" + + amqp "github.com/rabbitmq/amqp091-go" + "go.uber.org/zap" + + "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" +) + +var errTemplateAssistConsumerClosed = errors.New("template assist consumer channel closed") + +type TemplateAssistWorker struct { + rabbitMQ *rabbitmq.Client + templateService *tenantapp.TemplateService + logger *zap.Logger + retryInterval time.Duration + processTimeout time.Duration + consumerPrefix string + workerConcurrency int +} + +func NewTemplateAssistWorker( + rabbitMQClient *rabbitmq.Client, + templateService *tenantapp.TemplateService, + logger *zap.Logger, + workerConcurrency int, + processTimeout time.Duration, +) *TemplateAssistWorker { + if workerConcurrency <= 0 { + workerConcurrency = 1 + } + if processTimeout <= 0 { + processTimeout = 5 * time.Minute + } + + return &TemplateAssistWorker{ + rabbitMQ: rabbitMQClient, + templateService: templateService, + logger: logger, + retryInterval: 5 * time.Second, + processTimeout: processTimeout, + consumerPrefix: fmt.Sprintf("template-assist-%d", os.Getpid()), + workerConcurrency: workerConcurrency, + } +} + +func (w *TemplateAssistWorker) Start(ctx context.Context) { + if w == nil || w.rabbitMQ == nil || w.templateService == nil { + return + } + + for i := 0; i < w.workerConcurrency; i++ { + consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1) + go w.run(ctx, consumerName) + } +} + +func (w *TemplateAssistWorker) run(ctx context.Context, consumerName string) { + for { + if err := w.consumeOnce(ctx, consumerName); err != nil && ctx.Err() == nil && w.logger != nil { + if errors.Is(err, errTemplateAssistConsumerClosed) { + w.logger.Info("template assist consumer channel closed, retrying", + zap.String("consumer", consumerName), + ) + } else { + w.logger.Warn("template assist worker stopped unexpectedly", + zap.Error(err), + zap.String("consumer", consumerName), + ) + } + } + + select { + case <-ctx.Done(): + return + case <-time.After(w.retryInterval): + } + } +} + +func (w *TemplateAssistWorker) consumeOnce(ctx context.Context, consumerName string) error { + deliveries, ch, err := w.rabbitMQ.ConsumeTemplateAssistTask(consumerName) + if err != nil { + return err + } + defer ch.Close() + + for { + select { + case <-ctx.Done(): + return nil + case delivery, ok := <-deliveries: + if !ok { + return errTemplateAssistConsumerClosed + } + w.handleDelivery(ctx, delivery) + } + } +} + +func (w *TemplateAssistWorker) handleDelivery(parent context.Context, delivery amqp.Delivery) { + job, err := tenantapp.DecodeTemplateAssistTask(delivery.Body) + if err != nil { + w.rejectDelivery(delivery, false, err, "", "") + return + } + + ctx, cancel := context.WithTimeout(parent, w.processTimeout) + defer cancel() + + terminal, err := w.templateService.IsAssistTaskTerminal(ctx, *job) + if err != nil { + requeue := !delivery.Redelivered + w.rejectDelivery(delivery, requeue, err, job.TaskID, job.TaskType) + return + } + if terminal { + if err := delivery.Ack(false); err != nil && w.logger != nil { + w.logger.Warn("template assist ack failed for terminal task", + zap.Error(err), + zap.String("task_id", job.TaskID), + zap.String("task_type", job.TaskType), + ) + } + return + } + + w.templateService.ProcessQueuedAssistJob(ctx, *job) + + if err := delivery.Ack(false); err != nil && w.logger != nil { + w.logger.Warn("template assist ack failed", + zap.Error(err), + zap.String("task_id", job.TaskID), + zap.String("task_type", job.TaskType), + ) + } +} + +func (w *TemplateAssistWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, taskID, taskType string) { + if nackErr := delivery.Nack(false, requeue); nackErr != nil && w.logger != nil { + w.logger.Warn("template assist nack failed", + zap.Error(nackErr), + zap.String("task_id", taskID), + zap.String("task_type", taskType), + ) + } + if w.logger != nil { + w.logger.Warn("template assist processing failed", + zap.Error(err), + zap.Bool("requeue", requeue), + zap.Bool("redelivered", delivery.Redelivered), + zap.String("task_id", taskID), + zap.String("task_type", taskType), + ) + } +}