root/geo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(infra): add scheduler and worker-generate standalone processes

Browse Source 
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:
liangxu
2026-04-15 14:19:21 +08:00
parent b46cfaaa61
commit 3f5b5c1e01
11 changed files with 1588 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/worker/generate/article_generation_worker.go
+261
View File
@@ -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
}
server/internal/worker/generate/template_assist_worker.go
+163
View File
@@ -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),
)
}
}