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:
@@ -0,0 +1,721 @@
|
||||
package generate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultGenerationTaskRecoveryInterval = time.Minute
|
||||
defaultGenerationTaskRecoveryTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
type generationTaskRecoveryConfig struct {
|
||||
LeaseTTL time.Duration
|
||||
RecoveryInterval time.Duration
|
||||
RecoveryTimeout time.Duration
|
||||
BatchSize int
|
||||
QueuedStaleAfter time.Duration
|
||||
MaxAttempts int
|
||||
}
|
||||
|
||||
type claimedGenerationTaskLease struct {
|
||||
task generationTaskRecoveryRecord
|
||||
token string
|
||||
}
|
||||
|
||||
type generationTaskRecoveryRecord struct {
|
||||
ID int64
|
||||
TenantID int64
|
||||
OperatorID int64
|
||||
ArticleID int64
|
||||
QuotaReservationID int64
|
||||
TaskType string
|
||||
InputParams map[string]interface{}
|
||||
AttemptCount int
|
||||
}
|
||||
|
||||
func (r generationTaskRecoveryRecord) toClaimedTask() tenantapp.ClaimedGenerationTask {
|
||||
return tenantapp.ClaimedGenerationTask{
|
||||
ID: r.ID,
|
||||
TenantID: r.TenantID,
|
||||
OperatorID: r.OperatorID,
|
||||
ArticleID: r.ArticleID,
|
||||
QuotaReservationID: r.QuotaReservationID,
|
||||
TaskType: r.TaskType,
|
||||
InputParams: r.InputParams,
|
||||
}
|
||||
}
|
||||
|
||||
type recoveredGenerationTask struct {
|
||||
ID int64
|
||||
TenantID int64
|
||||
ArticleID int64
|
||||
TaskType string
|
||||
AttemptCount int
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) startRecovery(ctx context.Context) {
|
||||
if w == nil || w.pool == nil {
|
||||
return
|
||||
}
|
||||
go w.runRecovery(ctx)
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) runRecovery(ctx context.Context) {
|
||||
w.runRecoveryOnce(context.Background())
|
||||
|
||||
for {
|
||||
interval := w.runtimeRecoveryConfig().RecoveryInterval
|
||||
timer := time.NewTimer(interval)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
return
|
||||
case <-timer.C:
|
||||
w.runRecoveryOnce(context.Background())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) runRecoveryOnce(parent context.Context) {
|
||||
cfg := w.runtimeRecoveryConfig()
|
||||
ctx, cancel := context.WithTimeout(parent, cfg.RecoveryTimeout)
|
||||
defer cancel()
|
||||
|
||||
queuedRecovered, requeueErr := w.requeueStaleGenerationTasks(ctx, cfg)
|
||||
if requeueErr != nil && w.logger != nil {
|
||||
w.logger.Warn("article generation queued task recovery failed", zap.Error(requeueErr))
|
||||
}
|
||||
|
||||
expired, expireErr := w.expireStaleRunningGenerationTasks(ctx, cfg)
|
||||
if expireErr != nil && w.logger != nil {
|
||||
w.logger.Warn("article generation running task recovery failed", zap.Error(expireErr))
|
||||
}
|
||||
|
||||
if w.logger != nil && (queuedRecovered > 0 || expired > 0) {
|
||||
w.logger.Info("article generation task recovery completed",
|
||||
zap.Int("queued_recovered_task_count", queuedRecovered),
|
||||
zap.Int("running_recovered_task_count", expired),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) runtimeRecoveryConfig() generationTaskRecoveryConfig {
|
||||
if w != nil && w.configProvider != nil {
|
||||
if cfg := w.configProvider.Current(); cfg != nil {
|
||||
return normalizeGenerationTaskRecoveryConfig(cfg.Generation)
|
||||
}
|
||||
}
|
||||
if w != nil {
|
||||
return normalizeGenerationTaskRecoveryConfig(w.generationCfg)
|
||||
}
|
||||
return normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{})
|
||||
}
|
||||
|
||||
func normalizeGenerationTaskRecoveryConfig(cfg config.GenerationConfig) generationTaskRecoveryConfig {
|
||||
normalized := cfg
|
||||
config.NormalizeGenerationConfig(&normalized)
|
||||
return generationTaskRecoveryConfig{
|
||||
LeaseTTL: normalized.TaskLeaseTTL,
|
||||
RecoveryInterval: normalized.TaskRecoveryInterval,
|
||||
RecoveryTimeout: normalized.TaskRecoveryTimeout,
|
||||
BatchSize: normalized.TaskRecoveryBatchSize,
|
||||
QueuedStaleAfter: normalized.TaskQueuedStaleAfter,
|
||||
MaxAttempts: normalized.TaskMaxAttempts,
|
||||
}
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) claimGenerationTask(ctx context.Context, taskID int64, owner string, leaseTTL time.Duration) (*claimedGenerationTaskLease, bool, error) {
|
||||
if leaseTTL <= 0 {
|
||||
leaseTTL = normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{}).LeaseTTL
|
||||
}
|
||||
token, err := newGenerationTaskLeaseToken()
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
leaseExpiresAt := time.Now().UTC().Add(leaseTTL)
|
||||
|
||||
var (
|
||||
row generationTaskRecoveryRecord
|
||||
operatorID sql.NullInt64
|
||||
articleID sql.NullInt64
|
||||
reservationID sql.NullInt64
|
||||
inputParamsJSON []byte
|
||||
)
|
||||
err = w.pool.QueryRow(ctx, `
|
||||
UPDATE generation_tasks
|
||||
SET status = 'running',
|
||||
started_at = COALESCE(started_at, NOW()),
|
||||
lease_token = $2,
|
||||
lease_owner = $3,
|
||||
lease_expires_at = $4,
|
||||
last_heartbeat_at = NOW(),
|
||||
attempt_count = attempt_count + 1,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND status = 'queued'
|
||||
RETURNING id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, attempt_count
|
||||
`, taskID, token, owner, leaseExpiresAt).Scan(
|
||||
&row.ID,
|
||||
&row.TenantID,
|
||||
&operatorID,
|
||||
&articleID,
|
||||
&reservationID,
|
||||
&row.TaskType,
|
||||
&inputParamsJSON,
|
||||
&row.AttemptCount,
|
||||
)
|
||||
if err != nil {
|
||||
if err == pgx.ErrNoRows {
|
||||
return nil, true, nil
|
||||
}
|
||||
return nil, false, err
|
||||
}
|
||||
|
||||
row.OperatorID = operatorID.Int64
|
||||
row.ArticleID = articleID.Int64
|
||||
row.QuotaReservationID = reservationID.Int64
|
||||
params, err := parseJSONMap(inputParamsJSON)
|
||||
if err != nil {
|
||||
row.InputParams = map[string]interface{}{}
|
||||
return &claimedGenerationTaskLease{task: row, token: token}, false, fmt.Errorf("decode generation input params: %w", err)
|
||||
}
|
||||
row.InputParams = params
|
||||
|
||||
return &claimedGenerationTaskLease{task: row, token: token}, false, nil
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) renewGenerationTaskLease(ctx context.Context, taskID int64, token string, leaseTTL time.Duration) (bool, error) {
|
||||
if taskID <= 0 || token == "" {
|
||||
return false, nil
|
||||
}
|
||||
if leaseTTL <= 0 {
|
||||
leaseTTL = normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{}).LeaseTTL
|
||||
}
|
||||
tag, err := w.pool.Exec(ctx, `
|
||||
UPDATE generation_tasks
|
||||
SET lease_expires_at = $3,
|
||||
last_heartbeat_at = NOW(),
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND lease_token = $2
|
||||
AND status = 'running'
|
||||
`, taskID, token, time.Now().UTC().Add(leaseTTL))
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return tag.RowsAffected() > 0, nil
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) releaseGenerationTaskLease(ctx context.Context, taskID int64, token string) {
|
||||
if w == nil || w.pool == nil || taskID <= 0 || token == "" {
|
||||
return
|
||||
}
|
||||
releaseCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if _, err := w.pool.Exec(releaseCtx, `
|
||||
UPDATE generation_tasks
|
||||
SET lease_token = NULL,
|
||||
lease_owner = NULL,
|
||||
lease_expires_at = NULL,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND lease_token = $2
|
||||
AND status IN ('completed', 'failed')
|
||||
`, taskID, token); err != nil && w.logger != nil {
|
||||
w.logger.Warn("article generation task lease release failed",
|
||||
zap.Error(err),
|
||||
zap.Int64("task_id", taskID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) runGenerationLeaseHeartbeat(ctx context.Context, taskID int64, token string, leaseTTL time.Duration) {
|
||||
if taskID <= 0 || token == "" || leaseTTL <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
interval := leaseTTL / 3
|
||||
if interval < 10*time.Second {
|
||||
interval = 10 * time.Second
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
heartbeatCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
ok, err := w.renewGenerationTaskLease(heartbeatCtx, taskID, token, leaseTTL)
|
||||
cancel()
|
||||
if err != nil {
|
||||
if w.logger != nil {
|
||||
w.logger.Warn("article generation task lease heartbeat failed",
|
||||
zap.Error(err),
|
||||
zap.Int64("task_id", taskID),
|
||||
)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !ok {
|
||||
if w.logger != nil {
|
||||
w.logger.Warn("article generation task lease lost",
|
||||
zap.Int64("task_id", taskID),
|
||||
)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) requeueStaleGenerationTasks(ctx context.Context, cfg generationTaskRecoveryConfig) (int, error) {
|
||||
if w == nil || w.pool == nil || w.rabbitMQ == nil || cfg.BatchSize <= 0 {
|
||||
return 0, nil
|
||||
}
|
||||
staleBefore := time.Now().UTC().Add(-cfg.QueuedStaleAfter)
|
||||
|
||||
rows, err := w.pool.Query(ctx, `
|
||||
WITH candidates AS (
|
||||
SELECT id, tenant_id, article_id, task_type, attempt_count
|
||||
FROM generation_tasks
|
||||
WHERE status = 'queued'
|
||||
AND article_id IS NOT NULL
|
||||
AND updated_at < $1
|
||||
ORDER BY updated_at ASC, id ASC
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT $2
|
||||
),
|
||||
touched AS (
|
||||
UPDATE generation_tasks gt
|
||||
SET updated_at = NOW()
|
||||
FROM candidates c
|
||||
WHERE gt.id = c.id
|
||||
RETURNING c.id, c.tenant_id, c.article_id, c.task_type, c.attempt_count
|
||||
)
|
||||
SELECT id, tenant_id, article_id, task_type, attempt_count
|
||||
FROM touched
|
||||
`, staleBefore, cfg.BatchSize)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]recoveredGenerationTask, 0, cfg.BatchSize)
|
||||
for rows.Next() {
|
||||
var item recoveredGenerationTask
|
||||
if err := rows.Scan(&item.ID, &item.TenantID, &item.ArticleID, &item.TaskType, &item.AttemptCount); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
requeued := 0
|
||||
for _, item := range items {
|
||||
if item.AttemptCount >= cfg.MaxAttempts {
|
||||
ok, err := w.finalizeExpiredGenerationTask(ctx, generationTaskRecoveryRecord{
|
||||
ID: item.ID,
|
||||
TenantID: item.TenantID,
|
||||
ArticleID: item.ArticleID,
|
||||
TaskType: item.TaskType,
|
||||
AttemptCount: item.AttemptCount,
|
||||
})
|
||||
if err != nil && w.logger != nil {
|
||||
w.logger.Warn("article generation stale queued task finalization failed",
|
||||
zap.Error(err),
|
||||
zap.Int64("task_id", item.ID),
|
||||
zap.Int64("article_id", item.ArticleID),
|
||||
)
|
||||
}
|
||||
if ok {
|
||||
requeued++
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err := w.publishGenerationTask(ctx, item.ID, item.TaskType, item.TenantID, item.ArticleID); err != nil {
|
||||
if w.logger != nil {
|
||||
w.logger.Warn("article generation stale queued task republish failed",
|
||||
zap.Error(err),
|
||||
zap.Int64("task_id", item.ID),
|
||||
zap.Int64("article_id", item.ArticleID),
|
||||
)
|
||||
}
|
||||
continue
|
||||
}
|
||||
requeued++
|
||||
}
|
||||
|
||||
return requeued, nil
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) expireStaleRunningGenerationTasks(ctx context.Context, cfg generationTaskRecoveryConfig) (int, error) {
|
||||
if w == nil || w.pool == nil || cfg.BatchSize <= 0 {
|
||||
return 0, nil
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
|
||||
rows, err := w.pool.Query(ctx, `
|
||||
SELECT id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, attempt_count
|
||||
FROM generation_tasks
|
||||
WHERE status = 'running'
|
||||
AND article_id IS NOT NULL
|
||||
AND (
|
||||
lease_expires_at IS NULL
|
||||
OR lease_expires_at < $1
|
||||
)
|
||||
ORDER BY COALESCE(lease_expires_at, started_at, updated_at, created_at) ASC, id ASC
|
||||
LIMIT $2
|
||||
`, now, cfg.BatchSize)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
recovered := make([]generationTaskRecoveryRecord, 0, cfg.BatchSize)
|
||||
for rows.Next() {
|
||||
var (
|
||||
item generationTaskRecoveryRecord
|
||||
operatorID sql.NullInt64
|
||||
articleID sql.NullInt64
|
||||
reservationID sql.NullInt64
|
||||
inputParamsJSON []byte
|
||||
)
|
||||
if err := rows.Scan(
|
||||
&item.ID,
|
||||
&item.TenantID,
|
||||
&operatorID,
|
||||
&articleID,
|
||||
&reservationID,
|
||||
&item.TaskType,
|
||||
&inputParamsJSON,
|
||||
&item.AttemptCount,
|
||||
); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
item.OperatorID = operatorID.Int64
|
||||
item.ArticleID = articleID.Int64
|
||||
item.QuotaReservationID = reservationID.Int64
|
||||
params, err := parseJSONMap(inputParamsJSON)
|
||||
if err != nil {
|
||||
params = map[string]interface{}{}
|
||||
}
|
||||
item.InputParams = params
|
||||
recovered = append(recovered, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
handled := 0
|
||||
for _, item := range recovered {
|
||||
if item.AttemptCount < cfg.MaxAttempts {
|
||||
ok, err := w.requeueExpiredRunningGenerationTask(ctx, item)
|
||||
if err != nil && w.logger != nil {
|
||||
w.logger.Warn("article generation expired task requeue failed",
|
||||
zap.Error(err),
|
||||
zap.Int64("task_id", item.ID),
|
||||
zap.Int64("article_id", item.ArticleID),
|
||||
)
|
||||
}
|
||||
if ok {
|
||||
handled++
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
ok, err := w.finalizeExpiredGenerationTask(ctx, item)
|
||||
if err != nil && w.logger != nil {
|
||||
w.logger.Warn("article generation expired task finalization failed",
|
||||
zap.Error(err),
|
||||
zap.Int64("task_id", item.ID),
|
||||
zap.Int64("article_id", item.ArticleID),
|
||||
)
|
||||
}
|
||||
if ok {
|
||||
handled++
|
||||
}
|
||||
}
|
||||
|
||||
return handled, nil
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) requeueExpiredRunningGenerationTask(ctx context.Context, item generationTaskRecoveryRecord) (bool, error) {
|
||||
if w == nil || w.rabbitMQ == nil {
|
||||
return false, nil
|
||||
}
|
||||
tx, err := w.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer func() {
|
||||
_ = tx.Rollback(ctx)
|
||||
}()
|
||||
|
||||
var (
|
||||
tenantID int64
|
||||
articleID int64
|
||||
taskType string
|
||||
)
|
||||
if err := tx.QueryRow(ctx, `
|
||||
UPDATE generation_tasks
|
||||
SET status = 'queued',
|
||||
error_message = NULL,
|
||||
lease_token = NULL,
|
||||
lease_owner = NULL,
|
||||
lease_expires_at = NULL,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND status = 'running'
|
||||
AND (
|
||||
lease_expires_at IS NULL
|
||||
OR lease_expires_at < NOW()
|
||||
)
|
||||
RETURNING tenant_id, article_id, task_type
|
||||
`, item.ID).Scan(&tenantID, &articleID, &taskType); err != nil {
|
||||
if err == pgx.ErrNoRows {
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return false, err
|
||||
}
|
||||
if err := w.publishGenerationTask(ctx, item.ID, taskType, tenantID, articleID); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) finalizeExpiredGenerationTask(ctx context.Context, item generationTaskRecoveryRecord) (bool, error) {
|
||||
errMsg := formatGenerationTaskRecoveryFailure()
|
||||
|
||||
tx, err := w.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer func() {
|
||||
_ = tx.Rollback(ctx)
|
||||
}()
|
||||
|
||||
var (
|
||||
claimedID int64
|
||||
operatorID sql.NullInt64
|
||||
articleID sql.NullInt64
|
||||
quotaReservationID sql.NullInt64
|
||||
taskType string
|
||||
)
|
||||
if err := tx.QueryRow(ctx, `
|
||||
UPDATE generation_tasks
|
||||
SET status = 'failed',
|
||||
error_message = $2,
|
||||
completed_at = COALESCE(completed_at, NOW()),
|
||||
lease_token = NULL,
|
||||
lease_owner = NULL,
|
||||
lease_expires_at = NULL,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND status IN ('running', 'queued')
|
||||
AND (
|
||||
status = 'queued'
|
||||
OR lease_expires_at IS NULL
|
||||
OR lease_expires_at < NOW()
|
||||
)
|
||||
RETURNING id, operator_id, article_id, quota_reservation_id, task_type
|
||||
`, item.ID, errMsg).Scan(&claimedID, &operatorID, &articleID, "aReservationID, &taskType); err != nil {
|
||||
if err == pgx.ErrNoRows {
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
_ = claimedID
|
||||
item.OperatorID = operatorID.Int64
|
||||
item.ArticleID = articleID.Int64
|
||||
item.QuotaReservationID = quotaReservationID.Int64
|
||||
item.TaskType = strings.TrimSpace(taskType)
|
||||
|
||||
if item.ArticleID > 0 {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE articles
|
||||
SET generate_status = 'failed',
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND tenant_id = $2
|
||||
AND generate_status = 'generating'
|
||||
AND deleted_at IS NULL
|
||||
`, item.ArticleID, item.TenantID); err != nil {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
|
||||
if item.TaskType == kolGenerationTaskType {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE task_records
|
||||
SET status = 'failed',
|
||||
error_message = $1,
|
||||
updated_at = NOW()
|
||||
WHERE tenant_id = $2
|
||||
AND task_type = $3
|
||||
AND resource_type = $4
|
||||
AND resource_id = $5
|
||||
`, errMsg, item.TenantID, "article_generation", "generation_task", item.ID); err != nil {
|
||||
return false, err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE kol_usage_logs
|
||||
SET status = 'failed',
|
||||
error_message = $1,
|
||||
finished_at = COALESCE(finished_at, NOW())
|
||||
WHERE generation_task_id = $2
|
||||
AND tenant_id = $3
|
||||
AND status IN ('pending', 'running')
|
||||
`, errMsg, item.ID, item.TenantID); err != nil {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
|
||||
if item.QuotaReservationID > 0 {
|
||||
var reservationStatus string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
WITH refunded AS (
|
||||
UPDATE quota_reservations
|
||||
SET status = 'refunded',
|
||||
refunded_amount = reserved_amount,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND tenant_id = $2
|
||||
AND status = 'pending'
|
||||
RETURNING status
|
||||
)
|
||||
SELECT status FROM refunded
|
||||
`, item.QuotaReservationID, item.TenantID).Scan(&reservationStatus); err != nil && err != pgx.ErrNoRows {
|
||||
return false, err
|
||||
}
|
||||
if reservationStatus == "refunded" {
|
||||
if err := insertGenerationRecoveryRefundLedger(ctx, tx, item); err != nil {
|
||||
return false, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
cacheCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if item.ArticleID > 0 {
|
||||
tenantapp.InvalidateArticleCaches(cacheCtx, w.cache, item.TenantID, &item.ArticleID)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func insertGenerationRecoveryRefundLedger(ctx context.Context, tx pgx.Tx, item generationTaskRecoveryRecord) error {
|
||||
reason := "generation_refund"
|
||||
referenceType := "generation_task"
|
||||
var operatorID interface{}
|
||||
if item.OperatorID > 0 {
|
||||
operatorID = item.OperatorID
|
||||
}
|
||||
_, err := tx.Exec(ctx, `
|
||||
INSERT INTO tenant_quota_ledgers (
|
||||
tenant_id,
|
||||
operator_id,
|
||||
quota_type,
|
||||
delta,
|
||||
balance_after,
|
||||
reason,
|
||||
reference_type,
|
||||
reference_id
|
||||
)
|
||||
SELECT
|
||||
$1,
|
||||
$2,
|
||||
'article_generation',
|
||||
1,
|
||||
COALESCE((
|
||||
SELECT balance_after
|
||||
FROM tenant_quota_ledgers
|
||||
WHERE tenant_id = $1
|
||||
AND quota_type = 'article_generation'
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT 1
|
||||
), 0) + 1,
|
||||
$3,
|
||||
$4,
|
||||
$5
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM tenant_quota_ledgers
|
||||
WHERE tenant_id = $1
|
||||
AND quota_type = 'article_generation'
|
||||
AND reason = $3
|
||||
AND reference_type = $4
|
||||
AND reference_id = $5
|
||||
)
|
||||
`, item.TenantID, operatorID, reason, referenceType, item.ID)
|
||||
return err
|
||||
}
|
||||
|
||||
func formatGenerationTaskRecoveryFailure() string {
|
||||
return "文章生成任务执行超时，worker 重启恢复后自动标记失败 [worker_recovery]"
|
||||
}
|
||||
|
||||
func (w *ArticleGenerationWorker) publishGenerationTask(ctx context.Context, taskID int64, taskType string, tenantID, articleID int64) error {
|
||||
if w == nil || w.rabbitMQ == nil {
|
||||
return nil
|
||||
}
|
||||
payload, err := json.Marshal(tenantapp.ArticleGenerationTaskEnvelope{
|
||||
TaskID: taskID,
|
||||
TaskType: taskType,
|
||||
TenantID: tenantID,
|
||||
ArticleID: articleID,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal generation task envelope: %w", err)
|
||||
}
|
||||
if err := w.rabbitMQ.PublishGenerationTask(ctx, payload); err != nil {
|
||||
return fmt.Errorf("publish generation task: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func newGenerationTaskLeaseToken() (string, error) {
|
||||
var raw [16]byte
|
||||
if _, err := rand.Read(raw[:]); err != nil {
|
||||
return "", fmt.Errorf("create generation task lease token: %w", err)
|
||||
}
|
||||
return hex.EncodeToString(raw[:]), nil
|
||||
}
|
||||
|
||||
func generationWorkerOwner(consumerName string) string {
|
||||
host, _ := os.Hostname()
|
||||
owner := strings.TrimSpace(host + "/" + consumerName)
|
||||
if owner == "/" {
|
||||
return "worker-generate"
|
||||
}
|
||||
if len(owner) > 128 {
|
||||
return owner[:128]
|
||||
}
|
||||
return owner
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package generate
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
)
|
||||
|
||||
func TestNormalizeGenerationTaskRecoveryConfigUsesGenerationDefaults(t *testing.T) {
|
||||
cfg := normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{
|
||||
ArticleTimeout: 90 * time.Second,
|
||||
})
|
||||
|
||||
if cfg.LeaseTTL != 3*time.Minute+30*time.Second {
|
||||
t.Fatalf("expected lease ttl article_timeout+2m, got %s", cfg.LeaseTTL)
|
||||
}
|
||||
if cfg.RecoveryInterval != time.Minute {
|
||||
t.Fatalf("expected recovery interval 1m, got %s", cfg.RecoveryInterval)
|
||||
}
|
||||
if cfg.RecoveryTimeout != 30*time.Second {
|
||||
t.Fatalf("expected recovery timeout 30s, got %s", cfg.RecoveryTimeout)
|
||||
}
|
||||
if cfg.BatchSize != 100 {
|
||||
t.Fatalf("expected batch size 100, got %d", cfg.BatchSize)
|
||||
}
|
||||
if cfg.QueuedStaleAfter != 2*time.Minute {
|
||||
t.Fatalf("expected queued stale after 2m, got %s", cfg.QueuedStaleAfter)
|
||||
}
|
||||
if cfg.MaxAttempts != 3 {
|
||||
t.Fatalf("expected max attempts 3, got %d", cfg.MaxAttempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeGenerationTaskRecoveryConfigKeepsOverrides(t *testing.T) {
|
||||
cfg := normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{
|
||||
ArticleTimeout: 5 * time.Minute,
|
||||
TaskLeaseTTL: 7 * time.Minute,
|
||||
TaskRecoveryInterval: 15 * time.Second,
|
||||
TaskRecoveryTimeout: 10 * time.Second,
|
||||
TaskRecoveryBatchSize: 12,
|
||||
TaskQueuedStaleAfter: 45 * time.Second,
|
||||
TaskMaxAttempts: 5,
|
||||
})
|
||||
|
||||
if cfg.LeaseTTL != 7*time.Minute ||
|
||||
cfg.RecoveryInterval != 15*time.Second ||
|
||||
cfg.RecoveryTimeout != 10*time.Second ||
|
||||
cfg.BatchSize != 12 ||
|
||||
cfg.QueuedStaleAfter != 45*time.Second ||
|
||||
cfg.MaxAttempts != 5 {
|
||||
t.Fatalf("unexpected recovery config: %#v", cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerationWorkerOwnerIsBounded(t *testing.T) {
|
||||
owner := generationWorkerOwner(strings.Repeat("x", 256))
|
||||
if len(owner) > 128 {
|
||||
t.Fatalf("expected owner to be bounded to 128 chars, got %d", len(owner))
|
||||
}
|
||||
if owner == "" {
|
||||
t.Fatalf("expected non-empty owner")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerationTaskRecoveryFailureMessage(t *testing.T) {
|
||||
message := formatGenerationTaskRecoveryFailure()
|
||||
if !strings.Contains(message, "worker_recovery") {
|
||||
t.Fatalf("expected recovery marker in failure message, got %q", message)
|
||||
}
|
||||
}
|
||||
@@ -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