Files
sub2api/backend/internal/service/batch_image_worker_test.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

158 lines
5.1 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//go:build unit
package service
import (
"context"
"errors"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestBatchImageWorker_ProcessesJobOnce(t *testing.T) {
queue := newFakeBatchImageQueue("imgbatch_worker_once")
processor := &fakeBatchImageProcessor{}
worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{ReserveBlockTimeout: time.Millisecond})
require.NoError(t, worker.RunOnce(context.Background()))
require.Equal(t, []string{"imgbatch_worker_once"}, processor.processed)
require.Len(t, queue.requeued, 1)
require.Equal(t, defaultBatchImageWorkerRequeueDelay, queue.requeued[0].delay)
require.Equal(t, 1, queue.releaseCount)
}
func TestBatchImageWorker_RequeuesNonTerminalResultWithRequestedDelay(t *testing.T) {
queue := newFakeBatchImageQueue("imgbatch_worker_requeue")
processor := &fakeBatchImageProcessor{result: BatchImageProcessResult{RequeueAfter: 42 * time.Second}}
worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{})
require.NoError(t, worker.RunOnce(context.Background()))
require.Len(t, queue.requeued, 1)
require.Equal(t, "imgbatch_worker_requeue", queue.requeued[0].batchID)
require.Equal(t, 42*time.Second, queue.requeued[0].delay)
require.Empty(t, queue.acked)
}
func TestBatchImageWorker_AcksTerminalResult(t *testing.T) {
queue := newFakeBatchImageQueue("imgbatch_worker_terminal")
processor := &fakeBatchImageProcessor{result: BatchImageProcessResult{Terminal: true}}
worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{})
require.NoError(t, worker.RunOnce(context.Background()))
require.Equal(t, []string{"imgbatch_worker_terminal"}, queue.acked)
require.Empty(t, queue.requeued)
}
func TestBatchImageWorker_RequeuesOnProcessorError(t *testing.T) {
queue := newFakeBatchImageQueue("imgbatch_worker_error")
processor := &fakeBatchImageProcessor{err: errors.New("processor failed")}
worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{ErrorRetryDelay: 7 * time.Second})
require.NoError(t, worker.RunOnce(context.Background()))
require.Len(t, queue.requeued, 1)
require.Equal(t, 7*time.Second, queue.requeued[0].delay)
require.Empty(t, queue.acked)
}
func TestBatchImageWorker_RequeuesWhenJobLockNotAcquired(t *testing.T) {
queue := newFakeBatchImageQueue("imgbatch_worker_locked")
queue.lockAcquired = false
processor := &fakeBatchImageProcessor{}
worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{LockConflictDelay: 3 * time.Second})
// 锁冲突必须按冲突延迟重新入队;直接丢弃会让 job 滞留 active zset,
// 要等 StaleActiveAfter(默认 10 分钟)才被恢复。
require.NoError(t, worker.RunOnce(context.Background()))
require.Empty(t, processor.processed)
require.Len(t, queue.requeued, 1)
require.Equal(t, 3*time.Second, queue.requeued[0].delay)
require.Empty(t, queue.acked)
}
func TestNewBatchImageWorkerOptionsFromConfig_UsesFiniteReserveTimeout(t *testing.T) {
opts := NewBatchImageWorkerOptionsFromConfig(nil)
require.Equal(t, defaultBatchImageWorkerReserveBlockTimeout, opts.ReserveBlockTimeout)
require.Positive(t, opts.ReserveBlockTimeout)
}
type fakeBatchImageQueue struct {
reserved ReservedBatchImageJob
lockAcquired bool
acked []string
requeued []fakeBatchImageRequeue
releaseCount int
}
type fakeBatchImageRequeue struct {
batchID string
delay time.Duration
}
func newFakeBatchImageQueue(batchID string) *fakeBatchImageQueue {
return &fakeBatchImageQueue{
reserved: ReservedBatchImageJob{BatchID: batchID},
lockAcquired: true,
}
}
func (q *fakeBatchImageQueue) Enqueue(context.Context, string) error {
return nil
}
func (q *fakeBatchImageQueue) Reserve(context.Context, time.Duration) (ReservedBatchImageJob, error) {
return q.reserved, nil
}
func (q *fakeBatchImageQueue) RequeueAfter(_ context.Context, batchID string, delay time.Duration) error {
q.requeued = append(q.requeued, fakeBatchImageRequeue{batchID: batchID, delay: delay})
return nil
}
func (q *fakeBatchImageQueue) Ack(_ context.Context, batchID string) error {
q.acked = append(q.acked, batchID)
return nil
}
func (q *fakeBatchImageQueue) Heartbeat(context.Context, string) error {
return nil
}
func (q *fakeBatchImageQueue) MoveDueDelayedToReady(context.Context, int) (int, error) {
return 0, nil
}
func (q *fakeBatchImageQueue) RecoverStaleActive(context.Context, time.Duration, int) (int, error) {
return 0, nil
}
func (q *fakeBatchImageQueue) TryAcquireJobLock(context.Context, string, time.Duration) (BatchImageJobLock, bool, error) {
if !q.lockAcquired {
return nil, false, nil
}
return fakeBatchImageLock{release: func() { q.releaseCount++ }}, true, nil
}
type fakeBatchImageLock struct {
release func()
}
func (l fakeBatchImageLock) Release(context.Context) error {
if l.release != nil {
l.release()
}
return nil
}
type fakeBatchImageProcessor struct {
result BatchImageProcessResult
err error
processed []string
}
func (p *fakeBatchImageProcessor) Process(_ context.Context, batchID string) (BatchImageProcessResult, error) {
p.processed = append(p.processed, batchID)
return p.result, p.err
}