feat(billing): 完成异步结算与请求执行闭环

统一任务成功、Attempt 与结算 Outbox 的事务边界,增加多实例安全结算、人工复核、请求幂等与执行租约。钱包决策使用九位精确金额,并通过审计保护约束保留流水事实;同时补充管理接口、指标与 PostgreSQL 集成测试。
This commit is contained in:
2026-07-21 00:25:53 +08:00
parent 7cea21f765
commit 5b2b94b1bd
33 changed files with 2884 additions and 418 deletions
+330 -54
View File
@@ -3,6 +3,7 @@ package store
import (
"context"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
@@ -152,18 +153,71 @@ func nullableTaskListTime(value *time.Time) any {
return *value
}
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType string, normalizedRequest map[string]any) error {
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
_, err := s.pool.Exec(ctx, `
func (s *Store) ClaimTaskExecution(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) {
if leaseTTL <= 0 {
leaseTTL = 5 * time.Minute
}
task, err := scanGatewayTask(s.pool.QueryRow(ctx, `
UPDATE gateway_tasks
SET status = 'running',
model_type = NULLIF($2::text, ''),
normalized_request = $3::jsonb,
execution_token = $2::uuid,
execution_lease_expires_at = now() + ($3::int * interval '1 second'),
locked_at = now(),
heartbeat_at = now(),
updated_at = now()
WHERE id = $1::uuid`, taskID, modelType, string(normalizedJSON))
return err
WHERE id = $1::uuid
AND (
(status = 'queued' AND next_run_at <= now())
OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()))
)
RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
if errors.Is(err, pgx.ErrNoRows) {
return GatewayTask{}, ErrTaskExecutionLeaseUnavailable
}
return task, err
}
func (s *Store) RenewTaskExecutionLease(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) error {
if leaseTTL <= 0 {
leaseTTL = 5 * time.Minute
}
tag, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET execution_lease_expires_at = now() + ($3::int * interval '1 second'),
heartbeat_at = now(),
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $2::uuid`, taskID, executionToken, int(leaseTTL/time.Second))
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrTaskExecutionLeaseLost
}
return nil
}
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, normalizedRequest map[string]any) error {
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
tag, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
model_type = NULLIF($3::text, ''),
normalized_request = $4::jsonb,
locked_at = now(),
heartbeat_at = now(),
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $2::uuid`, taskID, executionToken, modelType, string(normalizedJSON))
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrTaskExecutionLeaseLost
}
return nil
}
func (s *Store) ClaimAsyncQueuedTask(ctx context.Context, workerID string) (GatewayTask, error) {
@@ -189,7 +243,7 @@ WHERE t.id = picked.task_id
RETURNING `+gatewayTaskColumns, workerID))
}
func (s *Store) RequeueTask(ctx context.Context, taskID string, delay time.Duration, queueKey string) (GatewayTask, error) {
func (s *Store) RequeueTask(ctx context.Context, taskID string, executionToken string, delay time.Duration, queueKey string) (GatewayTask, error) {
if delay < time.Second {
delay = time.Second
}
@@ -203,14 +257,18 @@ SET status = 'queued',
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
next_run_at = $2::timestamptz,
queue_key = COALESCE(NULLIF($3::text, ''), queue_key),
execution_token = NULL,
execution_lease_expires_at = NULL,
next_run_at = $3::timestamptz,
queue_key = COALESCE(NULLIF($4::text, ''), queue_key),
error = NULL,
error_code = NULL,
error_message = NULL,
updated_at = now()
WHERE id = $1::uuid
RETURNING `+gatewayTaskColumns, taskID, nextRunAt, strings.TrimSpace(queueKey)))
AND status = 'running'
AND execution_token = $2::uuid
RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt, strings.TrimSpace(queueKey)))
}
func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error {
@@ -225,25 +283,32 @@ WHERE id = $1::uuid`, taskID, riverJobID)
return err
}
func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, attemptID string, remoteTaskID string, payload map[string]any) error {
func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, executionToken string, attemptID string, remoteTaskID string, payload map[string]any) error {
payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload))
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET remote_task_id = NULLIF($2::text, ''),
remote_task_payload = $3::jsonb,
SET remote_task_id = NULLIF($3::text, ''),
remote_task_payload = $4::jsonb,
updated_at = now()
WHERE id = $1::uuid`,
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $2::uuid`,
taskID,
executionToken,
remoteTaskID,
string(payloadJSON),
); err != nil {
)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrTaskExecutionLeaseLost
}
if strings.TrimSpace(attemptID) == "" {
return nil
}
_, err := tx.Exec(ctx, `
_, err = tx.Exec(ctx, `
UPDATE gateway_task_attempts
SET remote_task_id = NULLIF($2::text, ''),
response_snapshot = COALESCE(response_snapshot, '{}'::jsonb) || jsonb_build_object('remote_task_payload', $3::jsonb)
@@ -256,6 +321,20 @@ WHERE id = $1::uuid`,
})
}
func (s *Store) SetAttemptUpstreamSubmissionStatus(ctx context.Context, attemptID string, status string) error {
switch status {
case "not_submitted", "submitting", "response_received":
default:
return fmt.Errorf("invalid upstream submission status %q", status)
}
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_attempts
SET upstream_submission_status = $2,
upstream_submission_updated_at = now()
WHERE id = $1::uuid`, attemptID, status)
return err
}
func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message string) (GatewayTask, bool, error) {
message = strings.TrimSpace(message)
if message == "" {
@@ -293,10 +372,13 @@ func (s *Store) ListRecoverableAsyncTasks(ctx context.Context, limit int) ([]Asy
limit = 500
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, priority
SELECT id::text, priority, next_run_at
FROM gateway_tasks
WHERE async_mode = true
AND status IN ('queued', 'running')
AND (
status = 'queued'
OR (status = 'running' AND (execution_lease_expires_at IS NULL OR execution_lease_expires_at <= now()))
)
ORDER BY priority ASC, created_at ASC
LIMIT $1`, limit)
if err != nil {
@@ -306,7 +388,7 @@ LIMIT $1`, limit)
items := make([]AsyncTaskQueueItem, 0)
for rows.Next() {
var item AsyncTaskQueueItem
if err := rows.Scan(&item.ID, &item.Priority); err != nil {
if err := rows.Scan(&item.ID, &item.Priority, &item.NextRunAt); err != nil {
return nil, err
}
items = append(items, item)
@@ -320,6 +402,7 @@ LIMIT $1`, limit)
func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) {
requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot))
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
pricingJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot))
tx, err := s.pool.Begin(ctx)
if err != nil {
return "", err
@@ -330,11 +413,13 @@ func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptIn
err = tx.QueryRow(ctx, `
INSERT INTO gateway_task_attempts (
task_id, attempt_no, platform_id, platform_model_id, client_id, queue_key,
status, simulated, request_snapshot, metrics
status, simulated, request_snapshot, metrics, pricing_snapshot, request_fingerprint,
upstream_submission_status, upstream_submission_updated_at
)
VALUES (
$1::uuid, $2::int, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, NULLIF($5::text, ''), $6,
$7, $8, $9::jsonb, $10::jsonb
$7, $8, $9::jsonb, $10::jsonb, $11::jsonb, NULLIF($12, ''),
'not_submitted', now()
)
RETURNING id::text`,
input.TaskID,
@@ -347,6 +432,8 @@ RETURNING id::text`,
input.Simulated,
string(requestJSON),
string(metricsJSON),
string(pricingJSON),
input.RequestFingerprint,
).Scan(&attemptID)
if err != nil {
return "", err
@@ -489,6 +576,8 @@ SELECT a.id::text, a.task_id::text, a.attempt_no,
a.request_snapshot, COALESCE(a.response_snapshot, '{}'::jsonb),
COALESCE(a.response_started_at::text, ''), COALESCE(a.response_finished_at::text, ''),
COALESCE(a.response_duration_ms, 0), COALESCE(a.error_code, ''), COALESCE(a.error_message, ''),
COALESCE(a.pricing_snapshot, '{}'::jsonb), COALESCE(a.request_fingerprint, ''),
a.upstream_submission_status, COALESCE(a.upstream_submission_updated_at::text, ''),
a.started_at, COALESCE(a.finished_at::text, '')
FROM gateway_task_attempts a
LEFT JOIN integration_platforms p ON p.id = a.platform_id
@@ -518,6 +607,7 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
var metricsBytes []byte
var requestBytes []byte
var responseBytes []byte
var pricingSnapshotBytes []byte
if err := scanner.Scan(
&item.ID,
&item.TaskID,
@@ -544,6 +634,10 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
&item.ResponseDurationMS,
&item.ErrorCode,
&item.ErrorMessage,
&pricingSnapshotBytes,
&item.RequestFingerprint,
&item.UpstreamSubmissionStatus,
&item.UpstreamSubmissionUpdatedAt,
&item.StartedAt,
&item.FinishedAt,
); err != nil {
@@ -553,6 +647,7 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
item.Metrics = decodeObject(metricsBytes)
item.RequestSnapshot = decodeObject(requestBytes)
item.ResponseSnapshot = decodeObject(responseBytes)
item.PricingSnapshot = decodeObject(pricingSnapshotBytes)
enrichTaskAttemptFromMetrics(&item)
return item, nil
}
@@ -670,7 +765,41 @@ func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessIn
usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage))
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
billingSummaryJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingSummary))
if _, err := s.pool.Exec(ctx, `
pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot))
finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText)
if finalChargeAmount == "" {
finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64)
}
currency := normalizeWalletCurrency(input.BillingCurrency)
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if strings.TrimSpace(input.AttemptID) != "" {
if _, err := tx.Exec(ctx, `
UPDATE gateway_task_attempts
SET status = 'succeeded',
retryable = false,
request_id = NULLIF($2, ''),
usage = $3::jsonb,
metrics = $4::jsonb,
response_snapshot = $5::jsonb,
pricing_snapshot = $6::jsonb,
request_fingerprint = NULLIF($7, ''),
upstream_submission_status = 'response_received',
upstream_submission_updated_at = now(),
response_started_at = $8::timestamptz,
response_finished_at = $9::timestamptz,
response_duration_ms = $10,
error_code = NULL,
error_message = NULL,
finished_at = now()
WHERE id = $1::uuid`,
input.AttemptID, input.RequestID, string(usageJSON), string(metricsJSON),
string(resultJSON), string(pricingSnapshotJSON), input.RequestFingerprint,
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS,
); err != nil {
return err
}
}
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'succeeded',
result = $2::jsonb,
@@ -680,32 +809,142 @@ SET status = 'succeeded',
usage = $6::jsonb,
metrics = $7::jsonb,
billing_summary = $8::jsonb,
final_charge_amount = $9,
response_started_at = $10::timestamptz,
response_finished_at = $11::timestamptz,
response_duration_ms = $12,
final_charge_amount = $9::numeric,
billing_version = 'effective-pricing-v2',
billing_status = CASE
WHEN run_mode = 'production' AND gateway_user_id IS NOT NULL THEN 'pending'
ELSE 'not_required'
END,
billing_currency = $10,
pricing_snapshot = $11::jsonb,
request_fingerprint = NULLIF($12, ''),
billing_updated_at = now(),
response_started_at = $13::timestamptz,
response_finished_at = $14::timestamptz,
response_duration_ms = $15,
error = NULL,
error_code = NULL,
error_message = NULL,
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`,
input.TaskID,
string(resultJSON),
string(billingsJSON),
input.RequestID,
input.ResolvedModel,
string(usageJSON),
string(metricsJSON),
string(billingSummaryJSON),
input.FinalChargeAmount,
nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt),
input.ResponseDurationMS,
); err != nil {
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $16::uuid`,
input.TaskID,
string(resultJSON),
string(billingsJSON),
input.RequestID,
input.ResolvedModel,
string(usageJSON),
string(metricsJSON),
string(billingSummaryJSON),
finalChargeAmount,
currency,
string(pricingSnapshotJSON),
input.RequestFingerprint,
nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt),
input.ResponseDurationMS,
input.ExecutionToken,
)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrTaskExecutionLeaseLost
}
payloadJSON, _ := json.Marshal(map[string]any{"taskId": input.TaskID, "pricingVersion": "effective-pricing-v2"})
_, err = tx.Exec(ctx, `
INSERT INTO settlement_outbox (
task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at
)
SELECT id, 'task.billing.settle', 'settle', $2::numeric, $3, $4::jsonb, $5::jsonb, 'pending', now()
FROM gateway_tasks
WHERE id = $1::uuid
AND run_mode = 'production'
AND gateway_user_id IS NOT NULL
ON CONFLICT (task_id, event_type) DO NOTHING`,
input.TaskID, finalChargeAmount, currency, string(pricingSnapshotJSON), string(payloadJSON))
return err
})
if err != nil {
return GatewayTask{}, err
}
return s.GetTask(ctx, input.TaskID)
}
func (s *Store) FinishTaskManualReview(ctx context.Context, input FinishTaskManualReviewInput) (GatewayTask, error) {
status := strings.TrimSpace(input.TaskStatus)
if status != "succeeded" && status != "failed" {
return GatewayTask{}, fmt.Errorf("manual review task status must be succeeded or failed")
}
resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result))
pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot))
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if strings.TrimSpace(input.AttemptID) != "" {
attemptStatus := "failed"
if status == "succeeded" {
attemptStatus = "succeeded"
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_task_attempts
SET status = $2,
request_id = NULLIF($3, ''),
error_code = CASE WHEN $2 = 'failed' THEN NULLIF($4, '') ELSE NULL END,
error_message = CASE WHEN $2 = 'failed' THEN NULLIF($5, '') ELSE NULL END,
upstream_submission_status = CASE WHEN $2 = 'succeeded' THEN 'response_received' ELSE upstream_submission_status END,
upstream_submission_updated_at = now(),
response_started_at = $6::timestamptz,
response_finished_at = $7::timestamptz,
response_duration_ms = $8,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, input.Message,
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS); err != nil {
return err
}
}
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = $2,
result = $3::jsonb,
request_id = NULLIF($4, ''),
error = CASE WHEN $2 = 'failed' THEN NULLIF($6, '') ELSE NULL END,
error_code = NULLIF($5, ''),
error_message = NULLIF($6, ''),
billing_status = 'manual_review',
pricing_snapshot = CASE WHEN $7::jsonb = '{}'::jsonb THEN pricing_snapshot ELSE $7::jsonb END,
request_fingerprint = COALESCE(NULLIF($8, ''), request_fingerprint),
billing_updated_at = now(),
response_started_at = $9::timestamptz,
response_finished_at = $10::timestamptz,
response_duration_ms = $11,
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $12::uuid`, input.TaskID, status, string(resultJSON), input.RequestID, input.Code, input.Message,
string(pricingSnapshotJSON), input.RequestFingerprint, nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ExecutionToken)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrTaskExecutionLeaseLost
}
return nil
})
if err != nil {
return GatewayTask{}, err
}
return s.GetTask(ctx, input.TaskID)
@@ -743,6 +982,9 @@ ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
task.GatewayTenantID, task.GatewayUserID, task.TenantID, task.TenantKey, task.UserID, currency); err != nil {
return err
}
if err := ensureWalletAccountAuditGuard(ctx, tx, task.GatewayUserID, currency); err != nil {
return err
}
var exists bool
var accountID string
var balanceBefore float64
@@ -859,7 +1101,8 @@ func taskBillingString(value any) string {
func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) {
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result))
if _, err := s.pool.Exec(ctx, `
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
error = NULLIF($2::text, ''),
@@ -871,22 +1114,55 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
response_finished_at = $7::timestamptz,
response_duration_ms = $8,
result = $9::jsonb,
billing_status = CASE
WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
WHEN reservation_amount > 0 THEN 'pending'
ELSE 'released'
END,
billing_updated_at = now(),
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`,
input.TaskID,
input.Message,
input.Code,
input.RequestID,
string(metricsJSON),
nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt),
input.ResponseDurationMS,
string(resultJSON),
); err != nil {
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $10::uuid`,
input.TaskID,
input.Message,
input.Code,
input.RequestID,
string(metricsJSON),
nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt),
input.ResponseDurationMS,
string(resultJSON),
input.ExecutionToken,
)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrTaskExecutionLeaseLost
}
payloadJSON, _ := json.Marshal(map[string]any{"taskId": input.TaskID, "reason": input.Code})
_, err = tx.Exec(ctx, `
INSERT INTO settlement_outbox (
task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at
)
SELECT id, 'task.billing.release', 'release', reservation_amount, billing_currency,
pricing_snapshot, $2::jsonb, 'pending', now()
FROM gateway_tasks
WHERE id = $1::uuid
AND run_mode = 'production'
AND gateway_user_id IS NOT NULL
AND reservation_amount > 0
ON CONFLICT (task_id, event_type) DO NOTHING`, input.TaskID, string(payloadJSON))
return err
})
if err != nil {
return GatewayTask{}, err
}
return s.GetTask(ctx, input.TaskID)