Merge pull request #3799 from Wei-Shaw/fix/batch-image-audit-hardening

fix(batch-image): 修复审计发现的计费死锁、状态机与队列原子性缺陷
This commit is contained in:
Wesley Liddick
2026-07-07 20:20:49 +08:00
committed by GitHub
32 changed files with 1190 additions and 90 deletions
@@ -8,10 +8,12 @@ import (
"strings"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
type BatchImageHandler struct {
@@ -160,7 +162,18 @@ func (h *BatchImageHandler) ItemContent(c *gin.Context) {
if _, err := io.Copy(c.Writer, stream.Reader); err != nil {
return
}
_ = h.service.MarkDownloaded(c.Request.Context(), owner, c.Param("id"))
h.markDownloadedBestEffort(c, owner)
}
// markDownloadedBestEffort 在响应体已写出后标记下载状态;
// 此时无法再向客户端返回错误,失败只能记日志(不能静默丢弃)。
func (h *BatchImageHandler) markDownloadedBestEffort(c *gin.Context, owner service.BatchImageOwner) {
if err := h.service.MarkDownloaded(c.Request.Context(), owner, c.Param("id")); err != nil {
logger.L().Warn("batch_image.mark_downloaded_failed",
zap.String("batch_id", c.Param("id")),
zap.Error(err),
)
}
}
func (h *BatchImageHandler) Download(c *gin.Context) {
@@ -186,7 +199,7 @@ func (h *BatchImageHandler) Download(c *gin.Context) {
}
return
}
_ = h.service.MarkDownloaded(c.Request.Context(), owner, c.Param("id"))
h.markDownloadedBestEffort(c, owner)
}
func (h *BatchImageHandler) DeleteRecord(c *gin.Context) {
@@ -20,6 +20,10 @@ const (
defaultBatchImageLockPrefix = "batch_image:queue:lock:"
defaultBatchImageInflightTTL = 7 * 24 * time.Hour
defaultBatchImageJobLockTTL = 5 * time.Minute
// batchImageReservePollInterval 是原子 Reserve 脚本空轮询的间隔。
// 用轮询替代 BRPop 是为了保证 "弹出 + 写 active" 的原子性。
batchImageReservePollInterval = time.Second
)
var batchImageMoveDueDelayedScript = redis.NewScript(`
@@ -47,6 +51,36 @@ end
return 0
`)
var batchImageRefreshLockScript = redis.NewScript(`
if redis.call("GET", KEYS[1]) == ARGV[1] then
return redis.call("PEXPIRE", KEYS[1], ARGV[2])
end
return 0
`)
// batchImageReserveScript 原子地从 ready 弹出并写入 active zset。
// BRPop + ZAdd 两步方案在两步之间进程崩溃时 job 会脱离所有队列结构,
// 且 inflight 去重键(默认 7 天)会挡住所有重新入队。
var batchImageReserveScript = redis.NewScript(`
local job = redis.call("RPOP", KEYS[1])
if not job then
return nil
end
redis.call("ZADD", KEYS[2], ARGV[1], job)
return job
`)
// batchImageEnqueueScript 原子地设置 inflight 去重键并推入 ready。
// SetNX + LPush 两步方案在两步之间进程崩溃时,inflight 键(默认 7 天)
// 会挡住所有后续入队,而 job 从未进入 ready。
var batchImageEnqueueScript = redis.NewScript(`
if redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) then
redis.call("LPUSH", KEYS[2], ARGV[1])
return 1
end
return 0
`)
type batchImageQueue struct {
rdb *redis.Client
readyKey string
@@ -131,40 +165,65 @@ func (q *batchImageQueue) Enqueue(ctx context.Context, batchID string) error {
return service.ErrInvalidBatchImageQueuePayload
}
ok, err := q.rdb.SetNX(ctx, q.inflightKey(batchID), batchID, q.inflightTTL).Result()
applied, err := batchImageEnqueueScript.Run(ctx, q.rdb,
[]string{q.inflightKey(batchID), q.readyKey},
batchID, q.inflightTTL.Milliseconds(),
).Int()
if err != nil {
return err
}
if !ok {
if applied == 0 {
return service.ErrBatchImageAlreadyQueued
}
if err := q.rdb.LPush(ctx, q.readyKey, batchID).Err(); err != nil {
_ = q.rdb.Del(ctx, q.inflightKey(batchID)).Err()
return err
}
return nil
}
func (q *batchImageQueue) Reserve(ctx context.Context, blockTimeout time.Duration) (service.ReservedBatchImageJob, error) {
result, err := q.rdb.BRPop(ctx, blockTimeout, q.readyKey).Result()
deadline := time.Now().Add(blockTimeout)
for {
batchID, err := q.reserveOnce(ctx)
if err == nil {
return service.ReservedBatchImageJob{BatchID: batchID}, nil
}
if !errors.Is(err, service.ErrBatchImageQueueEmpty) {
return service.ReservedBatchImageJob{}, err
}
remaining := time.Until(deadline)
if remaining <= 0 {
return service.ReservedBatchImageJob{}, service.ErrBatchImageQueueEmpty
}
wait := batchImageReservePollInterval
if remaining < wait {
wait = remaining
}
timer := time.NewTimer(wait)
select {
case <-ctx.Done():
timer.Stop()
return service.ReservedBatchImageJob{}, ctx.Err()
case <-timer.C:
}
}
}
func (q *batchImageQueue) reserveOnce(ctx context.Context) (string, error) {
raw, err := batchImageReserveScript.Run(ctx, q.rdb, []string{q.readyKey, q.activeKey}, time.Now().UnixMilli()).Result()
if errors.Is(err, redis.Nil) {
return service.ReservedBatchImageJob{}, service.ErrBatchImageQueueEmpty
return "", service.ErrBatchImageQueueEmpty
}
if err != nil {
return service.ReservedBatchImageJob{}, err
return "", err
}
if len(result) != 2 || !service.IsValidBatchImageID(result[1]) {
return service.ReservedBatchImageJob{}, service.ErrInvalidBatchImageQueuePayload
batchID, ok := raw.(string)
if !ok || !service.IsValidBatchImageID(batchID) {
// 非法 payload 已被脚本写入 active,必须移除,否则 stale 恢复会把它
// 无限重投回 ready。
if ok && batchID != "" {
_ = q.rdb.ZRem(ctx, q.activeKey, batchID).Err()
}
return "", service.ErrInvalidBatchImageQueuePayload
}
batchID := result[1]
if err := q.rdb.ZAdd(ctx, q.activeKey, redis.Z{
Score: float64(time.Now().UnixMilli()),
Member: batchID,
}).Err(); err != nil {
return service.ReservedBatchImageJob{}, err
}
return service.ReservedBatchImageJob{BatchID: batchID}, nil
return batchID, nil
}
func (q *batchImageQueue) RequeueAfter(ctx context.Context, batchID string, delay time.Duration) error {
@@ -202,7 +261,9 @@ func (q *batchImageQueue) Heartbeat(ctx context.Context, batchID string) error {
if !service.IsValidBatchImageID(batchID) {
return service.ErrInvalidBatchImageQueuePayload
}
return q.rdb.ZAdd(ctx, q.activeKey, redis.Z{
// XX:只刷新已存在的 active 成员。无条件 ZAdd 会在 Ack/Requeue 之后的
// 竞态心跳里把幽灵成员塞回 active zset。
return q.rdb.ZAddXX(ctx, q.activeKey, redis.Z{
Score: float64(time.Now().UnixMilli()),
Member: batchID,
}).Err()
@@ -269,6 +330,19 @@ func (l *batchImageRedisJobLock) Release(ctx context.Context) error {
return batchImageReleaseLockScript.Run(ctx, l.rdb, []string{l.key}, l.token).Err()
}
// Refresh 在仍持有锁(token 匹配)时续期 TTL,供长处理任务的心跳调用。
func (l *batchImageRedisJobLock) Refresh(ctx context.Context, ttl time.Duration) error {
if l == nil || l.rdb == nil || l.key == "" || l.token == "" {
return nil
}
if ttl <= 0 {
ttl = defaultBatchImageJobLockTTL
}
return batchImageRefreshLockScript.Run(ctx, l.rdb, []string{l.key}, l.token, ttl.Milliseconds()).Err()
}
var _ service.BatchImageJobLockRefresher = (*batchImageRedisJobLock)(nil)
func newBatchImageLockToken() (string, error) {
var b [16]byte
if _, err := rand.Read(b[:]); err != nil {
@@ -108,6 +108,82 @@ func TestBatchImageQueue_JobLockReleaseOnlyDeletesMatchingToken(t *testing.T) {
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)
@@ -152,6 +152,45 @@ func (r *batchImageRepository) TransitionBatchImageJobStatus(ctx context.Context
return tx.Commit()
}
func (r *batchImageRepository) TouchBatchImageJobSubmitting(ctx context.Context, batchID string) error {
_, err := r.sql.ExecContext(ctx, `
UPDATE batch_image_jobs
SET updated_at = $2
WHERE batch_id = $1
AND status IN ('created', 'uploading')`, batchID, time.Now())
return err
}
func (r *batchImageRepository) FailStaleUnsubmittedBatchImageJob(ctx context.Context, batchID string, cutoff time.Time, code, message string) (bool, error) {
now := time.Now()
res, err := r.sql.ExecContext(ctx, `
UPDATE batch_image_jobs
SET status = 'failed',
last_error_code = $2,
last_error_message = $3,
finished_at = CASE WHEN finished_at IS NULL THEN $4 ELSE finished_at END,
updated_at = $4,
version = version + 1
WHERE batch_id = $1
AND status IN ('created', 'uploading')
AND provider_job_name IS NULL
AND updated_at <= $5`, batchID, code, message, now, cutoff)
if err != nil {
return false, err
}
affected, err := res.RowsAffected()
if err != nil {
return false, err
}
if affected == 0 {
return false, nil
}
return true, appendBatchImageEventWithSQL(ctx, r.sql, batchID, "billing_hold_recovery_failed_unsubmitted", map[string]any{
"batch_id": batchID,
"error_code": code,
})
}
func (r *batchImageRepository) UpdateBatchImageJobProviderOutputRef(ctx context.Context, batchID, providerOutputRef string) error {
res, err := r.sql.ExecContext(ctx, `
UPDATE batch_image_jobs
@@ -410,9 +449,15 @@ func (r *batchImageRepository) ReplaceBatchImageItemsForJob(ctx context.Context,
func (r *batchImageRepository) replaceBatchImageItemsForJobWithSQL(ctx context.Context, sqlq batchImageSQLExecutor, batchID string, items []service.CreateBatchImageItemParams, counts service.BatchImageCounts) error {
var id int64
if err := sqlq.QueryRowContext(ctx, `SELECT id FROM batch_image_jobs WHERE batch_id = $1 FOR UPDATE`, batchID).Scan(&id); err != nil {
var status string
if err := sqlq.QueryRowContext(ctx, `SELECT id, status FROM batch_image_jobs WHERE batch_id = $1 FOR UPDATE`, batchID).Scan(&id, &status); err != nil {
return translatePersistenceError(err, service.ErrBatchImageJobNotFound, nil)
}
// 仅允许 indexing 状态重建 item 表:防止锁过期后掉队的 worker
// 重写已完成/已结算 job 的条目,造成账目与结果漂移。
if status != service.BatchImageJobStatusIndexing {
return service.ErrBatchImageIndexStateConflict
}
promptPreviews, err := r.batchImageItemPromptPreviews(ctx, sqlq, batchID)
if err != nil {
return err
@@ -208,6 +208,16 @@ func TestBatchImageRepository_ReplaceBatchImageItemsForJob(t *testing.T) {
})
require.NoError(t, err)
// 非 indexing 状态不允许重建 item 表:防止锁过期后掉队的 worker
// 重写已完成/已结算 job 的条目。
err = repo.ReplaceBatchImageItemsForJob(ctx, batchID, []service.CreateBatchImageItemParams{
{CustomID: "old", Status: service.BatchImageItemStatusSuccess, SourceLineNumber: &lineOne, ImageCount: 1},
}, service.BatchImageCounts{SuccessCount: 1})
require.ErrorIs(t, err, service.ErrBatchImageIndexStateConflict)
require.NoError(t, repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusSubmitted, service.BatchImageTransitionOptions{}))
require.NoError(t, repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusIndexing, service.BatchImageTransitionOptions{}))
err = repo.ReplaceBatchImageItemsForJob(ctx, batchID, []service.CreateBatchImageItemParams{
{CustomID: "old", Status: service.BatchImageItemStatusSuccess, SourceLineNumber: &lineOne, ImageCount: 1},
}, service.BatchImageCounts{SuccessCount: 1})
@@ -335,6 +335,16 @@ func releaseUsageBillingBatchImageBalance(ctx context.Context, tx *sql.Tx, cmd *
if cmd.HoldAmount <= 0 {
return &service.BatchImageBalanceHoldResult{}, nil
}
// 释放前校验该 job 确实预留过 hold(hold request id 已被 claim),
// 防止从未成功冻结的 job 触发"幻影释放",从其他用户的冻结资金池中凭空生成余额。
held, heldErr := batchImageHoldClaimExists(ctx, tx, service.BatchImageHoldRequestID(cmd.BatchID), cmd.APIKeyID)
if heldErr != nil {
return nil, heldErr
}
if !held {
logger.LegacyPrintf("repository.usage_billing", "[BatchImage] release skipped, hold was never reserved: batch=%s", cmd.BatchID)
return &service.BatchImageBalanceHoldResult{}, nil
}
var balance, frozen float64
err := tx.QueryRowContext(ctx, `
UPDATE users
@@ -358,6 +368,35 @@ func releaseUsageBillingBatchImageBalance(ctx context.Context, tx *sql.Tx, cmd *
return nil, errors.New("batch image frozen balance is insufficient")
}
// batchImageHoldClaimExists 检查 hold request id 是否已在 dedup(或归档)表中被 claim,
// 即该 batch 的冻结操作确实成功提交过。
func batchImageHoldClaimExists(ctx context.Context, tx *sql.Tx, holdRequestID string, apiKeyID int64) (bool, error) {
var exists int
err := tx.QueryRowContext(ctx, `
SELECT 1
FROM usage_billing_dedup
WHERE request_id = $1 AND api_key_id = $2
`, holdRequestID, apiKeyID).Scan(&exists)
if err == nil {
return true, nil
}
if !errors.Is(err, sql.ErrNoRows) {
return false, err
}
err = tx.QueryRowContext(ctx, `
SELECT 1
FROM usage_billing_dedup_archive
WHERE request_id = $1 AND api_key_id = $2
`, holdRequestID, apiKeyID).Scan(&exists)
if err == nil {
return true, nil
}
if errors.Is(err, sql.ErrNoRows) {
return false, nil
}
return false, err
}
func userExistsForBilling(ctx context.Context, tx *sql.Tx, userID int64) (bool, error) {
var exists int
err := tx.QueryRowContext(ctx, `
@@ -217,15 +217,45 @@ func TestReleaseUsageBillingBatchImageBalance_ReturnsFrozenToAvailable(t *testin
mock.ExpectBegin()
tx, err := db.BeginTx(ctx, nil)
require.NoError(t, err)
mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup\s+WHERE request_id = \$1 AND api_key_id = \$2`).
WithArgs(service.BatchImageHoldRequestID("imgbatch_release"), int64(7)).
WillReturnRows(sqlmock.NewRows([]string{"?column?"}).AddRow(1))
mock.ExpectQuery(releaseBatchImageHoldSQL).
WithArgs(1.0, int64(42)).
WillReturnRows(sqlmock.NewRows([]string{"balance", "frozen_balance"}).AddRow(10.0, 0.0))
mock.ExpectCommit()
result, err := releaseUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, HoldAmount: 1})
result, err := releaseUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, APIKeyID: 7, BatchID: "imgbatch_release", HoldAmount: 1})
require.NoError(t, err)
require.InDelta(t, 10.0, *result.NewBalance, 0.000001)
require.InDelta(t, 0.0, *result.FrozenBalance, 0.000001)
require.NoError(t, tx.Commit())
require.NoError(t, mock.ExpectationsWereMet())
}
func TestReleaseUsageBillingBatchImageBalance_SkipsWhenHoldNeverReserved(t *testing.T) {
ctx := context.Background()
db, mock, err := sqlmock.New()
require.NoError(t, err)
defer func() { _ = db.Close() }()
mock.ExpectBegin()
tx, err := db.BeginTx(ctx, nil)
require.NoError(t, err)
// dedup 与归档表均无 hold claim:说明该 job 从未成功冻结,
// 释放必须跳过,不得从他人冻结资金池中凭空生成余额。
mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup\s+WHERE request_id = \$1 AND api_key_id = \$2`).
WithArgs(service.BatchImageHoldRequestID("imgbatch_phantom"), int64(7)).
WillReturnError(sql.ErrNoRows)
mock.ExpectQuery(`SELECT 1\s+FROM usage_billing_dedup_archive\s+WHERE request_id = \$1 AND api_key_id = \$2`).
WithArgs(service.BatchImageHoldRequestID("imgbatch_phantom"), int64(7)).
WillReturnError(sql.ErrNoRows)
mock.ExpectCommit()
result, err := releaseUsageBillingBatchImageBalance(ctx, tx, &service.BatchImageBalanceHoldCommand{UserID: 42, APIKeyID: 7, BatchID: "imgbatch_phantom", HoldAmount: 1})
require.NoError(t, err)
require.Nil(t, result.NewBalance)
require.Nil(t, result.FrozenBalance)
require.NoError(t, tx.Commit())
require.NoError(t, mock.ExpectationsWereMet())
}
@@ -289,14 +289,11 @@ func setGroupContext(c *gin.Context, group *service.Group) {
c.Request = c.Request.WithContext(ctx)
}
func apiKeyBalanceBelowAuthThreshold(balance float64, cfg *config.Config) bool {
if balance <= 0 {
return true
}
if cfg == nil || cfg.Billing.MinimumBalanceReserve <= 0 {
return false
}
return balance < cfg.Billing.MinimumBalanceReserve
// apiKeyBalanceBelowAuthThreshold 保持鉴权层的历史语义:仅在余额耗尽(<=0)时拒绝。
// MinimumBalanceReserve 只作为 billing-cache 预检的保守下限,不得复用为鉴权硬门槛,
// 否则已配置该值的存量部署升级后,0 < balance < reserve 的用户会在所有端点被静默 403。
func apiKeyBalanceBelowAuthThreshold(balance float64, _ *config.Config) bool {
return balance <= 0
}
func abortIfAPIKeyGroupUnavailable(c *gin.Context, apiKey *service.APIKey) bool {
@@ -542,6 +542,8 @@ func TestApiKeyAuthWithSubscriptionGoogle_InsufficientBalance(t *testing.T) {
func TestApiKeyAuthWithSubscriptionGoogle_BalanceBelowMinimumReserve(t *testing.T) {
gin.SetMode(gin.TestMode)
// 鉴权层保持历史语义:MinimumBalanceReserve 只用于 billing-cache 预检,
// 0 < balance < reserve 的用户不得在鉴权中间件被硬 403。
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
@@ -567,6 +569,36 @@ func TestApiKeyAuthWithSubscriptionGoogle_BalanceBelowMinimumReserve(t *testing.
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
}
func TestApiKeyAuthWithSubscriptionGoogle_RejectsExhaustedBalance(t *testing.T) {
gin.SetMode(gin.TestMode)
r := gin.New()
apiKeyService := newTestAPIKeyService(fakeAPIKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
return &service.APIKey{
ID: 1,
Key: key,
Status: service.StatusActive,
User: &service.User{
ID: 123,
Status: service.StatusActive,
Balance: 0,
},
}, nil
},
})
cfg := &config.Config{}
r.Use(APIKeyAuthWithSubscriptionGoogle(apiKeyService, nil, cfg))
r.GET("/v1beta/test", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) })
req := httptest.NewRequest(http.MethodGet, "/v1beta/test", nil)
req.Header.Set("Authorization", "Bearer ok")
rec := httptest.NewRecorder()
r.ServeHTTP(rec, req)
require.Equal(t, http.StatusForbidden, rec.Code)
var resp googleErrorResponse
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
@@ -1000,7 +1000,7 @@ func TestAPIKeyAuthTouchesLastUsedInStandardMode(t *testing.T) {
require.Equal(t, 1, touchCalls)
}
func TestAPIKeyAuthRejectsBalanceBelowMinimumReserve(t *testing.T) {
func TestAPIKeyAuthAllowsBalanceBelowMinimumReserve(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{
@@ -1039,6 +1039,49 @@ func TestAPIKeyAuthRejectsBalanceBelowMinimumReserve(t *testing.T) {
req.Header.Set("x-api-key", apiKey.Key)
router.ServeHTTP(w, req)
// 鉴权层保持历史语义:MinimumBalanceReserve 只用于 billing-cache 预检,
// 0 < balance < reserve 不得被鉴权中间件硬 403(存量部署静默行为变更)。
require.Equal(t, http.StatusOK, w.Code)
}
func TestAPIKeyAuthRejectsExhaustedBalance(t *testing.T) {
gin.SetMode(gin.TestMode)
user := &service.User{
ID: 10,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 0,
Concurrency: 3,
}
apiKey := &service.APIKey{
ID: 104,
UserID: user.ID,
Key: "held-balance-zero",
Status: service.StatusActive,
User: user,
}
apiKeyRepo := &stubApiKeyRepo{
getByKey: func(ctx context.Context, key string) (*service.APIKey, error) {
if key != apiKey.Key {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
userClone := *user
clone.User = &userClone
return &clone, nil
},
}
cfg := &config.Config{RunMode: config.RunModeStandard}
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
router := newAuthTestRouter(apiKeyService, nil, cfg)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("x-api-key", apiKey.Key)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusForbidden, w.Code)
requireAPIKeyAuthError(t, w, "INSUFFICIENT_BALANCE", "Insufficient account balance")
}
+11
View File
@@ -1877,6 +1877,11 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
}
batchImageHoldMultiplier = *input.BatchImageHoldMultiplier
}
// 不变式:hold 比例 >= discount 比例。否则批量任务成功率足够高时
// 实际成本会超过冻结额,结算永远失败、用户冻结余额无法解冻。
if batchImageHoldMultiplier < batchImageDiscountMultiplier {
return nil, errors.New("batch_image_hold_multiplier must be >= batch_image_discount_multiplier")
}
peakRateMultiplier := 1.0
if input.PeakRateMultiplier != nil {
@@ -2174,6 +2179,12 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
}
group.BatchImageHoldMultiplier = *input.BatchImageHoldMultiplier
}
// 仅在本次更新显式触碰任一比例时校验合并后的不变式(hold >= discount),
// 避免存量脏数据阻塞其他字段的正常更新(提交侧另有钳制兜底)。
if (input.BatchImageDiscountMultiplier != nil || input.BatchImageHoldMultiplier != nil) &&
group.BatchImageHoldMultiplier < group.BatchImageDiscountMultiplier {
return nil, errors.New("batch_image_hold_multiplier must be >= batch_image_discount_multiplier")
}
if input.PeakRateEnabled != nil {
group.PeakRateEnabled = *input.PeakRateEnabled
}
@@ -475,7 +475,7 @@ func TestAdminService_CreateGroup_BatchImagePricingSettings(t *testing.T) {
repo := &groupRepoStubForAdmin{}
svc := &adminServiceImpl{groupRepo: repo}
discount := 0.8
hold := 0.6
hold := 0.9
group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
Name: "batch-image-pricing",
@@ -488,7 +488,26 @@ func TestAdminService_CreateGroup_BatchImagePricingSettings(t *testing.T) {
require.NotNil(t, group)
require.NotNil(t, repo.created)
require.InDelta(t, 0.8, repo.created.BatchImageDiscountMultiplier, 1e-12)
require.InDelta(t, 0.6, repo.created.BatchImageHoldMultiplier, 1e-12)
require.InDelta(t, 0.9, repo.created.BatchImageHoldMultiplier, 1e-12)
}
func TestAdminService_CreateGroup_RejectsHoldBelowDiscount(t *testing.T) {
repo := &groupRepoStubForAdmin{}
svc := &adminServiceImpl{groupRepo: repo}
discount := 0.8
hold := 0.6
// hold < discount 时,成功率足够高的批量任务实际成本会超过冻结额,
// 结算永远失败,必须在配置入口拒绝。
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
Name: "batch-image-pricing-invalid",
Platform: PlatformGemini,
RateMultiplier: 1,
BatchImageDiscountMultiplier: &discount,
BatchImageHoldMultiplier: &hold,
})
require.Error(t, err)
require.Nil(t, repo.created)
}
func TestAdminService_GroupBatchImagePricingValidation(t *testing.T) {
+8
View File
@@ -50,6 +50,7 @@ var (
ErrBatchImageIndexParseFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_INDEX_PARSE_FAILED", "batch image provider output parse failed")
ErrBatchImageIndexNoResultLines = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_INDEX_NO_RESULT_LINES", "batch image provider output has no result lines")
ErrBatchImageDuplicateCustomID = infraerrors.New(http.StatusBadGateway, "DUPLICATE_CUSTOM_ID_IN_OUTPUT", "batch image provider output contains duplicate custom id")
ErrBatchImageIndexStateConflict = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_INDEX_STATE_CONFLICT", "batch image job is no longer in indexing state")
ErrBatchImageSettlementInvalidStatus = infraerrors.New(http.StatusBadRequest, "BATCH_IMAGE_SETTLEMENT_INVALID_STATUS", "batch image job is not ready for settlement")
ErrBatchImageSettlementManifestConflict = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_SETTLEMENT_MANIFEST_CONFLICT", "batch image settlement manifest hash conflict")
@@ -307,6 +308,13 @@ type BatchImageRepository interface {
GetBatchImageJobByID(ctx context.Context, id int64) (*BatchImageJob, error)
ListBatchImageJobsForOwner(ctx context.Context, userID, apiKeyID int64, filter BatchImageJobFilter) ([]*BatchImageJob, error)
TransitionBatchImageJobStatus(ctx context.Context, batchID, toStatus string, opts BatchImageTransitionOptions) error
// TouchBatchImageJobSubmitting 刷新未提交(created/uploading)job 的 updated_at,
// 作为慢提交期间的心跳,防止被 stale 恢复扫描误杀。
TouchBatchImageJobSubmitting(ctx context.Context, batchID string) error
// FailStaleUnsubmittedBatchImageJob 原子地将仍处于 created/uploading 且
// provider_job_name 为空、updated_at 早于 cutoff 的 job 转为 failed。
// 返回 false 表示 job 已被并发推进(如已提交成功),调用方不得释放冻结。
FailStaleUnsubmittedBatchImageJob(ctx context.Context, batchID string, cutoff time.Time, code, message string) (bool, error)
UpdateBatchImageJobProviderOutputRef(ctx context.Context, batchID, providerOutputRef string) error
UpdateBatchImageJobProviderSubmit(ctx context.Context, params UpdateBatchImageJobProviderSubmitParams) error
RecordBatchImageJobSubmitFailure(ctx context.Context, batchID, code, message string, markFailed bool) error
@@ -4,6 +4,9 @@ import (
"context"
"errors"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"go.uber.org/zap"
)
const (
@@ -98,6 +101,15 @@ func releaseBatchImageBalanceHold(ctx context.Context, repo UsageBillingReposito
return nil
}
if _, err := repo.ReleaseBatchImageBalance(ctx, cmd); err != nil {
// 同一 release request id 出现指纹冲突,说明此前已有一次携带不同
// payloadHash 的释放成功提交(资金已归还)。视为幂等成功,
// 避免历史指纹不一致的 job 永远卡在释放失败的毒消息循环里。
if errors.Is(err, ErrUsageBillingRequestConflict) {
logger.L().Warn("batch_image.release_fingerprint_conflict_treated_as_released",
zap.String("batch_id", job.BatchID),
)
return nil
}
return ErrBatchImageBillingHoldFailed.WithCause(err)
}
return nil
@@ -4,6 +4,9 @@ import (
"context"
"errors"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"go.uber.org/zap"
)
const (
@@ -15,6 +18,7 @@ type BatchImageBillingRecoveryService struct {
Repo BatchImageRepository
Billing UsageBillingRepository
AuthCache APIKeyAuthCacheInvalidator
Queue BatchImageQueue
StaleAfter time.Duration
Limit int
}
@@ -31,32 +35,69 @@ func (s *BatchImageBillingRecoveryService) ReleaseStaleUnsubmittedOnce(ctx conte
if limit <= 0 {
limit = defaultBatchImageBillingRecoveryLimit
}
jobs, err := s.Repo.ListStaleUnsubmittedBatchImageJobs(ctx, time.Now().Add(-staleAfter), limit)
cutoff := time.Now().Add(-staleAfter)
jobs, err := s.Repo.ListStaleUnsubmittedBatchImageJobs(ctx, cutoff, limit)
if err != nil {
return 0, err
}
released := 0
var lastErr error
for _, job := range jobs {
if job == nil {
continue
}
msg := "batch image submission did not reach provider before recovery cutoff"
if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusFailed, BatchImageTransitionOptions{
EventType: "billing_hold_recovery_failed_unsubmitted",
EventPayload: map[string]any{"batch_id": job.BatchID},
ErrorCode: batchImageStringPtr("SUBMIT_STALE_BEFORE_PROVIDER"),
ErrorMessage: batchImageStringPtr(msg),
}); err != nil && !errors.Is(err, ErrBatchImageInvalidTransition) {
if err := ctx.Err(); err != nil {
return released, err
}
msg := "batch image submission did not reach provider before recovery cutoff"
// 原子转 failed 并复核 stale 条件:List 与转态之间 job 可能已被慢提交
// 心跳续期或提交成功(provider_job_name 已写入),此时绝不能退款,
// 否则上游任务照常产生成本而用户已拿回冻结余额。
applied, err := s.Repo.FailStaleUnsubmittedBatchImageJob(ctx, job.BatchID, cutoff, "SUBMIT_STALE_BEFORE_PROVIDER", msg)
if err != nil {
// applied=true 时 UPDATE 已提交(仅审计事件写入失败):必须继续释放,
// 否则 job 已转 failed、不再出现在 stale 列表,冻结余额会永久泄漏。
if !applied {
lastErr = err
continue
}
logger.L().Warn("batch_image.recovery_fail_event_append_failed",
zap.String("batch_id", job.BatchID),
zap.Error(err),
)
}
if !applied {
continue
}
job.Status = BatchImageJobStatusFailed
if err := releaseBatchImageBalanceHold(ctx, s.Billing, job, batchImageDerefString(job.RequestHash)); err != nil {
return released, err
// job 已转 failed、不会再进入 stale 列表:必须给释放失败留下
// 自动重试路径(入队后由 worker 的 releaseTerminalHold 兜底),
// 否则冻结余额永久泄漏。
logger.L().Warn("batch_image.recovery_release_failed",
zap.String("batch_id", job.BatchID),
zap.Error(err),
)
s.enqueueReleaseRetry(ctx, job.BatchID)
lastErr = err
continue
}
if s.AuthCache != nil && job.UserID > 0 {
s.AuthCache.InvalidateAuthCacheByUserID(ctx, job.UserID)
}
released++
}
return released, nil
return released, lastErr
}
func (s *BatchImageBillingRecoveryService) enqueueReleaseRetry(ctx context.Context, batchID string) {
if s == nil || s.Queue == nil {
return
}
if err := s.Queue.Enqueue(ctx, batchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
logger.L().Warn("batch_image.recovery_release_retry_enqueue_failed",
zap.String("batch_id", batchID),
zap.Error(err),
)
}
}
@@ -4,12 +4,23 @@ package service
import (
"context"
"errors"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type recordingBatchImageQueue struct {
*fakeBatchImageQueue
enqueued []string
}
func (q *recordingBatchImageQueue) Enqueue(_ context.Context, batchID string) error {
q.enqueued = append(q.enqueued, batchID)
return nil
}
func TestBatchImageBillingRecoveryService_ReleasesStaleUnsubmittedHold(t *testing.T) {
repo := newFakeBatchImageRepository()
apiKeyID := int64(22)
@@ -50,3 +61,57 @@ func TestBatchImageBillingRecoveryService_ReleasesStaleUnsubmittedHold(t *testin
require.Equal(t, BatchImageReleaseRequestID(stale.BatchID), billing.releases[0].RequestID)
require.Equal(t, BatchImageJobStatusSubmitted, repo.jobs[active.BatchID].Status)
}
func TestBatchImageBillingRecoveryService_SkipsJobRefreshedByHeartbeat(t *testing.T) {
repo := newFakeBatchImageRepository()
apiKeyID := int64(22)
holdAmount := 0.5
// updated_at 在 cutoff 之后(慢提交心跳持续续期):不得误杀退款。
fresh := &BatchImageJob{
BatchID: "imgbatch_fresh_uploading",
UserID: 11,
APIKeyID: &apiKeyID,
Status: BatchImageJobStatusUploading,
EstimatedCost: holdAmount,
HoldAmount: &holdAmount,
CreatedAt: time.Now().Add(-time.Hour),
UpdatedAt: time.Now(),
}
repo.jobs[fresh.BatchID] = fresh
billing := &fakeBatchImageBillingRepo{}
svc := &BatchImageBillingRecoveryService{Repo: repo, Billing: billing, StaleAfter: time.Minute, Limit: 10}
released, err := svc.ReleaseStaleUnsubmittedOnce(context.Background())
require.NoError(t, err)
require.Equal(t, 0, released)
require.Equal(t, BatchImageJobStatusUploading, repo.jobs[fresh.BatchID].Status)
require.Empty(t, billing.releases)
}
func TestBatchImageBillingRecoveryService_EnqueuesRetryWhenReleaseFails(t *testing.T) {
repo := newFakeBatchImageRepository()
apiKeyID := int64(22)
holdAmount := 0.5
stale := &BatchImageJob{
BatchID: "imgbatch_stale_release_fail",
UserID: 11,
APIKeyID: &apiKeyID,
Status: BatchImageJobStatusCreated,
EstimatedCost: holdAmount,
HoldAmount: &holdAmount,
CreatedAt: time.Now().Add(-time.Hour),
UpdatedAt: time.Now().Add(-time.Hour),
}
repo.jobs[stale.BatchID] = stale
billing := &fakeBatchImageBillingRepo{releaseErr: errors.New("billing db down")}
queue := &recordingBatchImageQueue{fakeBatchImageQueue: newFakeBatchImageQueue("")}
svc := &BatchImageBillingRecoveryService{Repo: repo, Billing: billing, Queue: queue, StaleAfter: time.Minute, Limit: 10}
released, err := svc.ReleaseStaleUnsubmittedOnce(context.Background())
// job 已转 failed、不会再出现在 stale 列表:释放失败必须入队重试
//(由 worker 的 releaseTerminalHold 兜底),否则冻结余额永久泄漏。
require.Error(t, err)
require.Equal(t, 0, released)
require.Equal(t, BatchImageJobStatusFailed, repo.jobs[stale.BatchID].Status)
require.Equal(t, []string{stale.BatchID}, queue.enqueued)
}
@@ -9,6 +9,8 @@ import (
"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 (
@@ -38,6 +40,17 @@ func NewBatchImageCleanupService(repo BatchImageRepository, accountRepo AccountR
}
}
// 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 {
@@ -49,7 +62,7 @@ func (s *BatchImageCleanupService) DeleteOutputsForOwner(ctx context.Context, ow
if job.Status != BatchImageJobStatusCompleted {
return nil, ErrBatchImageOutputDeleteNotReady
}
_ = s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "manual_output_delete_requested", map[string]any{
s.appendCleanupEvent(ctx, job.BatchID, "manual_output_delete_requested", map[string]any{
"batch_id": job.BatchID,
"cleanup_target": "output",
"reason": "manual",
@@ -178,7 +191,7 @@ func (s *BatchImageCleanupService) cleanupJob(ctx context.Context, job *BatchIma
if !IsTerminalBatchImageJobStatus(job.Status) {
return ErrBatchImageCleanupFailed
}
_ = s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "input_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil))
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
@@ -186,7 +199,7 @@ func (s *BatchImageCleanupService) cleanupJob(ctx context.Context, job *BatchIma
if job.Status != BatchImageJobStatusCompleted && job.Status != BatchImageJobStatusFailed && job.Status != BatchImageJobStatusCancelled {
return ErrBatchImageOutputDeleteNotReady
}
_ = s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "output_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil))
s.appendCleanupEvent(ctx, job.BatchID, "output_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil))
default:
return ErrUnsupportedCleanupTarget
}
@@ -194,9 +207,14 @@ func (s *BatchImageCleanupService) cleanupJob(ctx context.Context, job *BatchIma
if err := s.callProviderCleanup(ctx, job, target); err != nil {
code := cleanupFailureCode(err)
msg := sanitizeBatchImagePublicMessage(err.Error())
_ = s.Repo.RecordBatchImageCleanupFailure(ctx, job.BatchID, code, msg)
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.Repo.AppendBatchImageEvent(ctx, job.BatchID, event, map[string]any{"batch_id": job.BatchID, "cleanup_target": string(target), "reason": reason, "error_code": code})
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
}
@@ -184,8 +184,9 @@ func (s *BatchImageDownloadService) StreamZip(ctx context.Context, owner BatchIm
return nil, err
}
maxItems := opts.MaxItems
if maxItems <= 0 {
maxItems = s.maxZipItems()
if cap := s.maxZipItems(); maxItems <= 0 || maxItems > cap {
// 客户端传入的 max_items 不得放大管理员配置的 ZIP 上限。
maxItems = cap
}
if job.SuccessCount > maxItems {
return nil, ErrBatchImageZipTooManyItems
@@ -8,6 +8,7 @@ import (
"fmt"
"io"
"net/http"
"sort"
"strconv"
"strings"
"time"
@@ -188,6 +189,11 @@ func (p *BatchImageProviderProcessor) indexAndSettle(ctx context.Context, job *B
if errors.Is(err, ErrBatchImageIndexOutputMissing) {
return BatchImageProcessResult{}, err
}
// job 状态已被并发方推进(如已进入 settling/终态):不是索引数据问题,
// 短延迟 requeue 让下一轮按最新状态处理,不能误转 failed。
if errors.Is(err, ErrBatchImageIndexStateConflict) {
return BatchImageProcessResult{RequeueAfter: time.Millisecond}, nil
}
code := "INDEX_PARSE_FAILED"
if errors.Is(err, ErrBatchImageDuplicateCustomID) {
code = "DUPLICATE_CUSTOM_ID_IN_OUTPUT"
@@ -281,6 +287,11 @@ func (i *BatchImageResultIndexer) Index(ctx context.Context, job *BatchImageJob,
if i == nil || i.Repo == nil || job == nil || provider == nil {
return nil, ErrBatchImageIndexOutputMissing
}
expected, err := i.listExpectedCustomIDs(ctx, job.BatchID)
if err != nil {
return nil, err
}
r, _, err := provider.OpenResult(ctx, job, account)
if err != nil {
return nil, ErrBatchImageIndexOutputMissing.WithCause(err)
@@ -291,6 +302,7 @@ func (i *BatchImageResultIndexer) Index(ctx context.Context, job *BatchImageJob,
scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
seen := make(map[string]int)
unknownCount := 0
var items []CreateBatchImageItemParams
result := &BatchImageIndexResult{}
lineNumber := 0
@@ -310,6 +322,14 @@ func (i *BatchImageResultIndexer) Index(ctx context.Context, job *BatchImageJob,
if err != nil {
return nil, err
}
// 与提交时的 custom_id 集对账:provider 输出中未知/多余的行不能进入 item 表,
// 否则 success+fail > item_count 会让结算永远校验失败。
if len(expected) > 0 {
if _, ok := expected[parsed.CustomID]; !ok {
unknownCount++
continue
}
}
if firstLine, ok := seen[parsed.CustomID]; ok {
return nil, ErrBatchImageDuplicateCustomID.WithCause(fmt.Errorf("custom id %q duplicated at lines %d and %d", parsed.CustomID, firstLine, lineNumber))
}
@@ -344,9 +364,52 @@ func (i *BatchImageResultIndexer) Index(ctx context.Context, job *BatchImageJob,
}
return nil, err
}
// 输出中漏掉的已提交项必须补失败记录,而不是静默消失:
// 否则用户看不到该项,且只按成功数计费会掩盖 provider 的丢单。
missingCount := 0
if len(expected) > 0 {
missingIDs := make([]string, 0)
for customID := range expected {
if _, ok := seen[customID]; !ok {
missingIDs = append(missingIDs, customID)
}
}
sort.Strings(missingIDs)
for _, customID := range missingIDs {
items = append(items, CreateBatchImageItemParams{
JobID: job.BatchID,
CustomID: customID,
Status: BatchImageItemStatusFailed,
ProviderSourceObject: batchImageOptionalStringPtr(sourceObject),
ErrorCode: batchImageStringPtr("PROVIDER_RESULT_MISSING"),
ErrorMessage: batchImageStringPtr("provider output did not include a result for this item"),
IndexedAt: &now,
})
result.FailCount++
result.TotalCount++
}
missingCount = len(missingIDs)
}
if result.TotalCount == 0 {
return nil, ErrBatchImageIndexNoResultLines
}
if unknownCount > 0 || missingCount > 0 {
logger.L().Warn("batch_image.index_reconciled",
zap.String("batch_id", job.BatchID),
zap.Int("unknown_custom_ids", unknownCount),
zap.Int("missing_custom_ids", missingCount),
)
if err := i.Repo.AppendBatchImageEvent(ctx, job.BatchID, "index_reconciled", map[string]any{
"batch_id": job.BatchID,
"unknown_custom_ids": unknownCount,
"missing_custom_ids": missingCount,
}); err != nil {
logger.L().Warn("batch_image.index_reconcile_event_failed",
zap.String("batch_id", job.BatchID),
zap.Error(err),
)
}
}
if err := i.Repo.ReplaceBatchImageItemsForJob(ctx, job.BatchID, items, BatchImageCounts{
SuccessCount: result.SuccessCount,
FailCount: result.FailCount,
@@ -356,6 +419,29 @@ func (i *BatchImageResultIndexer) Index(ctx context.Context, job *BatchImageJob,
return result, nil
}
// listExpectedCustomIDs 返回该 job 当前 item 表中的全部 custom_id 集合,
// 即提交时预创建(或上一轮索引重建)的完整条目清单,用于与 provider 输出对账。
func (i *BatchImageResultIndexer) listExpectedCustomIDs(ctx context.Context, batchID string) (map[string]struct{}, error) {
const pageSize = 500
expected := make(map[string]struct{})
offset := 0
for {
page, err := i.Repo.ListBatchImageItems(ctx, batchID, BatchImageItemFilter{Limit: pageSize, Offset: offset})
if err != nil {
return nil, err
}
for _, item := range page {
if item != nil {
expected[item.CustomID] = struct{}{}
}
}
if len(page) < pageSize {
return expected, nil
}
offset += len(page)
}
}
type ParsedBatchImageResult struct {
CustomID string
Status string
@@ -114,12 +114,58 @@ func TestBatchImageResultIndexer_WritesCountsAndReplacesItems(t *testing.T) {
require.Equal(t, BatchImageCounts{SuccessCount: 1, FailCount: 1}, repo.counts[job.BatchID])
require.NotContains(t, fmt.Sprintf("%+v", repo.items[job.BatchID]), batchImageTestData)
// 重新索引时与现有 custom_id 集对账:未知的 "ok2" 被丢弃,
// 输出中缺失的 ok/bad 补为 PROVIDER_RESULT_MISSING 失败记录。
provider.result = `{"key":"ok2","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/webp","data":"` + batchImageTestData + `"}}]}}]}}` + "\n"
result, err = (&BatchImageResultIndexer{Repo: repo}).Index(context.Background(), job, provider, &Account{})
require.NoError(t, err)
require.Equal(t, 1, result.TotalCount)
require.Len(t, repo.items[job.BatchID], 1)
require.Equal(t, "ok2", repo.items[job.BatchID][0].CustomID)
require.Equal(t, 2, result.TotalCount)
require.Equal(t, 0, result.SuccessCount)
require.Equal(t, 2, result.FailCount)
require.Len(t, repo.items[job.BatchID], 2)
gotIDs := []string{repo.items[job.BatchID][0].CustomID, repo.items[job.BatchID][1].CustomID}
require.ElementsMatch(t, []string{"ok", "bad"}, gotIDs)
for _, item := range repo.items[job.BatchID] {
require.Equal(t, BatchImageItemStatusFailed, item.Status)
require.Equal(t, "PROVIDER_RESULT_MISSING", batchImageDerefString(item.ErrorCode))
}
}
func TestBatchImageResultIndexer_ReconcilesMissingAndUnknownCustomIDs(t *testing.T) {
repo := newFakeBatchImageRepository()
outputRef := "files/output"
job := &BatchImageJob{BatchID: "imgbatch_reconcile", ProviderOutputRef: &outputRef, ItemCount: 3}
// 预创建提交时的 pending 条目(提交流程的行为)。
require.NoError(t, repo.BulkCreateBatchImageItems(context.Background(), []CreateBatchImageItemParams{
{JobID: job.BatchID, CustomID: "a", Status: BatchImageItemStatusPending},
{JobID: job.BatchID, CustomID: "b", Status: BatchImageItemStatusPending},
{JobID: job.BatchID, CustomID: "c", Status: BatchImageItemStatusPending},
}))
// provider 输出:a 成功,b 失败,c 漏掉,多出未知的 x。
output := strings.Join([]string{
`{"key":"a","response":{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"` + batchImageTestData + `"}}]}}]}}`,
`{"key":"b","error":{"code":"SAFETY","message":"blocked"}}`,
`{"key":"x","error":{"code":"UNKNOWN","message":"not ours"}}`,
}, "\n") + "\n"
provider := &fakeProcessorProvider{result: output}
result, err := (&BatchImageResultIndexer{Repo: repo}).Index(context.Background(), job, provider, &Account{})
require.NoError(t, err)
require.Equal(t, 3, result.TotalCount)
require.Equal(t, 1, result.SuccessCount)
require.Equal(t, 2, result.FailCount)
require.Len(t, repo.items[job.BatchID], 3)
byID := make(map[string]CreateBatchImageItemParams)
for _, item := range repo.items[job.BatchID] {
byID[item.CustomID] = item
}
require.NotContains(t, byID, "x")
require.Equal(t, BatchImageItemStatusSuccess, byID["a"].Status)
require.Equal(t, BatchImageItemStatusFailed, byID["b"].Status)
require.Equal(t, BatchImageItemStatusFailed, byID["c"].Status)
require.Equal(t, "PROVIDER_RESULT_MISSING", batchImageDerefString(byID["c"].ErrorCode))
// 对账后 success+fail == item_count,结算计数校验可通过。
require.Equal(t, job.ItemCount, result.SuccessCount+result.FailCount)
}
func TestBatchImageResultIndexer_EmptyInvalidAndDuplicateOutput(t *testing.T) {
@@ -487,6 +533,37 @@ func (r *fakeBatchImageRepository) TransitionBatchImageJobStatus(_ context.Conte
return nil
}
func (r *fakeBatchImageRepository) TouchBatchImageJobSubmitting(_ context.Context, batchID string) error {
job, ok := r.jobs[batchID]
if !ok {
return ErrBatchImageJobNotFound
}
if job.Status == BatchImageJobStatusCreated || job.Status == BatchImageJobStatusUploading {
job.UpdatedAt = time.Now()
}
return nil
}
func (r *fakeBatchImageRepository) FailStaleUnsubmittedBatchImageJob(_ context.Context, batchID string, cutoff time.Time, code, message string) (bool, error) {
job, ok := r.jobs[batchID]
if !ok {
return false, ErrBatchImageJobNotFound
}
if job.Status != BatchImageJobStatusCreated && job.Status != BatchImageJobStatusUploading {
return false, nil
}
if batchImageDerefString(job.ProviderJobName) != "" || job.UpdatedAt.After(cutoff) {
return false, nil
}
job.Status = BatchImageJobStatusFailed
job.LastErrorCode = batchImageStringPtr(code)
job.LastErrorMessage = batchImageStringPtr(message)
job.UpdatedAt = time.Now()
r.transitions[batchID] = append(r.transitions[batchID], BatchImageJobStatusFailed)
r.events[batchID] = append(r.events[batchID], "billing_hold_recovery_failed_unsubmitted")
return true, nil
}
func (r *fakeBatchImageRepository) UpdateBatchImageJobProviderOutputRef(_ context.Context, batchID, providerOutputRef string) error {
job, ok := r.jobs[batchID]
if !ok {
@@ -589,6 +666,11 @@ func (r *fakeBatchImageRepository) BulkCreateBatchImageItems(ctx context.Context
}
func (r *fakeBatchImageRepository) ReplaceBatchImageItemsForJob(_ context.Context, batchID string, items []CreateBatchImageItemParams, counts BatchImageCounts) error {
// 与真实实现一致:仅 indexing 状态允许重建 item 表(未注册的 job 保持宽松,
// 供直接构造 job 的单测使用)。
if job, ok := r.jobs[batchID]; ok && job.Status != BatchImageJobStatusIndexing {
return ErrBatchImageIndexStateConflict
}
r.replaceCalls++
copied := append([]CreateBatchImageItemParams(nil), items...)
for idx := range copied {
@@ -17,6 +17,7 @@ import (
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
)
const defaultGeminiBatchRequeueAfter = 30 * time.Second
@@ -460,11 +461,24 @@ func NewGeminiBatchHTTPClient(baseURL string, client *http.Client) *GeminiBatchH
baseURL = geminicli.AIStudioBaseURL
}
if client == nil {
client = http.DefaultClient
client = batchImageDefaultHTTPClient()
}
return &GeminiBatchHTTPClient{baseURL: baseURL, client: client}
}
// batchImageDefaultHTTPClient 返回带连接/握手/响应头超时的共享客户端。
// 不设整体 Timeout:大文件上传与结果流式下载耗时不可预估,
// 但拨号、TLS、等待响应头必须有界,否则挂死的连接会无限占用提交路径。
func batchImageDefaultHTTPClient() *http.Client {
client, err := httpclient.GetClient(httpclient.Options{
ResponseHeaderTimeout: 60 * time.Second,
})
if err != nil {
return http.DefaultClient
}
return client
}
func (c *GeminiBatchHTTPClient) UploadJSONL(ctx context.Context, apiKey string, displayName string, r io.Reader) (*GeminiUploadedFile, error) {
var body bytes.Buffer
writer := multipart.NewWriter(&body)
@@ -567,6 +581,11 @@ func (c *GeminiBatchHTTPClient) DownloadFile(ctx context.Context, apiKey string,
if downloadURL == "" {
downloadURL = c.baseURL + "/v1beta/" + strings.TrimLeft(fileName, "/") + ":download"
}
// 纵深加固:downloadUri 来自上游响应,跟随前校验目标 host,
// 防止异常/被劫持的响应把带 api key 的请求带到任意主机。
if err := validateGeminiDownloadHost(downloadURL, c.baseURL); err != nil {
return nil, "", err
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
if err != nil {
return nil, "", err
@@ -643,6 +662,26 @@ func (c *GeminiBatchHTTPClient) newRequest(ctx context.Context, method, path, ap
return req, nil
}
// validateGeminiDownloadHost 只允许跟随到 googleapis.com(含子域)
// 或与配置的 baseURL 同 host 的下载地址。
func validateGeminiDownloadHost(downloadURL, baseURL string) error {
parsed, err := url.Parse(downloadURL)
if err != nil {
return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri is invalid", err)
}
if parsed.Scheme != "https" {
return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri must use https", nil)
}
host := strings.ToLower(parsed.Hostname())
if host == "googleapis.com" || strings.HasSuffix(host, ".googleapis.com") {
return nil
}
if base, err := url.Parse(baseURL); err == nil && strings.EqualFold(base.Hostname(), host) {
return nil
}
return geminiProviderError("GEMINI_INVALID_RESPONSE", "Gemini download uri host is not allowed", nil)
}
type GeminiAPIError struct {
StatusCode int
Code string
@@ -736,7 +736,7 @@ type VertexBatchHTTPClient struct {
func NewVertexBatchHTTPClient(baseURL string, client *http.Client) *VertexBatchHTTPClient {
if client == nil {
client = http.DefaultClient
client = batchImageDefaultHTTPClient()
}
return &VertexBatchHTTPClient{baseURL: strings.TrimRight(strings.TrimSpace(baseURL), "/"), client: client}
}
@@ -794,7 +794,7 @@ type VertexGCSObjectStore struct {
func NewVertexGCSObjectStore(baseURL string, client *http.Client) *VertexGCSObjectStore {
if client == nil {
client = http.DefaultClient
client = batchImageDefaultHTTPClient()
}
baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
if baseURL == "" {
+136 -4
View File
@@ -14,6 +14,8 @@ import (
"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 (
@@ -202,6 +204,11 @@ func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOw
if err != nil {
return nil, err
}
// 与 ListModels 使用同一鉴权谓词(AllowBatchImageGeneration + Platform==Gemini),
// 避免两个入口校验口径不一致留下防御纵深缺口。
if err := s.ensureGroupAllowsBatchImage(ctx, owner.GroupID); err != nil {
return nil, err
}
requestHash := HashBatchImageSubmitRequest(normalized)
idempotencyKey = strings.TrimSpace(idempotencyKey)
if idempotencyKey != "" {
@@ -319,7 +326,31 @@ func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOw
})
}
// 上游提交(上传参考图 + 创建批任务)可能长达数分钟且不刷新 updated_at,
// 会被 stale 恢复扫描误判为滞留并退款。提交前转入 uploading 刷新时间戳,
// 提交期间用心跳持续续期。
if err := s.Repo.TransitionBatchImageJobStatus(ctx, job.BatchID, BatchImageJobStatusUploading, BatchImageTransitionOptions{
EventType: "upload_started",
EventPayload: map[string]any{"batch_id": job.BatchID},
}); err != nil {
if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
return nil, releaseErr
}
// 并发 Cancel 等导致的非法转换:job 已处于终态,不再覆盖其状态。
if !errors.Is(err, ErrBatchImageInvalidTransition) {
_ = s.Repo.RecordBatchImageJobSubmitFailure(ctx, job.BatchID, "UPLOAD_TRANSITION_FAILED", sanitizeBatchImagePublicMessage(err.Error()), true)
s.hidePreUpstreamSubmitFailure(ctx, owner, job)
}
return nil, err
}
job.Status = BatchImageJobStatusUploading
hbCtx, hbCancel := context.WithCancel(ctx)
hbDone := make(chan struct{})
go s.runSubmitHeartbeat(hbCtx, job.BatchID, hbDone)
providerJob, err := provider.Submit(ctx, job, account, input)
hbCancel()
<-hbDone
if err != nil {
if releaseErr := s.releaseFailedSubmitHold(ctx, job, requestHash); releaseErr != nil {
return nil, releaseErr
@@ -348,6 +379,9 @@ func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOw
GCSOutputURI: batchImageGCSRef(provider.Name(), providerJob.ProviderOutputRef),
EventPayload: map[string]any{"provider": provider.Name()},
}); err != nil {
// job 可能已被恢复扫描转 failed 并退款:上游批任务已创建成功,
// 必须尽力取消并清理输入,否则上游照常产生成本(孤儿任务)。
s.abortOrphanProviderJob(ctx, provider, job, account, providerJob)
return nil, err
}
@@ -375,6 +409,75 @@ func (s *BatchImagePublicService) releaseFailedSubmitHold(ctx context.Context, j
return nil
}
// runSubmitHeartbeat 在 provider.Submit 期间周期性刷新 job 的 updated_at,
// 使 stale 恢复扫描能区分"仍在慢提交"与"进程死亡后的滞留"。
func (s *BatchImagePublicService) runSubmitHeartbeat(ctx context.Context, batchID string, done chan<- struct{}) {
defer close(done)
interval := s.submitHeartbeatInterval()
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
if err := s.Repo.TouchBatchImageJobSubmitting(ctx, batchID); err != nil && ctx.Err() == nil {
logger.L().Warn("batch_image.submit_heartbeat_failed",
zap.String("batch_id", batchID),
zap.Error(err),
)
}
}
}
}
func (s *BatchImagePublicService) submitHeartbeatInterval() time.Duration {
staleAfter := 10 * time.Minute
if s != nil && s.Config != nil && s.Config.BatchImage.StaleActiveAfterSeconds > 0 {
staleAfter = time.Duration(s.Config.BatchImage.StaleActiveAfterSeconds) * time.Second
}
interval := staleAfter / 3
if interval < 15*time.Second {
interval = 15 * time.Second
}
return interval
}
// abortOrphanProviderJob 在上游任务创建成功但本地状态推进失败时,
// 尽力取消上游批任务并清理已上传的输入文件,避免孤儿任务持续产生成本。
func (s *BatchImagePublicService) abortOrphanProviderJob(ctx context.Context, provider BatchImageProvider, job *BatchImageJob, account *Account, providerJob *BatchProviderJob) {
if s == nil || provider == nil || job == nil || providerJob == nil {
return
}
orphan := *job
orphan.ProviderJobName = batchImageOptionalStringPtr(providerJob.ProviderJobName)
orphan.ProviderInputRef = batchImageOptionalStringPtr(providerJob.ProviderInputRef)
orphan.GCSInputURI = batchImageOptionalStringPtr(batchImageGCSRef(provider.Name(), providerJob.ProviderInputRef))
if err := provider.Cancel(ctx, &orphan, account); err != nil {
logger.L().Warn("batch_image.orphan_provider_job_cancel_failed",
zap.String("batch_id", job.BatchID),
zap.String("provider", provider.Name()),
zap.Error(err),
)
}
if err := provider.Cleanup(ctx, &orphan, account, CleanupTargetInput); err != nil {
logger.L().Warn("batch_image.orphan_provider_job_cleanup_failed",
zap.String("batch_id", job.BatchID),
zap.String("provider", provider.Name()),
zap.Error(err),
)
}
if err := s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "provider_job_aborted_after_submit", map[string]any{
"batch_id": job.BatchID,
"provider": provider.Name(),
}); err != nil {
logger.L().Warn("batch_image.orphan_provider_job_event_failed",
zap.String("batch_id", job.BatchID),
zap.Error(err),
)
}
}
func (s *BatchImagePublicService) createPendingItems(ctx context.Context, batchID, requestHash string, items []BatchImageSubmitItem) error {
if s == nil || s.Repo == nil || len(items) == 0 {
return nil
@@ -399,10 +502,19 @@ func (s *BatchImagePublicService) enqueueBillingRetry(ctx context.Context, batch
return
}
if err := s.Queue.Enqueue(ctx, batchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
_ = s.Repo.AppendBatchImageEvent(ctx, batchID, "billing_retry_enqueue_failed", map[string]any{
logger.L().Warn("batch_image.billing_retry_enqueue_failed",
zap.String("batch_id", batchID),
zap.Error(err),
)
if eventErr := s.Repo.AppendBatchImageEvent(ctx, batchID, "billing_retry_enqueue_failed", map[string]any{
"batch_id": batchID,
"error": sanitizeBatchImagePublicMessage(err.Error()),
})
}); eventErr != nil {
logger.L().Warn("batch_image.billing_retry_event_failed",
zap.String("batch_id", batchID),
zap.Error(eventErr),
)
}
}
}
@@ -410,7 +522,12 @@ func (s *BatchImagePublicService) hidePreUpstreamSubmitFailure(ctx context.Conte
if s == nil || s.Repo == nil || job == nil || job.ProviderJobName != nil {
return
}
_ = s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now())
if err := s.Repo.MarkBatchImageJobUserDeleted(ctx, owner.UserID, owner.APIKeyID, job.BatchID, time.Now()); err != nil {
logger.L().Warn("batch_image.hide_pre_upstream_failure_failed",
zap.String("batch_id", job.BatchID),
zap.Error(err),
)
}
}
func (s *BatchImagePublicService) Get(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) {
@@ -615,7 +732,12 @@ func (s *BatchImagePublicService) Cancel(ctx context.Context, owner BatchImageOw
if err := provider.Cancel(ctx, job, account); err != nil {
return nil, ErrBatchImageCancelFailed
}
_ = s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "job_cancel_requested", map[string]any{"batch_id": job.BatchID})
if eventErr := s.Repo.AppendBatchImageEvent(ctx, job.BatchID, "job_cancel_requested", map[string]any{"batch_id": job.BatchID}); eventErr != nil {
logger.L().Warn("batch_image.cancel_event_failed",
zap.String("batch_id", job.BatchID),
zap.Error(eventErr),
)
}
if s.Queue != nil {
if err := s.Queue.Enqueue(ctx, job.BatchID); err != nil && !errors.Is(err, ErrBatchImageAlreadyQueued) {
return nil, ErrBatchImageCancelFailed
@@ -930,6 +1052,16 @@ func (s *BatchImagePublicService) resolvePricingSnapshot(ctx context.Context, ow
}
unit = resolvedUnit
}
// 定价不变式:hold 比例不得低于 discount 比例,否则成功率足够高时
// actualCost > holdAmount,结算永远失败、冻结余额无法解冻。
// 管理端已校验新配置,此处兜底钳制存量脏数据。
if holdMultiplier < discountMultiplier {
logger.L().Warn("batch_image.hold_multiplier_below_discount_clamped",
zap.Float64("hold_multiplier", holdMultiplier),
zap.Float64("discount_multiplier", discountMultiplier),
)
holdMultiplier = discountMultiplier
}
accountMultiplier := 1.0
if account != nil {
accountMultiplier = account.BillingRateMultiplier()
@@ -93,10 +93,12 @@ func TestBatchImagePublicService_Submit(t *testing.T) {
require.InDelta(t, 0.5, job.GroupRateMultiplier, 1e-12)
require.InDelta(t, 1.25, job.AccountRateMultiplier, 1e-12)
require.InDelta(t, 0.8, job.BatchDiscountMultiplier, 1e-12)
require.InDelta(t, 0.6, job.HoldMultiplier, 1e-12)
// 配置的 hold(0.6) < discount(0.8) 属于会导致结算死锁的脏数据,
// 快照时被钳制为 discount,保证 holdAmount >= 实际成本上限。
require.InDelta(t, 0.8, job.HoldMultiplier, 1e-12)
require.InDelta(t, 0.125, job.BillableUnitPrice, 1e-12)
require.InDelta(t, 0.09375, job.HoldUnitPrice, 1e-12)
require.InDelta(t, 0.1875, *job.HoldAmount, 1e-12)
require.InDelta(t, 0.125, job.HoldUnitPrice, 1e-12)
require.InDelta(t, 0.25, *job.HoldAmount, 1e-12)
})
t.Run("uses configured group 1k image price for batch image base price", func(t *testing.T) {
@@ -11,6 +11,8 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"go.uber.org/zap"
)
const (
@@ -98,28 +100,40 @@ func (s *BatchImageSettlementService) Settle(ctx context.Context, batchID string
if job.Status != BatchImageJobStatusSettling {
return nil, ErrBatchImageSettlementInvalidStatus
}
if job.SuccessCount < 0 || job.FailCount < 0 || job.ItemCount < 0 || job.SuccessCount+job.FailCount > job.ItemCount {
return nil, ErrBatchImageSettlementInvalidCounts
}
if strings.TrimSpace(batchImageDerefString(job.ManifestHash)) != "" && batchImageDerefString(job.ManifestHash) != manifestHash {
return nil, ErrBatchImageSettlementManifestConflict
}
if job.APIKeyID == nil || *job.APIKeyID <= 0 {
return nil, ErrBatchImageSettlementMissingAPIKeyID
}
if job.AccountID == nil || *job.AccountID <= 0 {
return nil, ErrBatchImageSettlementMissingAccountID
}
// 重试耗尽检查必须先于各类可重复失败的校验(counts/manifest/定价/超冻结),
// 否则这些错误路径会绕过耗尽出口,settling job 无限 requeue、冻结余额永不释放。
if isBatchImageSettlementRetryExhausted(job) {
return nil, s.failExhaustedSettlement(ctx, job, manifestHash, "settlement billing retry limit reached")
return nil, s.failExhaustedSettlement(ctx, job, "settlement retry limit reached: "+batchImageDerefString(job.LastErrorCode))
}
if job.SuccessCount < 0 || job.FailCount < 0 || job.ItemCount < 0 || job.SuccessCount+job.FailCount > job.ItemCount {
if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_INVALID_COUNTS",
fmt.Sprintf("success=%d fail=%d item_count=%d", job.SuccessCount, job.FailCount, job.ItemCount)); failErr != nil {
return nil, failErr
}
return nil, ErrBatchImageSettlementInvalidCounts
}
if strings.TrimSpace(batchImageDerefString(job.ManifestHash)) != "" && batchImageDerefString(job.ManifestHash) != manifestHash {
if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_MANIFEST_CONFLICT", "manifest hash conflict"); failErr != nil {
return nil, failErr
}
return nil, ErrBatchImageSettlementManifestConflict
}
unitPrice, err := s.settlementUnitPrice(ctx, job)
if err != nil {
return nil, err
if err == nil && unitPrice < 0 {
err = ErrBatchImageSettlementPricingMissing
}
if unitPrice < 0 {
return nil, ErrBatchImageSettlementPricingMissing
if err != nil {
if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_PRICING_MISSING", err.Error()); failErr != nil {
return nil, failErr
}
return nil, err
}
actualCost := float64(job.SuccessCount) * unitPrice
result.ActualCost = actualCost
@@ -129,16 +143,16 @@ func (s *BatchImageSettlementService) Settle(ctx context.Context, batchID string
}
if actualCost-holdAmount > batchImageCostEpsilon {
msg := fmt.Sprintf("actual cost %.10f exceeds held amount %.10f", actualCost, holdAmount)
_, _ = s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, "SETTLEMENT_COST_EXCEEDS_HOLD", msg)
if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_COST_EXCEEDS_HOLD", msg); failErr != nil {
return nil, failErr
}
return nil, ErrBatchImageSettlementCostExceedsHold
}
if err := captureBatchImageBalanceHold(ctx, s.BillingRepo, job, actualCost, manifestHash); err != nil {
msg := truncateBatchImageMessage(err.Error(), batchImageMaxErrorMessageLength)
retryCount, recordErr := s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, "SETTLEMENT_BILLING_FAILED", msg)
if recordErr == nil && retryCount >= batchImageSettlementMaxRetries {
job.RetryCount = retryCount
return nil, s.failExhaustedSettlement(ctx, job, manifestHash, msg)
if failErr := s.recordSettlementFailure(ctx, job, "SETTLEMENT_BILLING_FAILED", msg); failErr != nil {
return nil, failErr
}
return nil, err
}
@@ -168,20 +182,52 @@ func (s *BatchImageSettlementService) Settle(ctx context.Context, batchID string
return result, nil
}
// isBatchImageSettlementRetryExhausted 判断 settling job 是否已达重试上限。
// 必须覆盖所有 SETTLEMENT_* 失败码(而非仅 SETTLEMENT_BILLING_FAILED),
// 否则 SETTLEMENT_COST_EXCEEDS_HOLD / SETTLEMENT_INVALID_COUNTS 等错误会无限 requeue。
func isBatchImageSettlementRetryExhausted(job *BatchImageJob) bool {
return job != nil &&
job.Status == BatchImageJobStatusSettling &&
job.RetryCount >= batchImageSettlementMaxRetries &&
batchImageDerefString(job.LastErrorCode) == "SETTLEMENT_BILLING_FAILED"
strings.HasPrefix(batchImageDerefString(job.LastErrorCode), "SETTLEMENT_")
}
func (s *BatchImageSettlementService) failExhaustedSettlement(ctx context.Context, job *BatchImageJob, manifestHash, message string) error {
// recordSettlementFailure 记录一次结算失败并递增 retry_count。
// 重试达到上限时立即走耗尽出口(释放冻结余额并转 failed);
// 返回非 nil 时调用方应直接返回该错误。
func (s *BatchImageSettlementService) recordSettlementFailure(ctx context.Context, job *BatchImageJob, code, message string) error {
retryCount, recordErr := s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, code, truncateBatchImageMessage(message, batchImageMaxErrorMessageLength))
if recordErr != nil {
logger.L().Warn("batch_image.settlement_failure_record_failed",
zap.String("batch_id", job.BatchID),
zap.String("code", code),
zap.Error(recordErr),
)
return nil
}
job.RetryCount = retryCount
job.LastErrorCode = &code
if retryCount >= batchImageSettlementMaxRetries {
return s.failExhaustedSettlement(ctx, job, message)
}
return nil
}
func (s *BatchImageSettlementService) failExhaustedSettlement(ctx context.Context, job *BatchImageJob, message string) error {
if s == nil || s.Repo == nil {
return ErrBatchImageSettlementBillingFailed
}
if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, manifestHash); err != nil {
// 释放指纹必须与其余所有释放点(processor/Cancel/recovery)一致地使用 RequestHash:
// 它们共享同一 request id,payloadHash 不同会触发 ErrUsageBillingRequestConflict,
// 导致后续 Cancel/重试永远失败、terminal job 变成毒消息。
if err := releaseBatchImageBalanceHold(ctx, s.BillingRepo, job, batchImageDerefString(job.RequestHash)); err != nil {
msg := truncateBatchImageMessage(err.Error(), batchImageMaxErrorMessageLength)
_, _ = s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, "SETTLEMENT_RELEASE_FAILED", msg)
if _, recordErr := s.Repo.SetBatchImageJobSettlementFailed(ctx, job.BatchID, "SETTLEMENT_RELEASE_FAILED", msg); recordErr != nil {
logger.L().Warn("batch_image.settlement_release_failure_record_failed",
zap.String("batch_id", job.BatchID),
zap.Error(recordErr),
)
}
return ErrBatchImageSettlementBillingFailed.WithCause(err)
}
s.invalidateAuthCache(ctx, job.UserID)
@@ -278,6 +278,73 @@ func TestBatchImageSettlementRetryExhaustedReleaseIsIdempotentAfterTransitionFai
require.Len(t, billing.seen, 1)
}
func TestBatchImageSettlementService_CostExceedsHoldExhaustsAndReleases(t *testing.T) {
repo := newFakeBatchImageRepository()
job := testSettlingBatchImageJob("imgbatch_over_hold_exhausted")
job.SuccessCount = 2
job.FailCount = 0
job.ItemCount = 2
holdAmount := 0.5
job.HoldAmount = &holdAmount
job.EstimatedCost = holdAmount
requestHash := "request-hash-over-hold"
job.RequestHash = &requestHash
repo.jobs[job.BatchID] = job
billing := &fakeBatchImageBillingRepo{}
svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.50}}
// 前 N-1 次:记录失败并返回错误(等待 worker 重试)。
for i := 0; i < batchImageSettlementMaxRetries-1; i++ {
_, err := svc.Settle(context.Background(), job.BatchID)
require.ErrorIs(t, err, ErrBatchImageSettlementCostExceedsHold)
require.Equal(t, BatchImageJobStatusSettling, repo.jobs[job.BatchID].Status)
}
// 达到上限:必须走耗尽出口释放冻结并转 failed,而不是无限 requeue。
_, err := svc.Settle(context.Background(), job.BatchID)
require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed)
require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status)
require.Empty(t, billing.captures)
require.Len(t, billing.releases, 1)
require.Equal(t, BatchImageReleaseRequestID(job.BatchID), billing.releases[0].RequestID)
// 释放指纹必须与 processor/Cancel/recovery 一致地使用 RequestHash,
// 否则共享同一 request id 的后续释放会命中指纹冲突(毒消息)。
require.Equal(t, requestHash, billing.releases[0].RequestPayloadHash)
}
func TestBatchImageSettlementService_InvalidCountsExhaustsAndReleases(t *testing.T) {
repo := newFakeBatchImageRepository()
job := testSettlingBatchImageJob("imgbatch_bad_counts_exhausted")
job.SuccessCount = 2
job.FailCount = 2
job.ItemCount = 3
requestHash := "request-hash-bad-counts"
job.RequestHash = &requestHash
repo.jobs[job.BatchID] = job
billing := &fakeBatchImageBillingRepo{}
svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
for i := 0; i < batchImageSettlementMaxRetries-1; i++ {
_, err := svc.Settle(context.Background(), job.BatchID)
require.ErrorIs(t, err, ErrBatchImageSettlementInvalidCounts)
}
_, err := svc.Settle(context.Background(), job.BatchID)
require.ErrorIs(t, err, ErrBatchImageSettlementBillingFailed)
require.Equal(t, BatchImageJobStatusFailed, repo.jobs[job.BatchID].Status)
require.Empty(t, billing.captures)
require.Len(t, billing.releases, 1)
require.Equal(t, requestHash, billing.releases[0].RequestPayloadHash)
}
func TestReleaseBatchImageBalanceHold_TreatsFingerprintConflictAsReleased(t *testing.T) {
job := testSettlingBatchImageJob("imgbatch_release_conflict")
// 历史版本用 manifestHash 释放过一次:同一 request id 再以 RequestHash
// 释放会命中指纹冲突。资金已归还,必须视为幂等成功而非毒消息。
billing := &fakeBatchImageBillingRepo{releaseErr: ErrUsageBillingRequestConflict}
err := releaseBatchImageBalanceHold(context.Background(), billing, job, "request-hash")
require.NoError(t, err)
require.Len(t, billing.releases, 1)
}
func TestBatchImageSettlementManifestHash(t *testing.T) {
job := testSettlingBatchImageJob("imgbatch_hash")
first := BuildBatchImageSettlementManifestHash(job)
+57 -1
View File
@@ -150,13 +150,23 @@ func (w *BatchImageWorker) RunOnce(ctx context.Context) error {
return err
}
if !ok {
return nil
// 锁被其他实例持有:按冲突延迟重新入队。直接丢弃会让 job 滞留在
// active zset,最早要等 StaleActiveAfter 才被恢复,造成分钟级停摆。
return w.queue.RequeueAfter(ctx, reserved.BatchID, w.opts.LockConflictDelay)
}
defer func() {
_ = lock.Release(ctx)
}()
// 处理期间持续心跳:刷新 active zset 时间戳防止 stale 恢复把在处理的
// job 重投给其他 worker,并对支持续期的锁实现延长锁 TTL。
hbStop := make(chan struct{})
hbDone := make(chan struct{})
go w.runJobHeartbeat(ctx, reserved.BatchID, lock, hbStop, hbDone)
result, err := w.processor.Process(ctx, reserved.BatchID)
close(hbStop)
<-hbDone
if err != nil {
logger.L().Warn("batch_image.worker_process_failed",
zap.String("batch_id", reserved.BatchID),
@@ -174,6 +184,52 @@ func (w *BatchImageWorker) RunOnce(ctx context.Context) error {
return w.queue.RequeueAfter(ctx, reserved.BatchID, delay)
}
// BatchImageJobLockRefresher 是可选的锁续期能力;由具体锁实现按需提供。
type BatchImageJobLockRefresher interface {
Refresh(ctx context.Context, ttl time.Duration) error
}
func (w *BatchImageWorker) heartbeatInterval() time.Duration {
interval := w.opts.JobLockTTL
if w.opts.StaleActiveAfter < interval {
interval = w.opts.StaleActiveAfter
}
interval /= 3
if interval < time.Second {
interval = time.Second
}
return interval
}
func (w *BatchImageWorker) runJobHeartbeat(ctx context.Context, batchID string, lock BatchImageJobLock, stop <-chan struct{}, done chan<- struct{}) {
defer close(done)
ticker := time.NewTicker(w.heartbeatInterval())
defer ticker.Stop()
for {
select {
case <-stop:
return
case <-ctx.Done():
return
case <-ticker.C:
if err := w.queue.Heartbeat(ctx, batchID); err != nil && ctx.Err() == nil {
logger.L().Warn("batch_image.worker_heartbeat_failed",
zap.String("batch_id", batchID),
zap.Error(err),
)
}
if refresher, ok := lock.(BatchImageJobLockRefresher); ok {
if err := refresher.Refresh(ctx, w.opts.JobLockTTL); err != nil && ctx.Err() == nil {
logger.L().Warn("batch_image.worker_lock_refresh_failed",
zap.String("batch_id", batchID),
zap.Error(err),
)
}
}
}
}
}
func (w *BatchImageWorker) MoveDueDelayedOnce(ctx context.Context) (int, error) {
if w == nil || w.queue == nil {
return 0, nil
@@ -53,6 +53,7 @@ func ProvideBatchImageWorkerRuntime(
Repo: repo,
Billing: billingRepo,
AuthCache: authCache,
Queue: queue,
StaleAfter: NewBatchImageWorkerOptionsFromConfig(cfg).StaleActiveAfter,
Limit: NewBatchImageWorkerOptionsFromConfig(cfg).RecoverLimit,
}
@@ -56,15 +56,18 @@ func TestBatchImageWorker_RequeuesOnProcessorError(t *testing.T) {
require.Empty(t, queue.acked)
}
func TestBatchImageWorker_SkipsWhenJobLockNotAcquired(t *testing.T) {
func TestBatchImageWorker_RequeuesWhenJobLockNotAcquired(t *testing.T) {
queue := newFakeBatchImageQueue("imgbatch_worker_locked")
queue.lockAcquired = false
processor := &fakeBatchImageProcessor{}
worker := NewBatchImageWorker(queue, processor, BatchImageWorkerOptions{})
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.Empty(t, queue.requeued)
require.Len(t, queue.requeued, 1)
require.Equal(t, 3*time.Second, queue.requeued[0].delay)
require.Empty(t, queue.acked)
}
@@ -713,6 +713,14 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) {
// 1. 优先从动态价格服务获取
if s.pricingService != nil {
litellmPricing := s.pricingService.GetModelPricing(model)
// 仅有图片价、无 token 价的条目(如 LiteLLM 的 imagen 类模型)不能用于
// token 计费:直接返回会把 token 流量按 $0 计费。跳过后走 fallback,
// 无 fallback 则 fail-closed(ErrModelPricingUnavailable)。
// 图片计费路径(getDefaultImagePrice / getImageUnitPrice)直接读
// PricingService,不受影响。
if litellmPricing != nil && litellmPricing.TokenPricingAbsent {
litellmPricing = nil
}
if litellmPricing != nil {
// 启用 5m/1h 分类计费的条件:
// 1. 存在 1h 价格
@@ -73,6 +73,11 @@ type LiteLLMModelPricing struct {
SupportsPromptCaching bool `json:"supports_prompt_caching"`
OutputCostPerImage float64 `json:"output_cost_per_image"` // 图片生成模型每张图片价格
OutputCostPerImageToken float64 `json:"output_cost_per_image_token"` // 图片输出 token 价格
// TokenPricingAbsent 表示源数据中 input/output token 价格均缺失(仅有图片价)。
// 此类条目只可用于图片计费,token 计费必须回退到 fallback 或 fail-closed,
// 否则 token 流量会被按 $0 计费。零值(false)表示条目具备 token 价格。
TokenPricingAbsent bool `json:"-"`
}
// PricingRemoteClient 远程价格数据获取接口
@@ -383,6 +388,7 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel
Mode: entry.Mode,
SupportsPromptCaching: entry.SupportsPromptCaching,
SupportsServiceTier: entry.SupportsServiceTier,
TokenPricingAbsent: entry.InputCostPerToken == nil && entry.OutputCostPerToken == nil,
}
if entry.InputCostPerToken != nil {
@@ -54,6 +54,44 @@ func TestParsePricingData_KeepsImageOnlyPricing(t *testing.T) {
require.NotNil(t, pricing)
require.InDelta(t, 0.034, pricing.OutputCostPerImage, 1e-12)
require.Equal(t, "image_generation", pricing.Mode)
// 仅有图片价的条目必须标记 token 价缺失,供 token 计费路径 fail-closed。
require.True(t, pricing.TokenPricingAbsent)
}
func TestBillingService_GetModelPricing_FailsClosedForImageOnlyEntries(t *testing.T) {
pricingSvc := &PricingService{}
data, err := pricingSvc.parsePricingData([]byte(`{
"imagen-9.0-generate": {
"output_cost_per_image": 0.04,
"litellm_provider": "vertex_ai-image-models",
"mode": "image_generation"
},
"gemini-image-with-token-price": {
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
"output_cost_per_image": 0.034,
"litellm_provider": "vertex_ai-language-models",
"mode": "image_generation"
}
}`))
require.NoError(t, err)
pricingSvc.pricingData = data
billingSvc := NewBillingService(&config.Config{}, pricingSvc)
// image-only 条目不得进入 token 计费(否则 token 流量按 $0 计费),
// 必须落到 fallback / ErrModelPricingUnavailable 的 fail-closed 路径。
_, err = billingSvc.GetModelPricing("imagen-9.0-generate")
require.ErrorIs(t, err, ErrModelPricingUnavailable)
// 显式 0 token 价的免费条目保持历史行为:正常返回。
pricing, err := billingSvc.GetModelPricing("gemini-image-with-token-price")
require.NoError(t, err)
require.Zero(t, pricing.InputPricePerToken)
// 图片计费路径不受影响:仍能读到 image-only 条目的图片单价。
raw := pricingSvc.GetModelPricing("imagen-9.0-generate")
require.NotNil(t, raw)
require.InDelta(t, 0.04, raw.OutputCostPerImage, 1e-12)
}
func TestPricingService_MergesFallbackOnlyModels(t *testing.T) {