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:
+7
-1
@@ -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
|
||||
|
||||
|
||||
@@ -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...")
|
||||
}
|
||||
@@ -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...")
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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())),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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),
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user