Files
sub2api/backend/internal/service/batch_image_cleanup.go
T
shaw 80a229bce5 fix(batch-image): 修复审计发现的计费死锁、状态机与队列原子性缺陷
修复 PR #3768 批量图像 MVP 合并后审计报告中的全部问题:

结算与计费(高危):
- 所有 SETTLEMENT_* 失败(超冻结/计数非法/manifest 冲突/定价缺失/扣费失败)
  统一计入 retry_count 并在耗尽时释放冻结转 failed,消灭 settling 无限
  requeue 导致的冻结余额永久锁死
- 耗尽出口的释放指纹统一为 RequestHash,与 processor/Cancel/recovery 一致;
  release 遇同 request id 指纹冲突视为幂等成功,治愈历史毒消息
- 管理端校验 hold_multiplier >= discount_multiplier,定价快照对存量脏数据钳制
- 释放前校验 per-job hold claim(dedup+归档表),杜绝幻影释放

索引对账(高危):
- provider 输出与提交 custom_id 集对账:未知条目丢弃并记事件,
  漏项补 PROVIDER_RESULT_MISSING 失败行,保证 success+fail == item_count

提交与恢复(高危):
- 提交前转 uploading 并在 provider.Submit 期间心跳刷新 updated_at;
  恢复扫描改为原子复核(FailStaleUnsubmittedBatchImageJob),
  消灭慢提交被误杀退款而上游任务照常计费的孤儿场景
- 上游任务创建成功但本地状态推进失败时,尽力取消上游并清理输入
- recovery 释放失败时入队交由 worker releaseTerminalHold 兜底重试

队列与并发(中危):
- Enqueue(SetNX+LPush)与 Reserve(BRPop+ZAdd)均改为 Lua 原子脚本,
  消灭崩溃窗口导致 job 脱离队列、被 7 天 inflight 键锁死
- 锁冲突按 LockConflictDelay 重新入队(原直接丢弃需等 10 分钟 stale 恢复)
- 处理期间心跳:active zset 续期(ZAddXX 防幽灵成员)+ 锁 TTL 续期
- ReplaceBatchImageItemsForJob 增加 indexing 状态守卫,防掉队 worker 重写账目

存量回归(中危):
- image-only 定价条目(仅图片价无 token 价)恢复 token 计费 fail-closed,
  不再按 $0 计费;图片计费路径不受影响
- 鉴权余额门槛恢复 balance <= 0 语义,MinimumBalanceReserve 不再作硬 403

加固:
- ZIP max_items 钳制到管理员上限;Submit 补齐 Platform==Gemini 校验;
  gemini downloadUri 跟随前做 host 白名单校验
- 批量客户端改用共享 httpclient(拨号/TLS/响应头超时有界)
- 审计点名的忽略错误(MarkDownloaded/SettlementFailed/AppendEvent 等)改为记日志
2026-07-07 18:53:56 +08:00

313 lines
9.3 KiB
Go

package service
import (
"context"
"errors"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"go.uber.org/zap"
)
const (
defaultBatchImageInputRetentionAfterTerminal = 24 * time.Hour
defaultBatchImageOutputRetentionAfterTerminal = 72 * time.Hour
defaultBatchImageCleanupInterval = 30 * time.Minute
defaultBatchImageCleanupBatchSize = 100
)
type BatchImageCleanupService struct {
Repo BatchImageRepository
ProviderRegistry *BatchImageProviderRegistry
AccountResolver BatchImageAccountResolver
Config *config.Config
cancel context.CancelFunc
done chan struct{}
mu sync.Mutex
}
func NewBatchImageCleanupService(repo BatchImageRepository, accountRepo AccountRepository, cfg *config.Config) *BatchImageCleanupService {
return &BatchImageCleanupService{
Repo: repo,
ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg),
AccountResolver: &BatchImageAccountRepositoryResolver{Repo: accountRepo},
Config: cfg,
}
}
// appendCleanupEvent 追加清理审计事件;事件写入失败不阻断清理流程,但必须留痕。
func (s *BatchImageCleanupService) appendCleanupEvent(ctx context.Context, batchID, eventType string, payload any) {
if err := s.Repo.AppendBatchImageEvent(ctx, batchID, eventType, payload); err != nil {
logger.L().Warn("batch_image.cleanup_event_failed",
zap.String("batch_id", batchID),
zap.String("event_type", eventType),
zap.Error(err),
)
}
}
func (s *BatchImageCleanupService) DeleteOutputsForOwner(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) {
job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
if err != nil {
return nil, err
}
if job.Status == BatchImageJobStatusOutputDeleted || job.OutputDeletedAt != nil {
return BatchImageJobToPublic(job), nil
}
if job.Status != BatchImageJobStatusCompleted {
return nil, ErrBatchImageOutputDeleteNotReady
}
s.appendCleanupEvent(ctx, job.BatchID, "manual_output_delete_requested", map[string]any{
"batch_id": job.BatchID,
"cleanup_target": "output",
"reason": "manual",
})
if err := s.cleanupJob(ctx, job, CleanupTargetOutput, "manual"); err != nil {
return nil, err
}
updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
if err != nil {
return nil, err
}
return BatchImageJobToPublic(updated), nil
}
func (s *BatchImageCleanupService) CleanupInput(ctx context.Context, batchID string) error {
job, err := s.Repo.GetBatchImageJobByBatchID(ctx, batchID)
if err != nil {
return err
}
return s.cleanupJob(ctx, job, CleanupTargetInput, "ttl")
}
func (s *BatchImageCleanupService) CleanupOutput(ctx context.Context, batchID string, reason string) error {
job, err := s.Repo.GetBatchImageJobByBatchID(ctx, batchID)
if err != nil {
return err
}
return s.cleanupJob(ctx, job, CleanupTargetOutput, reason)
}
func (s *BatchImageCleanupService) RunOnce(ctx context.Context, now time.Time) (BatchImageCleanupRunResult, error) {
if s == nil || s.Repo == nil {
return BatchImageCleanupRunResult{}, ErrBatchImageCleanupFailed
}
if now.IsZero() {
now = time.Now()
}
limit := s.cleanupBatchSize()
result := BatchImageCleanupRunResult{}
inputCutoff := now.Add(-s.inputRetentionAfterTerminal())
inputJobs, err := s.Repo.ListBatchImageJobsDueForInputCleanup(ctx, inputCutoff, limit)
if err != nil {
return result, err
}
for _, job := range inputJobs {
if job == nil {
continue
}
if err := s.cleanupJob(ctx, job, CleanupTargetInput, "ttl"); err != nil {
result.Failures++
continue
}
result.InputCleaned++
}
outputJobs, err := s.Repo.ListBatchImageJobsDueForOutputCleanup(ctx, now, limit)
if err != nil {
return result, err
}
for _, job := range outputJobs {
if job == nil {
continue
}
if err := s.cleanupJob(ctx, job, CleanupTargetOutput, "expired"); err != nil {
result.Failures++
continue
}
result.OutputCleaned++
}
return result, nil
}
func (s *BatchImageCleanupService) Start() {
if s == nil || s.Repo == nil || s.Config == nil || !s.Config.BatchImage.Enabled || s.cleanupInterval() <= 0 {
return
}
s.mu.Lock()
defer s.mu.Unlock()
if s.cancel != nil {
return
}
ctx, cancel := context.WithCancel(context.Background())
s.cancel = cancel
s.done = make(chan struct{})
go func() {
defer close(s.done)
ticker := time.NewTicker(s.cleanupInterval())
defer ticker.Stop()
for {
_, _ = s.RunOnce(ctx, time.Now())
select {
case <-ctx.Done():
return
case <-ticker.C:
}
}
}()
}
func (s *BatchImageCleanupService) Stop() {
if s == nil {
return
}
s.mu.Lock()
cancel := s.cancel
done := s.done
s.cancel = nil
s.done = nil
s.mu.Unlock()
if cancel != nil {
cancel()
}
if done != nil {
<-done
}
}
func (s *BatchImageCleanupService) cleanupJob(ctx context.Context, job *BatchImageJob, target CleanupTarget, reason string) error {
if job == nil {
return ErrBatchImageJobNotFound
}
switch target {
case CleanupTargetInput:
if job.InputDeletedAt != nil {
return nil
}
if !IsTerminalBatchImageJobStatus(job.Status) {
return ErrBatchImageCleanupFailed
}
s.appendCleanupEvent(ctx, job.BatchID, "input_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil))
case CleanupTargetOutput:
if job.OutputDeletedAt != nil || job.Status == BatchImageJobStatusOutputDeleted {
return nil
}
if job.Status != BatchImageJobStatusCompleted && job.Status != BatchImageJobStatusFailed && job.Status != BatchImageJobStatusCancelled {
return ErrBatchImageOutputDeleteNotReady
}
s.appendCleanupEvent(ctx, job.BatchID, "output_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil))
default:
return ErrUnsupportedCleanupTarget
}
if err := s.callProviderCleanup(ctx, job, target); err != nil {
code := cleanupFailureCode(err)
msg := sanitizeBatchImagePublicMessage(err.Error())
if recordErr := s.Repo.RecordBatchImageCleanupFailure(ctx, job.BatchID, code, msg); recordErr != nil {
logger.L().Warn("batch_image.cleanup_failure_record_failed",
zap.String("batch_id", job.BatchID),
zap.Error(recordErr),
)
}
event := string(target) + "_cleanup_failed"
s.appendCleanupEvent(ctx, job.BatchID, event, map[string]any{"batch_id": job.BatchID, "cleanup_target": string(target), "reason": reason, "error_code": code})
if errors.Is(err, ErrBatchImageProviderUnsafeCleanupPath) {
return ErrBatchImageCleanupUnsafePath
}
return ErrBatchImageProviderCleanupFailed
}
deletedAt := time.Now()
if target == CleanupTargetInput {
return s.Repo.MarkBatchImageInputDeleted(ctx, job.BatchID, deletedAt)
}
return s.Repo.MarkBatchImageOutputDeleted(ctx, job.BatchID, deletedAt)
}
func (s *BatchImageCleanupService) callProviderCleanup(ctx context.Context, job *BatchImageJob, target CleanupTarget) error {
if s == nil || s.ProviderRegistry == nil || s.AccountResolver == nil {
return ErrBatchImageCleanupFailed
}
provider, ok := s.ProviderRegistry.Get(job.Provider)
if !ok || provider == nil {
return ErrBatchImageUnsupportedProvider
}
if job.AccountID == nil || *job.AccountID <= 0 {
return ErrBatchImageMissingAccountID
}
account, err := s.AccountResolver.ResolveBatchImageAccount(ctx, *job.AccountID)
if err != nil {
return err
}
if err := provider.Cleanup(ctx, job, account, target); err != nil {
if cleanupErrorIsNotFound(err) {
return nil
}
return err
}
return nil
}
func (s *BatchImageCleanupService) inputRetentionAfterTerminal() time.Duration {
if s != nil && s.Config != nil && s.Config.BatchImage.InputRetentionAfterTerminalHours > 0 {
return time.Duration(s.Config.BatchImage.InputRetentionAfterTerminalHours) * time.Hour
}
return defaultBatchImageInputRetentionAfterTerminal
}
func (s *BatchImageCleanupService) cleanupInterval() time.Duration {
if s != nil && s.Config != nil && s.Config.BatchImage.CleanupIntervalMinutes > 0 {
return time.Duration(s.Config.BatchImage.CleanupIntervalMinutes) * time.Minute
}
return defaultBatchImageCleanupInterval
}
func (s *BatchImageCleanupService) cleanupBatchSize() int {
if s != nil && s.Config != nil && s.Config.BatchImage.CleanupBatchSize > 0 {
return s.Config.BatchImage.CleanupBatchSize
}
return defaultBatchImageCleanupBatchSize
}
type BatchImageCleanupRunResult struct {
InputCleaned int
OutputCleaned int
Failures int
}
func cleanupEventPayload(batchID string, target CleanupTarget, reason string, deletedAt *time.Time) map[string]any {
payload := map[string]any{
"batch_id": batchID,
"cleanup_target": string(target),
"reason": reason,
}
if deletedAt != nil {
payload["deleted_at"] = deletedAt.UTC().Format(time.RFC3339)
}
return payload
}
func cleanupErrorIsNotFound(err error) bool {
if err == nil {
return false
}
reason := strings.ToUpper(infraerrors.Reason(err))
msg := strings.ToUpper(err.Error())
return strings.Contains(reason, "NOT_FOUND") || strings.Contains(msg, "NOT FOUND") || strings.Contains(msg, "404")
}
func cleanupFailureCode(err error) string {
if errors.Is(err, ErrBatchImageProviderUnsafeCleanupPath) {
return "BATCH_IMAGE_CLEANUP_UNSAFE_PATH"
}
reason := strings.TrimSpace(infraerrors.Reason(err))
if reason != "" {
return reason
}
return "BATCH_IMAGE_PROVIDER_CLEANUP_FAILED"
}