feat(cache,worker): overhaul cache layer and add generation task lease recovery
- shared/cache: add Options/L1/async/metrics/prefix decorators, multi-key ops, Redis pool tuning, and JSON readthrough metrics - worker-generate: claim tasks via DB lease + heartbeat, requeue stale queued tasks, expire dead leases with refund/cache invalidation - tenant: version article cache keys so worker recovery invalidations propagate cleanly - shared/config: expand Redis (pool/timeouts/TLS) and Generation (lease/recovery) configs with defaults Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -2,7 +2,6 @@ package generate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
@@ -10,11 +9,11 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
amqp "github.com/rabbitmq/amqp091-go"
|
||||
"go.uber.org/zap"
|
||||
|
||||
sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache"
|
||||
"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"
|
||||
@@ -31,9 +30,11 @@ type ArticleGenerationWorker struct {
|
||||
promptRuleService *tenantapp.PromptRuleGenerationService
|
||||
imitationService *tenantapp.ArticleImitationService
|
||||
kolWorker *KolGenerationWorker
|
||||
cache sharedcache.Cache
|
||||
logger *zap.Logger
|
||||
retryInterval time.Duration
|
||||
processTimeout time.Duration
|
||||
generationCfg config.GenerationConfig
|
||||
configProvider config.Provider
|
||||
consumerPrefix string
|
||||
workerConcurrency int
|
||||
@@ -64,6 +65,7 @@ func NewArticleGenerationWorker(
|
||||
logger: logger,
|
||||
retryInterval: defaultArticleGenerationWorkerRetryInterval,
|
||||
processTimeout: articleGenerationProcessTimeout(cfg),
|
||||
generationCfg: cfg,
|
||||
consumerPrefix: fmt.Sprintf("worker-generate-%d", os.Getpid()),
|
||||
workerConcurrency: workerConcurrency,
|
||||
}
|
||||
@@ -74,6 +76,11 @@ func (w *ArticleGenerationWorker) WithConfigProvider(provider config.Provider) *
|
||||
return w
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) WithCache(c sharedcache.Cache) *ArticleGenerationWorker {
|
||||
w.cache = c
|
||||
return w
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) Start(ctx context.Context) {
|
||||
if w == nil || w.pool == nil || w.rabbitMQ == nil {
|
||||
return
|
||||
@@ -83,6 +90,7 @@ func (w *ArticleGenerationWorker) Start(ctx context.Context) {
|
||||
consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1)
|
||||
go w.run(ctx, consumerName)
|
||||
}
|
||||
w.startRecovery(ctx)
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) run(ctx context.Context, consumerName string) {
|
||||
@@ -135,31 +143,25 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver
|
||||
return
|
||||
}
|
||||
|
||||
recoveryCfg := w.runtimeRecoveryConfig()
|
||||
ctx, cancel := context.WithTimeout(parent, w.runtimeProcessTimeout())
|
||||
defer cancel()
|
||||
|
||||
task, terminal, err := w.loadTask(ctx, envelope.TaskID)
|
||||
lease, terminal, err := w.claimGenerationTask(ctx, envelope.TaskID, generationWorkerOwner(delivery.ConsumerTag), recoveryCfg.LeaseTTL)
|
||||
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 {
|
||||
if lease != nil {
|
||||
task := lease.task.toClaimedTask()
|
||||
w.handleTaskFailure(ctx, task, "task_load", err)
|
||||
w.releaseGenerationTaskLease(context.Background(), task.ID, lease.token)
|
||||
if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil {
|
||||
w.logger.Warn("article generation ack failed after task load failure",
|
||||
w.logger.Warn("article generation ack failed after task claim failure",
|
||||
zap.Error(ackErr),
|
||||
zap.Int64("task_id", envelope.TaskID),
|
||||
)
|
||||
}
|
||||
return
|
||||
}
|
||||
requeue := !delivery.Redelivered
|
||||
requeue := !delivery.Redelivered && !isPermanentGenerationTaskDecodeError(err)
|
||||
w.rejectDelivery(delivery, requeue, err, envelope.TaskID)
|
||||
return
|
||||
}
|
||||
@@ -172,6 +174,18 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver
|
||||
}
|
||||
return
|
||||
}
|
||||
if lease == nil {
|
||||
w.rejectDelivery(delivery, !delivery.Redelivered, fmt.Errorf("generation task claim returned empty lease"), envelope.TaskID)
|
||||
return
|
||||
}
|
||||
|
||||
task := lease.task.toClaimedTask()
|
||||
heartbeatCtx, stopHeartbeat := context.WithCancel(parent)
|
||||
go w.runGenerationLeaseHeartbeat(heartbeatCtx, task.ID, lease.token, recoveryCfg.LeaseTTL)
|
||||
defer func() {
|
||||
stopHeartbeat()
|
||||
w.releaseGenerationTaskLease(context.Background(), task.ID, lease.token)
|
||||
}()
|
||||
|
||||
task.TaskType = strings.TrimSpace(task.TaskType)
|
||||
switch {
|
||||
@@ -215,10 +229,13 @@ func (w *ArticleGenerationWorker) runtimeProcessTimeout() time.Duration {
|
||||
if w != nil && w.processTimeout > 0 {
|
||||
return w.processTimeout
|
||||
}
|
||||
return articleGenerationProcessTimeout(config.GenerationConfig{})
|
||||
cfg := config.GenerationConfig{}
|
||||
config.NormalizeGenerationConfig(&cfg)
|
||||
return articleGenerationProcessTimeout(cfg)
|
||||
}
|
||||
|
||||
func articleGenerationProcessTimeout(cfg config.GenerationConfig) time.Duration {
|
||||
config.NormalizeGenerationConfig(&cfg)
|
||||
processTimeout := cfg.ArticleTimeout + 2*time.Minute
|
||||
if processTimeout <= 2*time.Minute {
|
||||
return 10 * time.Minute
|
||||
@@ -238,50 +255,11 @@ func (w *ArticleGenerationWorker) handleTaskFailure(ctx context.Context, task te
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
func isPermanentGenerationTaskDecodeError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
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
|
||||
return strings.Contains(err.Error(), "decode generation input params")
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, taskID int64) {
|
||||
|
||||
Reference in New Issue
Block a user