767 lines
20 KiB
Go
767 lines
20 KiB
Go
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
|
|
shouldRequeue bool
|
|
)
|
|
if err := tx.QueryRow(ctx, `
|
|
WITH existing AS (
|
|
SELECT gt.id,
|
|
gt.tenant_id,
|
|
gt.article_id,
|
|
gt.task_type,
|
|
a.generate_status
|
|
FROM generation_tasks gt
|
|
JOIN articles a ON a.id = gt.article_id
|
|
AND a.tenant_id = gt.tenant_id
|
|
WHERE gt.id = $1
|
|
AND gt.status = 'running'
|
|
AND (
|
|
gt.lease_expires_at IS NULL
|
|
OR gt.lease_expires_at < NOW()
|
|
)
|
|
FOR UPDATE OF gt
|
|
),
|
|
finalized_failed_article AS (
|
|
UPDATE generation_tasks gt
|
|
SET status = 'failed',
|
|
error_message = COALESCE(NULLIF(gt.error_message, ''), '文章已标记生成失败，停止后台恢复重试 [article_failed_stop_retry]'),
|
|
completed_at = COALESCE(gt.completed_at, NOW()),
|
|
lease_token = NULL,
|
|
lease_owner = NULL,
|
|
lease_expires_at = NULL,
|
|
updated_at = NOW()
|
|
FROM existing e
|
|
WHERE gt.id = e.id
|
|
AND e.generate_status = 'failed'
|
|
RETURNING e.tenant_id, e.article_id, e.task_type, false AS should_requeue
|
|
),
|
|
requeued AS (
|
|
UPDATE generation_tasks
|
|
SET status = 'queued',
|
|
error_message = NULL,
|
|
lease_token = NULL,
|
|
lease_owner = NULL,
|
|
lease_expires_at = NULL,
|
|
updated_at = NOW()
|
|
FROM existing e
|
|
WHERE generation_tasks.id = e.id
|
|
AND e.generate_status <> 'failed'
|
|
RETURNING e.tenant_id, e.article_id, e.task_type, true AS should_requeue
|
|
)
|
|
SELECT tenant_id, article_id, task_type, should_requeue
|
|
FROM finalized_failed_article
|
|
UNION ALL
|
|
SELECT tenant_id, article_id, task_type, should_requeue
|
|
FROM requeued
|
|
`, item.ID).Scan(&tenantID, &articleID, &taskType, &shouldRequeue); err != nil {
|
|
if err == pgx.ErrNoRows {
|
|
return false, nil
|
|
}
|
|
return false, err
|
|
}
|
|
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return false, err
|
|
}
|
|
if !shouldRequeue {
|
|
if w.logger != nil {
|
|
w.logger.Info("article generation expired task finalized because article is already failed",
|
|
zap.Int64("task_id", item.ID),
|
|
zap.Int64("article_id", articleID),
|
|
)
|
|
}
|
|
return true, nil
|
|
}
|
|
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
|
|
}
|