Files
sub2api/backend/internal/repository/batch_image_queue_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

200 lines
7.0 KiB
Go

//go:build unit
package repository
import (
"context"
"errors"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
)
func TestBatchImageQueue_DuplicateEnqueueReturnsAlreadyQueued(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
batchID := "imgbatch_duplicate"
require.NoError(t, queue.Enqueue(ctx, batchID))
err := queue.Enqueue(ctx, batchID)
require.Error(t, err)
require.True(t, errors.Is(err, service.ErrBatchImageAlreadyQueued))
}
func TestBatchImageQueue_RequeueAfterMovesJobFromActiveToDelayed(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
batchID := "imgbatch_requeue_after"
require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey, redis.Z{
Score: float64(time.Now().UnixMilli()),
Member: batchID,
}).Err())
require.NoError(t, queue.RequeueAfter(ctx, batchID, time.Minute))
require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, batchID).Err(), redis.Nil)
score, err := queue.rdb.ZScore(ctx, queue.delayedKey, batchID).Result()
require.NoError(t, err)
require.Greater(t, score, float64(time.Now().UnixMilli()))
}
func TestBatchImageQueue_MoveDueDelayedToReadyMovesDueJobs(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
dueBatchID := "imgbatch_due"
futureBatchID := "imgbatch_future"
now := time.Now()
require.NoError(t, queue.rdb.ZAdd(ctx, queue.delayedKey,
redis.Z{Score: float64(now.Add(-time.Second).UnixMilli()), Member: dueBatchID},
redis.Z{Score: float64(now.Add(time.Hour).UnixMilli()), Member: futureBatchID},
).Err())
moved, err := queue.MoveDueDelayedToReady(ctx, 10)
require.NoError(t, err)
require.Equal(t, 1, moved)
require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.delayedKey, dueBatchID).Err(), redis.Nil)
require.NoError(t, queue.rdb.ZScore(ctx, queue.delayedKey, futureBatchID).Err())
reserved, err := queue.Reserve(ctx, time.Millisecond)
require.NoError(t, err)
require.Equal(t, dueBatchID, reserved.BatchID)
}
func TestBatchImageQueue_RecoverStaleActiveMovesStaleJobsToReady(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
staleBatchID := "imgbatch_stale"
recentBatchID := "imgbatch_recent"
now := time.Now()
require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey,
redis.Z{Score: float64(now.Add(-time.Hour).UnixMilli()), Member: staleBatchID},
redis.Z{Score: float64(now.UnixMilli()), Member: recentBatchID},
).Err())
moved, err := queue.RecoverStaleActive(ctx, 10*time.Minute, 10)
require.NoError(t, err)
require.Equal(t, 1, moved)
require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, staleBatchID).Err(), redis.Nil)
require.NoError(t, queue.rdb.ZScore(ctx, queue.activeKey, recentBatchID).Err())
reserved, err := queue.Reserve(ctx, time.Millisecond)
require.NoError(t, err)
require.Equal(t, staleBatchID, reserved.BatchID)
}
func TestBatchImageQueue_JobLockReleaseOnlyDeletesMatchingToken(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
batchID := "imgbatch_lock"
lock, ok, err := queue.TryAcquireJobLock(ctx, batchID, time.Minute)
require.NoError(t, err)
require.True(t, ok)
require.NoError(t, queue.rdb.Set(ctx, queue.lockKey(batchID), "other-token", time.Minute).Err())
require.NoError(t, lock.Release(ctx))
got, err := queue.rdb.Get(ctx, queue.lockKey(batchID)).Result()
require.NoError(t, err)
require.Equal(t, "other-token", got)
require.NoError(t, queue.rdb.Del(ctx, queue.lockKey(batchID)).Err())
lock, ok, err = queue.TryAcquireJobLock(ctx, batchID, time.Minute)
require.NoError(t, err)
require.True(t, ok)
require.NoError(t, lock.Release(ctx))
require.ErrorIs(t, queue.rdb.Get(ctx, queue.lockKey(batchID)).Err(), redis.Nil)
}
func TestBatchImageQueue_ReserveAtomicallyMovesJobToActive(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
batchID := "imgbatch_reserve"
require.NoError(t, queue.Enqueue(ctx, batchID))
reserved, err := queue.Reserve(ctx, time.Second)
require.NoError(t, err)
require.Equal(t, batchID, reserved.BatchID)
// 弹出与写入 active 必须原子完成:ready 已空,active 中有该 job。
require.Equal(t, int64(0), queue.rdb.LLen(ctx, queue.readyKey).Val())
score, err := queue.rdb.ZScore(ctx, queue.activeKey, batchID).Result()
require.NoError(t, err)
require.Positive(t, score)
}
func TestBatchImageQueue_ReserveReturnsEmptyAfterTimeout(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
start := time.Now()
_, err := queue.Reserve(ctx, 50*time.Millisecond)
require.ErrorIs(t, err, service.ErrBatchImageQueueEmpty)
require.Less(t, time.Since(start), 5*time.Second)
}
func TestBatchImageQueue_ReserveDropsInvalidPayload(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
require.NoError(t, queue.rdb.LPush(ctx, queue.readyKey, "not-a-batch-id").Err())
_, err := queue.Reserve(ctx, 10*time.Millisecond)
require.ErrorIs(t, err, service.ErrInvalidBatchImageQueuePayload)
// 非法 payload 不得残留在 active zset,否则 stale 恢复会无限重投。
require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, "not-a-batch-id").Err(), redis.Nil)
}
func TestBatchImageQueue_HeartbeatOnlyRefreshesExistingActiveMember(t *testing.T) {
ctx := context.Background()
queue, _ := newBatchImageQueueTest(t)
batchID := "imgbatch_heartbeat"
// 不在 active 中:心跳不得创建幽灵成员。
require.NoError(t, queue.Heartbeat(ctx, batchID))
require.ErrorIs(t, queue.rdb.ZScore(ctx, queue.activeKey, batchID).Err(), redis.Nil)
require.NoError(t, queue.rdb.ZAdd(ctx, queue.activeKey, redis.Z{Score: 1, Member: batchID}).Err())
require.NoError(t, queue.Heartbeat(ctx, batchID))
score, err := queue.rdb.ZScore(ctx, queue.activeKey, batchID).Result()
require.NoError(t, err)
require.Greater(t, score, float64(1))
}
func TestBatchImageQueue_JobLockRefreshExtendsTTLOnlyForHolder(t *testing.T) {
ctx := context.Background()
queue, mr := newBatchImageQueueTest(t)
batchID := "imgbatch_lock_refresh"
lock, ok, err := queue.TryAcquireJobLock(ctx, batchID, time.Minute)
require.NoError(t, err)
require.True(t, ok)
refresher, isRefresher := lock.(service.BatchImageJobLockRefresher)
require.True(t, isRefresher)
require.NoError(t, refresher.Refresh(ctx, 10*time.Minute))
ttl := mr.TTL(queue.lockKey(batchID))
require.Greater(t, ttl, 5*time.Minute)
// token 不匹配时不得续期他人持有的锁。
require.NoError(t, queue.rdb.Set(ctx, queue.lockKey(batchID), "other-token", time.Minute).Err())
require.NoError(t, refresher.Refresh(ctx, 10*time.Minute))
ttl = mr.TTL(queue.lockKey(batchID))
require.LessOrEqual(t, ttl, time.Minute)
}
func newBatchImageQueueTest(t *testing.T) (*batchImageQueue, *miniredis.Miniredis) {
t.Helper()
mr := miniredis.RunT(t)
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() {
_ = rdb.Close()
})
queue := newBatchImageQueueWithOptions(rdb, batchImageQueueOptions{
InflightTTL: time.Hour,
LockTTL: time.Minute,
})
return queue, mr
}