Compare commits
4 Commits
037e83e92f
...
f8f5da04db
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f8f5da04db | ||
|
|
910e60ab59 | ||
|
|
65694cee15 | ||
|
|
923ba495c4 |
@@ -23,6 +23,11 @@ WORKER_CONCURRENCY=4
|
||||
# 单进程图片处理并发上限(API 与 Worker 均生效,默认等于 CPU 线程数)
|
||||
IMAGE_PROCESSING_CONCURRENCY=4
|
||||
|
||||
# ZIP 使用任务租约做 single-flight;总大小按解压前源文件字节计算。
|
||||
ZIP_BUILD_CONCURRENCY=2
|
||||
ZIP_MAX_ENTRIES=200
|
||||
ZIP_MAX_UNCOMPRESSED_BYTES=2147483648
|
||||
|
||||
# 仅当后端只能由可信反向代理访问时启用,否则客户端可伪造来源 IP
|
||||
TRUST_PROXY_HEADERS=false
|
||||
|
||||
|
||||
@@ -27,6 +27,11 @@ WORKER_TASK_CONCURRENCY=4
|
||||
WORKER_CONCURRENCY=2
|
||||
IMAGE_PROCESSING_CONCURRENCY=4
|
||||
|
||||
# Two concurrent 2 GiB ZIP builds require roughly 8 GiB of temporary disk.
|
||||
ZIP_BUILD_CONCURRENCY=2
|
||||
ZIP_MAX_ENTRIES=200
|
||||
ZIP_MAX_UNCOMPRESSED_BYTES=2147483648
|
||||
|
||||
# Resource ceilings tuned for an 8-core / 16 GB application host.
|
||||
POSTGRES_MEMORY_LIMIT=2g
|
||||
REDIS_MEMORY_LIMIT=1g
|
||||
|
||||
@@ -13,6 +13,9 @@ x-imageforge-environment: &imageforge-environment
|
||||
WORKER_TASK_CONCURRENCY: ${WORKER_TASK_CONCURRENCY:-4}
|
||||
WORKER_CONCURRENCY: ${WORKER_CONCURRENCY:-2}
|
||||
IMAGE_PROCESSING_CONCURRENCY: ${IMAGE_PROCESSING_CONCURRENCY:-2}
|
||||
ZIP_BUILD_CONCURRENCY: ${ZIP_BUILD_CONCURRENCY:-2}
|
||||
ZIP_MAX_ENTRIES: ${ZIP_MAX_ENTRIES:-200}
|
||||
ZIP_MAX_UNCOMPRESSED_BYTES: ${ZIP_MAX_UNCOMPRESSED_BYTES:-2147483648}
|
||||
ALLOW_ANONYMOUS_UPLOAD: ${ALLOW_ANONYMOUS_UPLOAD:-true}
|
||||
ANON_MAX_FILE_SIZE_MB: ${ANON_MAX_FILE_SIZE_MB:-5}
|
||||
ANON_MAX_FILES_PER_BATCH: ${ANON_MAX_FILES_PER_BATCH:-5}
|
||||
|
||||
@@ -278,6 +278,7 @@ dotenvy = "0.15"
|
||||
- 图片压缩使用 `spawn_blocking` 避免阻塞异步线程
|
||||
- `WORKER_TASK_CONCURRENCY` 控制任务级并发,避免大批量任务独占 Worker
|
||||
- `WORKER_CONCURRENCY` 控制单任务内文件并发,`IMAGE_PROCESSING_CONCURRENCY` 作为进程级 CPU 闸门
|
||||
- `ZIP_BUILD_CONCURRENCY` 是 API 进程级 ZIP 闸门;数据库租约保证同一任务跨实例只构建一次,`ZIP_MAX_ENTRIES` 和 `ZIP_MAX_UNCOMPRESSED_BYTES` 在下载源对象前拒绝超预算任务
|
||||
|
||||
```rust
|
||||
// 在独立线程池中执行 CPU 密集型压缩
|
||||
@@ -290,6 +291,7 @@ let result = tokio::task::spawn_blocking(move || {
|
||||
- 流式处理大文件
|
||||
- 限制并发压缩任务数
|
||||
- 压缩完成后立即清理临时文件
|
||||
- ZIP attempt 使用独立临时目录和对象键;发布 CAS 失败时立即删除,两项默认并发且每项 2 GiB 上限时应至少预留约 8 GiB 临时磁盘余量
|
||||
|
||||
### 3. 缓存策略
|
||||
- Redis 缓存用户会话
|
||||
|
||||
@@ -101,6 +101,7 @@ RETURNING used_units;
|
||||
- 批量任务的计量仍以“成功文件数”为准;失败文件(含 `QUOTA_EXCEEDED`)不计费。
|
||||
- 前端建议在上传前调用 `GET /billing/usage`(登录)或读取配额头(API)做本地提示/拦截。
|
||||
- 匿名批量任务先按文件数预留当日额度,终态结算只退还失败或未完成文件。未提供 `compression_rate` 属于正常压缩并计量;只有显式 `compression_rate=100`、同格式且无缩放的原样请求免计量。
|
||||
- 匿名单文件同样先预留,但响应中的 `units_charged` 只由实际输出决定:原样请求或输出未缩小均为 0。预留日期、session/IP 和任务 ID 会持久化;失败、跨日及进程中断由 Redis marker 幂等退款,不能退到请求结束时的新日期。
|
||||
|
||||
---
|
||||
|
||||
@@ -143,7 +144,7 @@ RETURNING used_units;
|
||||
- **乱序容忍**:订阅对象按 `(event.created, 事件优先级)` 保存独立水位;`deleted` 即使先到也会保留 tombstone,旧 `created/updated` 不得恢复已取消订阅。
|
||||
- **同秒歧义**:两个不同事件具有相同 `(event.created, 事件优先级)` 时,不能用不透明的 Event ID 排序,必须从 Stripe 拉取当前订阅快照并以快照响应时间推进水位。
|
||||
- **迁移对账**:历史版本用本地 `subscriptions.updated_at` 播种的非终态水位会标记为待对账;API 后台任务持租约获取 Stripe 快照,成功后才清除标记。未映射 Customer 或 Price 的受管订阅事件返回失败并等待重试,不能标记为已处理。
|
||||
- **发票一致性**:`invoices(provider, provider_invoice_id)` 唯一,发票事件也使用对象水位;新 `invoice.paid` 不会被迟到的旧 `invoice.payment_failed` 回退。同秒同等级事件从 Stripe 获取权威发票快照,未映射 Customer 时返回失败重试。
|
||||
- **发票一致性**:`invoices(provider, provider_invoice_id)` 唯一,发票事件也使用对象水位;新 `invoice.paid` 不会被迟到的旧 `invoice.payment_failed` 回退。同秒同等级事件从 Stripe 获取权威发票快照,未映射 Customer 时返回失败重试。迁移前已有 Stripe 发票会播种为待对账哨兵,首个后续事件必须先取权威快照;非 `paid` 状态不允许保留 `paid_at`。
|
||||
- **并发一致性**:`subscriptions(provider, provider_subscription_id)` 唯一,订阅业务写入与 `webhook_events=processed` 在同一事务提交。
|
||||
- **可重放**:保存原始 payload(脱敏)用于排查。
|
||||
|
||||
|
||||
@@ -315,6 +315,8 @@ Stripe 运行时还通过迁移维护三组一致性结构:
|
||||
- `provider_object_event_watermarks` 以 Stripe `event.created` 和事件等级保存对象水位;同秒同等级的不同事件标记为歧义并触发权威快照,不能按 Event ID 字典序决定先后。
|
||||
- `stripe_subscription_reconciliations` 保存历史非因果水位的租约化对账任务,允许多 API 实例用 `FOR UPDATE SKIP LOCKED` 安全消费。
|
||||
|
||||
迁移 `020` 会为尚无水位的历史 Stripe 发票写入 `requires_reconciliation=true` 哨兵。首个后续发票事件必须从 Stripe 获取当前对象后才能覆盖本地记录;`invoices_paid_at_status_check` 同时保证只有 `paid` 状态可以携带 `paid_at`。
|
||||
|
||||
数据库唯一索引同时保证非空 `users.billing_customer_id` 全局唯一、`subscriptions(provider, provider_subscription_id)` 唯一、非空 `invoices(provider, provider_invoice_id)` 唯一,以及每用户最多一条未取消 Stripe 订阅。部署这些索引前必须先清理存量冲突,具体检查见 `docs/deployment.md`。
|
||||
|
||||
### 4.8 tasks - 压缩任务
|
||||
@@ -355,7 +357,10 @@ CREATE TABLE tasks (
|
||||
zip_storage_endpoint_id UUID REFERENCES storage_endpoints(id) ON DELETE RESTRICT,
|
||||
zip_storage_key TEXT,
|
||||
zip_storage_etag TEXT,
|
||||
zip_size BIGINT
|
||||
zip_size BIGINT,
|
||||
zip_build_token UUID,
|
||||
zip_build_lease_until TIMESTAMPTZ,
|
||||
zip_build_attempt BIGINT NOT NULL DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE INDEX idx_tasks_user_id ON tasks(user_id);
|
||||
@@ -365,6 +370,10 @@ CREATE INDEX idx_tasks_created_at ON tasks(created_at);
|
||||
CREATE INDEX idx_tasks_expires_at ON tasks(expires_at);
|
||||
```
|
||||
|
||||
`zip_build_token/zip_build_lease_until` 是跨 API 实例的任务级 single-flight 租约。每个构建 attempt 写入独立对象键,只有 token 匹配的 CAS 更新可以发布到 `zip_storage_*`;失败或失租 attempt 必须删除对象。
|
||||
|
||||
匿名单文件预留单独存入 `anonymous_single_reservations`,不依赖尚未创建的 `tasks` 外键。`pending` 超时或 `refund_pending` 记录由 Worker 维护循环使用任务级 Redis marker 补偿;`charged/refunded` 记录保留 7 天后清理。
|
||||
|
||||
### 4.9 task_files - 任务文件
|
||||
```sql
|
||||
CREATE TABLE task_files (
|
||||
|
||||
@@ -13,6 +13,8 @@
|
||||
|
||||
Debian 13、4 核 CPU、8GB 内存的起始值建议为 `WORKER_TASK_CONCURRENCY=2`、`WORKER_CONCURRENCY=2` 和 `IMAGE_PROCESSING_CONCURRENCY=2`;8 核应用服务器可从 `4/2/4` 开始。三者分别表示同时处理的任务数、单任务内文件数和单进程 CPU 图片处理上限。最后一项是全局 CPU 闸门,因此不要把它设置为 CPU 核数的数倍。任务并发提高后,数据库连接池建议至少为 `WORKER_TASK_CONCURRENCY * WORKER_CONCURRENCY + 4`,生产示例使用 16。
|
||||
|
||||
8 核 16GB 应用服务器的 ZIP 起始值为 `ZIP_BUILD_CONCURRENCY=2`、`ZIP_MAX_ENTRIES=200`、`ZIP_MAX_UNCOMPRESSED_BYTES=2147483648`。ZIP 使用 stored 模式,构建时同时存在下载源和归档文件,按两个 2 GiB 构建估算应至少保留约 8 GiB 临时磁盘余量;磁盘较小时应先降低总字节或并发,而不是提高 HTTP 超时。
|
||||
|
||||
生产 Compose 默认按 8 核 16GB 主机设置可覆盖的资源上限:API 3GB、Worker 8GB、PostgreSQL 2GB、Redis 1GB;对应变量为 `API_MEMORY_LIMIT`、`WORKER_MEMORY_LIMIT`、`POSTGRES_MEMORY_LIMIT` 和 `REDIS_MEMORY_LIMIT`。Redis 的 `REDIS_MAXMEMORY` 默认 768MB,达到上限后返回写入错误而不是继续挤占宿主机内存。
|
||||
|
||||
### 首次启动
|
||||
@@ -55,7 +57,7 @@ curl --fail http://127.0.0.1:8080/metrics
|
||||
|
||||
### 更新与回滚
|
||||
|
||||
更新代码后保留 `.env.production` 和命名卷。包含迁移 `017` 至 `019` 的版本不能直接让旧、新 Worker 并行滚动:先备份数据库并停止旧 Worker,再构建新镜像。
|
||||
更新代码后保留 `.env.production` 和命名卷。包含迁移 `017` 至 `022` 的版本不能让旧、新 API 或 Worker 并行滚动:旧 API 不理解 ZIP 构建租约,旧 Worker 不理解任务 attempt fencing。先备份数据库并停止旧 API/Worker,再构建新镜像。
|
||||
|
||||
迁移 `017` 会在发现重复 Customer 或同用户多条未取消 Stripe 订阅时主动失败,迁移 `019` 会在发现同一 Stripe 发票对应多行时主动失败。部署前先检查并人工对账,三个查询都必须返回 0 行:
|
||||
|
||||
@@ -83,26 +85,27 @@ HAVING COUNT(*) > 1;
|
||||
|
||||
```bash
|
||||
git pull --ff-only
|
||||
docker compose --env-file .env.production -f docker/docker-compose.prod.yml stop worker
|
||||
docker compose --env-file .env.production -f docker/docker-compose.prod.yml stop api worker
|
||||
docker compose --env-file .env.production -f docker/docker-compose.prod.yml build api
|
||||
docker compose --env-file .env.production -f docker/docker-compose.prod.yml up -d postgres redis api
|
||||
docker compose --env-file .env.production -f docker/docker-compose.prod.yml up -d worker
|
||||
```
|
||||
|
||||
新 API 启动后会消费迁移 `018` 创建的 Stripe 对账队列。启动 Worker 前应确认 API 健康、`STRIPE_SECRET_KEY` 可用且服务器能访问 `STRIPE_API_BASE_URL`;对账可以后台继续,但必须监控失败项:
|
||||
新 API 启动后会消费迁移 `018` 创建的订阅对账队列。迁移 `020` 为历史发票写入待对账哨兵,发票不主动批量拉取,而是在首个后续事件到达时取 Stripe 快照。启动 Worker 前应确认 API 健康、`STRIPE_SECRET_KEY` 可用且服务器能访问 `STRIPE_API_BASE_URL`;订阅对账可以后台继续,但必须监控失败项:
|
||||
|
||||
```sql
|
||||
SELECT status, COUNT(*)
|
||||
FROM stripe_subscription_reconciliations
|
||||
GROUP BY status;
|
||||
|
||||
SELECT provider_object_id, reconciliation_reason, updated_at
|
||||
SELECT object_type, requires_reconciliation, COUNT(*)
|
||||
FROM provider_object_event_watermarks
|
||||
WHERE provider = 'stripe' AND requires_reconciliation = true
|
||||
ORDER BY updated_at;
|
||||
WHERE provider = 'stripe'
|
||||
GROUP BY object_type, requires_reconciliation
|
||||
ORDER BY object_type, requires_reconciliation;
|
||||
```
|
||||
|
||||
`failed` 会指数退避重试;持续失败通常表示 Stripe 凭据、网络、Customer/Price 映射不完整。上线验收要求 `pending/processing/failed` 最终归零,且 `requires_reconciliation=true` 为 0。生产镜像应使用不可变的 `IMAGEFORGE_TAG`。数据库迁移已应用后,不能只回滚旧二进制;应保留新 schema,并使用兼容该 schema 的修复镜像。
|
||||
`failed` 会指数退避重试;持续失败通常表示 Stripe 凭据、网络、Customer/Price 映射不完整。上线验收要求订阅队列的 `pending/processing/failed` 最终归零,且 subscription 水位不再待对账;invoice 水位在对应发票首个后续事件到达前保持 `requires_reconciliation=true` 属于预期状态。生产镜像应使用不可变的 `IMAGEFORGE_TAG`。数据库迁移已应用后,不能只回滚旧二进制;应保留新 schema,并使用兼容该 schema 的修复镜像。
|
||||
|
||||
### 反向代理
|
||||
|
||||
|
||||
@@ -66,7 +66,7 @@ flowchart LR
|
||||
| 低级会员 Pro | 7 天 | `results/7d/`、`archives/7d/` | 9 天 |
|
||||
| 高级会员 Business | 15 天 | `results/15d/`、`archives/15d/` | 17 天 |
|
||||
|
||||
Worker 每 5 分钟按 `expires_at` 精确删除对象,删除成功后才删除数据库任务。S3 生命周期多保留 2 天,只负责处理数据库故障、进程崩溃或上传后未能落库的孤儿对象,不能作为精确会员权限判断。未完成的分片上传 1 天后由生命周期中止。
|
||||
Worker 每 5 分钟按 `expires_at` 精确删除对象,删除成功后才删除数据库任务。ZIP 构建 attempt 位于对应 `archives/<retention>/.../attempts/` 前缀,发布失败会立即删除,进程崩溃遗留项仍由同一前缀生命周期兜底。S3 生命周期多保留 2 天,只负责处理数据库故障、进程崩溃或上传后未能落库的孤儿对象,不能作为精确会员权限判断。未完成的分片上传 1 天后由生命周期中止。
|
||||
|
||||
## 5. 119 首期容量
|
||||
|
||||
|
||||
31
migrations/020_existing_invoice_watermarks.sql
Normal file
31
migrations/020_existing_invoice_watermarks.sql
Normal file
@@ -0,0 +1,31 @@
|
||||
-- Existing Stripe invoices predate object watermarks. Force the first later
|
||||
-- event to reconcile against Stripe instead of treating it as authoritative.
|
||||
INSERT INTO provider_object_event_watermarks (
|
||||
provider, object_type, provider_object_id,
|
||||
last_event_created, last_event_rank, last_event_id,
|
||||
is_deleted, requires_reconciliation, reconciliation_reason,
|
||||
updated_at
|
||||
)
|
||||
SELECT
|
||||
'stripe', 'invoice', provider_invoice_id,
|
||||
0, 0, 'reconcile:migration:020',
|
||||
false, true, 'migration_020_existing_invoice',
|
||||
NOW()
|
||||
FROM invoices
|
||||
WHERE provider = 'stripe'
|
||||
AND provider_invoice_id IS NOT NULL
|
||||
ON CONFLICT (provider, object_type, provider_object_id) DO NOTHING;
|
||||
|
||||
-- A non-paid invoice must never retain the timestamp from an older paid
|
||||
-- payload. Clean historical contradictions before enforcing the invariant.
|
||||
UPDATE invoices
|
||||
SET paid_at = NULL
|
||||
WHERE status <> 'paid'
|
||||
AND paid_at IS NOT NULL;
|
||||
|
||||
ALTER TABLE invoices
|
||||
ADD CONSTRAINT invoices_paid_at_status_check
|
||||
CHECK (paid_at IS NULL OR status = 'paid') NOT VALID;
|
||||
|
||||
ALTER TABLE invoices
|
||||
VALIDATE CONSTRAINT invoices_paid_at_status_check;
|
||||
12
migrations/021_zip_build_leases.sql
Normal file
12
migrations/021_zip_build_leases.sql
Normal file
@@ -0,0 +1,12 @@
|
||||
ALTER TABLE tasks
|
||||
ADD COLUMN zip_build_token UUID,
|
||||
ADD COLUMN zip_build_lease_until TIMESTAMPTZ,
|
||||
ADD COLUMN zip_build_attempt BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
ALTER TABLE tasks
|
||||
ADD CONSTRAINT tasks_zip_build_lease_pair_check
|
||||
CHECK ((zip_build_token IS NULL) = (zip_build_lease_until IS NULL));
|
||||
|
||||
CREATE INDEX idx_tasks_zip_build_lease
|
||||
ON tasks(zip_build_lease_until)
|
||||
WHERE zip_storage_key IS NULL AND zip_build_token IS NOT NULL;
|
||||
20
migrations/022_anonymous_single_reservations.sql
Normal file
20
migrations/022_anonymous_single_reservations.sql
Normal file
@@ -0,0 +1,20 @@
|
||||
CREATE TABLE anonymous_single_reservations (
|
||||
task_id UUID PRIMARY KEY,
|
||||
session_id VARCHAR(100) NOT NULL,
|
||||
client_ip INET NOT NULL,
|
||||
quota_date DATE NOT NULL,
|
||||
units INTEGER NOT NULL,
|
||||
status VARCHAR(20) NOT NULL DEFAULT 'pending',
|
||||
refund_after TIMESTAMPTZ NOT NULL DEFAULT (NOW() + INTERVAL '15 minutes'),
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
settled_at TIMESTAMPTZ,
|
||||
CONSTRAINT anonymous_single_reservations_units_check
|
||||
CHECK (units > 0),
|
||||
CONSTRAINT anonymous_single_reservations_status_check
|
||||
CHECK (status IN ('pending', 'charged', 'refund_pending', 'refunded'))
|
||||
);
|
||||
|
||||
CREATE INDEX anonymous_single_reservations_unsettled
|
||||
ON anonymous_single_reservations(refund_after, created_at)
|
||||
WHERE status IN ('pending', 'refund_pending');
|
||||
@@ -1091,6 +1091,7 @@ mod tests {
|
||||
AppState {
|
||||
mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")),
|
||||
image_processing_semaphore: Arc::new(Semaphore::new(2)),
|
||||
zip_build_semaphore: Arc::new(Semaphore::new(2)),
|
||||
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
|
||||
storage_cache: crate::services::storage::StorageCache::new(),
|
||||
config,
|
||||
|
||||
@@ -79,6 +79,10 @@ fn default_units_charged() -> i32 {
|
||||
1
|
||||
}
|
||||
|
||||
fn metered_units(charged: bool) -> i32 {
|
||||
i32::from(charged)
|
||||
}
|
||||
|
||||
fn direct_response<B: IntoResponse>(
|
||||
body: B,
|
||||
format: ImageFmt,
|
||||
@@ -264,14 +268,16 @@ async fn compress_json(
|
||||
}
|
||||
}
|
||||
|
||||
let mut anonymous_reserved = false;
|
||||
let task_id = Uuid::new_v4();
|
||||
let mut anonymous_reservation_date = None;
|
||||
let op: Result<CompressResponse, AppError> = (async {
|
||||
match "a_ctx {
|
||||
QuotaContext::User(billing) => ensure_quota_available(&state, billing, 1).await?,
|
||||
QuotaContext::ApiKey(billing, _) => ensure_quota_available(&state, billing, 1).await?,
|
||||
QuotaContext::Anonymous { session_id, ip } => {
|
||||
quota::consume_anonymous_units(&state, session_id, *ip, 1).await?;
|
||||
anonymous_reserved = true;
|
||||
anonymous_reservation_date = Some(
|
||||
quota::reserve_anonymous_single_unit(&state, task_id, session_id, *ip).await?,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -299,8 +305,7 @@ async fn compress_json(
|
||||
} else {
|
||||
(saved_bytes as f64) * 100.0 / (original_size as f64)
|
||||
};
|
||||
let charge_units = anonymous_reserved
|
||||
|| quota::output_consumes_unit(
|
||||
let charge_units = quota::output_consumes_unit(
|
||||
req.compression_rate,
|
||||
format_in == format_out,
|
||||
req.max_width.is_some() || req.max_height.is_some(),
|
||||
@@ -309,7 +314,6 @@ async fn compress_json(
|
||||
compressed_size,
|
||||
);
|
||||
|
||||
let task_id = Uuid::new_v4();
|
||||
let file_id = Uuid::new_v4();
|
||||
let retention_hours = retention.num_hours();
|
||||
let object_key =
|
||||
@@ -356,6 +360,17 @@ async fn compress_json(
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
if anonymous_reservation_date.is_some() {
|
||||
if let Err(err) =
|
||||
quota::finalize_anonymous_single_reservation(&state, task_id, charge_units).await
|
||||
{
|
||||
// The durable reservation remains visible to maintenance, so a
|
||||
// transient Redis failure must not turn a successful image into
|
||||
// a failed, non-idempotent request.
|
||||
tracing::warn!(task_id = %task_id, charged = charge_units, error = %err, "anonymous single reservation finalization deferred");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(CompressResponse {
|
||||
task_id,
|
||||
file_id,
|
||||
@@ -368,7 +383,7 @@ async fn compress_json(
|
||||
download_url: format!("/downloads/{file_id}"),
|
||||
expires_at,
|
||||
billing: BillingView {
|
||||
units_charged: if charge_units { 1 } else { 0 },
|
||||
units_charged: metered_units(charge_units),
|
||||
},
|
||||
})
|
||||
})
|
||||
@@ -402,9 +417,11 @@ async fn compress_json(
|
||||
))
|
||||
}
|
||||
Err(err) => {
|
||||
if anonymous_reserved {
|
||||
if let QuotaContext::Anonymous { session_id, ip } = "a_ctx {
|
||||
let _ = quota::refund_anonymous_units(&state, session_id, *ip, 1).await;
|
||||
if anonymous_reservation_date.is_some() {
|
||||
if let Err(refund_err) =
|
||||
quota::refund_anonymous_single_reservation(&state, task_id).await
|
||||
{
|
||||
tracing::warn!(task_id = %task_id, error = %refund_err, "anonymous single reservation refund deferred");
|
||||
}
|
||||
}
|
||||
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
|
||||
@@ -653,7 +670,7 @@ async fn compress_direct(
|
||||
compressed_size,
|
||||
saved_bytes,
|
||||
saved_percent,
|
||||
units_charged: if charge_units { 1 } else { 0 },
|
||||
units_charged: metered_units(charge_units),
|
||||
};
|
||||
let response = direct_response(compressed, format_out, &idem_data);
|
||||
Ok((response, idem_data))
|
||||
@@ -1137,7 +1154,9 @@ async fn record_task_and_metering(
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建文件记录失败").with_source(err))?;
|
||||
|
||||
match quota_ctx {
|
||||
QuotaContext::Anonymous { .. } => {}
|
||||
QuotaContext::Anonymous { .. } => {
|
||||
quota::mark_anonymous_single_result(&mut tx, task_id, charge_units).await?;
|
||||
}
|
||||
QuotaContext::User(billing) => {
|
||||
if charge_units {
|
||||
charge_one_unit(
|
||||
@@ -1247,4 +1266,34 @@ mod tests {
|
||||
assert_eq!(response.headers()["imageforge-saved-percent"], "37.50");
|
||||
assert_eq!(response.headers()["imageforge-units-charged"], "1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn anonymous_response_units_follow_actual_output_metering() {
|
||||
let cases = [
|
||||
(Some(100), true, false, false, 100, 50, 0),
|
||||
(None, true, false, false, 100, 100, 0),
|
||||
(None, true, false, false, 100, 50, 1),
|
||||
(Some(100), true, true, false, 100, 50, 1),
|
||||
];
|
||||
for (
|
||||
compression_rate,
|
||||
same_format,
|
||||
has_resize,
|
||||
has_target_size,
|
||||
original_size,
|
||||
output_size,
|
||||
expected_units,
|
||||
) in cases
|
||||
{
|
||||
let charged = quota::output_consumes_unit(
|
||||
compression_rate,
|
||||
same_format,
|
||||
has_resize,
|
||||
has_target_size,
|
||||
original_size,
|
||||
output_size,
|
||||
);
|
||||
assert_eq!(metered_units(charged), expected_units);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -241,8 +241,27 @@ struct TaskZipFileRow {
|
||||
storage_key: Option<String>,
|
||||
original_name: String,
|
||||
output_format: String,
|
||||
compressed_size: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct ZipBuildStateRow {
|
||||
zip_storage_backend: Option<String>,
|
||||
zip_storage_endpoint_id: Option<Uuid>,
|
||||
zip_storage_key: Option<String>,
|
||||
expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum ZipBuildClaim {
|
||||
Acquired { token: Uuid },
|
||||
Cached(storage::ObjectLocator),
|
||||
Busy,
|
||||
}
|
||||
|
||||
const ZIP_BUILD_LEASE_SECONDS: i64 = 15 * 60;
|
||||
const ZIP_BUILD_WAIT_SECONDS: u64 = 30;
|
||||
|
||||
async fn download_task_zip(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
@@ -315,11 +334,31 @@ async fn download_task_zip(
|
||||
.await;
|
||||
}
|
||||
|
||||
let object = resolve_task_zip(&state, task_id, task.retention_hours as i64).await?;
|
||||
|
||||
respond_object(
|
||||
&state,
|
||||
jar,
|
||||
&object,
|
||||
&format!("task_{task_id}.zip"),
|
||||
"application/zip",
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_task_zip(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
retention_hours: i64,
|
||||
) -> Result<storage::ObjectLocator, AppError> {
|
||||
if let Some(object) = load_published_zip(state, task_id).await? {
|
||||
return Ok(object);
|
||||
}
|
||||
let rows = sqlx::query_as::<_, TaskZipFileRow>(
|
||||
r#"
|
||||
SELECT storage_backend, storage_endpoint_id,
|
||||
COALESCE(storage_key, storage_path) AS storage_key,
|
||||
original_name, output_format
|
||||
original_name, output_format, compressed_size
|
||||
FROM task_files
|
||||
WHERE task_id = $1 AND status = 'completed'
|
||||
ORDER BY created_at ASC
|
||||
@@ -333,44 +372,273 @@ async fn download_task_zip(
|
||||
if rows.is_empty() {
|
||||
return Err(AppError::new(ErrorCode::NotFound, "没有可打包的文件"));
|
||||
}
|
||||
validate_zip_budget(state, &rows)?;
|
||||
|
||||
let deadline =
|
||||
tokio::time::Instant::now() + std::time::Duration::from_secs(ZIP_BUILD_WAIT_SECONDS);
|
||||
loop {
|
||||
match claim_zip_build(state, task_id).await? {
|
||||
ZipBuildClaim::Cached(object) => return Ok(object),
|
||||
ZipBuildClaim::Acquired { token } => {
|
||||
return build_claimed_zip(state, task_id, retention_hours, token, &rows).await;
|
||||
}
|
||||
ZipBuildClaim::Busy if tokio::time::Instant::now() < deadline => {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
}
|
||||
ZipBuildClaim::Busy => {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::StorageUnavailable,
|
||||
"ZIP 正在生成,请稍后重试",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_zip_budget(state: &AppState, rows: &[TaskZipFileRow]) -> Result<u64, AppError> {
|
||||
if rows.len() > state.config.zip_max_entries as usize {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::FileTooLarge,
|
||||
format!("ZIP 文件数量超过 {} 个上限", state.config.zip_max_entries),
|
||||
));
|
||||
}
|
||||
rows.iter().try_fold(0_u64, |total, row| {
|
||||
if row.storage_key.is_none() {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::StorageUnavailable,
|
||||
"ZIP 源文件存储信息不完整",
|
||||
));
|
||||
}
|
||||
let size = row
|
||||
.compressed_size
|
||||
.and_then(|value| u64::try_from(value).ok())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::StorageUnavailable, "ZIP 源文件大小无效"))?;
|
||||
let next = total
|
||||
.checked_add(size)
|
||||
.ok_or_else(|| AppError::new(ErrorCode::FileTooLarge, "ZIP 源文件总大小超出限制"))?;
|
||||
if next > state.config.zip_max_uncompressed_bytes {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::FileTooLarge,
|
||||
format!(
|
||||
"ZIP 源文件总大小超过 {} 字节上限",
|
||||
state.config.zip_max_uncompressed_bytes
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(next)
|
||||
})
|
||||
}
|
||||
|
||||
async fn claim_zip_build(state: &AppState, task_id: Uuid) -> Result<ZipBuildClaim, AppError> {
|
||||
let token = Uuid::new_v4();
|
||||
let acquired: Option<i64> = sqlx::query_scalar(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET zip_build_token = $2,
|
||||
zip_build_lease_until = NOW() + ($3 * INTERVAL '1 second'),
|
||||
zip_build_attempt = zip_build_attempt + 1
|
||||
WHERE id = $1
|
||||
AND zip_storage_key IS NULL
|
||||
AND completed_at IS NOT NULL
|
||||
AND expires_at > NOW()
|
||||
AND (
|
||||
zip_build_token IS NULL
|
||||
OR zip_build_lease_until <= NOW()
|
||||
)
|
||||
RETURNING zip_build_attempt
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(token)
|
||||
.bind(ZIP_BUILD_LEASE_SECONDS)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "获取 ZIP 构建租约失败").with_source(err))?;
|
||||
if acquired.is_some() {
|
||||
return Ok(ZipBuildClaim::Acquired { token });
|
||||
}
|
||||
|
||||
let current = sqlx::query_as::<_, ZipBuildStateRow>(
|
||||
r#"
|
||||
SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key, expires_at
|
||||
FROM tasks
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询 ZIP 构建状态失败").with_source(err))?
|
||||
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "任务不存在"))?;
|
||||
if current.expires_at <= Utc::now() {
|
||||
return Err(AppError::new(ErrorCode::NotFound, "任务已过期或不存在"));
|
||||
}
|
||||
if let (Some(backend), Some(key)) = (current.zip_storage_backend, current.zip_storage_key) {
|
||||
return Ok(ZipBuildClaim::Cached(storage::ObjectLocator {
|
||||
backend,
|
||||
endpoint_id: current.zip_storage_endpoint_id,
|
||||
key,
|
||||
}));
|
||||
}
|
||||
// A lease may have been released between the UPDATE and this read. The
|
||||
// caller retries the atomic claim after a short wait in either case.
|
||||
Ok(ZipBuildClaim::Busy)
|
||||
}
|
||||
|
||||
async fn renew_zip_build(state: &AppState, task_id: Uuid, token: Uuid) -> Result<(), AppError> {
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET zip_build_lease_until = NOW() + ($3 * INTERVAL '1 second')
|
||||
WHERE id = $1
|
||||
AND zip_build_token = $2
|
||||
AND zip_storage_key IS NULL
|
||||
AND expires_at > NOW()
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(token)
|
||||
.bind(ZIP_BUILD_LEASE_SECONDS)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "续租 ZIP 构建失败").with_source(err))?;
|
||||
if updated.rows_affected() != 1 {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::StorageUnavailable,
|
||||
"ZIP 构建租约已失效,请重试",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn release_zip_build(state: &AppState, task_id: Uuid, token: Uuid) {
|
||||
if let Err(err) = sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET zip_build_token = NULL,
|
||||
zip_build_lease_until = NULL
|
||||
WHERE id = $1
|
||||
AND zip_build_token = $2
|
||||
AND zip_storage_key IS NULL
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(token)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(task_id = %task_id, zip_build_token = %token, error = %err, "failed to release ZIP build lease");
|
||||
}
|
||||
}
|
||||
|
||||
async fn build_claimed_zip(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
retention_hours: i64,
|
||||
token: Uuid,
|
||||
rows: &[TaskZipFileRow],
|
||||
) -> Result<storage::ObjectLocator, AppError> {
|
||||
let permit = match tokio::time::timeout(
|
||||
std::time::Duration::from_secs(ZIP_BUILD_WAIT_SECONDS),
|
||||
state.zip_build_semaphore.clone().acquire_owned(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => permit,
|
||||
Ok(Err(err)) => {
|
||||
release_zip_build(state, task_id, token).await;
|
||||
return Err(AppError::new(ErrorCode::Internal, "ZIP 并发闸门已关闭").with_source(err));
|
||||
}
|
||||
Err(_) => {
|
||||
release_zip_build(state, task_id, token).await;
|
||||
return Err(AppError::new(
|
||||
ErrorCode::StorageUnavailable,
|
||||
"ZIP 生成繁忙,请稍后重试",
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let temp_dir = PathBuf::from(format!(
|
||||
"{}/tmp/zips/{task_id}-{}",
|
||||
state.config.storage_path,
|
||||
Uuid::new_v4()
|
||||
"{}/tmp/zips/{task_id}-{token}",
|
||||
state.config.storage_path
|
||||
));
|
||||
tokio::fs::create_dir_all(&temp_dir).await.map_err(|err| {
|
||||
let zip_path = temp_dir.join(format!("task_{task_id}.zip"));
|
||||
let build_result = build_zip_attempt(
|
||||
state,
|
||||
task_id,
|
||||
token,
|
||||
retention_hours,
|
||||
rows,
|
||||
&temp_dir,
|
||||
&zip_path,
|
||||
)
|
||||
.await;
|
||||
drop(permit);
|
||||
let _ = tokio::fs::remove_dir_all(&temp_dir).await;
|
||||
|
||||
let stored = match build_result {
|
||||
Ok(stored) => stored,
|
||||
Err(err) => {
|
||||
release_zip_build(state, task_id, token).await;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
publish_zip_attempt(state, task_id, token, stored).await
|
||||
}
|
||||
|
||||
async fn build_zip_attempt(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
token: Uuid,
|
||||
retention_hours: i64,
|
||||
rows: &[TaskZipFileRow],
|
||||
temp_dir: &std::path::Path,
|
||||
zip_path: &std::path::Path,
|
||||
) -> Result<storage::StoredObject, AppError> {
|
||||
tokio::fs::create_dir_all(temp_dir).await.map_err(|err| {
|
||||
AppError::new(ErrorCode::StorageUnavailable, "创建 ZIP 临时目录失败").with_source(err)
|
||||
})?;
|
||||
let zip_path = temp_dir.join(format!("task_{task_id}.zip"));
|
||||
|
||||
let build_result: Result<storage::StoredObject, AppError> = async {
|
||||
let mut used_names: HashMap<String, usize> = HashMap::new();
|
||||
let mut entries: Vec<(String, String)> = Vec::new();
|
||||
for (index, row) in rows.into_iter().enumerate() {
|
||||
let Some(key) = row.storage_key else {
|
||||
continue;
|
||||
};
|
||||
let mut entries: Vec<(String, String)> = Vec::with_capacity(rows.len());
|
||||
let mut actual_bytes = 0_u64;
|
||||
for (index, row) in rows.iter().enumerate() {
|
||||
renew_zip_build(state, task_id, token).await?;
|
||||
let key = row.storage_key.as_ref().ok_or_else(|| {
|
||||
AppError::new(ErrorCode::StorageUnavailable, "ZIP 源文件存储信息不完整")
|
||||
})?;
|
||||
let path = temp_dir.join(format!("entry-{index}"));
|
||||
storage::download_to_file(
|
||||
&state,
|
||||
state,
|
||||
&storage::ObjectLocator {
|
||||
backend: row.storage_backend,
|
||||
backend: row.storage_backend.clone(),
|
||||
endpoint_id: row.storage_endpoint_id,
|
||||
key,
|
||||
key: key.clone(),
|
||||
},
|
||||
&path,
|
||||
)
|
||||
.await?;
|
||||
let name =
|
||||
build_zip_entry_name(&row.original_name, &row.output_format, &mut used_names);
|
||||
let size = tokio::fs::metadata(&path)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
AppError::new(ErrorCode::StorageUnavailable, "读取 ZIP 临时文件失败")
|
||||
.with_source(err)
|
||||
})?
|
||||
.len();
|
||||
actual_bytes = actual_bytes
|
||||
.checked_add(size)
|
||||
.ok_or_else(|| AppError::new(ErrorCode::FileTooLarge, "ZIP 实际文件总大小超出限制"))?;
|
||||
if actual_bytes > state.config.zip_max_uncompressed_bytes {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::FileTooLarge,
|
||||
"ZIP 实际文件总大小超出限制",
|
||||
));
|
||||
}
|
||||
let name = build_zip_entry_name(&row.original_name, &row.output_format, &mut used_names);
|
||||
entries.push((name, path.to_string_lossy().to_string()));
|
||||
}
|
||||
if entries.is_empty() {
|
||||
return Err(AppError::new(ErrorCode::NotFound, "没有可打包的文件"));
|
||||
}
|
||||
|
||||
let zip_path_cloned = zip_path.clone();
|
||||
renew_zip_build(state, task_id, token).await?;
|
||||
let zip_path_cloned = zip_path.to_path_buf();
|
||||
let task_id_str = task_id.to_string();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
generate_zip_file(&zip_path_cloned, &task_id_str, &entries)
|
||||
@@ -379,47 +647,111 @@ async fn download_task_zip(
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))?
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))?;
|
||||
|
||||
let object_key = storage::archive_key(task.retention_hours as i64, task_id);
|
||||
storage::store_file(&state, &object_key, &zip_path, "application/zip").await
|
||||
}
|
||||
.await;
|
||||
renew_zip_build(state, task_id, token).await?;
|
||||
let object_key = storage::archive_attempt_key(retention_hours, task_id, token);
|
||||
storage::store_file(state, &object_key, zip_path, "application/zip").await
|
||||
}
|
||||
|
||||
let _ = tokio::fs::remove_dir_all(&temp_dir).await;
|
||||
let stored = build_result?;
|
||||
|
||||
sqlx::query(
|
||||
async fn publish_zip_attempt(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
token: Uuid,
|
||||
stored: storage::StoredObject,
|
||||
) -> Result<storage::ObjectLocator, AppError> {
|
||||
let published = sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET zip_storage_backend = $2,
|
||||
zip_storage_endpoint_id = $3,
|
||||
zip_storage_key = $4,
|
||||
zip_storage_etag = $5,
|
||||
zip_size = $6
|
||||
WHERE id = $1 AND zip_storage_key IS NULL
|
||||
SET zip_storage_backend = $3,
|
||||
zip_storage_endpoint_id = $4,
|
||||
zip_storage_key = $5,
|
||||
zip_storage_etag = $6,
|
||||
zip_size = $7,
|
||||
zip_build_token = NULL,
|
||||
zip_build_lease_until = NULL
|
||||
WHERE id = $1
|
||||
AND zip_build_token = $2
|
||||
AND zip_storage_key IS NULL
|
||||
AND expires_at > NOW()
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(token)
|
||||
.bind(&stored.backend)
|
||||
.bind(stored.endpoint_id)
|
||||
.bind(&stored.key)
|
||||
.bind(&stored.etag)
|
||||
.bind(stored.size as i64)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "记录 ZIP 对象失败").with_source(err))?;
|
||||
.await;
|
||||
|
||||
respond_object(
|
||||
&state,
|
||||
jar,
|
||||
&storage::ObjectLocator {
|
||||
match published {
|
||||
Ok(result) if result.rows_affected() == 1 => Ok(storage::ObjectLocator {
|
||||
backend: stored.backend,
|
||||
endpoint_id: stored.endpoint_id,
|
||||
key: stored.key,
|
||||
},
|
||||
&format!("task_{task_id}.zip"),
|
||||
"application/zip",
|
||||
}),
|
||||
Ok(_) => {
|
||||
delete_unpublished_zip(state, task_id, token, &stored).await;
|
||||
let current = load_published_zip(state, task_id).await?;
|
||||
current.ok_or_else(|| {
|
||||
AppError::new(ErrorCode::StorageUnavailable, "ZIP 发布租约已失效,请重试")
|
||||
})
|
||||
}
|
||||
Err(err) => {
|
||||
delete_unpublished_zip(state, task_id, token, &stored).await;
|
||||
Err(AppError::new(ErrorCode::Internal, "记录 ZIP 对象失败").with_source(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_unpublished_zip(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
token: Uuid,
|
||||
stored: &storage::StoredObject,
|
||||
) {
|
||||
let object = storage::ObjectLocator {
|
||||
backend: stored.backend.clone(),
|
||||
endpoint_id: stored.endpoint_id,
|
||||
key: stored.key.clone(),
|
||||
};
|
||||
let mut last_error = None;
|
||||
for attempt in 1..=3_u64 {
|
||||
match storage::delete_object(state, &object).await {
|
||||
Ok(()) => {
|
||||
last_error = None;
|
||||
break;
|
||||
}
|
||||
Err(err) => {
|
||||
last_error = Some(err);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100 * attempt)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(err) = last_error {
|
||||
tracing::error!(task_id = %task_id, zip_build_token = %token, object_key = %stored.key, error = %err, "failed to delete unpublished ZIP attempt after retries");
|
||||
}
|
||||
release_zip_build(state, task_id, token).await;
|
||||
}
|
||||
|
||||
async fn load_published_zip(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
) -> Result<Option<storage::ObjectLocator>, AppError> {
|
||||
let row: Option<(Option<String>, Option<Uuid>, Option<String>)> = sqlx::query_as(
|
||||
"SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1",
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询已发布 ZIP 失败").with_source(err))?;
|
||||
Ok(row.and_then(|(backend, endpoint_id, key)| {
|
||||
Some(storage::ObjectLocator {
|
||||
backend: backend?,
|
||||
endpoint_id,
|
||||
key: key?,
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
fn build_zip_entry_name(
|
||||
@@ -497,6 +829,12 @@ fn generate_zip_file(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::Config;
|
||||
use crate::services::mail::Mailer;
|
||||
use crate::services::settings;
|
||||
use sqlx::postgres::PgPoolOptions;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{Barrier, Semaphore};
|
||||
|
||||
#[test]
|
||||
fn content_disposition_supports_unicode_names() {
|
||||
@@ -525,4 +863,321 @@ mod tests {
|
||||
assert_eq!(output_file_name("photo.png", "webp"), "photo.webp");
|
||||
assert_eq!(output_file_name("没有扩展名", "jpeg"), "没有扩展名.jpg");
|
||||
}
|
||||
|
||||
async fn build_zip_test_state(
|
||||
pool: sqlx::PgPool,
|
||||
database_url: String,
|
||||
redis_url: String,
|
||||
storage_path: String,
|
||||
) -> AppState {
|
||||
let mut config = Config::from_env().expect("load ZIP test config");
|
||||
config.database_url = database_url;
|
||||
config.redis_url = redis_url;
|
||||
config.storage_path = storage_path;
|
||||
config.zip_build_concurrency = 2;
|
||||
config.zip_max_entries = 200;
|
||||
config.zip_max_uncompressed_bytes = 2 * 1024 * 1024;
|
||||
config.mail_enabled = false;
|
||||
config.mail_log_links_when_disabled = false;
|
||||
let redis = redis::Client::open(config.redis_url.clone())
|
||||
.expect("create ZIP test Redis client")
|
||||
.get_connection_manager()
|
||||
.await
|
||||
.expect("connect ZIP test Redis");
|
||||
AppState {
|
||||
mailer: Arc::new(Mailer::new(&config).expect("create disabled ZIP test mailer")),
|
||||
image_processing_semaphore: Arc::new(Semaphore::new(2)),
|
||||
zip_build_semaphore: Arc::new(Semaphore::new(2)),
|
||||
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
|
||||
storage_cache: storage::StorageCache::new(),
|
||||
config,
|
||||
db: pool,
|
||||
redis,
|
||||
}
|
||||
}
|
||||
|
||||
async fn insert_zip_task(
|
||||
pool: &sqlx::PgPool,
|
||||
task_id: Uuid,
|
||||
marker: &str,
|
||||
input_path: &std::path::Path,
|
||||
recorded_size: i64,
|
||||
) {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO tasks (
|
||||
id, session_id, status, total_files, completed_files,
|
||||
total_original_size, total_compressed_size,
|
||||
started_at, completed_at, expires_at, retention_hours
|
||||
) VALUES (
|
||||
$1, $2, 'completed', 1, 1,
|
||||
$3, $3, NOW(), NOW(), NOW() + INTERVAL '1 day', 24
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(format!("zip-session-{marker}"))
|
||||
.bind(recorded_size)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("insert ZIP test task");
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO task_files (
|
||||
id, task_id, original_name, original_format, output_format,
|
||||
original_size, compressed_size, saved_percent,
|
||||
storage_path, storage_backend, storage_key,
|
||||
status, completed_at
|
||||
) VALUES (
|
||||
$1, $2, $3, 'png', 'png',
|
||||
$4, $4, 0,
|
||||
$5, 'local', $5,
|
||||
'completed', NOW()
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(Uuid::new_v4())
|
||||
.bind(task_id)
|
||||
.bind(format!("{marker}.png"))
|
||||
.bind(recorded_size)
|
||||
.bind(input_path.to_string_lossy().to_string())
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("insert ZIP test file");
|
||||
}
|
||||
|
||||
async fn configure_test_s3(state: &AppState, marker: &str) -> Option<Uuid> {
|
||||
let endpoint = std::env::var("IMAGEFORGE_TEST_S3_ENDPOINT").ok()?;
|
||||
let bucket = std::env::var("IMAGEFORGE_TEST_S3_BUCKET").ok()?;
|
||||
let access_key = std::env::var("IMAGEFORGE_TEST_S3_ACCESS_KEY").ok()?;
|
||||
let secret_key = std::env::var("IMAGEFORGE_TEST_S3_SECRET_KEY").ok()?;
|
||||
let endpoint_id = Uuid::new_v4();
|
||||
let encrypted_access =
|
||||
settings::encrypt_secret(state, &access_key).expect("encrypt ZIP test S3 access key");
|
||||
let encrypted_secret =
|
||||
settings::encrypt_secret(state, &secret_key).expect("encrypt ZIP test S3 secret key");
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO storage_endpoints (
|
||||
id, name, internal_endpoint, public_endpoint,
|
||||
bucket, region,
|
||||
access_key_encrypted, secret_key_encrypted, access_key_hint,
|
||||
force_path_style, is_active
|
||||
) VALUES (
|
||||
$1, $2, $3, $3,
|
||||
$4, 'us-east-1',
|
||||
$5, $6, 'test',
|
||||
true, true
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(endpoint_id)
|
||||
.bind(format!("zip-test-{marker}"))
|
||||
.bind(endpoint)
|
||||
.bind(bucket)
|
||||
.bind(encrypted_access)
|
||||
.bind(encrypted_secret)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.expect("insert ZIP test S3 endpoint");
|
||||
Some(endpoint_id)
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
#[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL; optional IMAGEFORGE_TEST_S3_* uses MinIO"]
|
||||
async fn zip_build_is_single_flight_bounded_and_fenced() {
|
||||
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
|
||||
.expect("IMAGEFORGE_TEST_DATABASE_URL must be set");
|
||||
assert!(
|
||||
database_url.to_ascii_lowercase().contains("test"),
|
||||
"refusing to run destructive integration test outside a test database"
|
||||
);
|
||||
let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL")
|
||||
.expect("IMAGEFORGE_TEST_REDIS_URL must be set");
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(32)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("connect ZIP test database");
|
||||
sqlx::migrate!().run(&pool).await.expect("run migrations");
|
||||
let marker = Uuid::new_v4().simple().to_string();
|
||||
let storage_root = std::env::temp_dir().join(format!("imageforge-zip-test-{marker}"));
|
||||
tokio::fs::create_dir_all(&storage_root)
|
||||
.await
|
||||
.expect("create ZIP test storage root");
|
||||
let state = build_zip_test_state(
|
||||
pool.clone(),
|
||||
database_url,
|
||||
redis_url,
|
||||
storage_root.to_string_lossy().to_string(),
|
||||
)
|
||||
.await;
|
||||
let endpoint_id = configure_test_s3(&state, &marker).await;
|
||||
|
||||
let input_path = storage_root.join("single-flight-input.png");
|
||||
tokio::fs::write(&input_path, b"single-flight-payload")
|
||||
.await
|
||||
.expect("write ZIP input");
|
||||
let task_id = Uuid::new_v4();
|
||||
insert_zip_task(&pool, task_id, &marker, &input_path, 21).await;
|
||||
let barrier = Arc::new(Barrier::new(20));
|
||||
let mut joins = Vec::new();
|
||||
for _ in 0..20 {
|
||||
let state = state.clone();
|
||||
let barrier = barrier.clone();
|
||||
joins.push(tokio::spawn(async move {
|
||||
barrier.wait().await;
|
||||
resolve_task_zip(&state, task_id, 24).await
|
||||
}));
|
||||
}
|
||||
let mut locators = Vec::new();
|
||||
for join in joins {
|
||||
locators.push(
|
||||
join.await
|
||||
.expect("join concurrent ZIP request")
|
||||
.expect("resolve concurrent ZIP request"),
|
||||
);
|
||||
}
|
||||
assert!(locators
|
||||
.iter()
|
||||
.all(|locator| locator.key == locators[0].key));
|
||||
assert!(locators[0]
|
||||
.key
|
||||
.replace('\\', "/")
|
||||
.contains(&format!("/attempts/{task_id}/")));
|
||||
if endpoint_id.is_some() {
|
||||
assert_eq!(locators[0].backend, "s3");
|
||||
}
|
||||
let attempts: i64 = sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("query ZIP build attempts");
|
||||
assert_eq!(attempts, 1, "concurrent ZIP requests built more than once");
|
||||
|
||||
let over_budget_task = Uuid::new_v4();
|
||||
let missing_path = storage_root.join("must-not-be-downloaded.png");
|
||||
insert_zip_task(
|
||||
&pool,
|
||||
over_budget_task,
|
||||
&format!("{marker}-over-budget"),
|
||||
&missing_path,
|
||||
state.config.zip_max_uncompressed_bytes as i64 + 1,
|
||||
)
|
||||
.await;
|
||||
let over_budget = resolve_task_zip(&state, over_budget_task, 24)
|
||||
.await
|
||||
.expect_err("over-budget ZIP reached the download phase");
|
||||
assert_eq!(over_budget.code, ErrorCode::FileTooLarge);
|
||||
let over_budget_attempts: i64 =
|
||||
sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1")
|
||||
.bind(over_budget_task)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("query over-budget ZIP attempts");
|
||||
assert_eq!(over_budget_attempts, 0);
|
||||
|
||||
let takeover_task = Uuid::new_v4();
|
||||
insert_zip_task(
|
||||
&pool,
|
||||
takeover_task,
|
||||
&format!("{marker}-takeover"),
|
||||
&input_path,
|
||||
21,
|
||||
)
|
||||
.await;
|
||||
let first_token = match claim_zip_build(&state, takeover_task)
|
||||
.await
|
||||
.expect("claim simulated failing ZIP builder")
|
||||
{
|
||||
ZipBuildClaim::Acquired { token } => token,
|
||||
other => panic!("unexpected initial ZIP claim: {other:?}"),
|
||||
};
|
||||
let waiter_state = state.clone();
|
||||
let waiter =
|
||||
tokio::spawn(async move { resolve_task_zip(&waiter_state, takeover_task, 24).await });
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
release_zip_build(&state, takeover_task, first_token).await;
|
||||
let takeover_locator = waiter
|
||||
.await
|
||||
.expect("join ZIP takeover waiter")
|
||||
.expect("waiter safely took over ZIP build");
|
||||
let takeover_attempts: i64 =
|
||||
sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1")
|
||||
.bind(takeover_task)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("query takeover attempts");
|
||||
assert_eq!(takeover_attempts, 2);
|
||||
|
||||
let deleted_task = Uuid::new_v4();
|
||||
insert_zip_task(
|
||||
&pool,
|
||||
deleted_task,
|
||||
&format!("{marker}-deleted"),
|
||||
&input_path,
|
||||
21,
|
||||
)
|
||||
.await;
|
||||
let deleted_token = match claim_zip_build(&state, deleted_task)
|
||||
.await
|
||||
.expect("claim deleted-task ZIP builder")
|
||||
{
|
||||
ZipBuildClaim::Acquired { token } => token,
|
||||
other => panic!("unexpected deleted-task ZIP claim: {other:?}"),
|
||||
};
|
||||
let unpublished_path = storage_root.join("unpublished.zip");
|
||||
tokio::fs::write(&unpublished_path, b"unpublished-zip")
|
||||
.await
|
||||
.expect("write unpublished ZIP fixture");
|
||||
let unpublished = storage::store_file(
|
||||
&state,
|
||||
&storage::archive_attempt_key(24, deleted_task, deleted_token),
|
||||
&unpublished_path,
|
||||
"application/zip",
|
||||
)
|
||||
.await
|
||||
.expect("store unpublished ZIP attempt");
|
||||
sqlx::query("DELETE FROM tasks WHERE id = $1")
|
||||
.bind(deleted_task)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("delete task before ZIP publish");
|
||||
let publish_error =
|
||||
publish_zip_attempt(&state, deleted_task, deleted_token, unpublished.clone())
|
||||
.await
|
||||
.expect_err("published ZIP after task deletion");
|
||||
assert_eq!(publish_error.code, ErrorCode::StorageUnavailable);
|
||||
let orphan_read = storage::read_bytes(
|
||||
&state,
|
||||
&storage::ObjectLocator {
|
||||
backend: unpublished.backend,
|
||||
endpoint_id: unpublished.endpoint_id,
|
||||
key: unpublished.key,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
assert!(orphan_read.is_err(), "unpublished ZIP object was orphaned");
|
||||
|
||||
for locator in [&locators[0], &takeover_locator] {
|
||||
storage::delete_object(&state, locator)
|
||||
.await
|
||||
.expect("delete published ZIP test object");
|
||||
}
|
||||
sqlx::query("DELETE FROM tasks WHERE id = ANY($1)")
|
||||
.bind(&[task_id, over_budget_task, takeover_task][..])
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("delete ZIP test tasks");
|
||||
if let Some(endpoint_id) = endpoint_id {
|
||||
sqlx::query("DELETE FROM storage_endpoints WHERE id = $1")
|
||||
.bind(endpoint_id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("delete ZIP test S3 endpoint");
|
||||
}
|
||||
tokio::fs::remove_dir_all(&storage_root)
|
||||
.await
|
||||
.expect("remove ZIP test storage root");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1187,6 +1187,7 @@ mod tests {
|
||||
let state = AppState {
|
||||
mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")),
|
||||
image_processing_semaphore: Arc::new(Semaphore::new(2)),
|
||||
zip_build_semaphore: Arc::new(Semaphore::new(2)),
|
||||
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
|
||||
storage_cache: crate::services::storage::StorageCache::new(),
|
||||
config,
|
||||
|
||||
@@ -1150,10 +1150,13 @@ async fn resolve_invoice(
|
||||
.get("period_end")
|
||||
.and_then(|value| value.as_i64())
|
||||
.and_then(|timestamp| Utc.timestamp_opt(timestamp, 0).single());
|
||||
let paid_at = object
|
||||
let paid_at = (status == "paid").then(|| {
|
||||
object
|
||||
.pointer("/status_transitions/paid_at")
|
||||
.and_then(|value| value.as_i64())
|
||||
.and_then(|timestamp| Utc.timestamp_opt(timestamp, 0).single());
|
||||
.and_then(|timestamp| Utc.timestamp_opt(timestamp, 0).single())
|
||||
});
|
||||
let paid_at = paid_at.flatten();
|
||||
|
||||
Ok(ResolvedInvoice {
|
||||
provider_invoice_id: provider_invoice_id.to_string(),
|
||||
@@ -1200,7 +1203,11 @@ async fn write_invoice(
|
||||
pdf_url = COALESCE(EXCLUDED.pdf_url, invoices.pdf_url),
|
||||
period_start = COALESCE(EXCLUDED.period_start, invoices.period_start),
|
||||
period_end = COALESCE(EXCLUDED.period_end, invoices.period_end),
|
||||
paid_at = COALESCE(EXCLUDED.paid_at, invoices.paid_at)
|
||||
paid_at = CASE
|
||||
WHEN EXCLUDED.status = 'paid'
|
||||
THEN COALESCE(EXCLUDED.paid_at, invoices.paid_at)
|
||||
ELSE NULL
|
||||
END
|
||||
"#,
|
||||
)
|
||||
.bind(invoice.user_id)
|
||||
@@ -1338,6 +1345,7 @@ mod tests {
|
||||
AppState {
|
||||
mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")),
|
||||
image_processing_semaphore: Arc::new(Semaphore::new(2)),
|
||||
zip_build_semaphore: Arc::new(Semaphore::new(2)),
|
||||
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
|
||||
storage_cache: crate::services::storage::StorageCache::new(),
|
||||
config,
|
||||
@@ -2215,4 +2223,177 @@ mod tests {
|
||||
);
|
||||
assert!(index_exists.is_none(), "unique index was partially applied");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL with CREATE DATABASE"]
|
||||
async fn existing_paid_invoice_reconciles_before_accepting_a_later_event() {
|
||||
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
|
||||
.expect("IMAGEFORGE_TEST_DATABASE_URL must be set");
|
||||
assert!(
|
||||
database_url.to_ascii_lowercase().contains("test"),
|
||||
"refusing to run destructive integration test outside a test database"
|
||||
);
|
||||
let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL")
|
||||
.expect("IMAGEFORGE_TEST_REDIS_URL must be set");
|
||||
let mut admin_url = url::Url::parse(&database_url).expect("parse test database URL");
|
||||
admin_url.set_path("/postgres");
|
||||
let admin_pool = PgPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(admin_url.as_str())
|
||||
.await
|
||||
.expect("connect PostgreSQL admin database");
|
||||
let marker = Uuid::new_v4().simple().to_string();
|
||||
let child_database = format!("imageforge_test_invoice_seed_{marker}");
|
||||
sqlx::query(&format!("CREATE DATABASE {child_database}"))
|
||||
.execute(&admin_pool)
|
||||
.await
|
||||
.expect("create invoice seed test database");
|
||||
|
||||
let mut child_url = url::Url::parse(&database_url).expect("parse child database URL");
|
||||
child_url.set_path(&format!("/{child_database}"));
|
||||
let child_pool = PgPoolOptions::new()
|
||||
.max_connections(8)
|
||||
.connect(child_url.as_str())
|
||||
.await
|
||||
.expect("connect invoice seed test database");
|
||||
let all_migrations = sqlx::migrate!();
|
||||
let through_019 = sqlx::migrate::Migrator {
|
||||
migrations: std::borrow::Cow::Owned(
|
||||
all_migrations
|
||||
.iter()
|
||||
.filter(|migration| migration.version <= 19)
|
||||
.cloned()
|
||||
.collect(),
|
||||
),
|
||||
ignore_missing: false,
|
||||
locking: true,
|
||||
no_tx: false,
|
||||
};
|
||||
through_019
|
||||
.run(&child_pool)
|
||||
.await
|
||||
.expect("run migrations through 019");
|
||||
|
||||
let user_id = Uuid::new_v4();
|
||||
let customer_id = format!("cus_{marker}_existing_invoice");
|
||||
let invoice_id = format!("inv_{marker}_existing_paid");
|
||||
let invoice_number = format!("INV-{marker}-EXISTING");
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO users (
|
||||
id, email, username, password_hash, billing_customer_id
|
||||
) VALUES ($1, $2, $3, 'test-only', $4)
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(format!("invoice-seed-{marker}@example.test"))
|
||||
.bind(format!("invoice_seed_{marker}"))
|
||||
.bind(&customer_id)
|
||||
.execute(&child_pool)
|
||||
.await
|
||||
.expect("insert existing invoice user");
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO invoices (
|
||||
user_id, invoice_number, status, provider, provider_invoice_id, paid_at
|
||||
) VALUES ($1, $2, 'paid', 'stripe', $3, to_timestamp($4))
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(&invoice_number)
|
||||
.bind(&invoice_id)
|
||||
.bind(1_700_020_000_i64)
|
||||
.execute(&child_pool)
|
||||
.await
|
||||
.expect("insert existing paid invoice");
|
||||
|
||||
let pre_migration_watermarks: i64 = sqlx::query_scalar(
|
||||
"SELECT COUNT(*) FROM provider_object_event_watermarks WHERE object_type = 'invoice' AND provider_object_id = $1",
|
||||
)
|
||||
.bind(&invoice_id)
|
||||
.fetch_one(&child_pool)
|
||||
.await
|
||||
.expect("count pre-migration invoice watermarks");
|
||||
assert_eq!(pre_migration_watermarks, 0);
|
||||
|
||||
all_migrations
|
||||
.run(&child_pool)
|
||||
.await
|
||||
.expect("run invoice seed correction migration");
|
||||
let seeded: (bool, String) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT requires_reconciliation, last_event_id
|
||||
FROM provider_object_event_watermarks
|
||||
WHERE provider = 'stripe'
|
||||
AND object_type = 'invoice'
|
||||
AND provider_object_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&invoice_id)
|
||||
.fetch_one(&child_pool)
|
||||
.await
|
||||
.expect("query seeded invoice watermark");
|
||||
assert!(seeded.0);
|
||||
assert_eq!(seeded.1, "reconcile:migration:020");
|
||||
|
||||
let (stripe_base_url, stripe_mock, stripe_mock_task) = spawn_stripe_snapshot_mock().await;
|
||||
let authoritative = invoice_event(
|
||||
&format!("evt_{marker}_snapshot"),
|
||||
"invoice.paid",
|
||||
1_700_020_000,
|
||||
&invoice_id,
|
||||
&customer_id,
|
||||
Some(&invoice_number),
|
||||
);
|
||||
stripe_mock
|
||||
.objects
|
||||
.write()
|
||||
.await
|
||||
.insert(invoice_id.clone(), authoritative.data.object);
|
||||
let state = build_test_state(
|
||||
child_pool.clone(),
|
||||
child_url.to_string(),
|
||||
redis_url,
|
||||
stripe_base_url,
|
||||
)
|
||||
.await;
|
||||
let stale_failure = invoice_event(
|
||||
&format!("evt_{marker}_stale_failure"),
|
||||
"invoice.payment_failed",
|
||||
1_700_010_000,
|
||||
&invoice_id,
|
||||
&customer_id,
|
||||
Some(&invoice_number),
|
||||
);
|
||||
apply_test_event(&state, &stale_failure)
|
||||
.await
|
||||
.expect("reconcile existing invoice before applying stale event");
|
||||
assert_invoice_once(&child_pool, &invoice_id, "paid", &invoice_number).await;
|
||||
let watermark: (i16, bool, String) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT last_event_rank, requires_reconciliation, last_event_id
|
||||
FROM provider_object_event_watermarks
|
||||
WHERE provider = 'stripe'
|
||||
AND object_type = 'invoice'
|
||||
AND provider_object_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&invoice_id)
|
||||
.fetch_one(&child_pool)
|
||||
.await
|
||||
.expect("query reconciled invoice watermark");
|
||||
assert_eq!(watermark.0, 100);
|
||||
assert!(!watermark.1);
|
||||
assert!(watermark.2.starts_with("snapshot:"));
|
||||
assert_eq!(stripe_mock.calls.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(state);
|
||||
stripe_mock_task.abort();
|
||||
child_pool.close().await;
|
||||
sqlx::query(&format!("DROP DATABASE {child_database} WITH (FORCE)"))
|
||||
.execute(&admin_pool)
|
||||
.await
|
||||
.expect("drop invoice seed test database");
|
||||
admin_pool.close().await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,6 +17,9 @@ pub struct Config {
|
||||
pub worker_task_concurrency: u32,
|
||||
pub worker_concurrency: u32,
|
||||
pub image_processing_concurrency: u32,
|
||||
pub zip_build_concurrency: u32,
|
||||
pub zip_max_entries: u32,
|
||||
pub zip_max_uncompressed_bytes: u64,
|
||||
|
||||
pub jwt_secret: String,
|
||||
pub jwt_expiry_hours: i64,
|
||||
@@ -81,6 +84,15 @@ impl Config {
|
||||
.map(|v| v.get() as u32)
|
||||
.unwrap_or(4)
|
||||
});
|
||||
let zip_build_concurrency = env_u32("ZIP_BUILD_CONCURRENCY")
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(2);
|
||||
let zip_max_entries = env_u32("ZIP_MAX_ENTRIES")
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(200);
|
||||
let zip_max_uncompressed_bytes = env_u64("ZIP_MAX_UNCOMPRESSED_BYTES")
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(2 * 1024 * 1024 * 1024);
|
||||
|
||||
let jwt_secret = env_string("JWT_SECRET")
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 JWT_SECRET"))?;
|
||||
@@ -140,6 +152,9 @@ impl Config {
|
||||
worker_task_concurrency,
|
||||
worker_concurrency,
|
||||
image_processing_concurrency,
|
||||
zip_build_concurrency,
|
||||
zip_max_entries,
|
||||
zip_max_uncompressed_bytes,
|
||||
jwt_secret,
|
||||
jwt_expiry_hours,
|
||||
api_key_pepper,
|
||||
|
||||
@@ -36,6 +36,9 @@ async fn main() -> Result<(), AppError> {
|
||||
let image_processing_semaphore = std::sync::Arc::new(tokio::sync::Semaphore::new(
|
||||
config.image_processing_concurrency as usize,
|
||||
));
|
||||
let zip_build_semaphore = std::sync::Arc::new(tokio::sync::Semaphore::new(
|
||||
config.zip_build_concurrency as usize,
|
||||
));
|
||||
|
||||
let state = AppState {
|
||||
config,
|
||||
@@ -43,6 +46,7 @@ async fn main() -> Result<(), AppError> {
|
||||
redis,
|
||||
mailer: std::sync::Arc::new(mailer),
|
||||
image_processing_semaphore,
|
||||
zip_build_semaphore,
|
||||
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
|
||||
storage_cache: crate::services::storage::StorageCache::new(),
|
||||
};
|
||||
|
||||
@@ -320,55 +320,257 @@ pub async fn reserve_anonymous_units(
|
||||
Ok(date)
|
||||
}
|
||||
|
||||
pub async fn refund_anonymous_units(
|
||||
pub async fn reserve_anonymous_single_unit(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
session_id: &str,
|
||||
ip: IpAddr,
|
||||
units: u32,
|
||||
) -> Result<(), AppError> {
|
||||
refund_anonymous_units_for_date(state, session_id, ip, utc8_date(), units).await
|
||||
) -> Result<NaiveDate, AppError> {
|
||||
reserve_anonymous_single_unit_for_date(state, task_id, session_id, ip, utc8_date()).await
|
||||
}
|
||||
|
||||
async fn refund_anonymous_units_for_date(
|
||||
async fn reserve_anonymous_single_unit_for_date(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
session_id: &str,
|
||||
ip: IpAddr,
|
||||
date: NaiveDate,
|
||||
units: u32,
|
||||
) -> Result<(), AppError> {
|
||||
if units == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
) -> Result<NaiveDate, AppError> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO anonymous_single_reservations (
|
||||
task_id, session_id, client_ip, quota_date, units
|
||||
) VALUES ($1, $2, $3::inet, $4, 1)
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(session_id)
|
||||
.bind(ip.to_string())
|
||||
.bind(date)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建匿名单文件预留失败").with_source(err))?;
|
||||
|
||||
let limit = crate::services::settings::runtime_policy(state)
|
||||
.await?
|
||||
.rate_limits
|
||||
.anonymous_units_per_day as i64;
|
||||
let session_key = anonymous_session_key(session_id, date);
|
||||
let ip_key = anonymous_ip_key(ip, date);
|
||||
let reservation_key = anonymous_single_reservation_key(task_id);
|
||||
let mut conn = state.redis.clone();
|
||||
let script = redis::Script::new(
|
||||
r#"
|
||||
local limit = tonumber(ARGV[1])
|
||||
local ttl = tonumber(ARGV[2])
|
||||
|
||||
if redis.call('EXISTS', KEYS[3]) == 1 then
|
||||
return tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
end
|
||||
|
||||
local session_value = tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
local ip_value = tonumber(redis.call('GET', KEYS[2]) or '0')
|
||||
if session_value + 1 > limit or ip_value + 1 > limit then
|
||||
return -1
|
||||
end
|
||||
|
||||
session_value = redis.call('INCRBY', KEYS[1], 1)
|
||||
ip_value = redis.call('INCRBY', KEYS[2], 1)
|
||||
if session_value == 1 then redis.call('EXPIRE', KEYS[1], ttl) end
|
||||
if ip_value == 1 then redis.call('EXPIRE', KEYS[2], ttl) end
|
||||
redis.call('SET', KEYS[3], '1', 'EX', ttl)
|
||||
return session_value
|
||||
"#,
|
||||
);
|
||||
let reserved: i64 = script
|
||||
.key(session_key)
|
||||
.key(ip_key)
|
||||
.key(reservation_key)
|
||||
.arg(limit)
|
||||
.arg(48 * 60 * 60)
|
||||
.invoke_async(&mut conn)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
// Keep the durable pending row: the script may have committed even
|
||||
// if the response was lost, and maintenance can safely reconcile it.
|
||||
AppError::new(ErrorCode::Internal, "匿名配额检查失败").with_source(err)
|
||||
})?;
|
||||
if reserved < 0 {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE anonymous_single_reservations
|
||||
SET status = 'refunded', settled_at = NOW(), updated_at = NOW()
|
||||
WHERE task_id = $1 AND status = 'pending'
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "关闭匿名单文件预留失败").with_source(err)
|
||||
})?;
|
||||
return Err(AppError::new(
|
||||
ErrorCode::QuotaExceeded,
|
||||
format!("匿名试用次数已用完(每日 {limit} 次)"),
|
||||
));
|
||||
}
|
||||
Ok(date)
|
||||
}
|
||||
|
||||
pub async fn mark_anonymous_single_result(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
task_id: Uuid,
|
||||
charged: bool,
|
||||
) -> Result<(), AppError> {
|
||||
let status = if charged { "charged" } else { "refund_pending" };
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE anonymous_single_reservations
|
||||
SET status = $2,
|
||||
settled_at = CASE WHEN $2 = 'charged' THEN NOW() ELSE settled_at END,
|
||||
updated_at = NOW()
|
||||
WHERE task_id = $1 AND status = 'pending'
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(status)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新匿名单文件预留失败").with_source(err))?;
|
||||
if updated.rows_affected() != 1 {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::StorageUnavailable,
|
||||
"匿名单文件预留已失效,请重试",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn finalize_anonymous_single_reservation(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
charged: bool,
|
||||
) -> Result<(), AppError> {
|
||||
if charged {
|
||||
let mut conn = state.redis.clone();
|
||||
let _: i64 = redis::cmd("DEL")
|
||||
.arg(anonymous_single_reservation_key(task_id))
|
||||
.query_async(&mut conn)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "完成匿名单文件计费失败").with_source(err)
|
||||
})?;
|
||||
return Ok(());
|
||||
}
|
||||
refund_anonymous_single_reservation(state, task_id).await
|
||||
}
|
||||
|
||||
pub async fn refund_anonymous_single_reservation(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
) -> Result<(), AppError> {
|
||||
let row: Option<(String, String, NaiveDate, i32)> = sqlx::query_as(
|
||||
r#"
|
||||
SELECT session_id, host(client_ip), quota_date, units
|
||||
FROM anonymous_single_reservations
|
||||
WHERE task_id = $1
|
||||
AND status IN ('pending', 'refund_pending', 'refunded')
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询匿名单文件预留失败").with_source(err))?;
|
||||
let Some((session_id, ip, date, units)) = row else {
|
||||
return Ok(());
|
||||
};
|
||||
let ip: IpAddr = ip.parse().map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "匿名单文件预留 IP 无效").with_source(err)
|
||||
})?;
|
||||
let units = u32::try_from(units).unwrap_or(0);
|
||||
let session_key = anonymous_session_key(&session_id, date);
|
||||
let ip_key = anonymous_ip_key(ip, date);
|
||||
let reservation_key = anonymous_single_reservation_key(task_id);
|
||||
let refund_key = format!("anon_quota_refund:{task_id}");
|
||||
let mut conn = state.redis.clone();
|
||||
let script = redis::Script::new(
|
||||
r#"
|
||||
local dec = tonumber(ARGV[1])
|
||||
local ttl = tonumber(ARGV[2])
|
||||
if redis.call('EXISTS', KEYS[4]) == 1 then return 0 end
|
||||
|
||||
if redis.call('EXISTS', KEYS[3]) == 1 then
|
||||
local function refund(key)
|
||||
local current = tonumber(redis.call('GET', key) or '0')
|
||||
if current <= 0 then return 0 end
|
||||
return redis.call('DECRBY', key, math.min(current, dec))
|
||||
end
|
||||
|
||||
refund(KEYS[1])
|
||||
refund(KEYS[2])
|
||||
redis.call('DEL', KEYS[3])
|
||||
end
|
||||
redis.call('SET', KEYS[4], '1', 'EX', ttl)
|
||||
return 1
|
||||
"#,
|
||||
);
|
||||
|
||||
let _: i64 = script
|
||||
.key(session_key)
|
||||
.key(ip_key)
|
||||
.key(reservation_key)
|
||||
.key(refund_key)
|
||||
.arg(units as i64)
|
||||
.arg(48 * 60 * 60)
|
||||
.invoke_async(&mut conn)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "退还匿名配额失败").with_source(err))?;
|
||||
.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "退还匿名单文件配额失败").with_source(err)
|
||||
})?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE anonymous_single_reservations
|
||||
SET status = 'refunded', settled_at = NOW(), updated_at = NOW()
|
||||
WHERE task_id = $1 AND status IN ('pending', 'refund_pending', 'refunded')
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "记录匿名单文件退款失败").with_source(err))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn settle_stale_anonymous_single_reservations(
|
||||
state: &AppState,
|
||||
limit: i64,
|
||||
) -> Result<usize, AppError> {
|
||||
let task_ids: Vec<Uuid> = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT task_id
|
||||
FROM anonymous_single_reservations
|
||||
WHERE status = 'refund_pending'
|
||||
OR (status = 'pending' AND refund_after <= NOW())
|
||||
ORDER BY refund_after ASC
|
||||
LIMIT $1
|
||||
"#,
|
||||
)
|
||||
.bind(limit)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "查询待补偿匿名单文件预留失败").with_source(err)
|
||||
})?;
|
||||
let mut settled = 0;
|
||||
for task_id in task_ids {
|
||||
match refund_anonymous_single_reservation(state, task_id).await {
|
||||
Ok(()) => settled += 1,
|
||||
Err(err) => {
|
||||
tracing::warn!(task_id = %task_id, error = %err, "anonymous single reservation refund deferred")
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(settled)
|
||||
}
|
||||
|
||||
pub async fn refund_anonymous_reservation_once(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
@@ -541,6 +743,10 @@ fn anonymous_session_key(session_id: &str, date: NaiveDate) -> String {
|
||||
format!("anon_quota:{session_id}:{}", date.format("%Y-%m-%d"))
|
||||
}
|
||||
|
||||
fn anonymous_single_reservation_key(task_id: Uuid) -> String {
|
||||
format!("anon_quota_reservation:{task_id}")
|
||||
}
|
||||
|
||||
pub(crate) fn anonymous_ip_scope(ip: IpAddr) -> String {
|
||||
match ip {
|
||||
IpAddr::V4(ip) => ip.to_string(),
|
||||
@@ -605,6 +811,7 @@ mod tests {
|
||||
AppState {
|
||||
mailer: Arc::new(Mailer::new(&config).expect("create disabled quota test mailer")),
|
||||
image_processing_semaphore: Arc::new(Semaphore::new(2)),
|
||||
zip_build_semaphore: Arc::new(Semaphore::new(2)),
|
||||
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
|
||||
storage_cache: crate::services::storage::StorageCache::new(),
|
||||
config,
|
||||
@@ -979,4 +1186,230 @@ mod tests {
|
||||
.expect("delete quota settlement Redis keys");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
#[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL"]
|
||||
async fn anonymous_single_reservations_charge_actual_work_and_refund_original_date() {
|
||||
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
|
||||
.expect("IMAGEFORGE_TEST_DATABASE_URL must be set");
|
||||
assert!(
|
||||
database_url.to_ascii_lowercase().contains("test"),
|
||||
"refusing to run destructive integration test outside a test database"
|
||||
);
|
||||
let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL")
|
||||
.expect("IMAGEFORGE_TEST_REDIS_URL must be set");
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(16)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("connect anonymous single test database");
|
||||
sqlx::migrate!().run(&pool).await.expect("run migrations");
|
||||
let state = build_test_state(pool.clone(), database_url, redis_url).await;
|
||||
let marker = Uuid::new_v4().simple().to_string();
|
||||
let current_date = utc8_date();
|
||||
let previous_date = current_date.pred_opt().expect("previous quota date");
|
||||
let mut cleanup = Vec::new();
|
||||
|
||||
async fn mark_and_finalize(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
charged: bool,
|
||||
) -> Result<(), AppError> {
|
||||
let mut tx = state.db.begin().await.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "begin anonymous single test").with_source(err)
|
||||
})?;
|
||||
mark_anonymous_single_result(&mut tx, task_id, charged).await?;
|
||||
tx.commit().await.map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "commit anonymous single test").with_source(err)
|
||||
})?;
|
||||
finalize_anonymous_single_reservation(state, task_id, charged).await
|
||||
}
|
||||
|
||||
let passthrough_task = Uuid::new_v4();
|
||||
let passthrough_session = format!("single-passthrough-{marker}");
|
||||
let passthrough_ip: IpAddr = "198.51.100.31".parse().expect("parse passthrough IP");
|
||||
reserve_anonymous_single_unit_for_date(
|
||||
&state,
|
||||
passthrough_task,
|
||||
&passthrough_session,
|
||||
passthrough_ip,
|
||||
current_date,
|
||||
)
|
||||
.await
|
||||
.expect("reserve passthrough unit");
|
||||
let passthrough_charged = output_consumes_unit(Some(100), true, false, false, 100, 50);
|
||||
assert!(!passthrough_charged);
|
||||
mark_and_finalize(&state, passthrough_task, passthrough_charged)
|
||||
.await
|
||||
.expect("refund passthrough unit");
|
||||
assert_eq!(
|
||||
anonymous_quota_counts(&state, &passthrough_session, passthrough_ip, current_date)
|
||||
.await,
|
||||
(0, 0)
|
||||
);
|
||||
cleanup.push((
|
||||
passthrough_task,
|
||||
passthrough_session,
|
||||
passthrough_ip,
|
||||
current_date,
|
||||
));
|
||||
|
||||
let unchanged_task = Uuid::new_v4();
|
||||
let unchanged_session = format!("single-unchanged-{marker}");
|
||||
let unchanged_ip: IpAddr = "198.51.100.32".parse().expect("parse unchanged IP");
|
||||
reserve_anonymous_single_unit_for_date(
|
||||
&state,
|
||||
unchanged_task,
|
||||
&unchanged_session,
|
||||
unchanged_ip,
|
||||
current_date,
|
||||
)
|
||||
.await
|
||||
.expect("reserve unchanged-output unit");
|
||||
let unchanged_charged = output_consumes_unit(None, true, false, false, 100, 100);
|
||||
assert!(!unchanged_charged);
|
||||
mark_and_finalize(&state, unchanged_task, unchanged_charged)
|
||||
.await
|
||||
.expect("refund unchanged-output unit");
|
||||
assert_eq!(
|
||||
anonymous_quota_counts(&state, &unchanged_session, unchanged_ip, current_date).await,
|
||||
(0, 0)
|
||||
);
|
||||
cleanup.push((
|
||||
unchanged_task,
|
||||
unchanged_session,
|
||||
unchanged_ip,
|
||||
current_date,
|
||||
));
|
||||
|
||||
let compressed_task = Uuid::new_v4();
|
||||
let compressed_session = format!("single-compressed-{marker}");
|
||||
let compressed_ip: IpAddr = "198.51.100.33".parse().expect("parse compressed IP");
|
||||
reserve_anonymous_single_unit_for_date(
|
||||
&state,
|
||||
compressed_task,
|
||||
&compressed_session,
|
||||
compressed_ip,
|
||||
current_date,
|
||||
)
|
||||
.await
|
||||
.expect("reserve compressed unit");
|
||||
let compressed_charged = output_consumes_unit(None, true, false, false, 100, 50);
|
||||
assert!(compressed_charged);
|
||||
mark_and_finalize(&state, compressed_task, compressed_charged)
|
||||
.await
|
||||
.expect("finalize compressed unit");
|
||||
assert_eq!(
|
||||
anonymous_quota_counts(&state, &compressed_session, compressed_ip, current_date).await,
|
||||
(1, 1)
|
||||
);
|
||||
let charged_status: String = sqlx::query_scalar(
|
||||
"SELECT status FROM anonymous_single_reservations WHERE task_id = $1",
|
||||
)
|
||||
.bind(compressed_task)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("query charged reservation");
|
||||
assert_eq!(charged_status, "charged");
|
||||
cleanup.push((
|
||||
compressed_task,
|
||||
compressed_session,
|
||||
compressed_ip,
|
||||
current_date,
|
||||
));
|
||||
|
||||
let cross_day_task = Uuid::new_v4();
|
||||
let cross_day_session = format!("single-cross-day-{marker}");
|
||||
let cross_day_ip: IpAddr = "198.51.100.34".parse().expect("parse cross-day IP");
|
||||
reserve_anonymous_units(&state, &cross_day_session, cross_day_ip, 2)
|
||||
.await
|
||||
.expect("seed current-day quota");
|
||||
reserve_anonymous_single_unit_for_date(
|
||||
&state,
|
||||
cross_day_task,
|
||||
&cross_day_session,
|
||||
cross_day_ip,
|
||||
previous_date,
|
||||
)
|
||||
.await
|
||||
.expect("reserve previous-day unit");
|
||||
refund_anonymous_single_reservation(&state, cross_day_task)
|
||||
.await
|
||||
.expect("refund previous-day failure");
|
||||
assert_eq!(
|
||||
anonymous_quota_counts(&state, &cross_day_session, cross_day_ip, previous_date).await,
|
||||
(0, 0)
|
||||
);
|
||||
assert_eq!(
|
||||
anonymous_quota_counts(&state, &cross_day_session, cross_day_ip, current_date).await,
|
||||
(2, 2),
|
||||
"cross-day refund changed the current quota bucket"
|
||||
);
|
||||
cleanup.push((
|
||||
cross_day_task,
|
||||
cross_day_session.clone(),
|
||||
cross_day_ip,
|
||||
previous_date,
|
||||
));
|
||||
|
||||
let interrupted_task = Uuid::new_v4();
|
||||
let interrupted_session = format!("single-interrupted-{marker}");
|
||||
let interrupted_ip: IpAddr = "198.51.100.35".parse().expect("parse interrupted IP");
|
||||
reserve_anonymous_single_unit_for_date(
|
||||
&state,
|
||||
interrupted_task,
|
||||
&interrupted_session,
|
||||
interrupted_ip,
|
||||
current_date,
|
||||
)
|
||||
.await
|
||||
.expect("reserve interrupted unit");
|
||||
sqlx::query(
|
||||
"UPDATE anonymous_single_reservations SET refund_after = NOW() - INTERVAL '1 second' WHERE task_id = $1",
|
||||
)
|
||||
.bind(interrupted_task)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("expire interrupted reservation");
|
||||
assert_eq!(
|
||||
settle_stale_anonymous_single_reservations(&state, 10)
|
||||
.await
|
||||
.expect("settle interrupted reservation"),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
anonymous_quota_counts(&state, &interrupted_session, interrupted_ip, current_date)
|
||||
.await,
|
||||
(0, 0)
|
||||
);
|
||||
cleanup.push((
|
||||
interrupted_task,
|
||||
interrupted_session,
|
||||
interrupted_ip,
|
||||
current_date,
|
||||
));
|
||||
|
||||
let mut redis = state.redis.clone();
|
||||
for (task_id, session_id, ip, date) in cleanup {
|
||||
sqlx::query("DELETE FROM anonymous_single_reservations WHERE task_id = $1")
|
||||
.bind(task_id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("delete anonymous single reservation");
|
||||
let _: i64 = redis::cmd("DEL")
|
||||
.arg(anonymous_session_key(&session_id, date))
|
||||
.arg(anonymous_ip_key(ip, date))
|
||||
.arg(anonymous_single_reservation_key(task_id))
|
||||
.arg(format!("anon_quota_refund:{task_id}"))
|
||||
.query_async(&mut redis)
|
||||
.await
|
||||
.expect("delete anonymous single Redis keys");
|
||||
}
|
||||
let _: i64 = redis::cmd("DEL")
|
||||
.arg(anonymous_session_key(&cross_day_session, current_date))
|
||||
.arg(anonymous_ip_key(cross_day_ip, current_date))
|
||||
.query_async(&mut redis)
|
||||
.await
|
||||
.expect("delete cross-day current Redis keys");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -311,10 +311,10 @@ pub fn result_attempt_key(
|
||||
)
|
||||
}
|
||||
|
||||
pub fn archive_key(retention_hours: i64, task_id: Uuid) -> String {
|
||||
pub fn archive_attempt_key(retention_hours: i64, task_id: Uuid, token: Uuid) -> String {
|
||||
let now = Utc::now();
|
||||
format!(
|
||||
"archives/{}/{:04}/{:02}/{task_id}.zip",
|
||||
"archives/{}/{:04}/{:02}/attempts/{task_id}/{token}.zip",
|
||||
retention_prefix(retention_hours),
|
||||
now.year(),
|
||||
now.month()
|
||||
@@ -983,7 +983,7 @@ mod tests {
|
||||
let key = result_key(168, task_id, file_id, "webp");
|
||||
assert!(key.starts_with("results/7d/"));
|
||||
assert!(key.ends_with("/00000000-0000-0000-0000-000000000001.webp"));
|
||||
assert!(archive_key(360, task_id).starts_with("archives/15d/"));
|
||||
assert!(archive_attempt_key(360, task_id, Uuid::new_v4()).starts_with("archives/15d/"));
|
||||
let attempt_key = result_attempt_key(24, task_id, file_id, 2, 3, "avif");
|
||||
assert!(attempt_key.contains("-t2-f3.avif"));
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ pub struct AppState {
|
||||
pub redis: redis::aio::ConnectionManager,
|
||||
pub mailer: std::sync::Arc<Mailer>,
|
||||
pub image_processing_semaphore: std::sync::Arc<tokio::sync::Semaphore>,
|
||||
pub zip_build_semaphore: std::sync::Arc<tokio::sync::Semaphore>,
|
||||
pub runtime_policy_cache: RuntimePolicyCache,
|
||||
pub storage_cache: StorageCache,
|
||||
}
|
||||
|
||||
@@ -1715,6 +1715,7 @@ async fn charge_one_unit(
|
||||
}
|
||||
|
||||
async fn maintenance(state: &AppState) -> Result<(), AppError> {
|
||||
settle_stale_anonymous_single_reservations(state).await?;
|
||||
settle_finished_anonymous_reservations(state).await?;
|
||||
cleanup_expired_tasks(state).await?;
|
||||
cleanup_stale_zip_temp(state).await?;
|
||||
@@ -1722,6 +1723,19 @@ async fn maintenance(state: &AppState) -> Result<(), AppError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn settle_stale_anonymous_single_reservations(state: &AppState) -> Result<(), AppError> {
|
||||
for _ in 0..MAX_MAINTENANCE_BATCHES {
|
||||
let settled =
|
||||
quota::settle_stale_anonymous_single_reservations(state, MAINTENANCE_BATCH_SIZE)
|
||||
.await?;
|
||||
if settled < MAINTENANCE_BATCH_SIZE as usize {
|
||||
break;
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn settle_finished_anonymous_reservations(state: &AppState) -> Result<(), AppError> {
|
||||
for _ in 0..MAX_MAINTENANCE_BATCHES {
|
||||
let task_ids: Vec<Uuid> = sqlx::query_scalar(
|
||||
@@ -1816,6 +1830,16 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
let _ = sqlx::query(
|
||||
r#"
|
||||
DELETE FROM anonymous_single_reservations
|
||||
WHERE status IN ('charged', 'refunded')
|
||||
AND settled_at < NOW() - INTERVAL '7 days'
|
||||
"#,
|
||||
)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
let _ =
|
||||
sqlx::query("DELETE FROM webhook_events WHERE received_at < NOW() - INTERVAL '90 days'")
|
||||
.execute(&state.db)
|
||||
@@ -2026,6 +2050,7 @@ mod tests {
|
||||
let state = AppState {
|
||||
mailer: Arc::new(Mailer::new(&config).expect("create disabled test mailer")),
|
||||
image_processing_semaphore: Arc::new(Semaphore::new(2)),
|
||||
zip_build_semaphore: Arc::new(Semaphore::new(2)),
|
||||
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
|
||||
storage_cache: storage::StorageCache::new(),
|
||||
config,
|
||||
|
||||
Reference in New Issue
Block a user