Compare commits

...
12 Commits
Author SHA1 Message Date
ZyyandCodex 0fc25bfce9 chore(release): v3.26.0
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 22:32:21 +08:00
ZyyandCodex eed8c93ae2 refactor(store): consolidate unreleased database migrations
Merge migrations 0030 through 0032 after the v3.25.0 release baseline. Queue analytics bucket rebuilds once and verify upgrade preservation and repeated startup.

Co-authored-by: Codex <noreply@openai.com>
2026-10-13 22:13:26 +08:00
ZyyandCodex d865a00768 style(pricing): simplify model search placeholder
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 21:58:00 +08:00
ZyyandCodex a0671fb70b style(settings): refine privacy controls and pricing table
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 21:50:53 +08:00
ZyyandCodex fce543d4ce feat(settings): remove account request interval limits
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 21:50:53 +08:00
219888d5c2 perf(store): scale analytics and pagination for large request histories (#423)
Aggregate analytics in PostgreSQL, reconcile synchronous usage with pending source changes, and preserve exact pagination and forecast totals across large histories.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 21:29:25 +08:00
5b80340135 fix(privacy): avoid repeated full-request processing (#420)
Stage scoped body fragments and reuse context locations while preserving per-rule rollback. Account for encoded size incrementally and use bounded work plus cancellation instead of a fixed wall-clock cutoff.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 17:09:33 +08:00
b2dab66a80 fix(accounts): import account file entries independently (#411)
* fix(accounts): import valid account file entries independently

Co-authored-by: Codex <noreply@openai.com>

* test(openai): cover single-token account imports

* style(openai): format token import test assertion

* fix(accounts): preserve provider boundaries and bound import batches

Co-authored-by: Codex <noreply@openai.com>

* fix(ui): keep account dialogs within narrow viewports

Co-authored-by: Codex <noreply@openai.com>

* chore(deps): update UI v0.4.1 and latest plugin main

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
2026-10-13 12:16:17 +08:00
JiahaoandCodex f0cf8e7198 feat(usage): highlight latency by fixed thresholds (#412)
Color first-output and total latency independently using fixed thresholds.
Keep the existing layout and details popover without a side indicator.

Closes #407

Co-authored-by: Codex <noreply@openai.com>
2026-10-13 11:45:38 +08:00
a9a86071bf fix: preserve terminal failure and cleanup causes (#417)
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 11:31:10 +08:00
27cef052c9 fix(accounts): preserve quota forecast totals on same-window reversals (#414)
* fix(accounts): preserve quota forecast totals on same-window reversals

Share exhausted-window recovery evidence between account recovery and
forecast sampling. Keep recorded cycle usage on ordinary reversals and
restart totals only after the recovery evidence is satisfied.

Co-authored-by: Codex <noreply@openai.com>

* fix(accounts): separate forecast cycle boundaries from quota recovery

Keep same-window usage totals after quota reversals, including observations near exhaustion. Leave access recovery in the OpenAI provider and reset forecast totals only at an identified window or plan boundary.

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 11:20:04 +08:00
ef40088303 fix(deps): update quinn-proto to address security advisories (#415)
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-13 11:11:41 +08:00
173 changed files with 7076 additions and 2253 deletions
+2 -2
View File
@@ -3401,9 +3401,9 @@ dependencies = [
[[package]]
name = "quinn-proto"
version = "0.11.16"
version = "0.11.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560"
checksum = "a9746dbde176634f4f2f1faf2404e30a31b2bc1e9cafb5329c95d8177a18c9fc"
dependencies = [
"bytes",
"getrandom 0.4.3",
+42 -17
View File
@@ -491,11 +491,13 @@ impl BackupTask {
)
.await
{
let _ = store.delete_object(&config, &record.object_key).await;
self.fail_task(
self.fail_uploaded_task(
record,
error.code(),
format!("上传到对象存储失败:{}", error.message()),
&config,
BackupError::new(
error.code(),
format!("上传到对象存储失败:{}", error.message()),
),
)
.await?;
return Ok(());
@@ -533,16 +535,19 @@ impl BackupTask {
sha256_matches,
"备份远端对象校验不一致"
);
let _ = store.delete_object(&config, &record.object_key).await;
self.fail_task(
self.fail_uploaded_task(
record,
code::REMOTE_VERIFICATION_FAILED,
match (size_matches, sha256_matches) {
(false, true) => "远端对象大小不一致",
(true, false) => "远端对象 sha256 校验值不一致",
(false, false) => "远端对象大小和 sha256 校验值均不一致",
(true, true) => "远端对象校验状态异常",
},
&config,
BackupError::new(
code::REMOTE_VERIFICATION_FAILED,
match (size_matches, sha256_matches) {
(false, true) => "远端对象大小不一致",
(true, false) => "远端对象 sha256 校验值不一致",
(false, false) => "远端对象大小和 sha256 校验值均不一致",
(true, true) => "远端对象校验状态异常",
}
.to_owned(),
),
)
.await?;
}
@@ -556,11 +561,13 @@ impl BackupTask {
.await?;
}
Err(error) => {
let _ = store.delete_object(&config, &record.object_key).await;
self.fail_task(
self.fail_uploaded_task(
record,
error.code(),
format!("远端对象校验失败:{}", error.message()),
&config,
BackupError::new(
error.code(),
format!("远端对象校验失败:{}", error.message()),
),
)
.await?;
}
@@ -568,6 +575,24 @@ impl BackupTask {
Ok(())
}
/// 清理失败并入原终态,后续清理沿用备份记录的到期策略
async fn fail_uploaded_task(
&self,
record: &BackupRecord,
config: &BackupStorageConfig,
mut error: BackupError,
) -> Result<(), WorkerTaskError> {
if let Err(cleanup) = self
.object_store
.delete_object(config, &record.object_key)
.await
{
error =
error.with_cleanup_failure("DeleteObject", &cleanup, record.expires_at.is_some());
}
self.fail_task(record, error.code(), error.message()).await
}
/// 把活跃任务写入失败终态并清理暂存
async fn fail_task(
&self,
@@ -66,7 +66,7 @@ pub struct AccountGroupCapacity {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupUsage {
pub today_usd: DecimalAmount,
/// 当前 usage retention 窗口内、按请求发生时 routing group 快照归属的累计成本
/// 当前 usage retention 窗口内、当前组成员实际完成请求的累计成本
pub retained_total_usd: DecimalAmount,
}
@@ -474,6 +474,26 @@ impl BackupError {
pub fn message(&self) -> &str {
&self.message
}
/// 清理失败附加到原始终态,不覆盖上传失败或主动取消的分类
#[must_use]
pub fn with_cleanup_failure(
mut self,
stage: &'static str,
error: &Self,
retry_on_expiration: bool,
) -> Self {
let recovery = if retry_on_expiration {
"等待备份到期清理"
} else {
"未安排自动重试"
};
self.message.push_str(&format!(
";清理阶段 {stage} 失败({}),{recovery}",
error.code()
));
self
}
}
/// 构造稳定错误码常量
@@ -157,7 +157,6 @@ impl OtherRequestOutcome {
/// 用量记录过滤条件
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
pub struct UsageFilter {
pub routing_group_refs: Vec<String>,
pub client_api_key_ref: Option<String>,
pub request_id: Option<String>,
pub provider_account_ref: Option<String>,
@@ -211,7 +210,6 @@ impl DashboardQuery {
/// 运维错误过滤条件
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct OpsErrorFilter {
pub routing_group_refs: Vec<String>,
pub client_api_key_ref: Option<String>,
pub request_id: Option<String>,
pub provider_kind: Option<String>,
@@ -142,7 +142,7 @@ pub fn select_forecast_sample(
};
history.push(current);
// 重置边界不在查询起点时,从首个新观测重新计数,避免把旧周期累计带入总量
let mut cycle_baseline = interrupted.then(|| &history[0]);
let cycle_baseline = interrupted.then(|| &history[0]);
let mut anchors: Vec<&QuotaForecastPoint> = Vec::new();
let mut high_water: Option<&QuotaForecastPoint> = None;
for point in &history {
@@ -155,9 +155,9 @@ pub fn select_forecast_sample(
if previous.used_percent - point.used_percent <= 1.0 {
continue;
}
// 同窗口回落不能证明周期重置,只重建采样并保留已记录累计
sample.discontinuous = true;
anchors.clear();
cycle_baseline = Some(point);
}
}
high_water = Some(point);
@@ -51,7 +51,6 @@ pub struct RuntimeSettingsValues {
pub refresh_margin_seconds: u64,
pub refresh_concurrency: u32,
pub max_concurrent_per_account: u32,
pub request_interval_ms: u64,
pub max_waiting_per_key: u32,
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
@@ -212,10 +211,6 @@ impl RuntimeSettingsValues {
"refresh_margin_seconds",
),
(self.refresh_concurrency > 0, "refresh_concurrency"),
(
i64::try_from(self.request_interval_ms).is_ok(),
"request_interval_ms",
),
] {
if !valid {
return Err(field);
@@ -432,6 +432,9 @@ pub trait UsageBillingProjector: Send + Sync {
/// 用量、趋势、诊断与运维错误的只读能力
#[async_trait]
pub trait ObservabilityStore: Send + Sync {
/// 读取实现声明的总预算,缓存等待沿用同一配置,避免外层提前截断
fn query_timeout(&self) -> std::time::Duration;
/// 缓存必须与删除和后台投影提交的版本一致
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64>;
@@ -12,7 +12,6 @@ use tokio::{sync::Mutex, time::Instant};
const CAPACITY: usize = 64;
const TTL: Duration = Duration::from_secs(15);
const MAX_STALE: Duration = Duration::from_secs(60);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
type Slot<T> = Arc<Mutex<Option<Cached<T>>>>;
struct Cached<T> {
@@ -72,7 +71,7 @@ impl<K: Eq + Hash, T: AnalyticsValue> Cache<K, T> {
{
// 刷新失败后的同键等待也计入预算,避免串行重试把延迟放大为队列长度倍数
tokio::time::timeout(
REQUEST_TIMEOUT,
store.query_timeout(),
self.get_within_budget(key, store, ttl, load),
)
.await
@@ -174,12 +173,6 @@ fn with_age<T: AnalyticsValue>(mut value: T, stale: bool) -> T {
value
}
fn normalized(mut filter: UsageFilter) -> UsageFilter {
filter.routing_group_refs.sort_unstable();
filter.routing_group_refs.dedup();
filter
}
type UsageKey = (TimeRange, UsageFilter);
pub(crate) struct CachedObservabilityStore {
inner: Arc<dyn ObservabilityStore>,
@@ -202,6 +195,10 @@ impl CachedObservabilityStore {
#[async_trait]
impl ObservabilityStore for CachedObservabilityStore {
fn query_timeout(&self) -> Duration {
self.inner.query_timeout()
}
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64> {
self.inner.analytics_cache_generation().await
}
@@ -241,7 +238,6 @@ impl ObservabilityStore for CachedObservabilityStore {
filter: UsageFilter,
granularity: Granularity,
) -> AdminStoreResult<UsageInsightsObservation> {
let filter = normalized(filter);
self.insights
.get(
((range, filter.clone()), granularity),
@@ -259,7 +255,6 @@ impl ObservabilityStore for CachedObservabilityStore {
range: TimeRange,
filter: UsageFilter,
) -> AdminStoreResult<UsageOverview> {
let filter = normalized(filter);
self.summaries
.get((range, filter.clone()), self.inner.as_ref(), TTL, || {
self.inner.usage_summary(range, filter)
@@ -273,7 +268,6 @@ impl ObservabilityStore for CachedObservabilityStore {
dimension: DiagnosticDimension,
limit: u16,
) -> AdminStoreResult<DiagnosticsObservation> {
let filter = normalized(filter);
self.diagnostics
.get(
((range, filter.clone()), dimension, limit),
@@ -536,6 +536,9 @@ impl DatabaseDumpPort for FakeDumpPort {
pub(crate) struct FakeObjectStore {
pub(crate) objects: Mutex<std::collections::HashMap<String, Vec<u8>>>,
upload_error: Option<BackupError>,
head_error: Option<BackupError>,
delete_error: Option<BackupError>,
connection_error: Option<BackupError>,
}
impl FakeObjectStore {
@@ -543,6 +546,9 @@ impl FakeObjectStore {
Self {
objects: Mutex::new(std::collections::HashMap::new()),
upload_error: None,
head_error: None,
delete_error: None,
connection_error: None,
}
}
@@ -552,6 +558,21 @@ impl FakeObjectStore {
self
}
pub(crate) fn fail_head(mut self, code: &'static str) -> Self {
self.head_error = Some(BackupError::new(code, "HEAD failed".into()));
self
}
pub(crate) fn fail_delete(mut self, code: &'static str) -> Self {
self.delete_error = Some(BackupError::new(code, "DELETE failed".into()));
self
}
pub(crate) fn fail_connection(mut self, code: &'static str) -> Self {
self.connection_error = Some(BackupError::new(code, "DELETE failed".into()));
self
}
pub(crate) fn object(&self, key: &str) -> Option<Vec<u8>> {
self.objects.lock().expect("objects").get(key).cloned()
}
@@ -563,6 +584,14 @@ impl BackupObjectStorePort for FakeObjectStore {
&self,
_config: &BackupStorageConfig,
) -> Result<ConnectionTestResult, gateway_admin::model::backup::BackupError> {
if let Some(error) = &self.connection_error {
return Ok(ConnectionTestResult {
ok: false,
stage: "deleteObject",
code: Some(error.code()),
message: error.message().into(),
});
}
Ok(ConnectionTestResult {
ok: true,
stage: "deleteObject",
@@ -597,6 +626,9 @@ impl BackupObjectStorePort for FakeObjectStore {
_config: &BackupStorageConfig,
object_key: &str,
) -> Result<Option<BackupObjectMetadata>, gateway_admin::model::backup::BackupError> {
if let Some(error) = &self.head_error {
return Err(error.clone());
}
let Some(content) = self.object(object_key) else {
return Ok(None);
};
@@ -613,6 +645,9 @@ impl BackupObjectStorePort for FakeObjectStore {
_config: &BackupStorageConfig,
object_key: &str,
) -> Result<(), gateway_admin::model::backup::BackupError> {
if let Some(error) = &self.delete_error {
return Err(error.clone());
}
self.objects.lock().expect("objects").remove(object_key);
Ok(())
}
@@ -324,3 +324,58 @@ async fn completed_recovery_removes_staging_for_dumping_and_uploading_tasks() {
);
}
}
#[tokio::test]
async fn manual_backup_preserves_primary_and_final_delete_failures() {
for (upload_fails, expiring) in [(false, false), (true, false), (false, true)] {
let repository = Arc::new(FakeBackupRepository::new(configured_settings()));
let dump = Arc::new(FakeDumpPort::new());
let store = FakeObjectStore::new().fail_delete(code::S3_PERMISSION_DENIED);
let store = Arc::new(if upload_fails {
store.fail_upload(code::S3_UPLOAD_FAILED, "上传失败")
} else {
store.fail_head(code::REMOTE_VERIFICATION_FAILED)
});
let seed = BackupRecordSeed {
id: backup_id("cleanup-failure"),
trigger_kind: BackupTriggerKind::Manual,
scheduled_at: None,
object_key: "codex/cleanup-failure.dump".into(),
expires_at: expiring.then(|| Utc::now() + Duration::hours(1)),
};
repository.insert_backup_record(seed.clone()).await.unwrap();
let task = BackupTask::new(repository.clone(), dump.clone(), store.clone());
task.run_cycle(&CancellationToken::new()).await.unwrap();
let records = repository.all_records();
let record = &records[0];
assert_eq!(record.status, BackupStatus::Failed);
assert_eq!(
record.error_code.as_deref(),
Some(if upload_fails {
code::S3_UPLOAD_FAILED
} else {
code::REMOTE_VERIFICATION_FAILED
})
);
let message = record.error_message.as_deref().unwrap();
assert!(message.contains(if upload_fails {
"上传失败"
} else {
"HEAD failed"
}));
assert!(message.contains("DeleteObject"));
assert!(message.contains(code::S3_PERMISSION_DENIED));
assert!(message.contains(if expiring {
"等待备份到期清理"
} else {
"未安排自动重试"
}));
assert_eq!(store.object(&seed.object_key).is_some(), !upload_fails);
assert!(dump.inspect_staging(&seed.id).await.unwrap().is_none());
task.run_cycle(&CancellationToken::new()).await.unwrap();
assert_eq!(
repository.all_records()[0].error_message,
record.error_message
);
}
}
@@ -73,14 +73,14 @@ fn snapshot_at_a_block_boundary_keeps_only_three_recent_blocks() {
}
#[test]
fn large_percentage_reversal_restarts_sampling_without_crossing_old_consumption() {
fn large_percentage_reversal_restarts_only_sampling_without_crossing_old_consumption() {
let sample = select(
vec![point(1, 40.0, 100), point(2, 50.0, 200), point(3, 1.0, 210)],
point(4, 8.0, 280),
);
assert_eq!(sample.baseline_percent, 1.0);
assert_eq!(sample.usage.tokens, 70);
assert_eq!(sample.cycle_usage.tokens, 70);
assert_eq!(sample.cycle_usage.tokens, 280);
assert_eq!(sample.sampled_percent, 7.0);
assert!(!sample.discontinuous);
let waiting = select(
@@ -90,6 +90,30 @@ fn large_percentage_reversal_restarts_sampling_without_crossing_old_consumption(
assert!(waiting.discontinuous);
}
#[test]
fn same_window_reversal_preserves_all_recorded_cycle_usage() {
let current = point(3, 35.0, 140);
let recorded = current.usage.clone();
let sample = select(vec![point(1, 32.0, 100), point(2, 30.0, 120)], current);
assert_eq!(sample.cycle_usage, recorded);
assert_eq!(sample.usage.tokens, 20);
assert!((sample.usage.usd - 0.02).abs() < 1e-10);
assert!(!sample.discontinuous);
}
#[test]
fn small_reversals_after_full_usage_preserve_cycle_totals_and_cumulative_forecasts() {
let current = point(3, 99.5, 10_020);
let recorded = current.usage.clone();
let sample = select(
vec![point(1, 100.0, 10_000), point(2, 99.0, 10_010)],
current,
);
assert_eq!(sample.cycle_usage, recorded);
assert_eq!(sample.method, QuotaForecastMethod::Cumulative);
assert!(!sample.discontinuous);
}
#[test]
fn one_point_reversal_is_skipped_without_counting_an_extra_block() {
let sample = select(
@@ -128,17 +152,41 @@ fn counter_regressions_fail_closed_and_future_points_are_not_used() {
}
#[test]
fn exhausted_history_is_retained_to_detect_a_reset_with_an_unchanged_deadline() {
fn same_window_recovery_restarts_sampling_but_preserves_recorded_cycle_usage() {
let sample = select(
vec![point(1, 100.0, 10_000), point(2, 0.0, 10_000)],
point(3, 10.0, 10_100),
);
assert_eq!(sample.cycle_usage.tokens, 100);
assert_eq!(sample.cycle_usage.tokens, 10_100);
assert!((sample.cycle_usage.usd - 10.1).abs() < 1e-10);
assert_eq!(sample.usage.tokens, 100);
assert_eq!(sample.sampled_percent, 10.0);
assert!(!sample.discontinuous);
}
#[test]
fn a_single_non_exhausted_observation_does_not_restart_cycle_usage() {
let sample = select(vec![point(1, 100.0, 10_000)], point(2, 0.0, 10_000));
assert_eq!(sample.cycle_usage.tokens, 10_000);
assert!(sample.discontinuous);
}
#[test]
fn repeated_same_window_recoveries_preserve_recorded_cycle_usage() {
let sample = select(
vec![
point(1, 100.0, 10_000),
point(2, 0.0, 10_000),
point(3, 100.0, 10_100),
point(4, 0.0, 10_100),
],
point(5, 10.0, 10_200),
);
assert_eq!(sample.cycle_usage.tokens, 10_200);
assert_eq!(sample.usage.tokens, 100);
assert!(!sample.discontinuous);
}
#[test]
fn changed_window_restarts_cycle_totals_at_first_matching_observation() {
let sample = select_forecast_sample(
@@ -1121,7 +1121,6 @@ impl SettingsStore for StaticSettingsStore {
refresh_margin_seconds: 300,
refresh_concurrency: 2,
max_concurrent_per_account: 1,
request_interval_ms: 0,
max_waiting_per_key: 0,
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
@@ -2399,6 +2398,93 @@ async fn quota_forecast_only_queries_the_two_selected_source_windows() {
);
}
#[tokio::test]
async fn quota_forecast_same_window_reversal_keeps_recorded_usage_in_the_report() {
use gateway_admin::model::quota_forecast_sampling::QuotaForecastHistoryPoint;
for (percents, tokens, expected_tokens, expected_usd) in [
([32.0, 30.0, 35.0], [100, 120, 140], 400, 0.4),
(
[100.0, 99.0, 99.5],
[10_000, 10_010, 10_020],
10_070,
10.02 / 0.995,
),
] {
let now = Utc::now();
let observed = now - TimeDelta::minutes(1);
let reset = now + TimeDelta::days(1);
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
plan_type: Some("pro".to_owned()),
observed_at: Some(observed),
windows: vec![ProviderQuotaWindow {
key: "week".to_owned(),
group: "shortTerm".to_owned(),
label: "周额度".to_owned(),
limit_id: None,
limit_name: None,
role: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(604_800),
used_percent: Some(percents[2]),
reset_at: Some(reset),
limit_reached: false,
local_usage: None,
provider_data: None,
}],
..empty_quota()
});
let mut account = account_record("openai");
account.created_at = now - TimeDelta::days(7);
let store = FakeAccountStore::with_account(account, events());
let usage = |requests, tokens| QuotaForecastUsage {
request_count: requests,
tokens,
input_tokens: tokens,
known_cost_count: requests,
usd: tokens as f64 / 1_000.0,
..Default::default()
};
*store.quota_forecast_history.lock().unwrap() = QuotaForecastHistory {
points: [(3, percents[0], tokens[0]), (2, percents[1], tokens[1])]
.into_iter()
.enumerate()
.map(
|(index, (minutes, percent, tokens))| QuotaForecastHistoryPoint {
started_at: now - TimeDelta::minutes(5),
completed_at: now - TimeDelta::minutes(minutes),
usage: usage(index as u64 + 1, tokens),
provider_observation: ProviderDocument::new(OpaqueProviderData::new(
serde_json::json!({
"percent": percent, "reset": reset.to_rfc3339(), "plan": "pro",
})
.as_object()
.unwrap()
.clone(),
)),
},
)
.collect(),
usage: usage(3, tokens[2]),
pending_request_count: 0,
};
let services = accounts_service(provider, store).await;
let result = services
.accounts()
.quota_forecast(&ProviderAccountId::new("acct_test").unwrap())
.await
.unwrap();
let forecast = &result.forecasts[0];
let source = forecast.source.as_ref().unwrap();
assert_eq!(source.tokens, Some(tokens[2]));
assert!((source.usd.unwrap() - tokens[2] as f64 / 1_000.0).abs() < 1e-10);
assert_eq!(forecast.estimated_tokens, Some(expected_tokens));
assert!((forecast.estimated_usd.unwrap() - expected_usd).abs() < 1e-10);
assert!(forecast.unavailable_reason.is_none());
}
}
#[tokio::test]
async fn quota_forecast_mid_cycle_sampling_accepts_small_reset_jitter_but_not_a_changed_bucket() {
use gateway_admin::model::quota_forecast_sampling::QuotaForecastHistoryPoint;
@@ -102,3 +102,51 @@ async fn backup_service_rejects_download_of_non_completed_record() {
.expect_err("non-completed record must be rejected");
assert_eq!(error.kind(), gateway_admin::model::AdminErrorKind::Conflict);
}
#[tokio::test]
async fn failed_delete_probe_does_not_refresh_verification() {
use gateway_admin::{model::backup::code, ports::backup::BackupRepository as _};
for succeeds in [false, true] {
let mut settings = configured_settings();
settings.last_verified_at = None;
let repository = Arc::new(FakeBackupRepository::new(settings));
let auth = Arc::new(FakeAuthStore::new());
let store = FakeObjectStore::new();
let store = if succeeds {
store
} else {
store.fail_connection(code::S3_PERMISSION_DENIED)
};
let services = AdminHarness::new()
.auth(auth.clone())
.backup(BackupStorePorts::new(
repository.clone(),
Arc::new(FakeDumpPort::new()),
Arc::new(store),
))
.build()
.await;
let result = services
.backups()
.test_connection(&system_context())
.await
.unwrap();
assert_eq!(result.ok, succeeds);
assert_eq!(
repository
.load_settings()
.await
.unwrap()
.last_verified_at
.is_some(),
succeeds
);
assert_eq!(
auth.audit_actions()
.iter()
.filter(|action| action.as_str() == "backup.s3_connection_tested")
.count(),
1
);
}
}
@@ -38,7 +38,6 @@ fn runtime_settings(enabled: bool, probe_enabled: bool, adaptive: bool) -> Runti
refresh_margin_seconds: 300,
refresh_concurrency: 2,
max_concurrent_per_account: 5,
request_interval_ms: 0,
max_waiting_per_key: 0,
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
@@ -747,6 +747,10 @@ impl ClientKeyStore for UnavailableStore {
#[async_trait]
impl ObservabilityStore for UnavailableStore {
fn query_timeout(&self) -> std::time::Duration {
std::time::Duration::from_secs(10)
}
async fn dashboard_summary(
&self,
_: DashboardQuery,
@@ -945,6 +945,10 @@ impl FixtureObservabilityStore {
#[async_trait]
impl ObservabilityStore for FixtureObservabilityStore {
fn query_timeout(&self) -> std::time::Duration {
std::time::Duration::from_secs(10)
}
async fn dashboard_summary(
&self,
query: DashboardQuery,
@@ -1104,7 +1108,6 @@ impl SettingsStore for FixtureSettingsStore {
refresh_margin_seconds: 300,
refresh_concurrency: 2,
max_concurrent_per_account: self.max_concurrent_per_account,
request_interval_ms: 0,
max_waiting_per_key: 0,
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
@@ -1406,23 +1409,23 @@ async fn usage_summary_cache_should_coalesce_when_generation_advances_during_the
}
#[tokio::test]
async fn usage_summary_cache_should_isolate_key_identity_and_normalize_group_sets() {
async fn usage_summary_cache_should_isolate_key_identity_and_filters() {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
let services = observability_services(store.clone()).await;
let service = services.observability();
for (key, groups) in [
("key-a", vec!["deleted-b", "deleted-a", "deleted-b"]),
("key-a", vec!["deleted-a", "deleted-b"]),
("key-b", vec!["deleted-a", "deleted-b"]),
("key-a", vec!["deleted-a"]),
for (key, model) in [
("key-a", None),
("key-a", None),
("key-b", None),
("key-a", Some("model-a")),
] {
service
.usage_summary(
range,
UsageFilter {
client_api_key_ref: Some(key.to_owned()),
routing_group_refs: groups.into_iter().map(str::to_owned).collect(),
model: model.map(str::to_owned),
..Default::default()
},
)
@@ -1431,16 +1434,16 @@ async fn usage_summary_cache_should_isolate_key_identity_and_normalize_group_set
}
let calls = store.usage_summary_filters.lock().unwrap();
assert_eq!(calls.len(), 3);
assert_eq!(calls[0].routing_group_refs, ["deleted-a", "deleted-b"]);
assert_eq!(calls[0].client_api_key_ref.as_deref(), Some("key-a"));
assert_eq!(calls[1].client_api_key_ref.as_deref(), Some("key-b"));
assert_eq!(calls[2].routing_group_refs, ["deleted-a"]);
assert_eq!(calls[2].model.as_deref(), Some("model-a"));
}
#[tokio::test(start_paused = true)]
async fn usage_summary_cache_should_bound_waiting_after_failed_refresh_and_allow_recovery() {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
store.set_dashboard_delay(StdDuration::from_secs(7));
store.set_dashboard_delay(store.query_timeout().mul_f64(0.7));
store.summary_fail.store(true, Ordering::SeqCst);
let services = observability_services(store.clone()).await;
let service = services.observability();
@@ -1453,7 +1456,7 @@ async fn usage_summary_cache_should_bound_waiting_after_failed_refresh_and_allow
);
assert!(first.is_err() && second.is_err() && third.is_err() && fourth.is_err());
assert!(
started.elapsed() <= StdDuration::from_secs(10),
started.elapsed() <= store.query_timeout(),
"同键失败不能使后续请求逐个等待完整数据库超时: {:?}",
started.elapsed()
);
@@ -331,7 +331,6 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() {
refresh_margin_seconds: 0,
refresh_concurrency: 1,
max_concurrent_per_account: 1,
request_interval_ms: 0,
max_waiting_per_key: 0,
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
@@ -29,7 +29,6 @@ impl DashboardQuery {
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UsageQuery {
pub routing_group_refs: Option<String>,
pub current_page: Option<u32>,
pub page_size: Option<u16>,
pub kind: Option<String>,
@@ -77,7 +76,6 @@ impl DetailQuery {
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct DiagnosticsQuery {
pub routing_group_refs: Option<String>,
pub dimension: Option<String>,
pub period: Option<String>,
pub as_of: Option<i64>,
@@ -86,7 +84,6 @@ pub struct DiagnosticsQuery {
pub provider: Option<String>,
pub model: Option<String>,
pub status_code: Option<i64>,
pub search: Option<String>,
}
impl DiagnosticsQuery {
@@ -100,7 +97,6 @@ impl DiagnosticsQuery {
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct OpsQuery {
pub routing_group_refs: Option<String>,
pub current_page: Option<u32>,
pub page_size: Option<u16>,
pub kind: Option<String>,
@@ -301,7 +297,6 @@ pub(crate) fn usage_filter(query: &UsageQuery) -> Result<domain::UsageFilter, Wi
})
});
Ok(domain::UsageFilter {
routing_group_refs: routing_group_refs(query.routing_group_refs.as_deref())?,
client_api_key_ref: non_empty(query.client_api_key_id.clone()),
request_id: non_empty(query.request_id.clone()),
provider_account_ref: non_empty(query.account_id.clone()),
@@ -318,6 +313,15 @@ pub(crate) fn usage_filter(query: &UsageQuery) -> Result<domain::UsageFilter, Wi
})
}
pub(crate) fn analytics_filter(
query: &UsageQuery,
) -> Result<domain::UsageFilter, WireValidationError> {
if query.search.is_some() {
return Err(WireValidationError::new("search"));
}
usage_filter(query)
}
pub(crate) fn usage_command(
query: &UsageQuery,
timezone: gateway_core::time::DeploymentTimeZone,
@@ -361,7 +365,6 @@ pub(crate) fn ops_command(
timezone,
)?,
filter: domain::OpsErrorFilter {
routing_group_refs: routing_group_refs(query.routing_group_refs.as_deref())?,
client_api_key_ref: non_empty(query.client_api_key_id.clone()),
request_id: non_empty(query.request_id.clone()),
provider_kind: non_empty(query.provider.clone()),
@@ -400,25 +403,3 @@ fn validate_pagination(
}
Ok((current_page, page_size))
}
/// 逗号分隔的稳定分组身份按集合归一化,不要求分组当前仍存在
pub(crate) fn routing_group_refs(value: Option<&str>) -> Result<Vec<String>, WireValidationError> {
let Some(value) = value else {
return Ok(Vec::new());
};
let mut refs = Vec::new();
for item in value.split(',') {
let item = item.trim();
if item.is_empty()
|| item.len() > 128
|| item.chars().any(char::is_control)
|| refs.len() >= 32
{
return Err(WireValidationError::new("routingGroupRefs"));
}
refs.push(item.to_owned());
}
refs.sort();
refs.dedup();
Ok(refs)
}
@@ -151,7 +151,7 @@ where
state.admin_services().timezone(),
)
.map_err(map_wire_error)?;
let filter = usage_filter(&query).map_err(map_wire_error)?;
let filter = analytics_filter(&query).map_err(map_wire_error)?;
let result = state
.admin_services()
.observability()
@@ -184,7 +184,7 @@ where
state.admin_services().timezone(),
)
.map_err(map_wire_error)?;
let filter = usage_filter(&query).map_err(map_wire_error)?;
let filter = analytics_filter(&query).map_err(map_wire_error)?;
let result = state
.admin_services()
.observability()
@@ -215,12 +215,9 @@ where
)
.map_err(map_wire_error)?;
let filter = domain::UsageFilter {
routing_group_refs: routing_group_refs(query.routing_group_refs.as_deref())
.map_err(map_wire_error)?,
provider_kind: non_empty(query.provider),
model: non_empty(query.model),
status_code: parse_status(query.status_code).map_err(map_wire_error)?,
search: non_empty(query.search),
..domain::UsageFilter::default()
};
let result = state
@@ -42,7 +42,6 @@ pub struct RuntimeSettingsFields {
pub refresh_margin_seconds: u64,
pub refresh_concurrency: u64,
pub max_concurrent_per_account: u64,
pub request_interval_ms: u64,
pub max_waiting_per_key: u32,
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
@@ -735,7 +734,6 @@ impl RuntimeSettingsFields {
.map_err(|_| WireValidationError::new("settingsRefreshConcurrencyOverflow"))?,
max_concurrent_per_account: u32::try_from(self.max_concurrent_per_account)
.map_err(|_| WireValidationError::new("settingsMaxConcurrencyOverflow"))?,
request_interval_ms: self.request_interval_ms,
max_waiting_per_key: self.max_waiting_per_key,
max_waiting_per_account: self.max_waiting_per_account,
concurrency_wait_timeout_seconds: self.concurrency_wait_timeout_seconds,
@@ -779,7 +777,6 @@ impl From<gateway_admin::model::settings::RuntimeSettingsValues> for RuntimeSett
refresh_margin_seconds: settings.refresh_margin_seconds,
refresh_concurrency: u64::from(settings.refresh_concurrency),
max_concurrent_per_account: u64::from(settings.max_concurrent_per_account),
request_interval_ms: settings.request_interval_ms,
max_waiting_per_key: settings.max_waiting_per_key,
max_waiting_per_account: settings.max_waiting_per_account,
concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds,
@@ -835,7 +832,6 @@ fn settings_validation_error(field: &'static str) -> WireValidationError {
"ops_event_retention_days" => "opsEventRetentionDays",
"refresh_concurrency" => "refreshConcurrency",
"refresh_margin_seconds" => "refreshMarginSeconds",
"request_interval_ms" => "requestIntervalMs",
"request_location" => "requestLocation",
"request_location_enabled" => "requestLocationEnabled",
"responses_max_decompressed_body_bytes" => "responsesMaxDecompressedBodyBytes",
@@ -228,6 +228,8 @@ async fn handle(
Ok(response.map(|body| body.map_err(|error| Box::new(error) as _).boxed_unsync()))
})
});
let request_id = context.request_id.clone();
let call_id = context.call_id.clone();
let result = match plan {
Some(plan) => plan.handle_http(context, request, next).await,
None => next.run(request).await,
@@ -251,7 +253,18 @@ async fn handle(
tokio::select! {
biased;
() = cancellation.cancelled() => {},
_ = task => {},
result = task => {
if let Err(contract::upgrade::SessionError::Failed { kind, plugin_instance_id }) = result {
tracing::warn!(
event = "middleware.session.failed",
request_id,
call_id,
plugin_instance_id,
kind,
"升级后的中间件会话失败"
);
}
},
}
});
}
@@ -610,7 +610,6 @@ impl SettingsStore for MemorySettingsStore {
refresh_margin_seconds: command.values.refresh_margin_seconds,
refresh_concurrency: command.values.refresh_concurrency,
max_concurrent_per_account: command.values.max_concurrent_per_account,
request_interval_ms: command.values.request_interval_ms,
max_waiting_per_key: command.values.max_waiting_per_key,
max_waiting_per_account: command.values.max_waiting_per_account,
concurrency_wait_timeout_seconds: command.values.concurrency_wait_timeout_seconds,
@@ -1374,6 +1373,10 @@ impl AccountRuntimeStore for UnusedStore {
#[async_trait]
impl ObservabilityStore for UnusedStore {
fn query_timeout(&self) -> std::time::Duration {
std::time::Duration::from_secs(10)
}
async fn dashboard_summary(
&self,
query: DashboardQuery,
@@ -1726,7 +1729,6 @@ fn test_runtime_settings() -> RuntimeSettings {
refresh_margin_seconds: 3_600,
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
max_waiting_per_key: 0,
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
@@ -1245,7 +1245,7 @@ async fn calendar_query_rejects_an_anchor_outside_the_local_calendar() {
}
#[tokio::test]
async fn request_and_ops_queries_should_forward_deleted_group_id_sets() {
async fn observability_queries_reject_group_filters_and_analytics_search() {
use crate::admin::{AdminTestFixture, AdminTestState};
use axum::{
body::Body,
@@ -1256,39 +1256,50 @@ async fn request_and_ops_queries_should_forward_deleted_group_id_sets() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let router = observability::router::<AdminTestState>().with_state(fixture.state());
for path in ["/api/admin/usage/records", "/api/admin/operations/errors"] {
for (path, parameter) in [
("/api/admin/usage/records", "routingGroupRefs=group-a"),
("/api/admin/operations/errors", "routingGroupRefs=group-a"),
(
"/api/admin/usage/records/summary",
"routingGroupRefs=group-a",
),
(
"/api/admin/usage/insights/overview",
"routingGroupRefs=group-a",
),
(
"/api/admin/usage/insights/diagnostics",
"routingGroupRefs=group-a",
),
("/api/admin/usage/records/summary", "search=request-prefix"),
(
"/api/admin/usage/insights/overview",
"search=request-prefix",
),
(
"/api/admin/usage/insights/diagnostics",
"search=request-prefix",
),
] {
let response = router
.clone()
.oneshot(
Request::builder()
.uri(format!(
"{path}?routingGroupRefs=deleted-b,deleted-a,deleted-b"
))
.uri(format!("{path}?{parameter}"))
.header(header::COOKIE, "cpr_session=valid-session")
.header("x-request-id", "req_group_filter")
.header("x-request-id", "req_rejected_observability_filter")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = axum::body::to_bytes(response.into_body(), 1024 * 1024)
.await
.unwrap();
assert_eq!(
status,
StatusCode::OK,
"{path}: {}",
String::from_utf8_lossy(&body)
response.status(),
StatusCode::BAD_REQUEST,
"{path}?{parameter}"
);
}
let observations = fixture.observations.lock().unwrap();
assert_eq!(
observations.records[0].filter.routing_group_refs,
["deleted-a", "deleted-b"]
);
assert_eq!(
observations.errors[0].filter.routing_group_refs,
["deleted-a", "deleted-b"]
);
assert!(observations.records.is_empty());
assert!(observations.errors.is_empty());
}
@@ -395,14 +395,7 @@ async fn published_fixture() -> AdminTestFixture {
fixture.published_snapshot.publish(
RuntimeSnapshot::new(
ConfigRevision::new(7).unwrap(),
gateway_core::settings::SettingsValues::new(
1,
0,
"smart",
Default::default(),
None,
None,
),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![],
vec![],
vec![],
@@ -50,7 +50,6 @@ fn update_body() -> Value {
"refreshMarginSeconds": 1800,
"refreshConcurrency": 4,
"maxConcurrentPerAccount": 5,
"requestIntervalMs": 25,
"maxWaitingPerKey": 0,
"maxWaitingPerAccount": 0,
"openaiGuardianReservedConcurrency": 0,
@@ -287,7 +286,6 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
refresh_margin_seconds: 1800,
refresh_concurrency: 4,
max_concurrent_per_account: 5,
request_interval_ms: 25,
max_waiting_per_key: 0,
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
@@ -337,7 +335,6 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
"refreshMarginSeconds": 1800,
"refreshConcurrency": 4,
"maxConcurrentPerAccount": 5,
"requestIntervalMs": 25,
"maxWaitingPerKey": 0,
"maxWaitingPerAccount": 0,
"openaiGuardianReservedConcurrency": 0,
@@ -415,7 +412,6 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() {
refresh_concurrency: u32::try_from(request.values.refresh_concurrency).expect("u32"),
max_concurrent_per_account: u32::try_from(request.values.max_concurrent_per_account)
.expect("u32"),
request_interval_ms: request.values.request_interval_ms,
max_waiting_per_key: 0,
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
@@ -1,5 +1,6 @@
//! HTTP 与 WebSocket 中间件测试入口及请求生命周期替身
mod session;
mod websocket;
use std::{
@@ -40,6 +41,9 @@ enum Mode {
idle: bool,
block_messages: bool,
},
SessionEnd {
kind: Option<&'static str>,
},
Service,
Echo,
Slow {
@@ -120,7 +124,7 @@ impl MiddlewarePlan for Plan {
assert!(input.is_null());
Box::pin(async move {
let mut result = next.run(input).await?;
result["request_interval_ms"] = serde_json::json!(4321);
result["refresh_margin_seconds"] = serde_json::json!(4321);
Ok(result)
})
}
@@ -161,13 +165,43 @@ impl MiddlewarePlan for Plan {
session.send(Message { kind: Kind::Text, payload: Bytes::from_static(b"ready") }).await?;
let message = receive.await?.unwrap();
session.send(message).await?;
session.send(Message { kind: Kind::Close { code: Some(1000) }, payload: Bytes::from_static(b"done") }).await
session.send(Message { kind: Kind::Close { code: Some(1000) }, payload: Bytes::from_static(b"done") }).await.map_err(Into::into)
})));
if reject {
*response.status_mut() = StatusCode::FORBIDDEN;
}
Ok(response)
}
Mode::SessionEnd { kind } => {
let upgrade = request
.extensions_mut()
.remove::<Arc<dyn contract::upgrade::WebSocketUpgrade>>()
.unwrap();
let (parts, _) = request.into_parts();
let (mut response, session) = upgrade.accept(parts, Vec::new()).await?;
response
.extensions_mut()
.insert(contract::upgrade::Upgraded::new(
StatusCode::SWITCHING_PROTOCOLS,
Box::pin(async move {
let _session = session.await.map_err(|_| MiddlewareError::Fault)?;
match kind {
None => Ok(()),
Some("cancelled") => {
Err(contract::upgrade::SessionError::Cancelled)
}
Some("reported") => {
Err(contract::upgrade::SessionError::Reported)
}
Some(kind) => Err(contract::upgrade::SessionError::Failed {
kind,
plugin_instance_id: Some("plugin-fixture".into()),
}),
}
}),
));
Ok(response)
}
Mode::Rewrite => {
assert_eq!(request.uri(), "/unrecognized/client/path?raw=%2f%FF");
assert_eq!(request.headers()["authorization"], "Bearer fixture-secret");
@@ -389,7 +423,7 @@ async fn admin_http_does_not_invoke_service_middleware_for_native_business_metho
.await
.unwrap();
let result: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_ne!(result["data"]["requestIntervalMs"], 4321);
assert_ne!(result["data"]["refreshMarginSeconds"], 4321);
assert_eq!(plan.calls.load(Ordering::SeqCst), 0);
}
@@ -0,0 +1,79 @@
//! 验证升级后终态只记录一次安全分类,正常结束及已报告失败不重复输出
use super::*;
use tracing::{
Event, Metadata, Subscriber,
span::{Attributes, Id, Record},
};
#[derive(Clone, Default)]
struct Recorder(Arc<Mutex<Vec<String>>>);
impl Subscriber for Recorder {
fn enabled(&self, _: &Metadata<'_>) -> bool {
true
}
fn new_span(&self, _: &Attributes<'_>) -> Id {
Id::from_u64(1)
}
fn record(&self, _: &Id, _: &Record<'_>) {}
fn record_follows_from(&self, _: &Id, _: &Id) {}
fn enter(&self, _: &Id) {}
fn exit(&self, _: &Id) {}
fn event(&self, event: &Event<'_>) {
struct Fields(String);
impl tracing::field::Visit for Fields {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
use std::fmt::Write as _;
write!(&mut self.0, "{}={value:?} ", field.name()).unwrap();
}
}
let mut fields = Fields(String::new());
event.record(&mut fields);
if fields.0.contains("middleware.session.failed") {
self.0.lock().unwrap().push(fields.0);
}
}
}
#[tokio::test]
async fn upgraded_session_records_one_correlated_failure_and_ignores_expected_endings() {
use tokio_tungstenite::{connect_async, tungstenite::client::IntoClientRequest as _};
for kind in [
None,
Some("cancelled"),
Some("reported"),
Some("permission_denied"),
] {
let recorder = Recorder::default();
let _subscriber = tracing::subscriber::set_default(recorder.clone());
let (router, plan) = router(Mode::SessionEnd { kind }, None).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
axum::serve(listener, router).await.unwrap();
});
let mut request = format!("ws://{address}/plugin-owned/session")
.into_client_request()
.unwrap();
request
.headers_mut()
.insert("x-request-id", "session-regression".parse().unwrap());
let (_socket, response) = connect_async(request).await.unwrap();
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
let cancellation = plan.cancellation.lock().unwrap().clone().unwrap();
tokio::time::timeout(Duration::from_secs(3), cancellation.cancelled())
.await
.unwrap();
let events = recorder.0.lock().unwrap();
if kind == Some("permission_denied") {
assert_eq!(events.len(), 1, "{events:?}");
assert!(events[0].contains("session-regression"));
assert!(events[0].contains("plugin-fixture"));
assert!(events[0].contains("permission_denied"));
} else {
assert!(events.is_empty(), "{events:?}");
}
server.abort();
}
}
@@ -308,7 +308,7 @@ fn snapshot_with_accounts(
);
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
gateway_core::settings::SettingsValues::new(2, 1, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
vec![provider.clone()],
["model-a", "model-b"]
.into_iter()
@@ -221,7 +221,7 @@ async fn turns_refresh_host_settings_and_key_identity_while_preserving_handshake
.unwrap();
let baseline = settings.runtime.as_ref().unwrap();
let mut values = serde_json::to_value(baseline.values()).unwrap();
values["request_interval_ms"] = json!(0);
values["max_account_rotations"] = json!(0);
settings.runtime = Some(
baseline
.replace(serde_json::from_value(values).unwrap(), "handshake")
@@ -292,7 +292,7 @@ async fn turns_refresh_host_settings_and_key_identity_while_preserving_handshake
for round in 0..3 {
if round == 1 {
let mut settings = serde_json::to_value(initial.settings()).unwrap();
settings["request_interval_ms"] = json!(50);
settings["max_account_rotations"] = json!(5);
settings["responses_max_decompressed_body_bytes"] = json!(2048);
snapshots.publish(
initial
@@ -337,8 +337,8 @@ async fn turns_refresh_host_settings_and_key_identity_while_preserving_handshake
observed[1]["runtime"]["responses_max_decompressed_body_bytes"],
2048
);
assert_eq!(observed[0]["runtime"]["request_interval_ms"], 0);
assert_eq!(observed[1]["runtime"]["request_interval_ms"], 0);
assert_eq!(observed[0]["runtime"]["max_account_rotations"], 0);
assert_eq!(observed[1]["runtime"]["max_account_rotations"], 0);
}
socket.close(None).await.unwrap();
server.abort();
@@ -89,7 +89,6 @@ pub struct AccountSelectionPolicy {
strategy: RotationStrategy,
smart_scheduling: SmartSchedulingConfig,
max_concurrent_per_account: AccountConcurrency,
request_interval: Duration,
queue_policy: ConcurrencyQueuePolicy,
openai_guardian_reserved_concurrency: u32,
openai_account_affinity: AccountAffinity,
@@ -102,13 +101,11 @@ impl AccountSelectionPolicy {
pub fn new(
strategy: RotationStrategy,
max_concurrent_per_account: impl Into<AccountConcurrency>,
request_interval: Duration,
) -> Self {
Self {
strategy,
smart_scheduling: SmartSchedulingConfig::default(),
max_concurrent_per_account: max_concurrent_per_account.into(),
request_interval,
openai_guardian_reserved_concurrency: 0,
openai_account_affinity: AccountAffinity::default(),
max_account_rotations: 3,
@@ -196,11 +193,6 @@ impl AccountSelectionPolicy {
pub const fn max_concurrent_per_account(self) -> AccountConcurrency {
self.max_concurrent_per_account
}
#[must_use]
pub const fn request_interval(self) -> Duration {
self.request_interval
}
}
/// Store 提供并发事实,Provider 叠加自己解释的额度事实;全部信号均可重建
@@ -543,8 +535,8 @@ impl AccountSelectionContext {
/// 选择账号时是否执行本地调度资格投影
///
/// 管理端对指定账号执行诊断时,会直接向上游确认实际状态;该模式跳过
/// `enabled`、可用性、冷却和 token 到期等本地投影,仍保留租约、并发和
/// 请求间隔约束,且只能与固定账号约束组合使用
/// `enabled`、可用性、冷却和 token 到期等本地投影,仍保留租约和并发
/// 约束,且只能与固定账号约束组合使用
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum AccountEligibilityPolicy {
#[default]
@@ -566,7 +558,6 @@ pub enum AccountSchedulingBlocker {
LocalAvailability,
Excluded,
ConcurrencyLimit,
RequestInterval,
LowerWeight,
}
@@ -791,7 +782,7 @@ impl AccountSelector {
}
}
/// 只有本地并发/调度间隔可等待;账号权限、失效、额度与上游冷却仍立即排除
/// 只有本地并发可等待;账号权限、失效、额度与上游冷却仍立即排除
#[must_use]
pub fn wait_candidates(
&self,
@@ -803,10 +794,7 @@ impl AccountSelector {
.filter(|candidate| {
matches!(
self.scheduling_blocker(candidate, context),
None | Some(
AccountSchedulingBlocker::ConcurrencyLimit
| AccountSchedulingBlocker::RequestInterval
)
None | Some(AccountSchedulingBlocker::ConcurrencyLimit)
)
})
.collect::<Vec<_>>();
@@ -879,18 +867,6 @@ impl AccountSelector {
{
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
}
if candidate
.signals
.last_started_at
.is_some_and(|last_started| {
!context
.now
.duration_since(last_started)
.is_ok_and(|elapsed| elapsed >= context.policy.request_interval())
})
{
return Some(AccountSchedulingBlocker::RequestInterval);
}
None
}
}
@@ -219,7 +219,7 @@ where
/// 对固定账号执行诊断请求
///
/// 仅跳过该账号的本地可用性投影;账号租约、并发和请求间隔仍必须满足
/// 仅跳过该账号的本地可用性投影;账号租约和并发仍必须满足
pub async fn start_diagnostic(
&self,
request: NewModelRequest,
@@ -2057,7 +2057,13 @@ where
let (outcome, gateway_error) = match error {
EngineError::Cancelled => (
ExecutionOutcome::Cancelled,
GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled"),
GatewayError::new(
GatewayErrorKind::Cancelled,
self.cancellation
.reason()
.unwrap_or(crate::lifecycle::CancellationReason::Requested)
.message(),
),
),
EngineError::Deadline => (
if self.response_commit.committed_at().is_some() {
@@ -8,7 +8,42 @@ use crate::{engine::middleware::MiddlewareError, engine::middleware::websocket::
pub type PendingSession =
Shared<BoxFuture<'static, Result<Arc<dyn Session>, Arc<MiddlewareError>>>>;
pub type SessionTask = BoxFuture<'static, Result<(), MiddlewareError>>;
pub type SessionTask = BoxFuture<'static, Result<(), SessionError>>;
/// 升级后的终态由传输 owner 消费,协议适配器只提供安全分类与诊断归属
#[derive(Debug)]
pub enum SessionError {
/// 预期取消,不作为异常终态输出
Cancelled,
/// 下层已记录同一失败,传输 owner 不重复输出
Reported,
Failed {
kind: &'static str,
plugin_instance_id: Option<String>,
},
}
impl From<MiddlewareError> for SessionError {
fn from(error: MiddlewareError) -> Self {
use crate::error::GatewayErrorKind;
if matches!(&error, MiddlewareError::Gateway(error) if error.kind() == GatewayErrorKind::Cancelled)
|| matches!(
error,
MiddlewareError::Engine(crate::engine::EngineError::Cancelled)
)
{
return Self::Cancelled;
}
Self::Failed {
kind: if error.is_rejected() {
"rejected"
} else {
"middleware"
},
plugin_instance_id: None,
}
}
}
/// 具体服务器负责握手与 socket;Core 不依赖 HTTP 服务器实现
pub trait WebSocketUpgrade: Send + Sync {
+76 -13
View File
@@ -1,16 +1,69 @@
//! 请求与连接的截止、租约、取消和 drain 契约
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::{Duration, SystemTime};
use event_listener::Event;
use futures::future::{BoxFuture, pending};
use futures_timer::Delay;
/// 参与业务并发约束的请求租约类别,观测续期不属于此合同
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RequestLeaseKind {
ClientAdmission,
Credential,
}
/// 续期最终结束的原因,不包含可恢复的中间失败
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LeaseFailure {
Lost,
RenewalTimedOut,
Expired,
}
/// 只携带稳定原因,不保存凭据、Redis key 或任意基础设施错误文本
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CancellationReason {
Requested,
Lease {
kind: RequestLeaseKind,
failure: LeaseFailure,
},
}
impl CancellationReason {
#[must_use]
pub const fn message(self) -> &'static str {
match self {
Self::Requested => "request was cancelled",
Self::Lease {
kind: RequestLeaseKind::ClientAdmission,
failure,
} => match failure {
LeaseFailure::Lost => "request was cancelled: client admission lease was lost",
LeaseFailure::RenewalTimedOut => {
"request was cancelled: client admission lease renewal timed out"
}
LeaseFailure::Expired => "request was cancelled: client admission lease expired",
},
Self::Lease {
kind: RequestLeaseKind::Credential,
failure,
} => match failure {
LeaseFailure::Lost => "request was cancelled: credential lease was lost",
LeaseFailure::RenewalTimedOut => {
"request was cancelled: credential lease renewal timed out"
}
LeaseFailure::Expired => "request was cancelled: credential lease expired",
},
}
}
}
struct CancellationState {
cancelled: AtomicBool,
reason: OnceLock<CancellationReason>,
event: Event,
}
@@ -41,7 +94,7 @@ impl CancellationToken {
pub fn new() -> Self {
Self {
state: Arc::new(CancellationState {
cancelled: AtomicBool::new(false),
reason: OnceLock::new(),
event: Event::new(),
}),
ancestors: Arc::from([]),
@@ -56,7 +109,7 @@ impl CancellationToken {
.collect::<Vec<_>>();
Self {
state: Arc::new(CancellationState {
cancelled: AtomicBool::new(false),
reason: OnceLock::new(),
event: Event::new(),
}),
ancestors: ancestors.into(),
@@ -64,19 +117,29 @@ impl CancellationToken {
}
pub fn cancel(&self) {
if self.state.cancelled.swap(true, Ordering::AcqRel) {
self.cancel_with_reason(CancellationReason::Requested);
}
/// 首次取消确定原因,后续资源回收不覆盖已观察到的取消事实
pub fn cancel_with_reason(&self, reason: CancellationReason) {
if self.is_cancelled() || self.state.reason.set(reason).is_err() {
return;
}
self.state.event.notify(usize::MAX);
}
#[must_use]
pub fn is_cancelled(&self) -> bool {
self.state.cancelled.load(Ordering::Acquire)
|| self
.ancestors
pub fn reason(&self) -> Option<CancellationReason> {
self.state.reason.get().copied().or_else(|| {
self.ancestors
.iter()
.any(|state| state.cancelled.load(Ordering::Acquire))
.find_map(|state| state.reason.get().copied())
})
}
#[must_use]
pub fn is_cancelled(&self) -> bool {
self.reason().is_some()
}
pub async fn cancelled(&self) {
@@ -92,13 +155,13 @@ impl CancellationToken {
fn wait_for_state(state: Arc<CancellationState>) -> BoxFuture<'static, ()> {
Box::pin(async move {
if state.cancelled.load(Ordering::Acquire) {
if state.reason.get().is_some() {
return;
}
// listener 随等待 future 丢弃而注销,避免 select 败选分支在自身及祖先状态中积累
let listener = state.event.listen();
// 注册后复查,覆盖取消发生在首次检查与注册之间的竞态
if state.cancelled.load(Ordering::Acquire) {
if state.reason.get().is_some() {
return;
}
listener.await;
@@ -55,7 +55,6 @@ pub struct ProviderSchedulingLeaseRequest {
credential_revision: CredentialRevision,
max_concurrent: AccountConcurrency,
concurrency_pool: ProviderConcurrencyPool,
request_interval: Duration,
deadline: crate::lifecycle::Deadline,
cancellation: crate::lifecycle::CancellationToken,
}
@@ -67,7 +66,6 @@ impl ProviderSchedulingLeaseRequest {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
max_concurrent: impl Into<AccountConcurrency>,
request_interval: Duration,
deadline: impl Into<crate::lifecycle::Deadline>,
) -> Self {
Self {
@@ -76,7 +74,6 @@ impl ProviderSchedulingLeaseRequest {
credential_revision,
max_concurrent: max_concurrent.into(),
concurrency_pool: ProviderConcurrencyPool::Shared,
request_interval,
deadline: deadline.into(),
cancellation: crate::lifecycle::CancellationToken::new(),
}
@@ -124,11 +121,6 @@ impl ProviderSchedulingLeaseRequest {
self.max_concurrent
}
#[must_use]
pub const fn request_interval(&self) -> Duration {
self.request_interval
}
#[must_use]
pub const fn deadline(&self) -> crate::lifecycle::Deadline {
self.deadline
@@ -166,7 +158,7 @@ pub enum ProviderLeaseRequest {
}
pub trait ProviderLeasePort: Send + Sync {
/// 按目标容量池读取在途数,账号最小请求间隔仍跨池共享
/// 按目标容量池读取在途数,最近调度时间仍跨池共享
fn load_state<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
@@ -73,7 +73,6 @@ impl CompiledSettings {
let account_selection_policy = AccountSelectionPolicy::new(
strategy,
AccountConcurrency::new(settings.max_concurrent_per_account),
Duration::from_millis(settings.request_interval_ms),
)
.with_openai_guardian_reserved_concurrency(settings.openai_guardian_reserved_concurrency)
.with_openai_account_affinity(settings.openai_account_affinity)
@@ -22,7 +22,6 @@ pub struct SettingsValues {
pub(crate) max_account_rotations: u32,
pub(crate) openai_session_affinity_ttl_hours: u32,
pub(crate) responses_max_decompressed_body_bytes: u64,
pub(crate) request_interval_ms: u64,
pub(crate) smart_scheduling: crate::account::SmartSchedulingConfig,
pub(crate) rotation_strategy: String,
pub(crate) model_mappings: Arc<BTreeMap<String, String>>,
@@ -129,7 +128,6 @@ impl SettingsValues {
#[must_use]
pub fn new(
max_concurrent_per_account: u32,
request_interval_ms: u64,
rotation_strategy: impl Into<String>,
model_mappings: BTreeMap<String, String>,
min_codex_desktop_version: Option<String>,
@@ -150,7 +148,6 @@ impl SettingsValues {
max_account_rotations: 3,
openai_session_affinity_ttl_hours: 24,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
request_interval_ms,
smart_scheduling: crate::account::SmartSchedulingConfig::default(),
rotation_strategy: rotation_strategy.into(),
model_mappings: Arc::new(model_mappings),
@@ -99,11 +99,7 @@ fn account_scheduling_should_default_and_resolve_concurrency_override() {
fn context(strategy: RotationStrategy) -> AccountSelectionContext {
AccountSelectionContext {
policy: AccountSelectionPolicy::new(
strategy,
NonZeroU32::new(3).expect("positive"),
Duration::ZERO,
),
policy: AccountSelectionPolicy::new(strategy, NonZeroU32::new(3).expect("positive")),
now: SystemTime::now(),
excluded_accounts: BTreeSet::new(),
preferred_account: None,
@@ -471,7 +467,6 @@ fn highest_priority_affinity_should_still_obey_existing_scheduling_constraints()
for blocker in [
AccountSchedulingBlocker::Excluded,
AccountSchedulingBlocker::ConcurrencyLimit,
AccountSchedulingBlocker::RequestInterval,
AccountSchedulingBlocker::LocalAvailability,
] {
let mut candidates = [
@@ -488,14 +483,6 @@ fn highest_priority_affinity_should_still_obey_existing_scheduling_constraints()
.insert(candidates[0].account.id().clone());
}
AccountSchedulingBlocker::ConcurrencyLimit => candidates[0].signals.in_flight = 3,
AccountSchedulingBlocker::RequestInterval => {
candidates[0].signals.last_started_at = Some(selection.now);
selection.policy = AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(3).expect("limit"),
Duration::from_secs(1),
);
}
AccountSchedulingBlocker::LocalAvailability => {
candidates[0].signals.cooldown =
Some((selection.now + Duration::from_secs(60)).into())
@@ -788,25 +775,25 @@ fn selector_should_reject_account_at_concurrency_limit() {
}
#[test]
fn selector_should_enforce_request_interval_until_boundary() {
fn selector_should_allow_immediate_reuse_regardless_of_last_started_time() {
let now = SystemTime::now();
let mut cooling = candidate("acct_cooling", 0, Some(100));
cooling.signals.last_started_at = Some(now - Duration::from_millis(9));
let mut ready = candidate("acct_ready", 0, Some(1));
ready.signals.last_started_at = Some(now - Duration::from_millis(10));
let candidates = vec![cooling, ready];
let mut context = context(RotationStrategy::Smart);
context.now = now;
context.policy = AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(3).expect("positive"),
Duration::from_millis(10),
);
let selected = AccountSelector
.select(&candidates, &context)
.expect("candidate available");
assert_eq!(selected.candidate().account.id().as_str(), "acct_ready");
for last_started in [
now - Duration::from_millis(1),
now,
now + Duration::from_secs(30),
] {
let mut candidate = candidate("acct_recent", 1, None);
candidate.signals.last_started_at = Some(last_started);
let candidates = [candidate];
let mut context = context(RotationStrategy::Sticky);
context.now = now;
context.preferred_account = Some(candidates[0].account.id().clone());
let selected = AccountSelector
.select(&candidates, &context)
.expect("available slot");
assert_eq!(selected.candidate().account.id().as_str(), "acct_recent");
assert_eq!(selected.preferred(), PreferredAccountSelection::Hit);
}
}
#[test]
@@ -1136,7 +1123,6 @@ fn reserved_capacity_stays_bounded_when_normal_concurrency_is_unlimited() {
normal.policy = AccountSelectionPolicy::new(
RotationStrategy::Smart,
gateway_core::account::AccountConcurrency::Unlimited,
Duration::ZERO,
);
assert!(AccountSelector.select(&candidates, &normal).is_some());
assert!(
@@ -13,7 +13,7 @@ use super::{candidate, candidate_with_concurrency, context};
const FAILURE_RATE_HALF_LIFE: Duration = Duration::from_secs(15 * 60);
#[test]
fn unlimited_account_concurrency_preserves_overrides_intervals_and_finite_scores() {
fn unlimited_account_concurrency_preserves_overrides_and_finite_scores() {
use gateway_core::account::{AccountConcurrency, AccountSelectionPolicy};
for strategy in [
@@ -23,8 +23,7 @@ fn unlimited_account_concurrency_preserves_overrides_intervals_and_finite_scores
RotationStrategy::QuotaResetPriority,
] {
let mut context = context(strategy);
context.policy =
AccountSelectionPolicy::new(strategy, AccountConcurrency::Unlimited, Duration::ZERO);
context.policy = AccountSelectionPolicy::new(strategy, AccountConcurrency::Unlimited);
let candidates = [
candidate("acct_unlimited", u32::MAX, None),
candidate_with_concurrency("acct_limited", 2, 2),
@@ -60,14 +59,9 @@ fn unlimited_account_concurrency_preserves_overrides_intervals_and_finite_scores
.is_finite()
);
}
context.policy = AccountSelectionPolicy::new(
strategy,
AccountConcurrency::Unlimited,
Duration::from_secs(1),
);
let mut recent = candidate("acct_unlimited", u32::MAX, None);
recent.signals.last_started_at = Some(context.now);
assert!(AccountSelector.select(&[recent], &context).is_none());
assert!(AccountSelector.select(&[recent], &context).is_some());
}
}
@@ -77,11 +71,8 @@ fn unavailable_unlimited_account_does_not_hide_the_finite_pool_capacity() {
AccountConcurrency, AccountSelectionPolicy, CredentialState, QuotaState,
};
let mut context = context(RotationStrategy::Smart);
context.policy = AccountSelectionPolicy::new(
RotationStrategy::Smart,
AccountConcurrency::Unlimited,
Duration::ZERO,
);
context.policy =
AccountSelectionPolicy::new(RotationStrategy::Smart, AccountConcurrency::Unlimited);
let mut unavailable = candidate("acct_disabled", 0, None);
unavailable.account = unavailable.account.with_account_facts(
false,
@@ -53,6 +53,7 @@ use serde_json::{Map, Value, json};
#[derive(Debug, Clone, PartialEq)]
struct FinalState {
error_message: Option<String>,
error_kind: Option<GatewayErrorKind>,
diagnostic_trace_json: Option<String>,
outcome: ExecutionOutcome,
@@ -228,6 +229,10 @@ impl ExecutionStore for FakeStore {
.expect("store lock")
.finalizations
.push(FinalState {
error_message: finalization
.error
.as_ref()
.map(|error| error.safe_message().to_owned()),
error_kind: finalization.error.as_ref().map(|error| error.kind()),
diagnostic_trace_json: finalization.diagnostic_trace_json,
outcome: finalization.outcome,
@@ -671,7 +676,6 @@ fn plan(operation: &Operation) -> RoutingPlan {
AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(2).expect("account concurrency"),
Duration::from_millis(50),
),
)
}
@@ -752,7 +756,6 @@ fn plan_with_profiles(
ConfigRevision::new(1).expect("config revision"),
gateway_core::settings::SettingsValues::new(
account_selection_policy.max_concurrent_per_account().get(),
u64::try_from(account_selection_policy.request_interval().as_millis()).unwrap(),
account_selection_policy.strategy().as_str(),
Default::default(),
None,
@@ -824,7 +827,7 @@ fn dual_provider_plan(operation: &Operation) -> RoutingPlan {
let directory = Arc::new(RuntimeAccountDirectory::new(entries));
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
gateway_core::settings::SettingsValues::new(2, 50, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
vec![openai.clone(), xai.clone()],
vec![
ProviderModel::new(
@@ -2156,7 +2159,6 @@ fn unavailable_native_continuation_uses_provider_recovery_with_the_configured_ac
let policy = AccountSelectionPolicy::new(
RotationStrategy::RoundRobin,
NonZeroU32::new(2).expect("account concurrency"),
Duration::from_millis(50),
);
let route_plan = plan_with_policy(&operation, policy);
let (coordinator, store, provider) = coordinator(vec![
@@ -5726,13 +5728,9 @@ fn configured_rotation_budget_limits_actual_account_changes() {
let operation = generate_operation();
let route_plan = plan_with_policy(
&operation,
AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(2).unwrap(),
Duration::ZERO,
)
.with_max_account_rotations(budget)
.with_openai_session_affinity_ttl(Duration::from_secs(168 * 3600)),
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
.with_max_account_rotations(budget)
.with_openai_session_affinity_ttl(Duration::from_secs(168 * 3600)),
);
let (coordinator, store, provider) = coordinator(
ROTATION_TEST_ACCOUNTS
@@ -5779,12 +5777,8 @@ fn larger_rotation_budget_reaches_a_healthy_account_after_four_failures() {
let operation = generate_operation();
let route_plan = plan_with_policy(
&operation,
AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(2).unwrap(),
Duration::ZERO,
)
.with_max_account_rotations(4),
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
.with_max_account_rotations(4),
);
let mut scripts = ROTATION_TEST_ACCOUNTS[..4]
.iter()
@@ -5820,12 +5814,8 @@ fn zero_rotation_budget_still_allows_retrying_the_same_account() {
let operation = generate_operation();
let route_plan = plan_with_policy(
&operation,
AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(2).unwrap(),
Duration::ZERO,
)
.with_max_account_rotations(0),
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
.with_max_account_rotations(0),
);
let (coordinator, _, provider) = coordinator(vec![
Script::Stream {
@@ -5995,3 +5985,40 @@ fn transport_keepalive_cannot_change_an_already_finalized_failure_to_http_200()
assert!(!state.finalizations[0].committed);
assert_eq!(state.recorded_statuses, vec![504]);
}
#[test]
fn lease_cancellation_is_preserved_in_the_single_request_terminal() {
use gateway_core::lifecycle::{CancellationReason, LeaseFailure, RequestLeaseKind};
block_on(async {
let operation = generate_operation();
let (coordinator, store, _) = coordinator(vec![Script::Pending]);
let cancellation = CancellationToken::new();
let mut session = coordinator
.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation.clone(),
plan(&operation),
None,
None,
cancellation.clone(),
)
.await
.unwrap();
let reason = CancellationReason::Lease {
kind: RequestLeaseKind::ClientAdmission,
failure: LeaseFailure::RenewalTimedOut,
};
let mut next = Box::pin(session.next_event());
assert!(futures::poll!(next.as_mut()).is_pending());
cancellation.cancel_with_reason(reason);
assert!(matches!(next.await, Err(EngineError::Cancelled)));
session.cancel_and_finalize().await.unwrap();
let state = store.state.lock().unwrap();
assert_eq!(state.finalizations.len(), 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Cancelled);
assert_eq!(
state.finalizations[0].error_message.as_deref(),
Some(reason.message())
);
});
}
@@ -1119,7 +1119,7 @@ fn request_policy_snapshot(
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(2, 1, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
providers.to_vec(),
providers
.iter()
@@ -1847,14 +1847,7 @@ fn model_routing_cannot_expand_the_frozen_key_model_scope() {
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(
2,
1,
"smart",
Default::default(),
None,
None,
),
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
vec![kind.clone()],
["gpt-start", "gpt-blocked"]
.into_iter()
@@ -2231,7 +2224,7 @@ fn fallback_observation_snapshot(reference: ExtensionSetReference) -> RuntimeSna
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(2, 1, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
providers.to_vec(),
providers
.iter()
@@ -3919,7 +3912,7 @@ fn probe_snapshot() -> RuntimeSnapshot {
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
gateway_core::settings::SettingsValues::new(1, 1, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![provider.clone()],
vec![ProviderModel::new(
provider,
@@ -3935,7 +3928,7 @@ fn client_snapshot() -> RuntimeSnapshot {
let provider = ProviderKind::new("openai").expect("provider kind");
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
gateway_core::settings::SettingsValues::new(1, 1, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![provider.clone()],
Vec::new(),
vec![ClientPolicy::new(
@@ -3964,7 +3957,7 @@ fn start_snapshot_with_policy(
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
RuntimeSnapshot::new(
ConfigRevision::new(revision).expect("config revision"),
gateway_core::settings::SettingsValues::new(1, 1, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![provider.clone()],
vec![ProviderModel::new(
provider.clone(),
@@ -4384,7 +4377,7 @@ fn request_settings_recompute_routing_and_admission_without_changing_sibling_or_
let mut settings =
serde_json::to_value(modified.request_settings().execution_values().unwrap()).unwrap();
settings["runtime"]["model_mappings"] = json!({"request-alias":"gpt-start"});
settings["runtime"]["request_interval_ms"] = json!(0);
settings["runtime"]["max_account_rotations"] = json!(0);
settings["runtime"]["request_profiles"] = json!({"openai":{"identity":"request-local"}});
settings["fast_mode"] = json!("default");
settings["client_limits"] = json!({"max_concurrency":0,"requests_per_minute":0});
@@ -4604,7 +4597,7 @@ fn entry_settings_freeze_authentication_and_rebase_only_explicit_overrides() {
let baseline = service.request_settings().unwrap();
let mut values = serde_json::to_value(baseline.values()).unwrap();
values["responses_max_decompressed_body_bytes"] = json!(1024);
values["request_interval_ms"] = json!(0);
values["max_account_rotations"] = json!(0);
let first = baseline
.replace(serde_json::from_value(values).unwrap(), "first")
.unwrap();
@@ -4621,7 +4614,7 @@ fn entry_settings_freeze_authentication_and_rebase_only_explicit_overrides() {
.unwrap();
let source = second.inspect();
assert_eq!(
source["overrides"]["request_interval_ms"]["instance_id"],
source["overrides"]["max_account_rotations"]["instance_id"],
"first"
);
assert_eq!(
@@ -4819,14 +4812,7 @@ fn child_settings_recompute_key_scope_without_inheriting_parent_defaults() {
.collect();
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(
1,
0,
"smart",
Default::default(),
None,
None,
),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![provider],
vec![],
policies,
@@ -1,11 +1,6 @@
//! 验证插件调度策略对软亲和的覆盖与宿主候选范围限制
use std::{
collections::BTreeSet,
num::NonZeroU32,
sync::Arc,
time::{Duration, SystemTime},
};
use std::{collections::BTreeSet, num::NonZeroU32, sync::Arc, time::SystemTime};
use futures::{executor::block_on, future::BoxFuture};
use gateway_core::{
@@ -106,11 +101,7 @@ fn policy(choice: Option<&'static str>) -> RequestPolicyContext {
fn context() -> AccountSelectionContext {
AccountSelectionContext {
policy: AccountSelectionPolicy::new(
RotationStrategy::Sticky,
NonZeroU32::new(2).unwrap(),
Duration::ZERO,
),
policy: AccountSelectionPolicy::new(RotationStrategy::Sticky, NonZeroU32::new(2).unwrap()),
now: SystemTime::now(),
excluded_accounts: BTreeSet::new(),
preferred_account: Some(account_id("acct_a")),
@@ -181,15 +172,15 @@ fn smart_switchback_does_not_override_explicit_plugin_choices_or_their_telemetry
.clone()
.with_scheduling(None, AccountWeight::new(10).unwrap());
let mut selection_context = context();
selection_context.policy = AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(2).unwrap(),
Duration::ZERO,
)
.with_smart_scheduling(
gateway_core::account::SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], true)
.unwrap(),
);
selection_context.policy =
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
.with_smart_scheduling(
gateway_core::account::SmartSchedulingConfig::new(
[1.0, 0.8, 1.0, 0.5, 0.0, 0.0],
true,
)
.unwrap(),
);
for chosen in [Some("acct_a"), Some("acct_b"), None] {
let selected = policy(chosen)
.select_account(
@@ -322,3 +322,25 @@ fn connection_lifecycle_contract_is_object_safe() {
assert!(!object.is_draining());
}
#[test]
fn cancellation_reason_survives_cleanup_and_propagates_only_to_descendants() {
use gateway_core::lifecycle::{CancellationReason, LeaseFailure, RequestLeaseKind};
let parent = CancellationToken::new();
let child = parent.child_token();
let descendant = child.child_token();
let reason = CancellationReason::Lease {
kind: RequestLeaseKind::Credential,
failure: LeaseFailure::Lost,
};
child.cancel_with_reason(reason);
child.cancel();
descendant.cancel();
assert_eq!(parent.reason(), None);
assert_eq!(child.reason(), Some(reason));
assert_eq!(descendant.reason(), Some(reason));
let requested = CancellationToken::new();
requested.cancel();
requested.cancel_with_reason(reason);
assert_eq!(requested.reason(), Some(CancellationReason::Requested));
}
@@ -23,7 +23,7 @@ mod request_settings;
mod snapshot;
fn settings() -> gateway_core::settings::SettingsValues {
gateway_core::settings::SettingsValues::new(3, 50, "smart", BTreeMap::new(), None, None)
gateway_core::settings::SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
}
fn capabilities() -> ModelCapabilities {
@@ -29,7 +29,7 @@ fn snapshot(revision: u64, name: &str) -> Arc<RuntimeSnapshot> {
Arc::new(
RuntimeSnapshot::new(
ConfigRevision::new(revision).unwrap(),
SettingsValues::new(3, 50, "smart", BTreeMap::new(), None, None)
SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
.with_request_profiles(profiles(&[("openai", name), ("xai", name)])),
vec![],
vec![],
@@ -136,7 +136,6 @@ fn compile_aliases(
revision(1),
SettingsValues::new(
3,
50,
"smart",
BTreeMap::from([("configured".into(), "upstream-model".into())]),
None,
@@ -411,7 +410,7 @@ fn compiler_accepts_unlimited_default_account_concurrency() {
let facts = SnapshotFacts::new(
revision(1),
revision(1),
SettingsValues::new(0, 50, "smart", BTreeMap::new(), None, None),
SettingsValues::new(0, "smart", BTreeMap::new(), None, None),
Vec::new(),
Vec::new(),
vec![SnapshotProviderAccountFacts::new(
@@ -437,10 +436,6 @@ fn compiler_accepts_unlimited_default_account_concurrency() {
plan.account_selection_policy().max_concurrent_per_account(),
gateway_core::account::AccountConcurrency::Unlimited
);
assert_eq!(
plan.account_selection_policy().request_interval(),
std::time::Duration::from_millis(50)
);
}
#[test]
@@ -455,7 +450,7 @@ fn compiled_plans_keep_their_scheduling_config_after_a_new_snapshot_is_built() {
let facts = SnapshotFacts::new(
revision(index as u64 + 1),
revision(index as u64 + 1),
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None)
SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
.with_smart_scheduling(config)
.with_openai_guardian_reserved_concurrency(index as u32),
vec![],
@@ -511,7 +506,6 @@ fn withdrawn_provider_keeps_accounts_without_becoming_a_route_or_expanding_group
revision(1),
SettingsValues::new(
3,
0,
"smart",
BTreeMap::from([("alias".into(), "listed-model".into())]),
None,
@@ -1063,7 +1057,6 @@ fn facts_with_min_versions(
revision(observed_current_revision),
SettingsValues::new(
3,
50,
"smart",
BTreeMap::from([("public-model".to_owned(), "upstream-model".to_owned())]),
desktop,
@@ -1105,7 +1098,7 @@ fn global_request_location_should_be_frozen_when_snapshot_is_published() {
let facts = SnapshotFacts::new(
revision(version),
revision(version),
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None)
SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
.with_request_location(location, enabled),
Vec::new(),
Vec::new(),
@@ -1182,7 +1175,7 @@ fn decompression_setting_should_validate_and_remain_frozen_across_publication()
let facts = SnapshotFacts::new(
revision(version),
revision(version),
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None)
SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
.with_responses_max_decompressed_body_bytes(bytes),
Vec::new(),
Vec::new(),
@@ -1245,7 +1238,7 @@ fn fast_mode_merges_only_bound_groups_without_changing_account_scope() {
let facts = SnapshotFacts::new(
revision(1),
revision(1),
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None),
SettingsValues::new(3, "smart", BTreeMap::new(), None, None),
vec![SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_fast_policy").unwrap(),
PlaintextClientApiKey::new("sk_fast_policy").unwrap(),
@@ -1326,7 +1319,7 @@ fn key_profiles_replace_whole_global_choice_and_previous_snapshot_stays_frozen()
};
let build = |global, overridden| {
let profiles = BTreeMap::from([(provider.clone(), document(global))]);
let settings = SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None)
let settings = SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
.with_request_profiles(profiles);
let inherited = SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_inherited").unwrap(),
@@ -823,7 +823,7 @@ fn scoped_facts(value: u64, allow_removed: bool) -> SnapshotFacts {
SnapshotFacts::new(
revision(value),
revision(value),
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None),
SettingsValues::new(3, "smart", BTreeMap::new(), None, None),
vec![SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_audit_synthetic").expect("key ID"),
PlaintextClientApiKey::new("sk_audit_synthetic_not_a_real_key").expect("synthetic key"),
@@ -858,7 +858,6 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
revision(observed_current_revision),
SettingsValues::new(
3,
50,
"smart",
BTreeMap::from([("public-model".to_owned(), "upstream-model".to_owned())]),
None,
@@ -883,7 +882,7 @@ fn compiler(store: Arc<dyn SnapshotStorePort>) -> RuntimeSnapshotCompiler {
fn empty_snapshot(value: u64) -> RuntimeSnapshot {
RuntimeSnapshot::new(
revision(value),
gateway_core::settings::SettingsValues::new(1, 0, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
Vec::new(),
Vec::new(),
Vec::new(),
@@ -11,6 +11,7 @@ mod process;
mod release;
mod state;
mod swap;
mod swap_paths;
pub use config::SystemUpdateConfig;
pub use operations::ProcessSystemOperations;
@@ -245,18 +245,7 @@ fn swap_file(current: &Path, backup: &Path) -> io::Result<()> {
if swap.exists() {
fs::remove_file(&swap)?;
}
move_file(current, &swap)?;
if let Err(error) = move_file(backup, current) {
let _ = move_file(&swap, current);
return Err(error);
}
if let Err(error) = move_file(&swap, backup) {
let _ = move_file(current, &swap);
let _ = move_file(backup, current);
let _ = move_file(&swap, backup);
return Err(error);
}
Ok(())
super::swap_paths::swap_paths(current, backup, &swap, move_file)
}
fn swap_dir(current: &Path, backup: &Path) -> io::Result<()> {
@@ -270,18 +259,7 @@ fn swap_dir(current: &Path, backup: &Path) -> io::Result<()> {
if swap.exists() {
remove_dir_all(&swap)?;
}
move_dir(current, &swap)?;
if let Err(error) = move_dir(backup, current) {
let _ = move_dir(&swap, current);
return Err(error);
}
if let Err(error) = move_dir(&swap, backup) {
let _ = move_dir(current, &swap);
let _ = move_dir(backup, current);
let _ = move_dir(&swap, backup);
return Err(error);
}
Ok(())
super::swap_paths::swap_paths(current, backup, &swap, move_dir)
}
fn restore_dir(current: &Path, backup: &Path) -> io::Result<()> {
@@ -0,0 +1,59 @@
//! 文件和目录交换共用的恢复顺序,保留主失败及补偿失败
use std::{io, path::Path};
pub(super) fn swap_paths(
current: &Path,
backup: &Path,
swap: &Path,
mut move_path: impl FnMut(&Path, &Path) -> io::Result<()>,
) -> io::Result<()> {
move_path(current, swap)?;
if let Err(error) = move_path(backup, current) {
return Err(with_recovery(
error,
"restore current from swap",
move_path(swap, current),
[current, backup, swap],
));
}
if let Err(error) = move_path(swap, backup) {
// 先把备份版本放回 backup,才能把原版本移回 current
// 补偿第一步失败时保留 swap,不能覆盖仍在 current 的备份版本
if let Err(recovery) = move_path(current, backup) {
return Err(with_recovery(
error,
"restore backup from current",
Err(recovery),
[current, backup, swap],
));
}
return Err(with_recovery(
error,
"restore current from swap",
move_path(swap, current),
[current, backup, swap],
));
}
Ok(())
}
fn with_recovery(
primary: io::Error,
action: &'static str,
recovery: io::Result<()>,
[current, backup, swap]: [&Path; 3],
) -> io::Error {
match recovery {
Ok(()) => primary,
Err(recovery) => io::Error::new(
primary.kind(),
format!(
"{primary}; rollback failed: {action}: {recovery}; manual recovery required (current={}, backup={}, swap={})",
current.display(),
backup.display(),
swap.display()
),
),
}
}
@@ -7,6 +7,7 @@ mod installation;
mod operations;
mod release;
mod state;
mod swap_paths;
use std::fs;
use std::path::PathBuf;
@@ -0,0 +1,76 @@
//! 对真实暂存文件注入移动双故障,验证恢复次序和人工恢复线索
#[path = "../../src/system_update/swap_paths.rs"]
mod implementation;
use std::{fs, io};
#[test]
fn failed_swap_preserves_both_versions_and_reports_compensation_failures() {
for directory in [false, true] {
for primary in [2, 3] {
let mut recoveries = vec![None, Some(primary + 1)];
if primary == 3 {
recoveries.push(Some(5));
}
for recovery in recoveries {
let temp = tempfile::tempdir().unwrap();
let current = temp.path().join("current");
let backup = temp.path().join("backup");
let swap = temp.path().join("swap");
for (path, content) in [(&current, "current version"), (&backup, "backup version")]
{
if directory {
fs::create_dir(path).unwrap();
fs::write(path.join("asset"), content).unwrap();
} else {
fs::write(path, content).unwrap();
}
}
let mut moves = 0;
let error = implementation::swap_paths(&current, &backup, &swap, |from, to| {
moves += 1;
if moves == primary {
return Err(io::Error::other("primary injected failure"));
}
if recovery == Some(moves) {
return Err(io::Error::other("recovery injected failure"));
}
fs::rename(from, to)
})
.unwrap_err();
let message = error.to_string();
assert!(message.contains("primary injected failure"));
if recovery.is_some() {
assert!(message.contains("recovery injected failure"));
assert!(message.contains("restore "));
assert!(message.contains("manual recovery required"));
assert!(message.contains(swap.to_str().unwrap()));
} else {
assert!(!message.contains("rollback failed"));
}
let mut contents = [&current, &backup, &swap]
.into_iter()
.filter_map(|path| {
fs::read_to_string(if directory {
path.join("asset")
} else {
path.clone()
})
.ok()
})
.collect::<Vec<_>>();
contents.sort();
assert_eq!(contents, ["backup version", "current version"]);
if recovery.is_none() {
let original = if directory {
current.join("asset")
} else {
current.clone()
};
assert_eq!(fs::read_to_string(original).unwrap(), "current version");
}
}
}
}
}
@@ -87,16 +87,16 @@ pub(super) async fn invoke(
let _binding = _binding;
tokio::select! {
biased;
() = context.cancellation.cancelled() => Err(MiddlewareError::Fault),
() = context.cancellation.cancelled() => Err(core::upgrade::SessionError::Cancelled),
result = async {
invocation.session_ready().await?;
match stream.next().await {
Ok(None) => Ok(()),
Ok(Some(_)) => {
session.invalid_response(Stage::Http);
Err(MiddlewareError::InvalidState)
let error = session.invalid_response(Stage::Http);
Err(session_error(error, &entry.instance_id))
}
Err(error) => Err(crate::callback::error::rpc_middleware(error)),
Err(error) => Err(session_error(error, &entry.instance_id)),
}
} => result,
}
@@ -143,3 +143,32 @@ pub(super) async fn invoke(
.await
.map_err(|_| invalid())
}
fn session_error(error: crate::RpcError, instance_id: &str) -> core::upgrade::SessionError {
use core::upgrade::SessionError;
use gateway_plugin_sdk::ErrorCode;
let kind = match &error {
crate::RpcError::Cancelled => return SessionError::Cancelled,
// 本地期限已经按调用关联记录,终态消费不能再记同一次超时
crate::RpcError::Timeout => return SessionError::Reported,
crate::RpcError::Remote(fault) => match fault.code {
ErrorCode::Cancelled => return SessionError::Cancelled,
ErrorCode::Unsupported => "unsupported",
ErrorCode::Rejected => "rejected",
ErrorCode::InvalidInput => "invalid_input",
ErrorCode::PermissionDenied => "permission_denied",
ErrorCode::Upstream => "upstream",
ErrorCode::Fault => "fault",
ErrorCode::Timeout => "timeout",
ErrorCode::Uncertain => "uncertain",
ErrorCode::Capacity => "capacity",
ErrorCode::Conflict => "conflict",
},
_ => error.diagnostic().0,
};
SessionError::Failed {
kind,
plugin_instance_id: Some(instance_id.to_owned()),
}
}
@@ -15,14 +15,7 @@ async fn middleware_preserves_three_state_fast_and_sources() {
] {
let snapshot = gateway_core::routing::RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(
1,
0,
"smart",
Default::default(),
None,
None,
),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![],
vec![],
vec![],
@@ -395,14 +395,7 @@ async fn active_http_children_inherit_settings_and_preserve_override_sources() {
let baseline = RequestSettings::new(Arc::new(
RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(
1,
10,
"smart",
Default::default(),
None,
None,
),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![],
vec![],
vec![],
@@ -439,7 +432,7 @@ async fn active_http_children_inherit_settings_and_preserve_override_sources() {
final_settings["responses_max_decompressed_body_bytes"],
8192
);
assert_eq!(final_settings["request_interval_ms"], 0);
assert_eq!(final_settings["max_account_rotations"], 0);
assert!(final_settings["min_codex_cli_version"].is_null());
let sources = effective.inspect();
assert_eq!(
@@ -447,14 +440,14 @@ async fn active_http_children_inherit_settings_and_preserve_override_sources() {
"dispatch-b"
);
assert_eq!(
sources["overrides"]["request_interval_ms"]["instance_id"],
sources["overrides"]["max_account_rotations"]["instance_id"],
"dispatch-a"
);
assert!(
sources["overrides"]["min_codex_cli_version"]["order"]
.as_u64()
.unwrap()
> sources["overrides"]["request_interval_ms"]["order"]
> sources["overrides"]["max_account_rotations"]["order"]
.as_u64()
.unwrap()
);
@@ -527,14 +527,7 @@ async fn http_onion_performance_baseline() {
let plan = runtime.execution_registry().middleware(&generation);
let snapshot = gateway_core::routing::RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(
1,
0,
"smart",
Default::default(),
None,
None,
),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![],
vec![],
vec![],
@@ -347,3 +347,65 @@ async fn dropping_an_uncommitted_upgrade_or_cancelling_a_session_reclaims_its_so
runtime.shutdown().await;
}
}
#[tokio::test]
async fn upgraded_plugin_business_failure_reaches_terminal_owner_with_safe_classification() {
for mode in ["session-error", "session-cancelled"] {
let (_cache, runtime, generation, plan) = http_plugin(mode).await;
let (_send, incoming) = tokio::sync::mpsc::channel(1);
let (outgoing, _receive) = tokio::sync::mpsc::channel(1);
let transport = Arc::new(Transport {
incoming: tokio::sync::Mutex::new(incoming),
outgoing,
dropped: Arc::default(),
});
let mut request = ::http::Request::builder()
.uri("/custom/session")
.body(core::empty_body())
.unwrap();
request
.extensions_mut()
.insert(Arc::new(Upgrade(transport)) as Arc<dyn upgrade::WebSocketUpgrade>);
let mut response = plan
.handle_http(
core::Context {
plugin_instance_id: None,
request_id: "session-failure".into(),
call_id: "http-call".into(),
parent_call_id: None,
extensions: Default::default(),
plan: Some(plan.clone()),
cancellation: CancellationToken::new(),
},
request,
compose(Vec::new(), |_| {
Box::pin(async { panic!("session owns its response") })
}),
)
.await
.unwrap();
assert_eq!(response.status(), 101);
let task = response
.extensions_mut()
.remove::<upgrade::Upgraded>()
.unwrap()
.take_task()
.unwrap();
let error = tokio::time::timeout(Duration::from_secs(5), task)
.await
.unwrap()
.unwrap_err();
if mode == "session-error" {
assert!(
matches!(error, upgrade::SessionError::Failed { kind: "permission_denied", plugin_instance_id: Some(ref id) } if id == "http-plugin")
);
} else {
assert!(matches!(error, upgrade::SessionError::Cancelled));
}
assert!(!format!("{error:?}").contains("private-plugin-error"));
drop(response);
drop(plan);
drop(generation);
runtime.shutdown().await;
}
}
@@ -553,7 +553,6 @@ fn selection_context() -> AccountSelectionContext {
policy: AccountSelectionPolicy::new(
RotationStrategy::RoundRobin,
NonZeroU32::new(4).unwrap(),
Duration::ZERO,
),
now: SystemTime::now(),
excluded_accounts: BTreeSet::new(),
@@ -1182,7 +1181,7 @@ async fn real_request_plugins_rewrite_settings_in_onion_order() {
)])
};
let baseline = ExecutionSettings {
runtime: SettingsValues::new(3, 50, "smart", BTreeMap::new(), None, None)
runtime: SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
.with_request_profiles(profiles("key-default")),
fast_mode: FastMode::Disabled,
client_limits: RateLimits {
@@ -1194,7 +1193,7 @@ async fn real_request_plugins_rewrite_settings_in_onion_order() {
let baseline_json = serde_json::to_value(&baseline).unwrap();
let host_snapshot = gateway_core::routing::RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(1, 0, "smart", Default::default(), None, None),
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
vec![],
vec![],
vec![],
@@ -1230,7 +1229,7 @@ async fn real_request_plugins_rewrite_settings_in_onion_order() {
);
let mut first = baseline_json.clone();
first["fast_mode"] = serde_json::json!("default");
first["runtime"]["request_interval_ms"] = serde_json::json!(0);
first["runtime"]["max_account_rotations"] = serde_json::json!(0);
first["runtime"]["model_mappings"] = serde_json::json!({"alias":"model-one"});
let mut last = first.clone();
last["runtime"]["model_mappings"] = serde_json::json!({"alias":"model-two"});
@@ -52,8 +52,8 @@ async fn explicit_service_dispatch_composes_once_and_preserves_native_transactio
.unwrap();
let visible = services.settings().load().await.unwrap();
assert_eq!(
visible.values.request_interval_ms,
original.values.request_interval_ms
visible.values.refresh_margin_seconds,
original.values.refresh_margin_seconds
);
assert_eq!(visible.config_revision, original.config_revision);
let through_registry = services
@@ -63,7 +63,7 @@ async fn explicit_service_dispatch_composes_once_and_preserves_native_transactio
.unwrap();
let typed: gateway_plugin_sdk::call::services::settings::RuntimeSettings =
serde_json::from_value(through_registry).unwrap();
assert_eq!(typed.request_interval_ms, 4321);
assert_eq!(typed.refresh_margin_seconds, 4321);
// 返回层改写不是持久化写入;原 owner 保留真实配置事实
assert_eq!(
environment
@@ -74,12 +74,12 @@ async fn explicit_service_dispatch_composes_once_and_preserves_native_transactio
.await
.unwrap()
.values
.request_interval_ms,
original.values.request_interval_ms
.refresh_margin_seconds,
original.values.refresh_margin_seconds
);
let mut command = ReplaceRuntimeSettings::from(original);
command.values.request_interval_ms = 999;
command.values.refresh_margin_seconds = 999;
let saved = services
.public_services()
.call(
@@ -90,7 +90,7 @@ async fn explicit_service_dispatch_composes_once_and_preserves_native_transactio
.unwrap();
let saved: gateway_admin::model::settings::RuntimeSettings =
serde_json::from_value(saved).unwrap();
assert_eq!(saved.values.request_interval_ms, 1234);
assert_eq!(saved.values.refresh_margin_seconds, 1234);
assert!(saved.config_revision > command.expected_revision);
let current = environment
.store
@@ -99,7 +99,7 @@ async fn explicit_service_dispatch_composes_once_and_preserves_native_transactio
.load_runtime_settings()
.await
.unwrap();
assert_eq!(current.values.request_interval_ms, 1234);
assert_eq!(current.values.refresh_margin_seconds, 1234);
assert_eq!(
core.snapshots()
.snapshot_for_diagnostics()
@@ -323,7 +323,7 @@ async fn cli_service_callback_uses_the_same_public_settings_and_plugin_chain() {
serde_json::from_str(&output.stdout).unwrap();
let returned: gateway_plugin_sdk::call::services::settings::RuntimeSettings =
serde_json::from_value(result.unwrap()).unwrap();
assert_eq!(returned.request_interval_ms, 4321);
assert_eq!(returned.refresh_margin_seconds, 4321);
assert_ne!(
environment
.store
@@ -333,7 +333,7 @@ async fn cli_service_callback_uses_the_same_public_settings_and_plugin_chain() {
.await
.unwrap()
.values
.request_interval_ms,
.refresh_margin_seconds,
4321
);
commands.shutdown().await;
@@ -161,7 +161,6 @@ fn context_with_fast_mode(
AccountSelectionPolicy::new(
gateway_core::account::RotationStrategy::RoundRobin,
std::num::NonZeroU32::new(1).unwrap(),
Duration::ZERO,
),
AccountAttemptContext::new(Default::default(), None, None),
None,
@@ -699,9 +699,6 @@ impl Environment {
} else {
gateway_store::initialize(config).await.unwrap()
};
// 迁移可能已创建默认行;该场景只测租约释放,明确关闭相邻请求的间隔限制
sqlx::raw_sql(sqlx::AssertSqlSafe(format!("insert into {schema}.runtime_settings(id, config_revision, request_interval_ms, updated_at) values (1,1,0,now()) on conflict (id) do update set request_interval_ms=0")))
.execute(&admin).await.unwrap();
Some(Self {
store,
directory,
@@ -189,13 +189,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut output = call.next.run(()).await.unwrap();
assert_eq!(output.config_revision, nested.config_revision);
if mode == "rewrite" {
output.request_interval_ms = 4321;
output.refresh_margin_seconds = 4321;
}
return ServiceResponse::from_result::<settings::Load>(Ok(output));
}
if call.is::<settings::Replace>() {
let mut call = call.into_typed::<settings::Replace>()?;
call.input.1.request_interval_ms = 1234;
call.input.1.refresh_margin_seconds = 1234;
return ServiceResponse::from_result::<settings::Replace>(
call.next.run(call.input).await,
);
@@ -311,7 +311,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
assert!(call.request.settings["min_codex_cli_version"].is_null());
call.request.settings["min_codex_cli_version"] =
serde_json::json!("0.40.0");
call.request.settings["request_interval_ms"] = serde_json::json!(0);
call.request.settings["max_account_rotations"] = serde_json::json!(0);
} else {
assert_eq!(call.request.timeout_ms, Some(90_000));
assert_eq!(
@@ -331,6 +331,22 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
}
return call.host.dispatch_http(call.request).await;
}
if mode == "session-error" || mode == "session-cancelled" {
use gateway_plugin_sdk::{ErrorCode, PluginFault};
let code = if mode == "session-error" {
ErrorCode::PermissionDenied
} else {
ErrorCode::Cancelled
};
return call
.upgrade(vec!["test-session".into()], move |_session| async move {
Err(PluginFault::new(
code,
"private-plugin-error-must-not-be-logged",
))
})
.await;
}
if mode == "callbacks" {
let host = call.host.clone();
return call.upgrade(vec!["test-session".into()], |mut session| async move {
@@ -95,7 +95,6 @@ pub struct RuntimeSettings {
pub refresh_margin_seconds: u64,
pub refresh_concurrency: u32,
pub max_concurrent_per_account: u32,
pub request_interval_ms: u64,
pub max_waiting_per_key: u32,
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
@@ -135,7 +134,6 @@ pub struct ReplaceRuntimeSettings {
pub refresh_margin_seconds: u64,
pub refresh_concurrency: u32,
pub max_concurrent_per_account: u32,
pub request_interval_ms: u64,
pub max_waiting_per_key: u32,
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
@@ -338,7 +336,6 @@ impl From<RuntimeSettings> for ReplaceRuntimeSettings {
refresh_margin_seconds: settings.refresh_margin_seconds,
refresh_concurrency: settings.refresh_concurrency,
max_concurrent_per_account: settings.max_concurrent_per_account,
request_interval_ms: settings.request_interval_ms,
max_waiting_per_key: settings.max_waiting_per_key,
max_waiting_per_account: settings.max_waiting_per_account,
concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds,
+1 -1
View File
@@ -38,7 +38,7 @@ uuid = { version = "1.23.3", features = ["v4", "v7"] }
[dev-dependencies]
tempfile = "3.23"
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt-multi-thread", "fs", "io-util", "net", "process", "time"] }
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt-multi-thread", "fs", "io-util", "net", "process", "time", "test-util"] }
[lints]
workspace = true
@@ -101,7 +101,6 @@ impl SettingsStore for AdminSettingsStoreAdapter {
"refresh_margin_seconds".to_owned(),
"refresh_concurrency".to_owned(),
"max_concurrent_per_account".to_owned(),
"request_interval_ms".to_owned(),
"max_waiting_per_key".to_owned(),
"max_waiting_per_account".to_owned(),
"concurrency_wait_timeout_seconds".to_owned(),
+20 -12
View File
@@ -192,7 +192,15 @@ impl BackupObjectStorePort for S3ObjectStoreAdapter {
}
}
let _ = best_effort_delete(&client, config, &probe_key).await;
if let Err(error) = client
.delete_object()
.bucket(&config.bucket)
.key(&probe_key)
.send()
.await
{
return Ok(failed(ConnectionTestStage::DeleteObject, &error));
}
Ok(ConnectionTestResult {
ok: true,
stage: ConnectionTestStage::DeleteObject.as_str(),
@@ -246,24 +254,24 @@ impl BackupObjectStorePort for S3ObjectStoreAdapter {
Err(BackupError::new(code::CANCELLED, "上传已取消".to_owned()))
}
};
if let Err(error) = &outcome {
let _ = client
if let Err(error) = outcome {
let cleanup = client
.abort_multipart_upload()
.bucket(&config.bucket)
.key(&request.object_key)
.upload_id(&upload_id)
.send()
.await;
if error.code() != code::CANCELLED {
let _ = client
.delete_object()
.bucket(&config.bucket)
.key(&request.object_key)
.send()
.await;
}
return Err(match cleanup {
Ok(_) => error,
Err(cleanup) => error.with_cleanup_failure(
"AbortMultipartUpload",
&map_s3_error(cleanup),
false,
),
});
}
outcome
Ok(())
}
async fn head_object(
+6 -2
View File
@@ -126,6 +126,7 @@ async fn connect(
let observability_query_budget = postgres::ObservabilityQueryBudget::try_new(
config.pool.observability_max_connections(),
config.pool.acquire_timeout(),
config.pool.observability_timeout(),
)?;
let redis_client = ::redis::Client::open(config.redis_url()?)
.map_err(|source| redis_unavailable("create Redis client", source))?;
@@ -180,8 +181,11 @@ async fn connect(
)),
admin_account_runtime,
Arc::new(
postgres::PgAccountGroupRepository::new(pool.clone())
.with_timezone(config.timezone),
postgres::PgAccountGroupRepository::new(
pool.clone(),
observability_query_budget.clone(),
)
.with_timezone(config.timezone),
),
Arc::new(postgres::PgProxyRepository::new(pool.clone())),
),
@@ -36,6 +36,7 @@ pub struct StoreConfig {
pub struct StorePoolConfig {
pub max_connections: u32,
pub acquire_timeout_seconds: u64,
pub observability_timeout_seconds: u64,
}
impl Default for StorePoolConfig {
@@ -43,6 +44,7 @@ impl Default for StorePoolConfig {
Self {
max_connections: 20,
acquire_timeout_seconds: 5,
observability_timeout_seconds: 120,
}
}
}
@@ -58,9 +60,22 @@ impl StorePoolConfig {
.to_owned(),
});
}
if self.observability_timeout_seconds == 0
|| self.observability_timeout().as_millis() > i32::MAX as u128
{
return Err(StoreError::InvalidData {
source: None,
entity: "store config",
message: "pool.observability_timeout_seconds must be positive and fit PostgreSQL statement_timeout".to_owned(),
});
}
Ok(())
}
pub(crate) const fn observability_timeout(self) -> Duration {
Duration::from_secs(self.observability_timeout_seconds)
}
/// 管理观测查询可并发占用的连接数;始终为数据面保留约 20% 的池容量
#[must_use]
pub const fn observability_max_connections(self) -> u32 {
@@ -3,7 +3,10 @@
use std::time::Duration;
use futures::future::BoxFuture;
use gateway_core::lifecycle::{CancellationToken, Deadline, REQUEST_LEASE_TTL};
use gateway_core::lifecycle::{
CancellationReason, CancellationToken, Deadline, LeaseFailure, REQUEST_LEASE_TTL,
RequestLeaseKind,
};
use crate::StoreResult;
@@ -12,22 +15,22 @@ pub(crate) struct LeaseRenewal(tokio::task::JoinHandle<()>);
impl LeaseRenewal {
pub(crate) fn spawn(
deadline: Deadline,
cancellation: Option<CancellationToken>,
cancellation: Option<(CancellationToken, RequestLeaseKind)>,
mut renew: impl FnMut(Duration) -> BoxFuture<'static, StoreResult<bool>> + Send + 'static,
) -> Self {
Self(tokio::spawn(async move {
let cancelled = async {
match &cancellation {
Some(cancellation) => cancellation.cancelled().await,
Some((cancellation, _)) => cancellation.cancelled().await,
None => std::future::pending().await,
}
};
tokio::pin!(cancelled);
let mut expires = tokio::time::Instant::now() + deadline.bounded(REQUEST_LEASE_TTL);
loop {
let failure = loop {
let ttl = deadline.bounded(REQUEST_LEASE_TTL);
if ttl.is_zero() {
break;
break LeaseFailure::Expired;
}
let renewed_until = tokio::time::Instant::now() + ttl;
let result = tokio::select! {
@@ -40,7 +43,8 @@ impl LeaseRenewal {
expires = renewed_until;
REQUEST_LEASE_TTL / 3
}
Ok(Ok(false)) | Err(_) => break,
Ok(Ok(false)) => break LeaseFailure::Lost,
Err(_) => break LeaseFailure::RenewalTimedOut,
Ok(Err(error)) => {
tracing::warn!(%error, "请求租约续期失败,将在租约有效期内重试");
Duration::from_secs(1)
@@ -49,15 +53,18 @@ impl LeaseRenewal {
tokio::select! {
biased;
() = &mut cancelled => return,
() = tokio::time::sleep_until(expires) => break,
() = tokio::time::sleep_until(expires) => break LeaseFailure::Expired,
() = tokio::time::sleep(delay) => {}
}
}
};
// 显式执行截止由 Core 归类为 timeout,租约不把它抢先改成 cancelled
if !deadline.is_elapsed()
&& let Some(cancellation) = &cancellation
&& let Some((cancellation, kind)) = &cancellation
{
cancellation.cancel();
cancellation.cancel_with_reason(CancellationReason::Lease {
kind: *kind,
failure,
});
}
}))
}
@@ -31,9 +31,9 @@ use crate::{
};
use super::{
PgRuntimeSettingsRepository, RuntimeSettingsRepository, account_summary_from_row,
append_admin_audit_event_in_transaction, bump_config_revision_in_transaction,
completed_usage_fact_predicate,
ObservabilityQueryBudget, PgRuntimeSettingsRepository, RuntimeSettingsRepository,
account_summary_from_row, append_admin_audit_event_in_transaction,
bump_config_revision_in_transaction,
};
const ENTITY: &str = "account group";
@@ -43,13 +43,15 @@ const ENTITY: &str = "account group";
pub struct PgAccountGroupRepository {
timezone: gateway_core::time::DeploymentTimeZone,
pool: PgPool,
query_budget: ObservabilityQueryBudget,
}
impl PgAccountGroupRepository {
#[must_use]
pub fn new(pool: PgPool) -> Self {
pub fn new(pool: PgPool, query_budget: ObservabilityQueryBudget) -> Self {
Self {
pool,
query_budget,
timezone: Default::default(),
}
}
@@ -74,9 +76,14 @@ impl PgAccountGroupRepository {
.await
.map_err(|error| admin_store_error(ENTITY, error))?
.ok_or_else(|| not_found(id.as_str()))?;
let costs = group_costs(&self.pool, &[id.as_str().to_owned()], self.timezone)
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
let costs = group_costs(
&self.pool,
&self.query_budget,
&[id.as_str().to_owned()],
self.timezone,
)
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
if let Some(usage) = costs.get(id.as_str()) {
record.usage = usage.clone();
}
@@ -152,7 +159,7 @@ impl AccountGroupStore for PgAccountGroupRepository {
.iter()
.map(|record| record.id.as_str().to_owned())
.collect::<Vec<_>>();
let costs = group_costs(&self.pool, &group_ids, self.timezone)
let costs = group_costs(&self.pool, &self.query_budget, &group_ids, self.timezone)
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
for record in &mut items {
@@ -601,71 +608,23 @@ fn group_ref(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupRef> {
async fn group_costs(
pool: &PgPool,
budget: &ObservabilityQueryBudget,
group_ids: &[String],
timezone: gateway_core::time::DeploymentTimeZone,
) -> StoreResult<BTreeMap<String, AccountGroupUsage>> {
if group_ids.is_empty() {
return Ok(BTreeMap::new());
}
let completed_usage = completed_usage_fact_predicate("mr");
// 用量按实际完成请求的账号统计:账号属于多个分组时计入每个所属分组,
// 不再按 Client Key 绑定分组快照归属,避免多分组 Key 的费用重复出现在未承接请求的分组上
let statement = format!(
"with requested_groups(group_id) as (
select unnest($1::text[])
)
select requested_groups.group_id,
coalesce(sum(mr.cost_amount) filter (
where mr.started_at >= $2
), 0)::text as today_usd,
coalesce(sum(mr.cost_amount), 0)::text as retained_total_usd
from requested_groups
cross join runtime_settings settings
left join account_group_accounts membership
on membership.account_group_id = requested_groups.group_id
left join model_requests mr
on mr.provider_account_ref = membership.provider_account_id
and mr.started_at >= now() - make_interval(days => settings.usage_retention_days::int)
and {completed_usage}
and mr.cost_currency = 'USD'
and mr.cost_amount is not null
where settings.id = 1
group by requested_groups.group_id"
);
// 动态片段仅为共享的固定 usage-fact predicate;group IDs 仍使用 bind
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(group_ids)
.bind(
timezone
.day_start(chrono::Utc::now())
.ok_or_else(|| invalid("invalid business date"))?,
budget
.run(
pool,
"load account group costs",
|mut connection| async move {
super::observability::account_group_usage(&mut connection, group_ids, timezone)
.await
},
)
.fetch_all(pool)
.await
.map_err(|source| unavailable("load account group costs", source))?;
rows.into_iter()
.map(|row| {
let group_id: String = row
.try_get("group_id")
.map_err(|source| invalid("invalid group ID").with_source(source))?;
let today = row
.try_get::<String, _>("today_usd")
.map_err(|source| invalid("invalid today cost").with_source(source))?;
let retained_total = row
.try_get::<String, _>("retained_total_usd")
.map_err(|source| invalid("invalid retained total cost").with_source(source))?;
Ok((
group_id,
AccountGroupUsage {
today_usd: DecimalAmount::from_str(&today)
.map_err(|source| invalid("invalid today cost").with_source(source))?,
retained_total_usd: DecimalAmount::from_str(&retained_total).map_err(
|source| invalid("invalid retained total cost").with_source(source),
)?,
},
))
})
.collect()
}
fn count_value(row: &sqlx::postgres::PgRow, field: &str) -> StoreResult<u64> {
@@ -197,7 +197,7 @@ impl PgObservabilityRepository {
)
.await?;
let account_usage =
provider_account_usage(&mut snapshot.tx, account_usage_query).await?;
provider_account_usage(&mut snapshot, account_usage_query).await?;
let metadata = snapshot.metadata();
snapshot.finish().await?;
Ok((totals, trend, metadata, account_usage))
@@ -290,12 +290,8 @@ impl PgObservabilityRepository {
"load usage insights snapshot",
|mut connection| async move {
let mut snapshot = analytics::AnalyticsSnapshot::new(&mut connection).await?;
let mut overview = snapshot.overview(range, &filter).await?;
let trend = snapshot
.trend(range, &filter, granularity, self.timezone, true)
.await?;
let costs = snapshot
.projected_costs(range, &filter, granularity, self.timezone)
let (mut overview, trend, costs) = snapshot
.insights(range, &filter, granularity, self.timezone)
.await?;
overview.analytics = snapshot.metadata();
snapshot.finish().await?;
@@ -318,12 +314,9 @@ impl PgObservabilityRepository {
&self.pool,
"load provider account usage",
|mut connection| async move {
let mut snapshot = analytics::AnalyticsSnapshot::new(&mut connection).await?;
snapshot.require_visible(Some(query.range)).await?;
if let Some(range) = query.request_bucket_range {
snapshot.require_visible(Some(range)).await?;
}
let usage = provider_account_usage(&mut snapshot.tx, query).await?;
let mut snapshot =
analytics::AnalyticsSnapshot::new_for_accounts(&mut connection).await?;
let usage = provider_account_usage(&mut snapshot, query).await?;
snapshot.finish().await?;
Ok(usage)
},
@@ -407,6 +400,10 @@ impl PgObservabilityRepository {
#[async_trait]
impl AdminObservabilityStore for PgAdminObservabilityStore {
fn query_timeout(&self) -> std::time::Duration {
self.repository.query_budget.query_timeout()
}
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64> {
self.repository
.query_budget
@@ -0,0 +1,655 @@
//! 账号用量复用汇总并补齐未投影源版本,保持同步查询的精确窗口口径
use super::query::push_summary_buckets;
use super::*;
use gateway_admin::model::{
account_groups::AccountGroupUsage,
accounts::{
AccountCost, AccountModelUsage, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult,
},
};
type AccountStates = BTreeMap<(String, String), MetricState>;
impl AnalyticsSnapshot<'_> {
pub(crate) async fn account_usage_states(
&mut self,
range: ObservabilityRange,
account_ids: &[String],
) -> StoreResult<AccountStates> {
self.account_states(range, account_ids, true).await
}
async fn account_states(
&mut self,
range: ObservabilityRange,
account_ids: &[String],
by_model: bool,
) -> StoreResult<AccountStates> {
if account_ids.is_empty() {
return Ok(BTreeMap::new());
}
// 原同步接口在首次回填期间仍可查询;只有完整覆盖的区间才允许用汇总替代源事实
if !self.projection_covers(range) {
return self
.account_fact_states("model_requests", range, account_ids, false, by_model)
.await;
}
let mut states = BTreeMap::new();
// 排序与分组费用只需要账号总量,不额外展开模型维度
let (collection, model) = if by_model {
(
"account_model",
"coalesce(nullif(dimensions[4],''),dimensions[3])",
)
} else {
("account", "''::text")
};
let (hours, quarters) = self.summary_buckets(range, true).await?;
for (table, seconds, buckets) in [
("request_analytics_hourly", 3600, hours),
("request_analytics_15m", 900, quarters),
] {
if buckets.is_empty() {
continue;
}
let mut query = QueryBuilder::<Postgres>::new(format!(
"with selected as materialized (select dimensions[1] as account_id,{model} as model,counts,costs,last_used_at from {table} a where collection='{collection}' and schema_version="
));
query
.push_bind(SCHEMA_VERSION)
.push(" and histogram_version=")
.push_bind(HISTOGRAM_VERSION)
.push(" and dimensions[1]=any(")
.push_bind(account_ids)
.push("::text[])");
push_summary_buckets(&mut query, &buckets, seconds);
let counters = usage_counter_names();
// 每行一次展开指标,避免对可能位于 TOAST 的同一 JSON 反复解码
query.push(format!(
"), totals as (select account_id,model,jsonb_build_object({}) as counts,max(last_used_at) as last_used_at from selected cross join lateral jsonb_to_record(counts) as metric({}) group by account_id,model),
currencies as (select account_id,model,c.key,sum(c.value::numeric)::text as amount from selected cross join lateral jsonb_each_text(costs) c where c.key not in ('standard_usd','no_cache_usd') group by account_id,model,c.key),
costs as (select account_id,model,jsonb_object_agg(key,amount) as costs from currencies group by account_id,model)
select totals.*,coalesce(costs.costs,'{{}}'::jsonb) as costs,'{{}}'::jsonb as distributions from totals left join costs using(account_id,model)",
counters.iter().map(|name| format!("'{name}',coalesce(sum(metric.{name}),0)")).collect::<Vec<_>>().join(","),
counters.iter().map(|name| format!("{name} numeric")).collect::<Vec<_>>().join(",")
));
let rows = query
.build()
.fetch_all(&mut *self.tx)
.await
.map_err(|error| {
postgres_unavailable("aggregate account usage summaries", error)
})?;
merge_account_rows(&mut states, rows)?;
}
for (start, end) in self.uncovered_spans(range).await? {
let boundary = self
.account_fact_states(
"request_analytics_facts",
ObservabilityRange { start, end },
account_ids,
false,
by_model,
)
.await?;
merge_account_states(&mut states, boundary)?;
}
if !self.has_pending_requests().await? {
states.retain(|_, state| state.count("completed_request_count") > 0);
return Ok(states);
}
let old = self
.account_fact_states(
"request_analytics_facts",
range,
account_ids,
true,
by_model,
)
.await?;
let current = self
.account_fact_states("model_requests", range, account_ids, true, by_model)
.await?;
let mut repair = BTreeSet::new();
for (key, previous) in old {
let state = states.entry(key.clone()).or_default();
for (name, count) in &previous.counts {
let value = state.counts.entry(name.clone()).or_default();
*value = value
.checked_sub(*count)
.ok_or_else(|| invalid("account usage pending counter underflow"))?;
}
for (currency, amount) in &previous.costs {
let before = state
.costs
.get(currency)
.ok_or_else(|| invalid("account usage pending cost is missing"))?;
let before = before
.parse::<gateway_core::metering::Decimal>()
.map_err(|_| invalid("invalid account cost"))?;
let removed = amount
.parse::<gateway_core::metering::Decimal>()
.map_err(|_| invalid("invalid pending account cost"))?;
let difference = before
.scaled()
.checked_sub(removed.scaled())
.ok_or_else(|| invalid("account usage pending cost underflow"))?;
let value = gateway_core::metering::Decimal::from_scaled(difference)
.map_err(|_| invalid("invalid account cost difference"))?;
state.costs.insert(currency.clone(), value.canonical());
}
repair.insert(key);
}
merge_account_states(&mut states, current)?;
// 撤销最大值不能用减法;只对发生源版本替换的账号模型做索引定位
let completed = completed_usage_fact_predicate("mr");
for (account, model) in repair {
let state = states
.get_mut(&(account.clone(), model.clone()))
.ok_or_else(|| invalid("pending account state disappeared"))?;
if state.count("completed_request_count") == 0 {
*state = MetricState::default();
continue;
}
let statement = format!(
"select started_at from model_requests mr where provider_account_ref=$1 and ($2::text is null or coalesce(upstream_model_id,requested_model_id,'')=$2) and started_at >= $3 and started_at < $4 and {completed} order by started_at desc,id desc limit 1"
);
state.last_used_at = sqlx::query_scalar(sqlx::AssertSqlSafe(statement))
.bind(&account)
.bind(by_model.then_some(&model))
.bind(range.start)
.bind(range.end)
.fetch_optional(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("repair account latest usage", error))?;
let zero_currencies = state
.costs
.iter()
.filter(|(_, amount)| amount.as_str() == "0")
.map(|(currency, _)| currency.clone())
.collect::<Vec<_>>();
for currency in zero_currencies {
let statement = format!(
"select exists(select 1 from model_requests mr where provider_account_ref=$1 and ($2::text is null or coalesce(upstream_model_id,requested_model_id,'')=$2) and started_at >= $3 and started_at < $4 and {completed} and cost_currency=$5 and cost_amount is not null)"
);
let exists: bool = sqlx::query_scalar(sqlx::AssertSqlSafe(statement))
.bind(&account)
.bind(by_model.then_some(&model))
.bind(range.start)
.bind(range.end)
.bind(&currency)
.fetch_one(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("repair account cost coverage", error))?;
if !exists {
state.costs.remove(&currency);
}
}
}
states.retain(|_, state| state.count("completed_request_count") > 0);
Ok(states)
}
pub(super) async fn has_pending_requests(&mut self) -> StoreResult<bool> {
if self.status.oldest_pending_at.is_none() {
return Ok(false);
}
sqlx::query_scalar("select exists(select 1 from request_analytics_pending_requests)")
.fetch_one(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("check pending account usage", error))
}
pub(super) fn projection_covers(&self, range: ObservabilityRange) -> bool {
self.status.schema_version == SCHEMA_VERSION as u32
&& self.status.histogram_version == HISTOGRAM_VERSION as u32
&& (!self.status.building
|| (self
.status
.coverage_start
.is_some_and(|start| range.start >= start)
&& self.status.coverage_end.is_some()))
}
pub(crate) async fn account_request_counts(
&mut self,
range: ObservabilityRange,
account_ids: &[String],
) -> StoreResult<BTreeMap<(String, i64), u64>> {
if account_ids.is_empty() {
return Ok(BTreeMap::new());
}
let mut counts = BTreeMap::new();
// 公开小时图采用 UTC 整点;非整刻内部窗口直接在源表按相对小时分组
if !self.projection_covers(range)
|| range.start.timestamp().rem_euclid(900) != 0
|| range.start.timestamp_subsec_nanos() != 0
{
let rows = self
.account_request_fact_counts(
"model_requests",
range,
range.start,
account_ids,
false,
)
.await?;
merge_request_counts(&mut counts, rows, false)?;
return Ok(counts);
}
let (hours, quarters) = self
.summary_buckets(range, range.start.timestamp().rem_euclid(3600) == 0)
.await?;
for (table, seconds, buckets) in [
("request_analytics_hourly", 3600, hours),
("request_analytics_15m", 900, quarters),
] {
if buckets.is_empty() {
continue;
}
let mut query = QueryBuilder::<Postgres>::new(
"select dimensions[1] as account_id,floor(extract(epoch from (bucket_start-",
);
query.push_bind(range.start).push("::timestamptz))/3600)::bigint as bucket_index,coalesce(sum((counts->>'completed_request_count')::bigint),0)::bigint as request_count from ")
.push(table).push(" a where collection='account' and schema_version=").push_bind(SCHEMA_VERSION)
.push(" and histogram_version=").push_bind(HISTOGRAM_VERSION).push(" and dimensions[1]=any(").push_bind(account_ids).push("::text[])");
push_summary_buckets(&mut query, &buckets, seconds);
query.push(" group by dimensions[1],bucket_index");
let rows = query
.build()
.fetch_all(&mut *self.tx)
.await
.map_err(|error| {
postgres_unavailable("read account hourly request summaries", error)
})?;
merge_request_counts(&mut counts, rows, false)?;
}
for (start, end) in self.uncovered_spans(range).await? {
let rows = self
.account_request_fact_counts(
"request_analytics_facts",
ObservabilityRange { start, end },
range.start,
account_ids,
false,
)
.await?;
merge_request_counts(&mut counts, rows, false)?;
}
if self.has_pending_requests().await? {
let old = self
.account_request_fact_counts(
"request_analytics_facts",
range,
range.start,
account_ids,
true,
)
.await?;
merge_request_counts(&mut counts, old, true)?;
let current = self
.account_request_fact_counts(
"model_requests",
range,
range.start,
account_ids,
true,
)
.await?;
merge_request_counts(&mut counts, current, false)?;
}
Ok(counts)
}
async fn account_request_fact_counts(
&mut self,
table: &'static str,
range: ObservabilityRange,
origin: DateTime<Utc>,
account_ids: &[String],
pending_only: bool,
) -> StoreResult<Vec<sqlx::postgres::PgRow>> {
let mut query = QueryBuilder::<Postgres>::new(
"select mr.provider_account_ref as account_id,floor(extract(epoch from (mr.started_at-",
);
query.push_bind(origin).push("::timestamptz))/3600)::bigint as bucket_index,count(*)::bigint as request_count from ").push(table).push(" mr");
if pending_only {
query.push(" join request_analytics_pending_requests p on p.request_id=mr.id");
}
query
.push(" where mr.provider_account_ref=any(")
.push_bind(account_ids)
.push("::text[]) and mr.started_at >= ")
.push_bind(range.start)
.push(" and mr.started_at < ")
.push_bind(range.end);
push_completed_usage_fact_filter(&mut query, "mr");
query.push(" group by mr.provider_account_ref,bucket_index");
query
.build()
.fetch_all(&mut *self.tx)
.await
.map_err(|error| {
postgres_unavailable("aggregate account hourly request boundaries", error)
})
}
async fn account_fact_states(
&mut self,
table: &'static str,
range: ObservabilityRange,
account_ids: &[String],
pending_only: bool,
by_model: bool,
) -> StoreResult<AccountStates> {
let mut query = QueryBuilder::<Postgres>::new(usage_fact_aggregate_sql(table, by_model));
if pending_only {
query.push(
" join request_analytics_pending_requests pending on pending.request_id=mr.id",
);
}
query
.push(" where mr.provider_account_ref=any(")
.push_bind(account_ids)
.push("::text[]) and mr.started_at >= ")
.push_bind(range.start)
.push(" and mr.started_at < ")
.push_bind(range.end);
push_completed_usage_fact_filter(&mut query, "mr");
query.push(" group by mr.provider_account_ref,mr.cost_currency");
if by_model {
query.push(",coalesce(mr.upstream_model_id,mr.requested_model_id,'')");
}
let rows = query
.build()
.fetch_all(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("aggregate account usage facts", error))?;
let mut states = BTreeMap::new();
merge_account_rows(&mut states, rows)?;
Ok(states)
}
}
fn merge_request_counts(
counts: &mut BTreeMap<(String, i64), u64>,
rows: Vec<sqlx::postgres::PgRow>,
subtract: bool,
) -> StoreResult<()> {
for row in rows {
let value = unsigned(&row, "request_count")?;
let total = counts
.entry((get(&row, "account_id")?, get(&row, "bucket_index")?))
.or_default();
*total = if subtract {
total.checked_sub(value)
} else {
total.checked_add(value)
}
.ok_or_else(|| invalid("invalid account request count delta"))?;
}
Ok(())
}
fn merge_account_rows(
states: &mut AccountStates,
rows: Vec<sqlx::postgres::PgRow>,
) -> StoreResult<()> {
for row in rows {
states
.entry((get(&row, "account_id")?, get(&row, "model")?))
.or_default()
.merge(&MetricState::decode(&row)?)?;
}
Ok(())
}
fn merge_account_states(states: &mut AccountStates, incoming: AccountStates) -> StoreResult<()> {
for (key, state) in incoming {
states.entry(key).or_default().merge(&state)?;
}
Ok(())
}
fn usage_counter_names() -> Vec<String> {
let mut names = [
"completed_request_count",
"image_request_count",
"image_request_failed_count",
"window_total_tokens",
"provider_reported_count",
"calculated_count",
"unavailable_count",
]
.map(str::to_owned)
.to_vec();
for name in super::metrics::TOKEN_METRICS {
names.push(name.to_owned());
names.push(format!("{name}_known_count"));
}
names
}
fn usage_fact_aggregate_sql(table: &str, by_model: bool) -> String {
let model = if by_model {
"coalesce(mr.upstream_model_id,mr.requested_model_id,'')"
} else {
"''::text"
};
let mut counters = vec![
"'completed_request_count',count(*)".to_owned(),
"'image_request_count',count(*) filter(where image_generation_succeeded is true)".to_owned(),
"'image_request_failed_count',count(*) filter(where image_generation_succeeded is false)".to_owned(),
"'window_total_tokens',coalesce(sum(coalesce(total_tokens,coalesce(input_tokens,0)+coalesce(output_tokens,0))),0)".to_owned(),
];
for name in super::metrics::TOKEN_METRICS {
counters.push(format!(
"'{name}',coalesce(sum({name}),0),'{name}_known_count',count({name})"
));
}
for source in ["provider_reported", "calculated", "unavailable"] {
counters.push(format!(
"'{source}_count',count(*) filter(where cost_source='{source}')"
));
}
format!("select mr.provider_account_ref as account_id,{model} as model,
jsonb_build_object({}) as counts,case when mr.cost_currency is not null and sum(mr.cost_amount) is not null
then jsonb_build_object(mr.cost_currency,sum(mr.cost_amount)::text) else '{{}}'::jsonb end as costs,
'{{}}'::jsonb as distributions,max(mr.started_at) as last_used_at from {table} mr",counters.join(","))
}
pub(crate) async fn account_usage_windows(
connection: &mut ObservabilityConnection,
windows: &[AccountUsageWindowQuery],
) -> StoreResult<Vec<AccountUsageWindowResult>> {
let mut snapshot = AnalyticsSnapshot::new_for_accounts(connection).await?;
// 相同窗口一次查询所有账号,避免列表逐账号重复读取同一套时间桶
let mut grouped = BTreeMap::<(DateTime<Utc>, DateTime<Utc>), BTreeSet<String>>::new();
for window in windows {
grouped
.entry((window.range.start, window.range.end))
.or_default()
.insert(window.account_id.clone());
}
let mut results = Vec::with_capacity(windows.len());
for ((start, end), accounts) in grouped {
let ids = accounts.into_iter().collect::<Vec<_>>();
let states = snapshot
.account_usage_states(ObservabilityRange { start, end }, &ids)
.await?;
for window in windows
.iter()
.filter(|window| window.range.start == start && window.range.end == end)
{
let mut total = MetricState::default();
let mut models = Vec::new();
for ((account, model), state) in &states {
if account != &window.account_id {
continue;
}
total.merge(state)?;
if !model.is_empty()
&& let Some(last_used_at) = state.last_used_at()
{
models.push(AccountModelUsage {
model: model.clone(),
request_count: state.count("completed_request_count"),
success_count: state.count("completed_request_count"),
input_tokens: state.optional_count("input_tokens"),
output_tokens: state.optional_count("output_tokens"),
cached_input_tokens: state.optional_count("cached_input_tokens"),
cache_write_input_tokens: state.optional_count("cache_write_input_tokens"),
reasoning_output_tokens: state.optional_count("reasoning_output_tokens"),
image_input_tokens: state.optional_count("image_input_tokens"),
image_output_tokens: state.optional_count("image_output_tokens"),
total_tokens: Some(state.count("window_total_tokens")),
image_request_count: state.count("image_request_count"),
image_request_failed_count: state.count("image_request_failed_count"),
cost_coverage: admin_cost_coverage(state.coverage()),
costs: account_costs(state)?,
last_used_at,
});
}
}
models.sort_by(|a, b| {
b.request_count
.cmp(&a.request_count)
.then(a.model.cmp(&b.model))
});
results.push(AccountUsageWindowResult {
account_id: window.account_id.clone(),
key: window.key.clone(),
usage: AccountUsage {
account_id: window.account_id.clone(),
request_count: total.count("completed_request_count"),
success_count: total.count("completed_request_count"),
input_tokens: total.optional_count("input_tokens"),
output_tokens: total.optional_count("output_tokens"),
cached_input_tokens: total.optional_count("cached_input_tokens"),
cache_write_input_tokens: total.optional_count("cache_write_input_tokens"),
reasoning_output_tokens: total.optional_count("reasoning_output_tokens"),
image_input_tokens: total.optional_count("image_input_tokens"),
image_output_tokens: total.optional_count("image_output_tokens"),
total_tokens: Some(total.count("window_total_tokens")),
image_request_count: total.count("image_request_count"),
image_request_failed_count: total.count("image_request_failed_count"),
cost_coverage: admin_cost_coverage(total.coverage()),
costs: account_costs(&total)?,
last_used_at: total.last_used_at(),
request_buckets: Vec::new(),
models,
},
});
}
}
snapshot.finish().await?;
results.sort_by(|a, b| a.account_id.cmp(&b.account_id).then(a.key.cmp(&b.key)));
Ok(results)
}
fn account_costs(state: &MetricState) -> StoreResult<Vec<AccountCost>> {
Ok(state
.costs()?
.into_iter()
.map(|cost| AccountCost {
currency: cost.currency,
amount: cost.amount,
})
.collect())
}
pub(crate) async fn account_usage_sort_values(
snapshot: &mut AnalyticsSnapshot<'_>,
now: DateTime<Utc>,
) -> StoreResult<(Vec<String>, Vec<i64>, Vec<Option<DateTime<Utc>>>)> {
let days: i32 =
sqlx::query_scalar("select usage_retention_days::integer from runtime_settings where id=1")
.fetch_one(&mut *snapshot.tx)
.await
.map_err(|error| postgres_unavailable("read account usage retention", error))?;
let ids: Vec<String> = sqlx::query_scalar("select id from provider_accounts")
.fetch_all(&mut *snapshot.tx)
.await
.map_err(|error| postgres_unavailable("read account usage sort identities", error))?;
let range = ObservabilityRange::new(now - TimeDelta::days(i64::from(days)), now)?;
let mut totals = BTreeMap::<String, MetricState>::new();
for ((account, _), state) in snapshot.account_states(range, &ids, false).await? {
totals.entry(account).or_default().merge(&state)?;
}
let mut tokens = Vec::with_capacity(ids.len());
let mut latest = Vec::with_capacity(ids.len());
for id in &ids {
let state = totals.remove(id).unwrap_or_default();
tokens.push(
i64::try_from(state.count("window_total_tokens"))
.map_err(|_| invalid("account usage sort token overflow"))?,
);
latest.push(state.last_used_at());
}
Ok((ids, tokens, latest))
}
pub(crate) async fn account_group_usage(
connection: &mut ObservabilityConnection,
group_ids: &[String],
timezone: gateway_core::time::DeploymentTimeZone,
) -> StoreResult<BTreeMap<String, AccountGroupUsage>> {
if group_ids.is_empty() {
return Ok(BTreeMap::new());
}
let mut snapshot = AnalyticsSnapshot::new_for_accounts(connection).await?;
let now = Utc::now();
let days: i32 =
sqlx::query_scalar("select usage_retention_days::integer from runtime_settings where id=1")
.fetch_one(&mut *snapshot.tx)
.await
.map_err(|error| postgres_unavailable("read group usage retention", error))?;
let memberships:Vec<(String,String)>=sqlx::query_as("select account_group_id,provider_account_id from account_group_accounts where account_group_id=any($1::text[])").bind(group_ids).fetch_all(&mut *snapshot.tx).await.map_err(|error|postgres_unavailable("read group usage membership",error))?;
let ids = memberships
.iter()
.map(|(_, account)| account.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let retained = ObservabilityRange::new(now - TimeDelta::days(i64::from(days)), now)?;
let today = ObservabilityRange::new(
timezone
.day_start(now)
.ok_or_else(|| invalid("invalid business date"))?
.max(retained.start),
now,
)?;
let retained_states = snapshot.account_states(retained, &ids, false).await?;
let today_states = snapshot.account_states(today, &ids, false).await?;
let mut results = BTreeMap::new();
for group in group_ids {
let mut usage = AccountGroupUsage {
today_usd: "0".parse().map_err(|_| invalid("invalid zero cost"))?,
retained_total_usd: "0".parse().map_err(|_| invalid("invalid zero cost"))?,
};
for (_, account) in memberships.iter().filter(|(id, _)| id == group) {
for (states, total) in [
(&retained_states, &mut usage.retained_total_usd),
(&today_states, &mut usage.today_usd),
] {
for ((id, _), state) in states {
if id == account
&& let Some(amount) = state.costs.get("USD")
{
let amount = amount
.parse()
.map_err(|_| invalid("invalid group usage cost"))?;
*total = total
.checked_add(&amount)
.ok_or_else(|| invalid("group usage cost overflow"))?;
}
}
}
}
results.insert(group.clone(), usage);
}
snapshot.finish().await?;
Ok(results)
}
@@ -0,0 +1,410 @@
//! 在 PostgreSQL 内完成事实统计与汇总合并,结果规模只随输出分组增长
use super::*;
/// 输入 matching 提供 output_bucket、name 和完整事实列及 is_completed_usage
/// 精确分位数在整个结果组内计算,不把局部分位数合并成新的分位数
pub(super) fn push_fact_aggregation(
query: &mut QueryBuilder<Postgres>,
exact_percentiles: bool,
read_costs: bool,
) {
push_fact_aggregation_selected(query, exact_percentiles, read_costs, true, None, None);
}
pub(super) fn push_fact_aggregation_selected(
query: &mut QueryBuilder<Postgres>,
exact_percentiles: bool,
read_costs: bool,
read_distributions: bool,
counter_names: Option<&[&str]>,
distribution_names: Option<&[&str]>,
) {
let counters = selected_counter_expressions(counter_names);
query.push(", counted as (select output_bucket,name,max(started_at) filter(where is_completed_usage) as last_used_at,");
push_counter_object(query, &counters);
query.push(" as counts");
if exact_percentiles && read_distributions {
query.push(",");
push_exact_distributions(query, distribution_names);
}
query.push(" from matching group by output_bucket,name)");
if read_costs {
query.push(", cost_values as (select output_bucket,name,currency,sum(amount)::text as amount from matching cross join lateral (values (cost_currency,cost_amount),('standard_usd',standard_cost_amount),('no_cache_usd',no_cache_cost_amount)) cost(currency,amount) where is_completed_usage and currency is not null and amount is not null group by output_bucket,name,currency), merged_costs as (select output_bucket,name,jsonb_object_agg(currency,amount) as costs from cost_values group by output_bucket,name)");
}
if !read_distributions {
query.push(" select c.output_bucket as bucket_start,c.name,c.counts,c.last_used_at,'{}'::jsonb as distributions,'{}'::jsonb as exact_percentiles,");
query.push(if read_costs {
"coalesce(k.costs,'{}'::jsonb)"
} else {
"'{}'::jsonb"
});
query.push(" as costs from counted c");
if read_costs {
query.push(" left join merged_costs k using(output_bucket,name)");
}
return;
}
if exact_percentiles {
query.push(" select c.output_bucket as bucket_start,c.name,c.counts,c.last_used_at,c.distributions,c.exact_percentiles,");
query.push(if read_costs {
"coalesce(k.costs,'{}'::jsonb)"
} else {
"'{}'::jsonb"
});
query.push(" as costs from counted c");
if read_costs {
query.push(" left join merged_costs k using(output_bucket,name)");
}
return;
}
query.push(", samples as (select output_bucket,name,metric,value from matching cross join lateral (values ");
for (index, (name, expression)) in selected_sample_expressions(distribution_names)
.iter()
.enumerate()
{
if index > 0 {
query.push(",");
}
query.push(format!("('{name}',{expression})"));
}
query.push(") sample(metric,value) where value is not null), sample_stats as (select output_bucket,name,metric,count(*) as count,sum(value) as sum,min(value) as minimum,max(value) as maximum,count(*) filter(where value=0) as zero");
query.push(" from samples group by output_bucket,name,metric)");
// 对数桶定义与 Histogram::add 相同,零值单独保留,不对零取对数
query.push(", bin_counts as (select output_bucket,name,metric,floor(log(2::numeric,value::numeric)*16)::integer as bin,count(*) as count from samples where value>0 group by output_bucket,name,metric,bin), merged_bins as (select output_bucket,name,metric,jsonb_object_agg(bin,count) as bins from bin_counts group by output_bucket,name,metric)");
query.push(", merged_distributions as (select s.output_bucket,s.name,jsonb_object_agg(s.metric,jsonb_build_object('count',s.count,'sum',s.sum,'minimum',s.minimum,'maximum',s.maximum,'zero',s.zero,'bins',");
query.push("coalesce(b.bins,'{}'::jsonb)");
query.push(")) as distributions");
query.push(" from sample_stats s");
query.push(" left join merged_bins b using(output_bucket,name,metric)");
query.push(" group by s.output_bucket,s.name) select c.output_bucket as bucket_start,c.name,c.counts,c.last_used_at,");
query.push(if read_costs {
"coalesce(k.costs,'{}'::jsonb)"
} else {
"'{}'::jsonb"
});
query.push(" as costs,coalesce(d.distributions,'{}'::jsonb) as distributions,");
query.push("'{}'::jsonb");
query.push(" as exact_percentiles from counted c left join merged_distributions d using(output_bucket,name)");
if read_costs {
query.push(" left join merged_costs k using(output_bucket,name)");
}
}
// 事实聚合与汇总读取共用指标清单,防止新增计数只在其中一条路径生效
fn counter_expressions() -> Vec<(String, String)> {
let mut counters = vec![
("request_count".to_owned(), "count(*)".to_owned()),
("completed_request_count".to_owned(), "count(*) filter(where is_completed_usage)".to_owned()),
("caller_error_count".to_owned(), "count(*) filter(where client_status_code between 400 and 499)".to_owned()),
("attempt_count".to_owned(), "coalesce(sum(attempt_count),0)".to_owned()),
("retry_count".to_owned(), "coalesce(sum(greatest(attempt_count-1,0)),0)".to_owned()),
("retried_request_count".to_owned(), "count(*) filter(where attempt_count>1)".to_owned()),
("capacity_sample_count".to_owned(), "count(capacity_total_slots)".to_owned()),
("window_total_tokens".to_owned(), "coalesce(sum(coalesce(total_tokens,coalesce(input_tokens,0)::numeric+coalesce(output_tokens,0))) filter(where is_completed_usage),0)".to_owned()),
("cache_eligible_request_count".to_owned(), "count(*) filter(where is_completed_usage and input_tokens is not null)".to_owned()),
("cache_hit_request_count".to_owned(), "count(*) filter(where is_completed_usage and input_tokens is not null and cached_input_tokens>0)".to_owned()),
("image_request_count".to_owned(), "count(*) filter(where is_completed_usage and image_generation_succeeded)".to_owned()),
("image_request_failed_count".to_owned(), "count(*) filter(where is_completed_usage and not image_generation_succeeded)".to_owned()),
];
for (outcome, name) in [
("succeeded", "success"),
("failed", "failure"),
("cancelled", "cancelled"),
("incomplete", "incomplete"),
] {
counters.push((
format!("{name}_count"),
format!("count(*) filter(where outcome='{outcome}')"),
));
let extra = if outcome == "failed" {
"+coalesce(sum(ops_failure_count),0)"
} else {
""
};
counters.push((
format!("attempt_{name}_count"),
format!("count(*) filter(where attempt_count>0 and outcome='{outcome}'){extra}"),
));
}
for (name, predicate, source) in [
(
"rate_limited_count",
"(error_kind in ('rate_limited','quota_exhausted') or coalesce(upstream_status_code,client_status_code)=429)",
"ops_rate_limited_count",
),
(
"auth_failure_count",
"(error_kind in ('authentication','authorization','invalid_credential') or coalesce(upstream_status_code,client_status_code) in (401,403))",
"ops_auth_failure_count",
),
(
"provider_5xx_count",
"coalesce(upstream_status_code,client_status_code) between 500 and 599",
"ops_provider_5xx_count",
),
] {
counters.push((name.to_owned(), format!("count(*) filter(where attempt_count>0 and outcome='failed' and {predicate})+coalesce(sum({source}),0)")));
}
for name in metrics::TOKEN_METRICS {
counters.push((
name.to_owned(),
format!("coalesce(sum({name}) filter(where is_completed_usage),0)"),
));
counters.push((
format!("{name}_known_count"),
format!("count({name}) filter(where is_completed_usage)"),
));
}
for name in ["latency_ms", "first_token_ms"] {
counters.push((
format!("{name}_sum"),
format!("coalesce(sum({name}) filter(where is_completed_usage),0)"),
));
}
for (name, column) in [
("standard_usd", "standard_cost_amount"),
("no_cache_usd", "no_cache_cost_amount"),
] {
counters.push((
format!("{name}_known_count"),
format!("count({column}) filter(where is_completed_usage)"),
));
}
for source in ["provider_reported", "calculated", "unavailable"] {
counters.push((
format!("{source}_count"),
format!("count(*) filter(where is_completed_usage and cost_source='{source}')"),
));
}
counters
}
fn selected_counter_expressions(names: Option<&[&str]>) -> Vec<(String, String)> {
let mut counters = counter_expressions();
if let Some(names) = names {
counters.retain(|(name, _)| names.contains(&name.as_str()));
}
counters
}
fn push_exact_distributions(query: &mut QueryBuilder<Postgres>, names: Option<&[&str]>) {
let samples = selected_sample_expressions(names);
// 独立 ordered-set 聚合只排序实际样本,避免把每条事实展开为六条记录再排序
query.push("jsonb_build_object(");
for (index, (name, expression)) in samples.iter().enumerate() {
if index > 0 {
query.push(",");
}
query.push(format!("'{name}',jsonb_build_object('count',count({expression}),'sum',coalesce(sum({expression}),0),'minimum',min({expression}),'maximum',max({expression}),'zero',count(*) filter(where ({expression})=0),'bins','{{}}'::jsonb)"));
}
query.push(") as distributions,jsonb_build_object(");
for (index, (name, expression)) in samples.iter().enumerate() {
if index > 0 {
query.push(",");
}
query.push(format!("'{name}',coalesce(to_jsonb(percentile_cont(array[0.1,0.5,0.9,0.95,0.99]::double precision[]) within group(order by ({expression}))),'[]'::jsonb)"));
}
query.push(") as exact_percentiles");
}
fn sample_expressions() -> [(&'static str, &'static str); 6] {
[
(
"admission_decision_ms",
"admission_decision_ms::double precision",
),
(
"account_selection_wait_ms",
"account_selection_wait_ms::double precision",
),
(
"capacity",
"case when capacity_total_slots>0 then capacity_used_slots::double precision/capacity_total_slots::double precision*10000.0 end",
),
(
"latency_ms",
"case when is_completed_usage then latency_ms::double precision end",
),
(
"first_token_ms",
"case when is_completed_usage then first_token_ms::double precision end",
),
(
"throughput",
"case when is_completed_usage and output_tokens>0 and upstream_response_ms>0 then output_tokens::double precision*1000.0/upstream_response_ms::double precision end",
),
]
}
fn selected_sample_expressions(names: Option<&[&str]>) -> Vec<(&'static str, &'static str)> {
sample_expressions()
.into_iter()
.filter(|(name, _)| names.is_none_or(|names| names.contains(name)))
.collect()
}
fn push_counter_object(query: &mut QueryBuilder<Postgres>, counters: &[(String, String)]) {
// PostgreSQL 单个函数调用最多接收一百个参数,分块构造不限制指标数量
for (index, chunk) in counters.chunks(40).enumerate() {
if index > 0 {
query.push(" || ");
}
query.push("jsonb_build_object(");
for (position, (name, expression)) in chunk.iter().enumerate() {
if position > 0 {
query.push(",");
}
query.push(format!("'{name}',{expression}"));
}
query.push(")");
}
}
/// 输入 matching 的行来自互斥时间桶,合并后每个最终输出组只返回一行
pub(super) fn push_summary_aggregation(
query: &mut QueryBuilder<Postgres>,
distributions: bool,
costs: bool,
counter_names: Option<&[&str]>,
distribution_names: Option<&[&str]>,
) {
// 每个源汇总只展开一次固定指标,避免按指标键放大行数再排序
let counters = selected_counter_expressions(counter_names)
.into_iter()
.map(|(name, _)| {
let expression = format!("coalesce(sum(metric.{name}),0)");
(name, expression)
})
.collect::<Vec<_>>();
query.push(", grouped as (select output_bucket,name,max(last_used_at) as last_used_at,");
push_counter_object(query, &counters);
query.push(" as counts from matching cross join lateral jsonb_to_record(counts) as metric(");
query.push(
counters
.iter()
.map(|(name, _)| format!("{name} numeric"))
.collect::<Vec<_>>()
.join(","),
);
query.push(") group by output_bucket,name)");
if costs {
query.push(", cost_values as (select output_bucket,name,c.key,sum(c.value::numeric)::text as value from matching cross join lateral jsonb_each_text(costs) c group by output_bucket,name,c.key), merged_costs as (select output_bucket,name,jsonb_object_agg(key,value) as costs from cost_values group by output_bucket,name)");
}
if distributions {
// 先按最终输出组归并分布,再在每组内合并直方图键,避免全局展开后大排序
query.push(", samples as (select output_bucket,name,d.key as metric,d.value from matching cross join lateral jsonb_each(distributions) d");
if let Some(names) = distribution_names {
query
.push(" where d.key=any(")
.push_bind(
names
.iter()
.map(|name| (*name).to_owned())
.collect::<Vec<_>>(),
)
.push("::text[])");
}
query.push("), sample_stats as (select output_bucket,name,metric,sum((value->>'count')::numeric) as count,sum((value->>'sum')::double precision) as sum,min((value->>'minimum')::double precision) as minimum,max((value->>'maximum')::double precision) as maximum,sum((value->>'zero')::numeric) as zero,array_agg(value->'bins') as bins from samples group by output_bucket,name,metric), merged_distributions as (select s.output_bucket,s.name,jsonb_object_agg(s.metric,jsonb_build_object('count',s.count,'sum',s.sum,'minimum',s.minimum,'maximum',s.maximum,'zero',s.zero,'bins',coalesce(b.bins,'{}'::jsonb))) as distributions from sample_stats s cross join lateral(select jsonb_object_agg(key,count) as bins from(select b.key,sum(b.value::numeric) as count from unnest(s.bins) input(value) cross join lateral jsonb_each_text(input.value) b group by b.key) merged)b group by s.output_bucket,s.name)");
}
query.push(" select g.output_bucket as bucket_start,g.name,g.last_used_at,g.counts,");
query.push(if costs {
"coalesce(k.costs,'{}'::jsonb)"
} else {
"'{}'::jsonb"
});
query.push(" as costs,");
query.push(if distributions {
"coalesce(d.distributions,'{}'::jsonb)"
} else {
"'{}'::jsonb"
});
query.push(" as distributions from grouped g");
if costs {
query.push(" left join merged_costs k using(output_bucket,name)");
}
if distributions {
query.push(" left join merged_distributions d using(output_bucket,name)");
}
}
pub(super) fn push_output_bucket(
query: &mut QueryBuilder<Postgres>,
column: &str,
range: ObservabilityRange,
granularity: Option<ObservationGranularity>,
boundaries: Option<&[DateTime<Utc>]>,
) {
match granularity {
None => {
query.push_bind(range.start).push("::timestamptz");
}
Some(ObservationGranularity::Day) => {
// 使用有序阈值二分定位,保留部署时区及夏令时日界,不逐行扫描所有日界
let boundaries = boundaries.unwrap_or_default();
query
.push("(")
.push_bind(boundaries.to_vec())
.push(format!("::timestamptz[])[width_bucket({column},"))
.push_bind(boundaries.to_vec())
.push("::timestamptz[])]");
}
Some(granularity) => {
query
.push("date_bin(make_interval(secs => ")
.push_bind(granularity.seconds() as f64)
.push(format!("),{column},timestamptz '2000-01-01 UTC')"));
}
}
}
pub(super) fn fact_dimension_sql(dimension: Option<DiagnosticDimension>) -> &'static str {
match dimension {
None => "''::text",
Some(DiagnosticDimension::Provider) => "coalesce(nullif(mr.provider_kind,''),'unrouted')",
Some(DiagnosticDimension::Model) => {
"coalesce(nullif(mr.upstream_model_id,''),mr.requested_model_id,'')"
}
Some(DiagnosticDimension::Account) => {
"coalesce(nullif(mr.provider_account_ref,''),'unrouted')"
}
Some(DiagnosticDimension::ApiKey) => "coalesce(mr.client_api_key_ref,'')",
Some(DiagnosticDimension::Transport) => {
"coalesce(nullif(mr.upstream_transport,''),nullif(mr.client_transport,''),'unknown')"
}
Some(DiagnosticDimension::Failure) => "coalesce(mr.error_kind,'')",
Some(DiagnosticDimension::Status) => {
"coalesce(mr.upstream_status_code::text,mr.client_status_code::text,'')"
}
}
}
pub(super) fn summary_dimension_sql(
collection: &str,
dimension: Option<DiagnosticDimension>,
) -> String {
let offset = usize::from(matches!(collection, "account_model" | "key_model"));
match dimension {
None => "''::text".to_owned(),
Some(DiagnosticDimension::Provider) => {
format!("coalesce(nullif(dimensions[{}],''),'unrouted')", 1 + offset)
}
Some(DiagnosticDimension::Model) => format!(
"coalesce(nullif(dimensions[{}],''),dimensions[{}],'')",
3 + offset,
2 + offset
),
Some(DiagnosticDimension::Account) => {
"coalesce(nullif(dimensions[1],''),'unrouted')".to_owned()
}
Some(DiagnosticDimension::Transport) => {
"coalesce(nullif(dimensions[2],''),nullif(dimensions[1],''),'unknown')".to_owned()
}
Some(DiagnosticDimension::Status) => {
"coalesce(nullif(dimensions[2],''),dimensions[1],'')".to_owned()
}
Some(DiagnosticDimension::ApiKey | DiagnosticDimension::Failure) => {
"coalesce(dimensions[1],'')".to_owned()
}
}
}
@@ -49,13 +49,13 @@ pub(super) async fn rebuild_bucket_batch(
0
};
if !same_version {
// 处理中事实发生变化时清掉未发布的部分状态,不能把新旧快照增量相加
// 显式重建版本变化时清掉旧前缀,普通事实变化由增量维护更新前缀
sqlx::query("update request_analytics_pending_buckets set build_version=change_version,build_phase=0,build_started_at=null,build_id=null,build_collection=null,build_dimensions=null,build_outcome=null where bucket_start=$1")
.bind(bucket).execute(&mut **tx).await.map_err(|error|postgres_unavailable("restart changed analytics bucket",error))?;
}
let hour = bucket_for(bucket, 3600)?;
if phase < 2 {
// 小时分段期间不改写其输入;源事实仍可投影,版本变化使小时任务先退回等待
// 小时分段期间暂停其他 quarter 的完整重建,普通增量同步修正已扫描前缀
let owned:bool=sqlx::query_scalar("select exists(select 1 from request_analytics_pending_buckets where bucket_start >= $1 and bucket_start < $1+interval '1 hour' and bucket_start <> $2 and build_phase >= 3)")
.bind(hour).bind(bucket).fetch_one(&mut **tx).await.map_err(|error|postgres_unavailable("claim analytics hour stage",error))?;
if owned {
@@ -119,8 +119,9 @@ pub(super) async fn rebuild_bucket_batch(
let limit = i64::from(batch_size) + 1;
let mut groups = GroupStates::new();
if phase == 1 {
// 增量前缀判断使用 Rust 字节序,SQL 游标必须使用同一排序规则
let statement = format!(
"{} where mr.started_at >= $1 and mr.started_at < $1+interval '15 minutes' and mr.recovered_at is null and (mr.started_at,mr.id)>($2,$3) order by mr.started_at,mr.id limit $4",
"{} where mr.started_at >= $1 and mr.started_at < $1+interval '15 minutes' and (mr.started_at,mr.id collate \"C\")>($2,$3) order by mr.started_at,mr.id collate \"C\" limit $4",
fact_select()
);
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
@@ -139,11 +140,16 @@ pub(super) async fn rebuild_bucket_batch(
let outcome: String = get(fact, "outcome")?;
let mut state = MetricState::default();
state.add_row(fact)?;
let forecast = forecast::state(fact)?;
for (collection, dimensions) in collections(fact)? {
groups
.entry((collection.to_owned(), dimensions, outcome.clone()))
.or_default()
.merge(&state)?;
.merge(if collection == forecast::COLLECTION {
&forecast
} else {
&state
})?;
}
}
merge_staged_groups(tx, "request_analytics_15m", bucket, groups).await?;
@@ -155,7 +161,7 @@ pub(super) async fn rebuild_bucket_batch(
.bind(bucket).bind(get::<DateTime<Utc>>(last,"started_at")?).bind(get::<String>(last,"id")?).execute(&mut **tx).await.map_err(|error|postgres_unavailable("save analytics fact cursor",error))?;
}
} else if phase == 4 {
let rows=sqlx::query("select bucket_start,collection,dimensions,outcome,counts,costs,distributions from request_analytics_15m where bucket_start >= $1 and bucket_start < $1+interval '1 hour' and schema_version=$2 and histogram_version=$3 and (bucket_start,collection,dimensions,outcome)>($4,$5,$6,$7) order by bucket_start,collection,dimensions,outcome limit $8")
let rows=sqlx::query("select bucket_start,collection,dimensions,outcome,counts,costs,distributions,last_used_at from request_analytics_15m where bucket_start >= $1 and bucket_start < $1+interval '1 hour' and schema_version=$2 and histogram_version=$3 and (bucket_start,collection collate \"C\",dimensions collate \"C\",outcome collate \"C\")>($4,$5,$6,$7) order by bucket_start,collection collate \"C\",dimensions collate \"C\",outcome collate \"C\" limit $8")
.bind(hour).bind(SCHEMA_VERSION).bind(HISTOGRAM_VERSION).bind(get::<Option<DateTime<Utc>>>(&row,"build_started_at")?.unwrap_or(hour-TimeDelta::microseconds(1)))
.bind(get::<Option<String>>(&row,"build_collection")?.unwrap_or_default()).bind(get::<Option<Vec<String>>>(&row,"build_dimensions")?.unwrap_or_default()).bind(get::<Option<String>>(&row,"build_outcome")?.unwrap_or_default()).bind(limit)
.fetch_all(&mut **tx).await.map_err(|error|postgres_unavailable("read analytics hour segment",error))?;
@@ -198,7 +204,7 @@ async fn merge_staged_groups(
}
let keys=groups.keys().map(|(collection,dimensions,outcome)|serde_json::json!({"collection":collection,"dimensions":dimensions,"outcome":outcome})).collect::<Vec<_>>();
let statement = format!(
"select stored.collection,stored.dimensions,stored.outcome,stored.counts,stored.costs,stored.distributions from {table} stored join jsonb_to_recordset($2::jsonb) requested(collection text,dimensions text[],outcome text) using(collection,dimensions,outcome) where bucket_start=$1"
"select stored.collection,stored.dimensions,stored.outcome,stored.counts,stored.costs,stored.distributions,stored.last_used_at from {table} stored join jsonb_to_recordset($2::jsonb) requested(collection text,dimensions text[],outcome text) using(collection,dimensions,outcome) where bucket_start=$1"
);
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(bucket)
@@ -220,37 +226,64 @@ async fn merge_staged_groups(
write_groups(tx, table, bucket, groups).await
}
async fn write_groups(
pub(super) async fn write_groups(
tx: &mut sqlx::Transaction<'_, Postgres>,
table: &str,
bucket: DateTime<Utc>,
groups: GroupStates,
) -> StoreResult<()> {
if groups.is_empty() {
return Ok(());
}
let mut statement = QueryBuilder::<Postgres>::new(format!(
"insert into {table}(bucket_start,collection,dimensions,outcome,schema_version,histogram_version,counts,costs,distributions) "
));
let rows = groups
.into_iter()
.map(|((collection, dimensions, outcome), state)| {
let (counts, costs, distributions) = state.encode()?;
Ok((
write_bucket_groups(tx, table, [(bucket, groups)]).await
}
pub(super) async fn write_bucket_groups(
tx: &mut sqlx::Transaction<'_, Postgres>,
table: &str,
buckets: impl IntoIterator<Item = (DateTime<Utc>, GroupStates)>,
) -> StoreResult<()> {
let mut groups = buckets.into_iter().flat_map(|(bucket, groups)| {
groups
.into_iter()
.map(move |(key, state)| (bucket, key, state))
});
loop {
// 跨桶合并 SQL 往返,但编码缓冲只保留当前批次,避免回填将所有状态展开为 JSON
let rows = groups
.by_ref()
.take(1000)
.map(|(bucket, (collection, dimensions, outcome), state)| {
let (counts, costs, distributions) = state.encode()?;
Ok((
bucket,
collection,
dimensions,
outcome,
counts,
costs,
distributions,
state.last_used_at(),
))
})
.collect::<StoreResult<Vec<_>>>()?;
if rows.is_empty() {
return Ok(());
}
let mut statement = QueryBuilder::<Postgres>::new(format!(
"insert into {table}(bucket_start,collection,dimensions,outcome,schema_version,histogram_version,counts,costs,distributions,last_used_at) "
));
// 十个绑定参数限制单语句批量,避免高基数时间桶越过 PostgreSQL 参数上限
statement.push_values(
&rows,
|mut row,
(
bucket,
collection,
dimensions,
outcome,
counts,
costs,
distributions,
))
})
.collect::<StoreResult<Vec<_>>>()?;
// 九个绑定参数限制单语句批量,避免高基数时间桶越过 PostgreSQL 参数上限
for chunk in rows.chunks(1000) {
statement.push_values(
chunk,
|mut row, (collection, dimensions, outcome, counts, costs, distributions)| {
last_used_at,
)| {
row.push_bind(bucket)
.push_bind(collection)
.push_bind(dimensions)
@@ -259,16 +292,15 @@ async fn write_groups(
.push_bind(HISTOGRAM_VERSION)
.push_bind(counts)
.push_bind(costs)
.push_bind(distributions);
.push_bind(distributions)
.push_bind(last_used_at);
},
);
statement.push(" on conflict(bucket_start,collection,dimensions,outcome) do update set counts=excluded.counts,costs=excluded.costs,distributions=excluded.distributions,schema_version=excluded.schema_version,histogram_version=excluded.histogram_version");
statement.push(" on conflict(bucket_start,collection,dimensions,outcome) do update set counts=excluded.counts,costs=excluded.costs,distributions=excluded.distributions,last_used_at=excluded.last_used_at,schema_version=excluded.schema_version,histogram_version=excluded.histogram_version");
statement
.build()
.execute(&mut **tx)
.await
.map_err(|error| postgres_unavailable("write analytics bucket", error))?;
statement.reset();
}
Ok(())
}
@@ -0,0 +1,113 @@
//! 用完整刻钟计数和同快照源差量定位精确页码,避免每页扫描整个保留期
use super::query::{collection_for, push_collection_filter, push_summary_buckets};
use super::*;
enum ReadVersion {
Published,
Before,
Current,
}
impl AnalyticsSnapshot<'_> {
pub(crate) async fn completed_request_counts(
&mut self,
range: ObservabilityRange,
filter: &UsageRecordFilter,
) -> StoreResult<Option<BTreeMap<DateTime<Utc>, u64>>> {
let Some(collection) = collection_for(filter, None) else {
return Ok(None);
};
if !self.projection_covers(range) {
return Ok(None);
}
let mut counts = BTreeMap::new();
let (_, buckets) = self.summary_buckets(range, false).await?;
if !buckets.is_empty() {
let mut query = QueryBuilder::<Postgres>::new(
"select bucket_start,coalesce(sum((counts->>'completed_request_count')::bigint),0)::bigint as request_count from request_analytics_15m a where collection=",
);
query
.push_bind(collection)
.push(" and schema_version=")
.push_bind(SCHEMA_VERSION)
.push(" and histogram_version=")
.push_bind(HISTOGRAM_VERSION);
push_summary_buckets(&mut query, &buckets, 900);
push_collection_filter(&mut query, collection, filter);
query.push(" group by bucket_start");
for row in query
.build()
.fetch_all(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("read usage page counts", error))?
{
counts.insert(
get(&row, "bucket_start")?,
to_u64(get(&row, "request_count")?)?,
);
}
}
for (start, end) in self.uncovered_spans(range).await? {
self.merge_completed_counts(
&mut counts,
ObservabilityRange { start, end },
filter,
ReadVersion::Published,
)
.await?;
}
if self.has_pending_requests().await? {
// 先撤销已投影版本,再加入源表当前版本;移动、删除和未确认的重放均只计算一次
self.merge_completed_counts(&mut counts, range, filter, ReadVersion::Before)
.await?;
self.merge_completed_counts(&mut counts, range, filter, ReadVersion::Current)
.await?;
}
Ok(Some(counts))
}
async fn merge_completed_counts(
&mut self,
counts: &mut BTreeMap<DateTime<Utc>, u64>,
range: ObservabilityRange,
filter: &UsageRecordFilter,
version: ReadVersion,
) -> StoreResult<()> {
let table = if matches!(version, ReadVersion::Current) {
"model_request_observations"
} else {
"request_analytics_facts"
};
let mut query = QueryBuilder::<Postgres>::new(format!(
"select date_bin(interval '15 minutes',mr.started_at,timestamptz '2000-01-01 UTC') as bucket_start,count(*)::bigint as request_count from {table} mr"
));
if !matches!(version, ReadVersion::Published) {
query.push(" join request_analytics_pending_requests p on p.request_id=mr.id");
}
query
.push(" where mr.started_at >= ")
.push_bind(range.start)
.push(" and mr.started_at < ")
.push_bind(range.end);
push_completed_usage_fact_filter(&mut query, "mr");
push_usage_filter(&mut query, filter, "mr");
query.push(" group by bucket_start");
let rows = query
.build()
.fetch_all(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("correct usage page counts", error))?;
for row in rows {
let value = to_u64(get(&row, "request_count")?)?;
let count = counts.entry(get(&row, "bucket_start")?).or_default();
*count = if matches!(version, ReadVersion::Before) {
count.checked_sub(value)
} else {
count.checked_add(value)
}
.ok_or_else(|| invalid("usage page count is inconsistent"))?;
}
Ok(())
}
}
@@ -0,0 +1,343 @@
//! 预测累计使用开始时间与完成时间的二维汇总,采样边界及待投影版本读取精确事实
use super::*;
use gateway_admin::model::provider_credentials::ProviderDocument;
use gateway_admin::model::quota_forecast_sampling::{
MAX_FORECAST_HISTORY_POINTS, QuotaForecastHistory, QuotaForecastHistoryPoint,
QuotaForecastUsage,
};
use gateway_core::account::OpaqueProviderData;
use sqlx::postgres::PgRow;
pub(super) const COLLECTION: &str = "account_completion";
const COMPLETION_SECONDS: i64 = 300;
const COUNTERS: [&str; 10] = [
"request_count",
"completed_request_count",
"window_total_tokens",
"input_tokens",
"output_tokens",
"cached_input_tokens",
"missing_token_count",
"known_cost_count",
"unavailable_cost_count",
"excluded_request_count",
];
pub(super) fn dimensions(row: &PgRow) -> StoreResult<Vec<String>> {
let completed: Option<DateTime<Utc>> = get(row, "completed_at")?;
let terminal = get::<String>(row, "outcome")? != "running";
Ok(vec![
get::<Option<String>>(row, "provider_account_ref")?.unwrap_or_default(),
completed
.filter(|_| terminal)
.map(|time| time.timestamp().div_euclid(COMPLETION_SECONDS).to_string())
.unwrap_or_default(),
])
}
pub(super) fn state(row: &PgRow) -> StoreResult<MetricState> {
let mut state = MetricState::default();
state.increment("request_count", 1)?;
state.add_cost("USD", "0")?;
if get::<Option<DateTime<Utc>>>(row, "completed_at")?.is_none()
|| get::<String>(row, "outcome")? == "running"
{
return Ok(state);
}
if !get::<bool>(row, "is_completed_usage")? {
state.increment("excluded_request_count", 1)?;
return Ok(state);
}
state.increment("completed_request_count", 1)?;
let input = optional_unsigned(row, "input_tokens")?;
let output = optional_unsigned(row, "output_tokens")?;
let total = optional_unsigned(row, "total_tokens")?;
state.increment(
"window_total_tokens",
total.unwrap_or(
input
.unwrap_or_default()
.checked_add(output.unwrap_or_default())
.ok_or_else(|| invalid("forecast token overflow"))?,
),
)?;
state.increment("input_tokens", input.unwrap_or_default())?;
state.increment("output_tokens", output.unwrap_or_default())?;
state.increment(
"cached_input_tokens",
optional_unsigned(row, "cached_input_tokens")?.unwrap_or_default(),
)?;
state.increment(
"missing_token_count",
u64::from(total.is_none() && (input.is_none() || output.is_none())),
)?;
let source: String = get(row, "cost_source")?;
let amount: Option<String> = get(row, "cost_amount_text")?;
let known = matches!(source.as_str(), "calculated" | "provider_reported")
&& get::<Option<String>>(row, "cost_currency")?.as_deref() == Some("USD")
&& amount.is_some();
state.increment("known_cost_count", u64::from(known))?;
state.increment("unavailable_cost_count", u64::from(!known))?;
if known {
state.add_cost(
"USD",
amount
.as_deref()
.ok_or_else(|| invalid("missing forecast amount"))?,
)?;
}
Ok(state)
}
impl AnalyticsSnapshot<'_> {
pub(crate) async fn forecast_history(
&mut self,
account: &str,
range: ObservabilityRange,
) -> StoreResult<QuotaForecastHistory> {
let documents = sqlx::query(DOCUMENTS)
.bind(account)
.bind(range.start)
.bind(range.end)
.bind(MAX_FORECAST_HISTORY_POINTS as i32)
.fetch_all(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("sample forecast documents", error))?;
let mut cutoffs = documents
.iter()
.map(|row| get(row, "completed_at"))
.collect::<StoreResult<Vec<DateTime<Utc>>>>()?;
// 最后一个状态专门代表整个窗口,文档恰好完成于窗口末端时复用同一截止值
if cutoffs.last() != Some(&range.end) {
cutoffs.push(range.end);
}
let boundaries = cutoffs
.iter()
.map(|time| time.timestamp().div_euclid(COMPLETION_SECONDS))
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let mut states = vec![MetricState::default(); cutoffs.len() + 1];
if self.projection_covers(range) {
let (_, quarters) = self.summary_buckets(range, false).await?;
if !quarters.is_empty() {
let sums = COUNTERS
.iter()
.map(|name| format!("'{name}',coalesce(sum((counts->>'{name}')::numeric),0)"))
.collect::<Vec<_>>()
.join(",");
let statement = format!(
"with selected as (select counts,costs,case when dimensions[2]='' then cardinality($3::timestamptz[]) else width_bucket(to_timestamp((dimensions[2]::bigint+1)*{COMPLETION_SECONDS})-interval '1 microsecond',$3::timestamptz[]) end as slot
from request_analytics_15m where collection='{COLLECTION}' and dimensions[1]=$1 and bucket_start=any($2::timestamptz[]) and schema_version=$5 and histogram_version=$6
and (dimensions[2]='' or not(nullif(dimensions[2],'')::bigint=any($4::bigint[]))))
select slot,jsonb_build_object({sums}) as counts,jsonb_build_object('USD',coalesce(sum((costs->>'USD')::numeric),0)::text) as costs,'{{}}'::jsonb as distributions,null::timestamptz as last_used_at from selected group by slot"
);
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(account)
.bind(&quarters)
.bind(&cutoffs)
.bind(&boundaries)
.bind(SCHEMA_VERSION)
.bind(HISTOGRAM_VERSION)
.fetch_all(&mut *self.tx)
.await
.map_err(|error| {
postgres_unavailable("read forecast completion summaries", error)
})?;
merge_rows(&mut states, rows, false)?;
// 只有跨过采样截止值的完成时间小桶需要回读,按完成时间索引逐段定位
let rows = self
.forecast_facts(
account,
range,
&cutoffs,
FactSelection::CompletionEdges {
quarters: &quarters,
boundaries: &boundaries,
},
)
.await?;
merge_rows(&mut states, rows, false)?;
}
for (start, end) in self.uncovered_spans(range).await? {
let rows = self
.forecast_facts(
account,
ObservabilityRange { start, end },
&cutoffs,
FactSelection::Projected,
)
.await?;
merge_rows(&mut states, rows, false)?;
}
if self.has_pending_requests().await? {
let old = self
.forecast_facts(account, range, &cutoffs, FactSelection::Previous)
.await?;
merge_rows(&mut states, old, true)?;
let current = self
.forecast_facts(account, range, &cutoffs, FactSelection::Current)
.await?;
merge_rows(&mut states, current, false)?;
}
} else {
let rows = self
.forecast_facts(account, range, &cutoffs, FactSelection::Source)
.await?;
merge_rows(&mut states, rows, false)?;
}
let pending_request_count = states
.last()
.map_or(0, |state| state.count("request_count"));
let mut cumulative = MetricState::default();
let mut usages = Vec::with_capacity(cutoffs.len());
for state in states.iter().take(cutoffs.len()) {
cumulative.merge(state)?;
usages.push(usage(&cumulative)?);
}
let mut points = Vec::with_capacity(documents.len());
for (row, usage) in documents.into_iter().zip(&usages) {
let document: sqlx::types::Json<serde_json::Map<String, serde_json::Value>> =
get(&row, "document")?;
points.push(QuotaForecastHistoryPoint {
started_at: get(&row, "started_at")?,
completed_at: get(&row, "completed_at")?,
usage: usage.clone(),
provider_observation: ProviderDocument::new(OpaqueProviderData::new(document.0)),
});
}
Ok(QuotaForecastHistory {
points,
usage: usage(&cumulative)?,
pending_request_count,
})
}
async fn forecast_facts(
&mut self,
account: &str,
range: ObservabilityRange,
cutoffs: &[DateTime<Utc>],
selection: FactSelection<'_>,
) -> StoreResult<Vec<PgRow>> {
let source = matches!(selection, FactSelection::Source | FactSelection::Current);
let table = if source {
"model_requests"
} else {
"request_analytics_facts"
};
let mut query = QueryBuilder::<Postgres>::new(
"with matching as (select mr.*,case when mr.completed_at is not null and mr.outcome <> 'running' then width_bucket(mr.completed_at-interval '1 microsecond',",
);
query.push_bind(cutoffs).push("::timestamptz[]) else ").push_bind(cutoffs.len() as i32)
.push(" end as slot,coalesce((").push(completed_usage_fact_predicate("mr"))
.push("),false) and mr.completed_at is not null and mr.outcome <> 'running' as included from ");
if let FactSelection::CompletionEdges { boundaries, .. } = selection {
// 保留按完成时间小桶的参数化索引扫描,避免展开后扫描整段账号事实再逐桶匹配
// OFFSET 0 不裁剪数据,只阻止此相关子查询被合并到外层连接
query.push("unnest(").push_bind(boundaries).push(format!("::bigint[]) edge(bucket) cross join lateral (select * from {table} f where f.provider_account_ref="))
.push_bind(account).push(format!(" and f.completed_at >= to_timestamp(edge.bucket*{COMPLETION_SECONDS}) and f.completed_at < to_timestamp((edge.bucket+1)*{COMPLETION_SECONDS}) and f.outcome <> 'running' offset 0) mr"));
} else {
query.push(format!("{table} mr"));
}
if matches!(selection, FactSelection::Previous | FactSelection::Current) {
query.push(" join request_analytics_pending_requests p on p.request_id=mr.id");
}
query
.push(" where mr.provider_account_ref=")
.push_bind(account)
.push(" and mr.started_at >= ")
.push_bind(range.start)
.push(" and mr.started_at < ")
.push_bind(range.end);
if let FactSelection::CompletionEdges { quarters, .. } = selection {
query.push(" and date_bin(interval '15 minutes',mr.started_at,timestamptz '2000-01-01 UTC')=any(").push_bind(quarters).push("::timestamptz[])");
}
// 多个累计指标复用同一分类结果,避免展开后重复判断完成用量并携带宽事实排序
// OFFSET 0 保留全部行,只让分类在聚合前计算一次
query.push("), classified as (select *,coalesce(cost_source in ('calculated','provider_reported') and cost_amount is not null and cost_currency='USD',false) as known_cost from matching offset 0)
select slot,jsonb_build_object(
'request_count',count(*),
'completed_request_count',count(*) filter(where included),
'window_total_tokens',coalesce(sum(coalesce(total_tokens,coalesce(input_tokens,0)+coalesce(output_tokens,0))) filter(where included),0),
'input_tokens',coalesce(sum(input_tokens) filter(where included),0),
'output_tokens',coalesce(sum(output_tokens) filter(where included),0),
'cached_input_tokens',coalesce(sum(cached_input_tokens) filter(where included),0),
'missing_token_count',count(*) filter(where included and total_tokens is null and (input_tokens is null or output_tokens is null)),
'known_cost_count',count(*) filter(where included and known_cost),
'unavailable_cost_count',count(*) filter(where included and not known_cost),
'excluded_request_count',count(*) filter(where not included and completed_at is not null and outcome <> 'running')) as counts,
jsonb_build_object('USD',coalesce(sum(cost_amount) filter(where included and known_cost),0)::text) as costs,
'{}'::jsonb as distributions,null::timestamptz as last_used_at from classified group by slot");
query
.build()
.fetch_all(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("read exact forecast boundaries", error))
}
}
#[derive(Clone, Copy)]
enum FactSelection<'a> {
Source,
Projected,
Previous,
Current,
CompletionEdges {
quarters: &'a [DateTime<Utc>],
boundaries: &'a [i64],
},
}
fn merge_rows(states: &mut [MetricState], rows: Vec<PgRow>, subtract: bool) -> StoreResult<()> {
for row in rows {
let index = usize::try_from(get::<i32>(&row, "slot")?)
.map_err(|_| invalid("negative forecast slot"))?;
let target = states
.get_mut(index)
.ok_or_else(|| invalid("forecast slot out of range"))?;
let state = MetricState::decode(&row)?;
if subtract {
incremental::subtract(target, &state, &MetricState::default())?;
} else {
target.merge(&state)?;
}
}
Ok(())
}
fn usage(state: &MetricState) -> StoreResult<QuotaForecastUsage> {
Ok(QuotaForecastUsage {
request_count: state.count("completed_request_count"),
tokens: state.count("window_total_tokens"),
input_tokens: state.count("input_tokens"),
output_tokens: state.count("output_tokens"),
cached_input_tokens: state.count("cached_input_tokens"),
missing_token_count: state.count("missing_token_count"),
known_cost_count: state.count("known_cost_count"),
unavailable_cost_count: state.count("unavailable_cost_count"),
excluded_request_count: state.count("excluded_request_count"),
usd: state
.costs
.get("USD")
.map(|value| {
value
.parse::<f64>()
.map_err(|_| invalid("invalid forecast dollars"))
})
.transpose()?
.unwrap_or_default(),
})
}
// 微秒边界向上取整,最后一槽包含观测时刻;文档采样与累计量补差共用读事务
const DOCUMENTS: &str = "select mr.started_at,mr.completed_at,mr.provider_observation_json as document
from generate_series(0,$4-1) b(bucket) cross join lateral (
select mr.started_at,mr.completed_at,mr.provider_observation_json
from model_requests mr where mr.provider_account_ref=$1 and mr.started_at >= $2 and mr.started_at < $3
and mr.completed_at >= $2 + ceil(b.bucket*greatest(1,extract(epoch from ($3::timestamptz-$2::timestamptz))/$4)*1000000)::bigint*interval '1 microsecond'
and mr.completed_at < case when b.bucket=$4-1 then $3+interval '1 microsecond' else least($3+interval '1 microsecond',$2+ceil((b.bucket+1)*greatest(1,extract(epoch from ($3::timestamptz-$2::timestamptz))/$4)*1000000)::bigint*interval '1 microsecond') end
and mr.outcome <> 'running' and mr.provider_observation_json is not null
order by mr.completed_at desc,mr.id desc limit 1
) mr order by mr.completed_at";
@@ -14,8 +14,8 @@ pub(super) struct Histogram {
pub sum: f64,
pub minimum: Option<f64>,
pub maximum: Option<f64>,
zero: u64,
bins: BTreeMap<i32, u64>,
pub(super) zero: u64,
pub(super) bins: BTreeMap<i32, u64>,
}
impl Histogram {
@@ -0,0 +1,371 @@
//! 同事务撤销旧事实并合并新事实,构建中的统计只更新已经扫描的前缀
use super::buckets::write_bucket_groups;
use super::worker::collections;
use super::*;
use sqlx::postgres::PgRow;
type GroupKey = (DateTime<Utc>, String, Vec<String>, String);
#[derive(Default)]
struct Change {
previous: MetricState,
current: MetricState,
}
struct Progress {
phase: i16,
started_at: Option<DateTime<Utc>>,
id: Option<String>,
collection: Option<String>,
dimensions: Option<Vec<String>>,
outcome: Option<String>,
}
impl Progress {
fn contains_fact(&self, time: DateTime<Utc>, id: &str) -> bool {
match self.phase {
0 => false,
1 => self
.started_at
.zip(self.id.as_deref())
.is_some_and(|cursor| (time, id) <= cursor),
_ => true,
}
}
fn contains_quarter(&self, key: &GroupKey) -> bool {
if self.phase != 4 {
return false;
}
match (
&self.started_at,
&self.collection,
&self.dimensions,
&self.outcome,
) {
(Some(time), Some(collection), Some(dimensions), Some(outcome)) => {
(&key.0, &key.1, &key.2, &key.3) <= (time, collection, dimensions, outcome)
}
_ => false,
}
}
}
pub(super) async fn update_summaries(
tx: &mut sqlx::Transaction<'_, Postgres>,
previous: &[PgRow],
current: &[PgRow],
) -> StoreResult<()> {
let mut hours = BTreeSet::new();
for row in previous.iter().chain(current) {
hours.insert(bucket_for(get(row, "started_at")?, 3600)?);
}
if hours.is_empty() {
return Ok(());
}
let progress_rows = sqlx::query("select p.* from unnest($1::timestamptz[]) h(start_at) join request_analytics_pending_buckets p on p.bucket_start >= h.start_at and p.bucket_start < h.start_at+interval '1 hour'")
.bind(hours.into_iter().collect::<Vec<_>>()).fetch_all(&mut **tx).await
.map_err(|error|postgres_unavailable("read incremental analytics stages",error))?;
let mut progress = BTreeMap::new();
for row in progress_rows {
let phase =
if get::<Option<i64>>(&row, "build_version")? == Some(get(&row, "change_version")?) {
get(&row, "build_phase")?
} else {
0
};
progress.insert(
get::<DateTime<Utc>>(&row, "bucket_start")?,
Progress {
phase,
started_at: get(&row, "build_started_at")?,
id: get(&row, "build_id")?,
collection: get(&row, "build_collection")?,
dimensions: get(&row, "build_dimensions")?,
outcome: get(&row, "build_outcome")?,
},
);
}
// 两种粒度共用一次事实解码,仍分别检查构建前缀,不能把未发布季度并入小时
let mut changes = [BTreeMap::<GroupKey, Change>::new(), BTreeMap::new()];
let clean_hours = progress.is_empty();
for (rows, adding) in [(previous, false), (current, true)] {
for row in rows {
let time = get(row, "started_at")?;
let quarter = bucket_for(time, 900)?;
let hour = bucket_for(time, 3600)?;
let id: String = get(row, "id")?;
let outcome: String = get(row, "outcome")?;
let mut state = MetricState::default();
state.add_row(row)?;
let forecast = forecast::state(row)?;
let mut hour_pending = progress.range(hour..hour + TimeDelta::hours(1)).peekable();
let hour_ready = hour_pending.peek().is_none();
let hour_owner =
hour_pending.find_map(|(_, stage)| (stage.phase >= 3).then_some(stage));
for (collection, dimensions) in collections(row)? {
let key = (quarter, collection.to_owned(), dimensions, outcome.clone());
let included = [
progress
.get(&quarter)
.is_none_or(|p| p.contains_fact(time, &id)),
!clean_hours
&& (hour_ready
|| hour_owner.is_some_and(|stage| stage.contains_quarter(&key))),
];
let hourly_key = (hour, key.1.clone(), key.2.clone(), key.3.clone());
for (index, key) in [key, hourly_key].into_iter().enumerate() {
if !included[index] {
continue;
}
let change = changes[index].entry(key).or_default();
let state = if collection == forecast::COLLECTION {
&forecast
} else {
&state
};
if adding {
change.current.merge(state)?;
} else {
change.previous.merge(state)?;
}
}
}
}
}
if clean_hours {
// 无重建中的时间桶时,小时贡献可直接合并季度贡献,避免逐事实再合并一遍全部指标
let [quarters, hours] = &mut changes;
for ((quarter, collection, dimensions, outcome), contribution) in quarters.iter() {
let key = (
bucket_for(*quarter, 3600)?,
collection.clone(),
dimensions.clone(),
outcome.clone(),
);
let hour = hours.entry(key).or_default();
hour.previous.merge(&contribution.previous)?;
hour.current.merge(&contribution.current)?;
}
}
for (index, changes) in changes.into_iter().enumerate() {
apply_changes(tx, index == 1, &progress, changes).await?;
}
Ok(())
}
async fn apply_changes(
tx: &mut sqlx::Transaction<'_, Postgres>,
hourly: bool,
progress: &BTreeMap<DateTime<Utc>, Progress>,
changes: BTreeMap<GroupKey, Change>,
) -> StoreResult<()> {
if changes.is_empty() {
return Ok(());
}
let table = if hourly {
"request_analytics_hourly"
} else {
"request_analytics_15m"
};
let keys = changes.keys().map(|(bucket,collection,dimensions,outcome)|serde_json::json!({"bucket_start":bucket,"collection":collection,"dimensions":dimensions,"outcome":outcome})).collect::<Vec<_>>();
let statement = format!(
"select a.* from {table} a join jsonb_to_recordset($1::jsonb) k(bucket_start timestamptz,collection text,dimensions text[],outcome text) using(bucket_start,collection,dimensions,outcome)"
);
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(serde_json::Value::Array(keys))
.fetch_all(&mut **tx)
.await
.map_err(|error| postgres_unavailable("read incremental analytics groups", error))?;
let mut stored = BTreeMap::new();
for row in rows {
stored.insert(
(
get(&row, "bucket_start")?,
get(&row, "collection")?,
get(&row, "dimensions")?,
get(&row, "outcome")?,
),
MetricState::decode(&row)?,
);
}
let mut buckets =
BTreeMap::<DateTime<Utc>, BTreeMap<(String, Vec<String>, String), MetricState>>::new();
for (key, change) in changes {
if change.previous.encode()? == change.current.encode()?
&& change.previous.last_used_at() == change.current.last_used_at()
{
continue;
}
let mut state = stored.remove(&key).unwrap_or_default();
let needs_rebuild = subtract(&mut state, &change.previous, &change.current)?;
state.merge(&change.current)?;
if state.count("request_count") == 0 {
let statement = format!(
"delete from {table} where bucket_start=$1 and collection=$2 and dimensions=$3 and outcome=$4"
);
sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(key.0)
.bind(&key.1)
.bind(&key.2)
.bind(&key.3)
.execute(&mut **tx)
.await
.map_err(|error| postgres_unavailable("remove empty analytics group", error))?;
} else {
// 预测状态只有可加计数和美元总额,零金额无需修复币种存在性
if needs_rebuild && key.1 != forecast::COLLECTION {
state = rebuild_group(tx, hourly, progress, &key).await?;
}
buckets
.entry(key.0)
.or_default()
.insert((key.1, key.2, key.3), state);
}
}
write_bucket_groups(tx, table, buckets).await
}
/// 返回是否撤销了不能从可加状态推导的极值或币种存在性
pub(super) fn subtract(
state: &mut MetricState,
previous: &MetricState,
current: &MetricState,
) -> StoreResult<bool> {
if previous.count("request_count") == 0 {
return Ok(false);
}
// 新贡献仍持有旧极值时无需重读事实,例如只补齐 Token、费用或重放同一请求
let mut rebuild = previous.last_used_at.is_some()
&& previous.last_used_at == state.last_used_at
&& current.last_used_at < state.last_used_at;
for (name, value) in &previous.counts {
let count = state.counts.entry(name.clone()).or_default();
*count = count
.checked_sub(*value)
.ok_or_else(|| invalid("analytics counter underflow"))?;
}
for (name, removed) in &previous.distributions {
let distribution = state
.distributions
.get_mut(name)
.ok_or_else(|| invalid("missing analytics distribution"))?;
distribution.count = distribution
.count
.checked_sub(removed.count)
.ok_or_else(|| invalid("analytics distribution underflow"))?;
distribution.zero = distribution
.zero
.checked_sub(removed.zero)
.ok_or_else(|| invalid("analytics zero count underflow"))?;
distribution.sum -= removed.sum;
for (bin, value) in &removed.bins {
let count = distribution
.bins
.get_mut(bin)
.ok_or_else(|| invalid("missing analytics bin"))?;
*count = count
.checked_sub(*value)
.ok_or_else(|| invalid("analytics bin underflow"))?;
}
distribution.bins.retain(|_, count| *count > 0);
if distribution.count == 0 {
*distribution = Default::default();
} else {
let replacement = current.distributions.get(name);
rebuild |= distribution.minimum == removed.minimum
&& replacement
.and_then(|value| value.minimum)
.is_none_or(|value| Some(value) > distribution.minimum);
rebuild |= distribution.maximum == removed.maximum
&& replacement
.and_then(|value| value.maximum)
.is_none_or(|value| Some(value) < distribution.maximum);
rebuild |= !distribution.sum.is_finite() || distribution.sum < 0.0;
}
}
for (currency, amount) in &previous.costs {
let stored_amount = state
.costs
.get(currency)
.ok_or_else(|| invalid("missing analytics currency"))?;
let stored_amount: gateway_core::metering::Decimal = stored_amount
.parse()
.map_err(|_| invalid("invalid analytics amount"))?;
let removed: gateway_core::metering::Decimal = amount
.parse()
.map_err(|_| invalid("invalid analytics amount"))?;
let value = stored_amount
.scaled()
.checked_sub(removed.scaled())
.ok_or_else(|| invalid("analytics amount underflow"))?;
rebuild |= value == 0 && !current.costs.contains_key(currency);
state.costs.insert(
currency.clone(),
gateway_core::metering::Decimal::from_scaled(value)
.map_err(|_| invalid("analytics amount overflow"))?
.canonical(),
);
}
Ok(rebuild)
}
async fn rebuild_group(
tx: &mut sqlx::Transaction<'_, Postgres>,
hourly: bool,
progress: &BTreeMap<DateTime<Utc>, Progress>,
key: &GroupKey,
) -> StoreResult<MetricState> {
let mut end = key.0 + TimeDelta::seconds(if hourly { 3600 } else { 900 });
let mut cursor = None;
if hourly {
if let Some((_, owner)) = progress.range(key.0..end).find(|(_, p)| p.phase == 4)
&& let Some(time) = owner.started_at
{
let candidate = (time, key.1.clone(), key.2.clone(), key.3.clone());
end = if owner.contains_quarter(&candidate) {
time + TimeDelta::minutes(15)
} else {
time
};
}
} else if let Some(stage) = progress.get(&key.0)
&& stage.phase == 1
{
cursor = stage.started_at.zip(stage.id.as_deref());
}
let mut query = QueryBuilder::<Postgres>::new("with matching as materialized (select mr.*,");
query.push_bind(key.0).push("::timestamptz as output_bucket,''::text as name,coalesce((")
.push(completed_usage_fact_predicate("mr")).push("),false) as is_completed_usage from request_analytics_facts mr where mr.started_at >= ")
.push_bind(key.0).push(" and mr.started_at < ").push_bind(end)
.push(" and mr.recovered_at is null and mr.outcome=").push_bind(&key.3);
if let Some((time, id)) = cursor {
query
.push(" and (mr.started_at,mr.id collate \"C\") <= (")
.push_bind(time)
.push(",")
.push_bind(id)
.push(")");
}
let columns = worker::COLLECTION_COLUMNS
.iter()
.find_map(|(collection, columns)| (*collection == key.1).then_some(*columns))
.ok_or_else(|| invalid("unknown incremental analytics collection"))?;
for (column, value) in columns.iter().zip(&key.2) {
query
.push(format!(" and coalesce(mr.{column}::text,'')="))
.push_bind(value);
}
query.push(")");
aggregate::push_fact_aggregation(&mut query, false, true);
let row = query
.build()
.fetch_optional(&mut **tx)
.await
.map_err(|error| postgres_unavailable("rebuild changed analytics extrema", error))?;
row.as_ref()
.map(MetricState::decode)
.transpose()
.map(Option::unwrap_or_default)
}
@@ -5,13 +5,25 @@ use super::histogram::Histogram;
use serde::{Deserialize, Serialize};
use sqlx::postgres::PgRow;
pub(super) const TOKEN_METRICS: [&str; 8] = [
"input_tokens",
"output_tokens",
"cached_input_tokens",
"cache_write_input_tokens",
"reasoning_output_tokens",
"total_tokens",
"image_input_tokens",
"image_output_tokens",
];
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub(super) struct MetricState {
counts: BTreeMap<String, u64>,
distributions: BTreeMap<String, Histogram>,
costs: BTreeMap<String, String>,
pub(crate) struct MetricState {
pub(super) counts: BTreeMap<String, u64>,
pub(super) distributions: BTreeMap<String, Histogram>,
pub(super) costs: BTreeMap<String, String>,
pub(super) last_used_at: Option<DateTime<Utc>>,
#[serde(skip)]
exact_samples: Option<BTreeMap<String, Vec<f64>>>,
exact_percentiles: Option<BTreeMap<String, Vec<Option<f64>>>>,
}
impl MetricState {
@@ -25,9 +37,23 @@ impl MetricState {
"distributions",
)?
.0,
exact_samples: None,
last_used_at: get(row, "last_used_at")?,
exact_percentiles: None,
})
}
pub(super) fn decode_aggregate(row: &PgRow, exact: bool) -> StoreResult<Self> {
let mut state = Self::decode(row)?;
if exact {
state.exact_percentiles = Some(
get::<sqlx::types::Json<BTreeMap<String, Vec<Option<f64>>>>>(
row,
"exact_percentiles",
)?
.0,
);
}
Ok(state)
}
pub fn encode(&self) -> StoreResult<(serde_json::Value, serde_json::Value, serde_json::Value)> {
Ok((
serde_json::to_value(&self.counts)
@@ -38,17 +64,24 @@ impl MetricState {
.map_err(|error| postgres_unavailable("encode analytics histograms", error))?,
))
}
pub fn exact() -> Self {
Self {
exact_samples: Some(BTreeMap::new()),
..Self::default()
}
}
pub fn count(&self, name: &str) -> u64 {
self.counts.get(name).copied().unwrap_or_default()
}
fn increment(&mut self, name: &str, amount: u64) -> StoreResult<()> {
let count = self.counts.entry(name.to_owned()).or_default();
pub fn optional_count(&self, name: &str) -> Option<u64> {
(self.count(&format!("{name}_known_count")) > 0).then(|| self.count(name))
}
pub fn last_used_at(&self) -> Option<DateTime<Utc>> {
self.last_used_at
}
pub(super) fn has_exact_percentiles(&self) -> bool {
self.exact_percentiles.is_some()
}
pub(super) fn increment(&mut self, name: &str, amount: u64) -> StoreResult<()> {
// 合并同一分组时键通常已经存在,避免每条事实都为已有指标分配字符串
let count = match self.counts.get_mut(name) {
Some(count) => count,
None => self.counts.entry(name.to_owned()).or_default(),
};
*count = count
.checked_add(amount)
.ok_or_else(|| invalid("analytics counter overflow"))?;
@@ -60,13 +93,10 @@ impl MetricState {
.entry(name.to_owned())
.or_default()
.add(sample)?;
if let Some(samples) = &mut self.exact_samples {
samples.entry(name.to_owned()).or_default().push(sample);
}
}
Ok(())
}
fn add_cost(&mut self, currency: &str, amount: &str) -> StoreResult<()> {
pub(super) fn add_cost(&mut self, currency: &str, amount: &str) -> StoreResult<()> {
let amount: DecimalAmount = amount
.parse()
.map_err(|_| invalid("analytics invalid cost"))?;
@@ -83,19 +113,28 @@ impl MetricState {
Ok(())
}
pub fn merge(&mut self, other: &Self) -> StoreResult<()> {
// 精确分位数只属于数据库已经完成分组的结果,不能把分位数再平均或合并
if self.count("request_count") > 0
&& other.count("request_count") > 0
&& (self.exact_percentiles.is_some() || other.exact_percentiles.is_some())
{
return Err(invalid(
"cannot merge finalized exact analytics percentiles",
));
}
if self.count("request_count") == 0 {
self.exact_percentiles.clone_from(&other.exact_percentiles);
}
self.last_used_at = self.last_used_at.max(other.last_used_at);
for (name, count) in &other.counts {
self.increment(name, *count)?;
}
for (name, distribution) in &other.distributions {
self.distributions
.entry(name.clone())
.or_default()
.merge(distribution)?;
}
if let (Some(ours), Some(theirs)) = (&mut self.exact_samples, &other.exact_samples) {
for (name, samples) in theirs {
ours.entry(name.clone()).or_default().extend(samples);
}
let stored = match self.distributions.get_mut(name) {
Some(stored) => stored,
None => self.distributions.entry(name.clone()).or_default(),
};
stored.merge(distribution)?;
}
for (currency, amount) in &other.costs {
self.add_cost(currency, amount)?;
@@ -185,18 +224,31 @@ impl MetricState {
self.sample("capacity", Some(used as f64 / total as f64 * 10_000.0))?;
}
if completed {
for name in [
"input_tokens",
"output_tokens",
"cached_input_tokens",
"cache_write_input_tokens",
"reasoning_output_tokens",
"total_tokens",
"image_input_tokens",
"image_output_tokens",
self.increment("completed_request_count", 1)?;
self.last_used_at = self.last_used_at.max(Some(get(row, "started_at")?));
for (value, name) in [
(Some(true), "image_request_count"),
(Some(false), "image_request_failed_count"),
] {
self.increment(name, optional_unsigned(row, name)?.unwrap_or_default())?;
if get::<Option<bool>>(row, "image_generation_succeeded")? == value {
self.increment(name, 1)?;
}
}
for name in TOKEN_METRICS {
if let Some(value) = optional_unsigned(row, name)? {
self.increment(name, value)?;
self.increment(&format!("{name}_known_count"), 1)?;
}
}
self.increment(
"window_total_tokens",
optional_unsigned(row, "total_tokens")?.unwrap_or(
optional_unsigned(row, "input_tokens")?
.unwrap_or_default()
.checked_add(optional_unsigned(row, "output_tokens")?.unwrap_or_default())
.ok_or_else(|| invalid("analytics token overflow"))?,
),
)?;
if get::<Option<i64>>(row, "input_tokens")?.is_some() {
self.increment("cache_eligible_request_count", 1)?;
if optional_unsigned(row, "cached_input_tokens")?.is_some_and(|value| value > 0) {
@@ -245,13 +297,11 @@ impl MetricState {
self.distributions.get(name).cloned().unwrap_or_default()
}
pub fn percentile(&self, name: &str, fraction: f64) -> Option<f64> {
if let Some(samples) = &self.exact_samples {
let mut samples = samples.get(name)?.clone();
samples.sort_by(f64::total_cmp);
let rank = samples.len().saturating_sub(1) as f64 * fraction;
let a = *samples.get(rank.floor() as usize)?;
let b = *samples.get(rank.ceil() as usize)?;
return Some(a + (b - a) * rank.fract());
if let Some(percentiles) = &self.exact_percentiles {
let index = [0.1, 0.5, 0.9, 0.95, 0.99]
.iter()
.position(|value| (*value - fraction).abs() < f64::EPSILON)?;
return percentiles.get(name)?.get(index).copied().flatten();
}
self.histogram(name).percentile(fraction)
}
@@ -1,7 +1,12 @@
//! 请求观测的持久化投影、可重建汇总与受预算约束的读取
pub(crate) mod accounts;
mod aggregate;
mod buckets;
mod counts;
mod forecast;
mod histogram;
mod incremental;
mod metrics;
mod query;
mod reset;
@@ -11,15 +16,15 @@ pub(crate) use query::*;
pub use worker::{AnalyticsMaintenance, AnalyticsStatus, PgRequestAnalyticsRepository};
use super::*;
use metrics::MetricState;
pub(crate) use metrics::MetricState;
const SCHEMA_VERSION: i32 = 1;
const HISTOGRAM_VERSION: i32 = histogram::HISTOGRAM_VERSION;
const MAX_FACT_ROWS: usize = 100_000;
pub(super) async fn begin_snapshot(
connection: &mut PgConnection,
connection: &mut ObservabilityConnection,
) -> StoreResult<sqlx::Transaction<'_, Postgres>> {
let query_timeout = connection.query_timeout();
let mut transaction = connection
.begin()
.await
@@ -28,7 +33,8 @@ pub(super) async fn begin_snapshot(
.execute(&mut *transaction)
.await
.map_err(|error| postgres_unavailable("isolate analytics snapshot", error))?;
sqlx::query("set local statement_timeout = '10s'")
sqlx::query("select set_config('statement_timeout',$1,true)")
.bind(format!("{}ms", query_timeout.as_millis()))
.execute(&mut *transaction)
.await
.map_err(|error| postgres_unavailable("set analytics query deadline", error))?;
@@ -5,7 +5,7 @@ use futures::TryStreamExt;
pub(crate) struct AnalyticsSnapshot<'a> {
pub tx: sqlx::Transaction<'a, Postgres>,
status: AnalyticsStatus,
pub(super) status: AnalyticsStatus,
approximate: bool,
}
@@ -13,10 +13,25 @@ struct ReadSelection<'a> {
distributions: bool,
costs: bool,
keys: Option<&'a [String]>,
counters: Option<&'a [&'a str]>,
distribution_names: Option<&'a [&'a str]>,
}
impl<'a> AnalyticsSnapshot<'a> {
pub async fn new(connection: &'a mut PgConnection) -> StoreResult<Self> {
pub async fn new(connection: &'a mut ObservabilityConnection) -> StoreResult<Self> {
Self::begin(connection, true).await
}
pub(crate) async fn new_for_accounts(
connection: &'a mut ObservabilityConnection,
) -> StoreResult<Self> {
Self::begin(connection, false).await
}
async fn begin(
connection: &'a mut ObservabilityConnection,
require_current_schema: bool,
) -> StoreResult<Self> {
let mut tx = begin_snapshot(connection).await?;
// 时间跨度和分组选择性差异很大,通用计划会把长区间误估成短区间
// 仅当前只读事务使用实际参数规划,不改变连接池后续事务的计划策略
@@ -25,8 +40,9 @@ impl<'a> AnalyticsSnapshot<'a> {
.await
.map_err(|error| postgres_unavailable("plan analytics with actual bounds", error))?;
let status = worker::query_status(&mut tx).await?;
if status.schema_version != SCHEMA_VERSION as u32
|| status.histogram_version != HISTOGRAM_VERSION as u32
if require_current_schema
&& (status.schema_version != SCHEMA_VERSION as u32
|| status.histogram_version != HISTOGRAM_VERSION as u32)
{
return Err(unavailable(
"request analytics is building; retry after the background projection completes",
@@ -163,14 +179,35 @@ impl<'a> AnalyticsSnapshot<'a> {
.load(range, filter, None, None, Default::default(), true)
.await?;
let state = merge_states(states.into_values())?;
self.overview_from_state(range, filter, state).await
}
async fn overview_from_state(
&mut self,
range: ObservabilityRange,
filter: &UsageRecordFilter,
state: MetricState,
) -> StoreResult<UsageOverview> {
// Provider 列表只消费计数,不重复合并费用与分布
let provider_states = self
.load(
.load_selected(
range,
filter,
Some(DiagnosticDimension::Provider),
None,
Default::default(),
false,
ReadSelection {
distributions: false,
costs: false,
keys: None,
counters: Some(&[
"request_count",
"attempt_count",
"failure_count",
"total_tokens",
]),
distribution_names: None,
},
)
.await?;
let mut providers = provider_states
@@ -216,47 +253,48 @@ impl<'a> AnalyticsSnapshot<'a> {
distributions: true,
costs: include_costs,
keys: None,
counters: None,
distribution_names: None,
},
)
.await?;
let mut points = BTreeMap::new();
for ((bucket, _), state) in states {
points.insert(
bucket,
RequestMetricPoint {
bucket_start: bucket,
granularity,
metrics: state.request_metrics()?,
cost_coverage: state.coverage(),
costs: if include_costs {
state.costs()?
} else {
Vec::new()
},
},
);
}
fill_metric_gaps(range, granularity, points, timezone)
metric_points(range, granularity, timezone, &states, include_costs)
}
pub async fn projected_costs(
pub async fn insights(
&mut self,
range: ObservabilityRange,
filter: &UsageRecordFilter,
granularity: ObservationGranularity,
timezone: gateway_core::time::DeploymentTimeZone,
) -> StoreResult<Vec<admin_observability::UsageProjectedCostPoint>> {
self.load(range, filter, None, Some(granularity), timezone, false)
.await?
.into_iter()
) -> StoreResult<(
UsageOverview,
Vec<RequestMetricPoint>,
Vec<admin_observability::UsageProjectedCostPoint>,
)> {
let states = self
.load(range, filter, None, Some(granularity), timezone, true)
.await?;
let trend = metric_points(range, granularity, timezone, &states, true)?;
let costs = states
.iter()
.map(|((bucket_start, _), state)| {
Ok(admin_observability::UsageProjectedCostPoint {
bucket_start,
bucket_start: *bucket_start,
standard_amount: state.projected_cost("standard_usd")?,
no_cache_amount: state.projected_cost("no_cache_usd")?,
})
})
.collect()
.collect::<StoreResult<Vec<_>>>()?;
// 趋势已包含可合并的完整指标和情景费用,概览直接复用同一份状态
// 精确分位数不可跨桶合并,明细回退时仍按完整区间单独计算概览
let overview = if states.values().any(MetricState::has_exact_percentiles) {
self.overview(range, filter).await?
} else {
self.overview_from_state(range, filter, merge_states(states.into_values())?)
.await?
};
Ok((overview, trend, costs))
}
pub async fn diagnostics(
@@ -273,8 +311,10 @@ impl<'a> AnalyticsSnapshot<'a> {
let ranked = if matches!(
dimension,
DiagnosticDimension::Account | DiagnosticDimension::ApiKey
) && collection_for(filter, Some(dimension)).is_some()
{
) && matches!(
collection_for(filter, Some(dimension)),
Some("account" | "key")
) {
Some(
self.diagnostic_top_keys(range, filter, dimension, limit)
.await?,
@@ -294,6 +334,22 @@ impl<'a> AnalyticsSnapshot<'a> {
distributions: true,
costs: true,
keys: selected_keys,
counters: Some(&[
"request_count",
"success_count",
"failure_count",
"attempt_count",
"total_tokens",
"latency_ms_sum",
"cancelled_count",
"incomplete_count",
"retry_count",
"retried_request_count",
"provider_reported_count",
"calculated_count",
"unavailable_count",
]),
distribution_names: Some(&["latency_ms", "first_token_ms"]),
},
)
.await?;
@@ -415,16 +471,18 @@ impl<'a> AnalyticsSnapshot<'a> {
.collect::<Vec<_>>();
let mut query = QueryBuilder::<Postgres>::new("with edge_facts as materialized (");
if starts.is_empty() {
query.push("select null::text as key where false");
query.push("select null::text as key,0::bigint as request_count where false");
} else {
query.push("select edge.key from unnest(").push_bind(starts).push("::timestamptz[],").push_bind(ends).push("::timestamptz[]) span(start_at,end_at) cross join lateral (select ")
.push(format!("coalesce(mr.{column},'') as key from request_analytics_facts mr where mr.started_at >= span.start_at and mr.started_at < span.end_at"));
query.push("select edge.key,sum(edge.request_count)::bigint as request_count from unnest(").push_bind(starts).push("::timestamptz[],").push_bind(ends).push("::timestamptz[]) span(start_at,end_at) cross join lateral (select ")
.push(format!("coalesce(mr.{column},'') as key,count(*)::bigint as request_count from request_analytics_facts mr where mr.started_at >= span.start_at and mr.started_at < span.end_at"));
push_unrecovered_request_filter(&mut query, "mr");
push_analytics_filter(&mut query, filter, "mr");
push_diagnostic_dimension_filter(&mut query, dimension);
query.push(" limit 100001) edge limit 100001");
query.push(format!(
" group by coalesce(mr.{column},'')) edge group by edge.key"
));
}
query.push("), contributions as (select key,count(*)::bigint as request_count from edge_facts group by key");
query.push("), contributions as (select key,request_count from edge_facts");
for (table, column, seconds, buckets) in [
("request_analytics_daily_counts", "dimension", 86400, days),
("request_analytics_hourly", "dimensions[1]", 3600, hours),
@@ -446,7 +504,7 @@ impl<'a> AnalyticsSnapshot<'a> {
}
query.push(format!(" group by {column}"));
}
query.push("), grouped as (select key,sum(request_count)::bigint as request_count from contributions group by key) select key,request_count,sum(request_count) over()::bigint as total_request_count,(select count(*) from edge_facts)::bigint as edge_count from grouped order by request_count desc,");
query.push("), grouped as (select key,sum(request_count)::bigint as request_count from contributions group by key) select key,request_count,sum(request_count) over()::bigint as total_request_count from grouped order by request_count desc,");
if dimension == DiagnosticDimension::Account {
query.push("coalesce(nullif(key,''),'unrouted')");
} else {
@@ -458,12 +516,6 @@ impl<'a> AnalyticsSnapshot<'a> {
.fetch_all(&mut *self.tx)
.await
.map_err(|error| postgres_unavailable("rank full-range analytics groups", error))?;
if rows.first().is_some_and(|row| {
row.try_get::<i64, _>("edge_count")
.is_ok_and(|count| count > MAX_FACT_ROWS as i64)
}) {
return Err(unavailable("analytics fact row budget exceeded"));
}
let total = rows
.first()
.map(|row| unsigned(row, "total_request_count"))
@@ -476,18 +528,15 @@ impl<'a> AnalyticsSnapshot<'a> {
Ok((total, keys))
}
async fn summary_buckets(
pub(super) async fn summary_buckets(
&mut self,
range: ObservabilityRange,
use_hourly: bool,
) -> StoreResult<(Vec<DateTime<Utc>>, Vec<DateTime<Utc>>)> {
// 先按时间确定互斥的完整桶,空的十五分钟集合不读取汇总表
let rows = sqlx::query("with quarters as materialized (select q.bucket_start,date_bin(interval '1 hour',q.bucket_start,timestamptz '2000-01-01 UTC') as hour_start from generate_series(date_bin(interval '15 minutes',$1::timestamptz,timestamptz '2000-01-01 UTC'),$2::timestamptz-interval '15 minutes',interval '15 minutes') q(bucket_start) where q.bucket_start >= $1 and not exists(select 1 from request_analytics_pending_buckets p where p.bucket_start=q.bucket_start)), hours as (select hour_start from quarters where $3 group by hour_start having count(*)=4) select q.bucket_start,h.hour_start from quarters q left join hours h on h.hour_start=q.hour_start order by q.bucket_start limit 100001")
let rows = sqlx::query("with quarters as materialized (select q.bucket_start,date_bin(interval '1 hour',q.bucket_start,timestamptz '2000-01-01 UTC') as hour_start from generate_series(date_bin(interval '15 minutes',$1::timestamptz,timestamptz '2000-01-01 UTC'),$2::timestamptz-interval '15 minutes',interval '15 minutes') q(bucket_start) where q.bucket_start >= $1 and not exists(select 1 from request_analytics_pending_buckets p where p.bucket_start=q.bucket_start)), hours as (select hour_start from quarters where $3 group by hour_start having count(*)=4) select q.bucket_start,h.hour_start from quarters q left join hours h on h.hour_start=q.hour_start order by q.bucket_start")
.bind(range.start).bind(range.end).bind(use_hourly)
.fetch_all(&mut *self.tx).await.map_err(|error|postgres_unavailable("select disjoint analytics buckets",error))?;
if rows.len() > MAX_FACT_ROWS {
return Err(unavailable("analytics time bucket budget exceeded"));
}
let mut hours = Vec::new();
let mut quarters = Vec::new();
for row in rows {
@@ -502,7 +551,7 @@ impl<'a> AnalyticsSnapshot<'a> {
Ok((hours, quarters))
}
async fn uncovered_spans(
pub(super) async fn uncovered_spans(
&mut self,
range: ObservabilityRange,
) -> StoreResult<Vec<(DateTime<Utc>, DateTime<Utc>)>> {
@@ -522,11 +571,8 @@ impl<'a> AnalyticsSnapshot<'a> {
spans.push((right_start, range.end));
}
if complete_start < complete_end {
let dirty:Vec<DateTime<Utc>>=sqlx::query_scalar("select bucket_start from request_analytics_pending_buckets where bucket_start >= $1 and bucket_start < $2 order by bucket_start limit 100001")
let dirty:Vec<DateTime<Utc>>=sqlx::query_scalar("select bucket_start from request_analytics_pending_buckets where bucket_start >= $1 and bucket_start < $2 order by bucket_start")
.bind(complete_start).bind(complete_end).fetch_all(&mut *self.tx).await.map_err(|error|postgres_unavailable("plan bounded analytics fact spans",error))?;
if dirty.len() > MAX_FACT_ROWS {
return Err(unavailable("analytics dirty interval budget exceeded"));
}
spans.extend(
dirty
.into_iter()
@@ -534,7 +580,17 @@ impl<'a> AnalyticsSnapshot<'a> {
);
}
spans.sort_unstable();
Ok(spans)
let mut merged: Vec<(DateTime<Utc>, DateTime<Utc>)> = Vec::new();
for (start, end) in spans {
if let Some(previous) = merged.last_mut()
&& previous.1 == start
{
previous.1 = end;
} else {
merged.push((start, end));
}
}
Ok(merged)
}
async fn load(
@@ -556,6 +612,8 @@ impl<'a> AnalyticsSnapshot<'a> {
distributions: read_distributions,
costs: true,
keys: None,
counters: None,
distribution_names: None,
},
)
.await
@@ -574,6 +632,8 @@ impl<'a> AnalyticsSnapshot<'a> {
distributions: read_distributions,
costs: read_costs,
keys: selected_keys,
counters,
distribution_names,
} = selection;
filter.validate()?;
self.require_visible(Some(range)).await?;
@@ -610,90 +670,63 @@ impl<'a> AnalyticsSnapshot<'a> {
if buckets.is_empty() {
continue;
}
let (dimension_filter, either_filter, either_positions) = match collection {
"provider_model" => {
(filter.provider_kind.clone(), filter.model.clone(), (2, 3))
}
"account" => (filter.provider_account_ref.clone(), None, (1, 2)),
"key" => (filter.client_api_key_ref.clone(), None, (1, 2)),
"operation" => (filter.operation.clone(), None, (1, 2)),
"transport" => (None, filter.transport.clone(), (1, 2)),
"status" => (
None,
filter.status_code.map(|value| value.to_string()),
(1, 2),
),
_ => (None, None, (1, 2)),
let mut query =
QueryBuilder::<Postgres>::new("with matching as not materialized (select ");
aggregate::push_output_bucket(
&mut query,
"bucket_start",
range,
granularity,
boundaries.as_deref(),
);
let name = aggregate::summary_dimension_sql(collection, dimension);
let costs = if read_costs {
"costs"
} else {
"'{}'::jsonb as costs"
};
let distributions = if read_distributions {
"distributions"
} else {
"'{}'::jsonb as distributions"
};
let costs = if read_costs {
"costs"
} else {
"'{}'::jsonb as costs"
};
let provider_counts = dimension == Some(DiagnosticDimension::Provider)
&& granularity.is_none()
&& !read_distributions;
let mut query = QueryBuilder::<Postgres>::new(if provider_counts {
"select min(bucket_start) as bucket_start, array[dimensions[1]] as dimensions, ''::text as outcome, jsonb_build_object('request_count',sum((counts->>'request_count')::bigint),'attempt_count',sum((counts->>'attempt_count')::bigint),'failure_count',sum(coalesce((counts->>'failure_count')::bigint,0)),'total_tokens',sum(coalesce((counts->>'total_tokens')::bigint,0))) as counts, '{}'::jsonb as costs, '{}'::jsonb as distributions from ("
} else {
""
});
query.push(format!("select bucket_start, dimensions, outcome, counts, {costs}, {distributions} from {table} a where collection="))
query.push(format!(" as output_bucket,{name} as name,counts,{costs},{distributions},last_used_at from {table} a where collection="))
.push_bind(collection).push(" and schema_version=").push_bind(SCHEMA_VERSION)
.push(" and histogram_version=").push_bind(HISTOGRAM_VERSION);
push_summary_buckets(&mut query, &buckets, seconds);
// 可选筛选仅在有值时生成谓词,保留维度与时间索引的选择空间
if let Some(outcome) = &filter.outcome {
query.push(" and outcome=").push_bind(outcome);
}
if let Some(value) = dimension_filter {
query.push(" and dimensions[1]=").push_bind(value);
}
if let Some(value) = either_filter {
query
.push(format!(" and (dimensions[{}]=", either_positions.0))
.push_bind(value.clone())
.push(format!(" or dimensions[{}]=", either_positions.1))
.push_bind(value)
.push(")");
}
push_collection_filter(&mut query, collection, filter);
if let Some(keys) = selected_keys {
query
.push(" and dimensions[1]=any(")
.push_bind(keys)
.push(")");
}
if provider_counts {
query.push(") matching group by dimensions[1]");
if matches!(
dimension,
Some(
DiagnosticDimension::Model
| DiagnosticDimension::Status
| DiagnosticDimension::Failure
)
) {
query.push(format!(" and ({name})<>''"));
}
query.push(" limit 100001");
query.push(")");
aggregate::push_summary_aggregation(
&mut query,
read_distributions,
read_costs,
counters,
distribution_names,
);
let mut rows = query.build().fetch(&mut *self.tx);
let mut count = 0;
while let Some(row) = rows
.try_next()
.await
.map_err(|error| postgres_unavailable("read analytics summaries", error))?
.map_err(|error| postgres_unavailable("merge analytics summaries", error))?
{
count += 1;
if count > MAX_FACT_ROWS {
return Err(unavailable("analytics aggregate row budget exceeded"));
}
let dimensions: Vec<String> = get(&row, "dimensions")?;
let name = dimension_name(collection, &dimensions, dimension);
let Some(name) = name else {
continue;
};
let bucket = output_bucket(
get(&row, "bucket_start")?,
range,
granularity,
boundaries.as_deref(),
)?;
let bucket = get(&row, "bucket_start")?;
let name = get(&row, "name")?;
let state = MetricState::decode(&row)?;
result.entry((bucket, name)).or_default().merge(&state)?;
}
@@ -707,28 +740,42 @@ impl<'a> AnalyticsSnapshot<'a> {
if spans.is_empty() {
return Ok(result);
}
// 单段直接绑定边界,让优化器看见真实跨度,避免把三十天误估成小范围 BitmapAnd
// 单段直接绑定边界,多个互斥区间用索引读取,所有统计都在数据库内完成
let single_span = spans.len() == 1;
let mut query = if single_span {
let mut query = QueryBuilder::<Postgres>::new(fact_select());
query
.push(" where mr.started_at >= ")
.push_bind(spans[0].0)
.push(" and mr.started_at < ")
.push_bind(spans[0].1);
query
} else {
let (starts, ends): (Vec<_>, Vec<_>) = spans.into_iter().unzip();
let mut query = QueryBuilder::<Postgres>::new("select fact.* from unnest(");
let mut query = QueryBuilder::<Postgres>::new("with matching as not materialized (");
if !single_span {
let (starts, ends): (Vec<_>, Vec<_>) = spans.iter().copied().unzip();
query
.push("select fact.* from unnest(")
.push_bind(starts)
.push("::timestamptz[],")
.push_bind(ends)
.push("::timestamptz[]) span(start_at,end_at) cross join lateral (")
.push(fact_select())
.push(" where mr.started_at >= span.start_at and mr.started_at < span.end_at");
.push("::timestamptz[]) span(start_at,end_at) cross join lateral (");
}
query.push(format!(
"select mr.*,coalesce(({}),false) as is_completed_usage,",
completed_usage_fact_predicate("mr")
));
aggregate::push_output_bucket(
&mut query,
"mr.started_at",
range,
granularity,
boundaries.as_deref(),
);
query.push(format!(
" as output_bucket,{} as name from request_analytics_facts mr where ",
aggregate::fact_dimension_sql(dimension)
));
if single_span {
query
};
.push("mr.started_at >= ")
.push_bind(spans[0].0)
.push(" and mr.started_at < ")
.push_bind(spans[0].1);
} else {
query.push("mr.started_at >= span.start_at and mr.started_at < span.end_at");
}
push_unrecovered_request_filter(&mut query, "mr");
push_analytics_filter(&mut query, filter, "mr");
if let Some(dimension) = dimension {
@@ -749,42 +796,40 @@ impl<'a> AnalyticsSnapshot<'a> {
}
query.push(")");
}
query.push(" limit 100001");
if !single_span {
query.push(") fact limit 100001");
// 保留逐区间的参数化索引扫描,避免展开 LATERAL 后全表读取再过滤
// OFFSET 0 不限制返回行数
query.push(" offset 0) fact");
}
query.push(")");
let exact = selected.is_none();
aggregate::push_fact_aggregation_selected(
&mut query,
exact,
read_costs,
read_distributions,
counters,
distribution_names,
);
let mut rows = query.build().fetch(&mut *self.tx);
let mut count = 0;
while let Some(row) = rows
.try_next()
.await
.map_err(|error| postgres_unavailable("read bounded analytics facts", error))?
.map_err(|error| postgres_unavailable("aggregate analytics facts", error))?
{
count += 1;
if count > MAX_FACT_ROWS {
return Err(unavailable("analytics fact row budget exceeded"));
let key = (get(&row, "bucket_start")?, get(&row, "name")?);
let state = MetricState::decode_aggregate(&row, exact)?;
if exact {
result.insert(key, state);
} else {
result.entry(key).or_default().merge(&state)?;
}
let bucket = output_bucket(
get(&row, "started_at")?,
range,
granularity,
boundaries.as_deref(),
)?;
let name = fact_dimension(&row, dimension)?;
let state = result.entry((bucket, name)).or_insert_with(|| {
if selected.is_some() {
MetricState::default()
} else {
MetricState::exact()
}
});
state.add_row(&row)?;
}
Ok(result)
}
}
fn push_summary_buckets(
pub(super) fn push_summary_buckets(
query: &mut QueryBuilder<Postgres>,
buckets: &[DateTime<Utc>],
seconds: i64,
@@ -810,6 +855,33 @@ fn push_summary_buckets(
}
}
fn metric_points(
range: ObservabilityRange,
granularity: ObservationGranularity,
timezone: gateway_core::time::DeploymentTimeZone,
states: &BTreeMap<(DateTime<Utc>, String), MetricState>,
include_costs: bool,
) -> StoreResult<Vec<RequestMetricPoint>> {
let mut points = BTreeMap::new();
for ((bucket, _), state) in states {
points.insert(
*bucket,
RequestMetricPoint {
bucket_start: *bucket,
granularity,
metrics: state.request_metrics()?,
cost_coverage: state.coverage(),
costs: if include_costs {
state.costs()?
} else {
Vec::new()
},
},
);
}
fill_metric_gaps(range, granularity, points, timezone)
}
fn merge_states(states: impl Iterator<Item = MetricState>) -> StoreResult<MetricState> {
let mut result: Option<MetricState> = None;
for state in states {
@@ -822,28 +894,7 @@ fn merge_states(states: impl Iterator<Item = MetricState>) -> StoreResult<Metric
Ok(result.unwrap_or_default())
}
fn output_bucket(
time: DateTime<Utc>,
range: ObservabilityRange,
granularity: Option<ObservationGranularity>,
boundaries: Option<&[DateTime<Utc>]>,
) -> StoreResult<DateTime<Utc>> {
match granularity {
None => Ok(range.start),
Some(ObservationGranularity::Day) => boundaries
.and_then(|values| {
values
.iter()
.rev()
.find(|boundary| **boundary <= time)
.copied()
})
.ok_or_else(|| invalid("missing analytics calendar boundary")),
Some(granularity) => bucket_for(time, granularity.seconds()),
}
}
fn collection_for(
pub(super) fn collection_for(
filter: &UsageRecordFilter,
dimension: Option<DiagnosticDimension>,
) -> Option<&'static str> {
@@ -852,7 +903,6 @@ fn collection_for(
|| filter.response_id.is_some()
|| filter.upstream_request_id.is_some()
|| filter.search.is_some()
|| !filter.routing_group_refs.is_empty()
{
return None;
}
@@ -885,85 +935,18 @@ fn collection_for(
DiagnosticDimension::Status => "status",
});
}
if required.len() > 1 {
if required.len() == 2 && required.contains("provider_model") && required.contains("account") {
Some("account_model")
} else if required.len() == 2 && required.contains("provider_model") && required.contains("key")
{
Some("key_model")
} else if required.len() > 1 {
None
} else {
Some(required.first().copied().unwrap_or("global"))
}
}
fn dimension_name(
collection: &str,
values: &[String],
dimension: Option<DiagnosticDimension>,
) -> Option<String> {
let Some(dimension) = dimension else {
return Some(String::new());
};
let value = |index| values.get(index).cloned().unwrap_or_default();
let coalesce = |preferred: usize, fallback: usize| {
let v = value(preferred);
if v.is_empty() { value(fallback) } else { v }
};
let name = match dimension {
DiagnosticDimension::Provider if collection == "provider_model" => {
let v = value(0);
if v.is_empty() {
"unrouted".to_owned()
} else {
v
}
}
DiagnosticDimension::Model => coalesce(2, 1),
DiagnosticDimension::Account => {
let v = value(0);
if v.is_empty() {
"unrouted".to_owned()
} else {
v
}
}
DiagnosticDimension::Transport => {
let v = coalesce(1, 0);
if v.is_empty() {
"unknown".to_owned()
} else {
v
}
}
DiagnosticDimension::Status => coalesce(1, 0),
_ => value(0),
};
if name.is_empty()
&& matches!(
dimension,
DiagnosticDimension::Model | DiagnosticDimension::Status | DiagnosticDimension::Failure
)
{
None
} else {
Some(name)
}
}
fn fact_dimension(
row: &sqlx::postgres::PgRow,
dimension: Option<DiagnosticDimension>,
) -> StoreResult<String> {
let Some(dimension) = dimension else {
return Ok(String::new());
};
let collection = collection_for(&UsageRecordFilter::default(), Some(dimension))
.ok_or_else(|| invalid("missing analytics dimension"))?;
let dimensions = worker::collections(row)?
.into_iter()
.find(|(name, _)| *name == collection)
.map(|(_, values)| values)
.ok_or_else(|| invalid("missing analytics dimension"))?;
dimension_name(collection, &dimensions, Some(dimension))
.ok_or_else(|| invalid("unexpected empty analytics dimension"))
}
pub(crate) fn push_analytics_filter(
query: &mut QueryBuilder<Postgres>,
filter: &UsageRecordFilter,
@@ -982,3 +965,53 @@ pub(crate) fn push_analytics_filter(
query.push(")");
}
}
pub(super) fn push_collection_filter(
query: &mut QueryBuilder<Postgres>,
collection: &str,
filter: &UsageRecordFilter,
) {
let (dimension_filter, either_filter, either_positions) = match collection {
"provider_model" => (filter.provider_kind.clone(), filter.model.clone(), (2, 3)),
"account_model" => (
filter.provider_account_ref.clone(),
filter.model.clone(),
(3, 4),
),
"key_model" => (
filter.client_api_key_ref.clone(),
filter.model.clone(),
(3, 4),
),
"account" => (filter.provider_account_ref.clone(), None, (1, 2)),
"key" => (filter.client_api_key_ref.clone(), None, (1, 2)),
"operation" => (filter.operation.clone(), None, (1, 2)),
"transport" => (None, filter.transport.clone(), (1, 2)),
"status" => (
None,
filter.status_code.map(|value| value.to_string()),
(1, 2),
),
_ => (None, None, (1, 2)),
};
// 可选筛选仅在有值时生成谓词,保留维度与时间索引的选择空间
if let Some(outcome) = &filter.outcome {
query.push(" and outcome=").push_bind(outcome);
}
if let Some(value) = dimension_filter {
query.push(" and dimensions[1]=").push_bind(value);
}
if matches!(collection, "account_model" | "key_model")
&& let Some(provider) = &filter.provider_kind
{
query.push(" and dimensions[2]=").push_bind(provider);
}
if let Some(value) = either_filter {
query
.push(format!(" and (dimensions[{}]=", either_positions.0))
.push_bind(value.clone())
.push(format!(" or dimensions[{}]=", either_positions.1))
.push_bind(value)
.push(")");
}
}
@@ -83,11 +83,11 @@ impl PgRequestAnalyticsRepository {
return Ok(AnalyticsMaintenance { reset_progress: true, ..Default::default() });
}
let mut result = AnalyticsMaintenance::default();
let rows = sqlx::query("select request_id, change_version, oldest_pending_at, previous_bucket from request_analytics_pending_requests order by oldest_pending_at, request_id limit $1")
let rows = sqlx::query("select request_id, change_version from request_analytics_pending_requests order by oldest_pending_at, request_id limit $1")
.bind(i64::from(batch_size)).fetch_all(&mut *transaction).await.map_err(|error| postgres_unavailable("load analytics request queue", error))?;
let mut jobs = Vec::with_capacity(rows.len());
for row in rows {
jobs.push(ProjectionJob { id: get(&row,"request_id")?, version: get(&row,"change_version")?, pending_at: get(&row,"oldest_pending_at")?, previous_bucket: get(&row,"previous_bucket")? });
jobs.push(ProjectionJob { id: get(&row,"request_id")?, version: get(&row,"change_version")? });
}
self.project_batch(&mut transaction, &jobs).await?;
result.projected_requests = jobs.len() as u64;
@@ -138,53 +138,36 @@ impl PgRequestAnalyticsRepository {
}
let ids = jobs.iter().map(|job| job.id.clone()).collect::<Vec<_>>();
let versions = jobs.iter().map(|job| job.version).collect::<Vec<_>>();
let pending = jobs
.iter()
.map(|job| (job.id.as_str(), job.pending_at))
.collect::<HashMap<_, _>>();
let mut buckets = BTreeMap::<DateTime<Utc>, DateTime<Utc>>::new();
let mut affect = |id: &str, start: DateTime<Utc>| -> StoreResult<()> {
let pending_at = pending
.get(id)
.copied()
.ok_or_else(|| invalid("analytics job disappeared"))?;
let bucket = bucket_for(start, 900)?;
buckets
.entry(bucket)
.and_modify(|old| *old = (*old).min(pending_at))
.or_insert(pending_at);
Ok(())
};
let rows = sqlx::query(
"delete from request_analytics_facts where id=any($1::text[]) returning id,started_at",
)
let previous = sqlx::query(sqlx::AssertSqlSafe(format!(
"{} where mr.id=any($1::text[])",
fact_select()
)))
.bind(&ids)
.fetch_all(&mut **tx)
.await
.map_err(|error| postgres_unavailable("replace analytics facts", error))?;
for row in rows {
affect(&get::<String>(&row, "id")?, get(&row, "started_at")?)?;
}
let rows = sqlx::query(PROJECT_REQUEST)
.map_err(|error| postgres_unavailable("read previous analytics contributions", error))?;
sqlx::query("delete from request_analytics_facts where id=any($1::text[])")
.bind(&ids)
.execute(&mut **tx)
.await
.map_err(|error| postgres_unavailable("replace analytics facts", error))?;
sqlx::query(PROJECT_REQUEST)
.bind(&ids)
.bind(versions)
.fetch_all(&mut **tx)
.execute(&mut **tx)
.await
.map_err(|error| postgres_unavailable("project analytics facts", error))?;
for row in rows {
affect(&get::<String>(&row, "id")?, get(&row, "started_at")?)?;
}
for job in jobs {
if let Some(bucket) = job.previous_bucket {
affect(&job.id, bucket)?;
}
}
self.project_billing_batch(tx, &ids).await?;
if !buckets.is_empty() {
let (starts, times): (Vec<_>, Vec<_>) = buckets.into_iter().unzip();
sqlx::query("insert into request_analytics_pending_buckets(bucket_start,oldest_pending_at) select * from unnest($1::timestamptz[],$2::timestamptz[]) on conflict(bucket_start) do update set change_version=nextval('request_analytics_change_version'),oldest_pending_at=least(request_analytics_pending_buckets.oldest_pending_at,excluded.oldest_pending_at)")
.bind(starts).bind(times).execute(&mut **tx).await.map_err(|error| postgres_unavailable("enqueue analytics buckets",error))?;
}
let current = sqlx::query(sqlx::AssertSqlSafe(format!(
"{} where mr.id=any($1::text[])",
fact_select()
)))
.bind(&ids)
.fetch_all(&mut **tx)
.await
.map_err(|error| postgres_unavailable("read current analytics contributions", error))?;
// 事实、已发布汇总与构建前缀在同一事务中撤销和添加,重试不会重复计数
super::incremental::update_summaries(tx, &previous, &current).await?;
Ok(())
}
@@ -199,7 +182,7 @@ impl PgRequestAnalyticsRepository {
return Ok(0);
}
let mut query = QueryBuilder::<Postgres>::new(
"select id,started_at,nextval('request_analytics_change_version') as version from model_requests",
"select id,started_at,nextval('request_analytics_change_version') as version,exists(select 1 from request_analytics_facts f where f.id=model_requests.id) as projected from model_requests",
);
// 首批与后续批次使用两种明确的计划,避免可空游标的 OR 使通用计划退化为全索引过滤
if let Some(start) = get::<Option<DateTime<Utc>>>(&state, "backfill_cursor_started_at")? {
@@ -222,11 +205,12 @@ impl PgRequestAnalyticsRepository {
.map_err(|error| postgres_unavailable("read analytics backfill batch", error))?;
let mut jobs = Vec::with_capacity(rows.len());
for row in &rows {
if get::<bool>(row, "projected")? {
continue;
}
jobs.push(ProjectionJob {
id: get(row, "id")?,
version: get(row, "version")?,
pending_at: Utc::now(),
previous_bucket: None,
});
}
self.project_batch(tx, &jobs).await?;
@@ -312,49 +296,68 @@ async fn read_status(
})
}
// 投影与极值重算共用同一维度顺序,避免两条路径对集合身份作不同解释
pub(super) const COLLECTION_COLUMNS: &[(&str, &[&str])] = &[
("global", &[]),
(
"provider_model",
&["provider_kind", "requested_model_id", "upstream_model_id"],
),
("account", &["provider_account_ref"]),
("key", &["client_api_key_ref"]),
(
"account_model",
&[
"provider_account_ref",
"provider_kind",
"requested_model_id",
"upstream_model_id",
],
),
(
"key_model",
&[
"client_api_key_ref",
"provider_kind",
"requested_model_id",
"upstream_model_id",
],
),
("operation", &["operation"]),
("transport", &["client_transport", "upstream_transport"]),
("status", &["client_status_code", "upstream_status_code"]),
("failure", &["error_kind"]),
];
pub(super) fn collections(
row: &sqlx::postgres::PgRow,
) -> StoreResult<Vec<(&'static str, Vec<String>)>> {
let text =
|name| -> StoreResult<String> { Ok(get::<Option<String>>(row, name)?.unwrap_or_default()) };
let status = |name| -> StoreResult<String> {
Ok(get::<Option<i32>>(row, name)?
.map(|value| value.to_string())
.unwrap_or_default())
};
Ok(vec![
("global", vec![]),
(
"provider_model",
vec![
text("provider_kind")?,
text("requested_model_id")?,
text("upstream_model_id")?,
],
),
("account", vec![text("provider_account_ref")?]),
("key", vec![text("client_api_key_ref")?]),
("operation", vec![text("operation")?]),
(
"transport",
vec![text("client_transport")?, text("upstream_transport")?],
),
(
"status",
vec![
status("client_status_code")?,
status("upstream_status_code")?,
],
),
("failure", vec![text("error_kind")?]),
])
let mut result = vec![(forecast::COLLECTION, forecast::dimensions(row)?)];
// 预测保留原始完成状态,包括被恢复的请求;其他统计继续排除恢复记录
if get::<Option<DateTime<Utc>>>(row, "recovered_at")?.is_some() {
return Ok(result);
}
for &(collection, columns) in COLLECTION_COLUMNS {
let dimensions = columns
.iter()
.map(|column| {
if matches!(*column, "client_status_code" | "upstream_status_code") {
Ok(get::<Option<i32>>(row, column)?
.map(|value| value.to_string())
.unwrap_or_default())
} else {
Ok(get::<Option<String>>(row, column)?.unwrap_or_default())
}
})
.collect::<StoreResult<Vec<_>>>()?;
result.push((collection, dimensions));
}
Ok(result)
}
struct ProjectionJob {
id: String,
version: i64,
pending_at: DateTime<Utc>,
previous_bucket: Option<DateTime<Utc>>,
}
async fn publish_generation(
@@ -17,7 +17,6 @@ pub(crate) fn store_range(
pub(crate) fn store_usage_filter(filter: admin_observability::UsageFilter) -> UsageRecordFilter {
UsageRecordFilter {
routing_group_refs: filter.routing_group_refs,
client_api_key_ref: filter.client_api_key_ref,
request_id: filter.request_id,
provider_account_ref: filter.provider_account_ref,
@@ -31,6 +31,10 @@ mod queries;
mod query_budget;
pub use admin_adapter::*;
pub(crate) use analytics::AnalyticsSnapshot;
pub(crate) use analytics::accounts::{
account_group_usage, account_usage_sort_values, account_usage_windows,
};
pub use analytics::{AnalyticsMaintenance, AnalyticsStatus, PgRequestAnalyticsRepository};
pub(crate) use mapping::*;
pub use model::*;
@@ -39,7 +39,6 @@ pub(crate) fn observability_page_offset(
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct UsageRecordFilter {
pub routing_group_refs: Vec<String>,
pub client_api_key_ref: Option<String>,
pub request_id: Option<String>,
pub provider_account_ref: Option<String>,
@@ -76,7 +75,6 @@ impl UsageRecordFilter {
] {
validate_optional_text(value, MAX_FILTER_BYTES, field)?;
}
validate_account_ids_if_present(&self.routing_group_refs)?;
validate_optional_text(self.search.as_deref(), MAX_SEARCH_BYTES, "search filter")?;
if self
.status_code
@@ -105,14 +103,6 @@ pub struct UsageRecordQuery {
pub use gateway_admin::model::observability::OpsErrorFilter;
pub(crate) fn validate_ops_error_filter(filter: &OpsErrorFilter) -> StoreResult<()> {
if filter.routing_group_refs.len() > 32
|| filter.routing_group_refs.iter().any(|value| {
value.is_empty() || value.len() > 128 || value.chars().any(char::is_control)
})
{
return Err(invalid("invalid routing group filter"));
}
for (value, field) in [
(
filter.client_api_key_ref.as_deref(),
@@ -384,10 +374,3 @@ pub struct DiagnosticObservation {
pub use gateway_admin::model::observability::OpsError as OpsErrorRecord;
pub use gateway_admin::model::observability::OpsErrorPage;
fn validate_account_ids_if_present(ids: &[String]) -> StoreResult<()> {
if ids.is_empty() {
return Ok(());
}
validate_account_ids(ids)
}
@@ -3,7 +3,7 @@
use super::super::*;
pub(crate) async fn provider_account_usage(
connection: &mut sqlx::PgConnection,
snapshot: &mut analytics::AnalyticsSnapshot<'_>,
query: ProviderAccountUsageQuery,
) -> StoreResult<Vec<ProviderAccountUsageObservation>> {
if let Some(account_ids) = &query.account_ids {
@@ -13,290 +13,133 @@ pub(crate) async fn provider_account_usage(
return Err(invalid("account usage limit must be between 1 and 200"));
}
// 账号卡片先确定有限候选,避免为只展示二十个账号聚合整个历史范围
let mut statement = QueryBuilder::<Postgres>::new(
"with candidate_accounts as materialized (select pa.* from provider_accounts pa",
);
// 账号卡片先确定有限候选,避免为少量账号聚合整个历史范围
let mut statement = QueryBuilder::<Postgres>::new("select pa.* from provider_accounts pa");
if let Some(account_ids) = &query.account_ids {
statement
.push(" where pa.id=any(")
.push_bind(account_ids.clone())
.push("::text[])");
} else {
statement.push(" join lateral (select mr.started_at from request_analytics_facts mr where mr.provider_account_ref=pa.id and mr.started_at >= ").push_bind(query.range.start)
statement.push(" join lateral (select mr.started_at from model_requests mr where mr.provider_account_ref=pa.id and mr.started_at >= ").push_bind(query.range.start)
.push(" and mr.started_at < ").push_bind(query.range.end);
push_completed_usage_fact_filter(&mut statement, "mr");
statement.push(" order by mr.started_at desc,mr.id desc limit 1) latest on true order by latest.started_at desc,pa.name,pa.id limit ").push_bind(i64::from(query.limit));
}
statement.push(
"), matched as materialized (
select pa.id, pa.provider_kind, pa.authentication_kind, pa.name, pa.email,
pa.plan_type, mr.id as request_id,
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
mr.cost_currency, mr.outcome, mr.input_tokens, mr.output_tokens,
mr.cached_input_tokens, mr.cache_write_input_tokens, mr.reasoning_output_tokens,
mr.image_input_tokens, mr.image_output_tokens,
mr.image_generation_succeeded, mr.total_tokens, mr.cost_source,
mr.cost_amount, mr.started_at
from candidate_accounts pa
left join request_analytics_facts mr
on mr.provider_account_ref = pa.id and mr.started_at >= ",
);
statement.push_bind(query.range.start);
statement.push(" and mr.started_at < ");
statement.push_bind(query.range.end);
push_completed_usage_fact_filter(&mut statement, "mr");
if let Some(account_ids) = &query.account_ids {
statement.push(" where pa.id = any(");
statement.push_bind(account_ids.clone());
statement.push("::text[])");
}
statement.push(
" limit 100001), aggregated as (
select id, provider_kind, authentication_kind, name, email, plan_type,
model, cost_currency,
grouping(model)::integer as model_grouping,
grouping(cost_currency)::integer as currency_grouping,
count(request_id)::bigint as request_count,
count(request_id) filter (where outcome = 'succeeded')::bigint
as success_count,
sum(input_tokens)::bigint as input_tokens,
sum(output_tokens)::bigint as output_tokens,
sum(cached_input_tokens)::bigint as cached_input_tokens,
sum(cache_write_input_tokens)::bigint as cache_write_input_tokens,
sum(reasoning_output_tokens)::bigint as reasoning_output_tokens,
sum(image_input_tokens)::bigint as image_input_tokens,
sum(image_output_tokens)::bigint as image_output_tokens,
count(request_id) filter (where image_generation_succeeded is true)::bigint
as image_request_count,
count(request_id) filter (where image_generation_succeeded is false)::bigint
as image_request_failed_count,
sum(total_tokens)::bigint as total_tokens,
count(request_id) filter (where cost_source = 'provider_reported')::bigint
as provider_reported_count,
count(request_id) filter (where cost_source = 'calculated')::bigint
as calculated_count,
count(request_id) filter (where cost_source = 'unavailable')::bigint
as unavailable_count,
max(started_at) as last_used_at,
sum(cost_amount)::text as amount
from matched
group by id, provider_kind, authentication_kind, name, email, plan_type,
grouping sets ((), (cost_currency), (model), (model, cost_currency))
), selected_accounts as (
select id,
row_number() over (order by last_used_at desc nulls last, name, id)
as sort_position
from aggregated
where model_grouping = 1 and currency_grouping = 1",
);
if query.account_ids.is_none() {
statement.push(" and request_count > 0");
}
statement.push(" order by last_used_at desc nulls last, name, id limit ");
statement.push_bind(i64::from(query.limit));
statement.push(
")
select aggregated.*, (select count(*) from matched)::bigint as matched_count
from aggregated
join selected_accounts selected using (id)
order by selected.sort_position, model_grouping desc, currency_grouping desc,
model nulls last, cost_currency nulls last",
);
let rows = statement
.build()
.fetch_all(&mut *connection)
.fetch_all(&mut *snapshot.tx)
.await
.map_err(|source| postgres_unavailable("load provider account usage", source))?;
if let Some(row) = rows.first()
&& unsigned(row, "matched_count")? > 100_000
.map_err(|error| postgres_unavailable("select usage accounts", error))?;
let account_ids = rows
.iter()
.map(|row| get(row, "id"))
.collect::<StoreResult<Vec<String>>>()?;
let mut models_by_account = HashMap::<String, Vec<ProviderAccountModelUsageObservation>>::new();
let mut totals = BTreeMap::<String, analytics::MetricState>::new();
for ((account, model), state) in snapshot
.account_usage_states(query.range, &account_ids)
.await?
{
return Err(crate::StoreError::Unavailable {
backend: crate::StoreBackend::PostgreSql,
message: "account analytics fact row budget exceeded".to_owned(),
source: None,
});
}
let mut observations = Vec::with_capacity(usize::from(query.limit));
let mut costs = HashMap::<String, Vec<CurrencyCostTotal>>::new();
let mut models = HashMap::<String, Vec<ProviderAccountModelUsageObservation>>::new();
let mut model_costs = HashMap::<(String, String), Vec<CurrencyCostTotal>>::new();
for row in &rows {
let model_grouping: i32 = get(row, "model_grouping")?;
let currency_grouping: i32 = get(row, "currency_grouping")?;
match (model_grouping, currency_grouping) {
(1, 1) => observations.push(provider_account_from_row(row)?),
(1, 0) if get::<Option<String>>(row, "cost_currency")?.is_some() => {
costs
.entry(get(row, "id")?)
.or_default()
.push(cost_from_row(row)?);
}
(0, 1) if get::<Option<String>>(row, "model")?.is_some() => {
models
.entry(get(row, "id")?)
.or_default()
.push(provider_account_model_from_row(row)?);
}
(0, 0)
if get::<Option<String>>(row, "model")?.is_some()
&& get::<Option<String>>(row, "cost_currency")?.is_some() =>
{
model_costs
.entry((get(row, "id")?, get(row, "model")?))
.or_default()
.push(cost_from_row(row)?);
}
(0 | 1, 0 | 1) => {}
_ => {
return Err(crate::StoreError::Unavailable {
backend: crate::StoreBackend::PostgreSql,
message: "decode provider account usage grouping".to_owned(),
source: None,
});
}
totals.entry(account.clone()).or_default().merge(&state)?;
if !model.is_empty()
&& let Some(last_used_at) = state.last_used_at()
{
models_by_account.entry(account).or_default().push(
ProviderAccountModelUsageObservation {
model,
request_count: state.count("completed_request_count"),
success_count: state.count("completed_request_count"),
input_tokens: state.optional_count("input_tokens"),
output_tokens: state.optional_count("output_tokens"),
cached_input_tokens: state.optional_count("cached_input_tokens"),
cache_write_input_tokens: state.optional_count("cache_write_input_tokens"),
reasoning_output_tokens: state.optional_count("reasoning_output_tokens"),
image_input_tokens: state.optional_count("image_input_tokens"),
image_output_tokens: state.optional_count("image_output_tokens"),
total_tokens: state.optional_count("total_tokens"),
image_request_count: state.count("image_request_count"),
image_request_failed_count: state.count("image_request_failed_count"),
cost_coverage: state.coverage(),
costs: state.costs()?,
last_used_at,
},
);
}
}
if observations.is_empty() {
return Ok(observations);
}
let account_ids = observations
.iter()
.map(|item| item.account_id.clone())
.collect::<Vec<_>>();
let mut request_buckets = if let Some(range) = query.request_bucket_range {
provider_account_request_buckets(connection, range, &account_ids).await?
provider_account_request_buckets(snapshot, range, &account_ids).await?
} else {
HashMap::new()
};
for observation in &mut observations {
observation.costs = costs.remove(&observation.account_id).unwrap_or_default();
observation.request_buckets = request_buckets
.remove(&observation.account_id)
.unwrap_or_default();
observation.models = models.remove(&observation.account_id).unwrap_or_default();
for model in &mut observation.models {
model.costs = model_costs
.remove(&(observation.account_id.clone(), model.model.clone()))
.unwrap_or_default();
}
observation.models.sort_by(|left, right| {
right
.request_count
.cmp(&left.request_count)
.then_with(|| left.model.cmp(&right.model))
let mut observations = Vec::with_capacity(rows.len());
for row in rows {
let account_id: String = get(&row, "id")?;
let state = totals.remove(&account_id).unwrap_or_default();
let mut models = models_by_account.remove(&account_id).unwrap_or_default();
models.sort_by(|a, b| {
b.request_count
.cmp(&a.request_count)
.then(a.model.cmp(&b.model))
});
observations.push(ProviderAccountUsageObservation {
provider_kind: get(&row, "provider_kind")?,
authentication_kind: get(&row, "authentication_kind")?,
name: get(&row, "name")?,
email: get(&row, "email")?,
plan_type: get(&row, "plan_type")?,
request_count: state.count("completed_request_count"),
success_count: state.count("completed_request_count"),
input_tokens: state.optional_count("input_tokens"),
output_tokens: state.optional_count("output_tokens"),
cached_input_tokens: state.optional_count("cached_input_tokens"),
cache_write_input_tokens: state.optional_count("cache_write_input_tokens"),
reasoning_output_tokens: state.optional_count("reasoning_output_tokens"),
image_input_tokens: state.optional_count("image_input_tokens"),
image_output_tokens: state.optional_count("image_output_tokens"),
total_tokens: state.optional_count("total_tokens"),
image_request_count: state.count("image_request_count"),
image_request_failed_count: state.count("image_request_failed_count"),
cost_coverage: state.coverage(),
costs: state.costs()?,
last_used_at: state.last_used_at(),
request_buckets: request_buckets.remove(&account_id).unwrap_or_default(),
models,
account_id,
});
}
observations.sort_by(|a, b| {
b.last_used_at
.cmp(&a.last_used_at)
.then(a.name.cmp(&b.name))
.then(a.account_id.cmp(&b.account_id))
});
observations.truncate(usize::from(query.limit));
Ok(observations)
}
pub(crate) async fn provider_account_request_buckets(
connection: &mut sqlx::PgConnection,
snapshot: &mut analytics::AnalyticsSnapshot<'_>,
range: ObservabilityRange,
account_ids: &[String],
) -> StoreResult<HashMap<String, Vec<ProviderAccountRequestBucket>>> {
let completed_usage = completed_usage_fact_predicate("mr");
let statement = format!(
"select provider_account_ref,
floor(extract(epoch from (started_at - $1)) / 3600)::bigint as bucket_index,
count(*)::bigint as request_count
from request_analytics_facts mr
where mr.provider_account_ref = any($2::text[])
and mr.started_at >= $1 and mr.started_at < $3
and {completed_usage}
group by mr.provider_account_ref, bucket_index
order by mr.provider_account_ref, bucket_index"
);
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(range.start)
.bind(account_ids)
.bind(range.end)
.fetch_all(&mut *connection)
.await
.map_err(|source| postgres_unavailable("load provider account request timeline", source))?;
let mut observed = HashMap::<String, BTreeMap<u64, u64>>::new();
for row in rows {
observed
.entry(get(&row, "provider_account_ref")?)
.or_default()
.insert(
unsigned(&row, "bucket_index")?,
unsigned(&row, "request_count")?,
);
}
let mut observed = snapshot.account_request_counts(range, account_ids).await?;
let mut timelines = HashMap::with_capacity(account_ids.len());
for account_id in account_ids {
let mut account_observed = observed.remove(account_id).unwrap_or_default();
let mut buckets = Vec::new();
for bucket_index in 0..ACCOUNT_USAGE_TIMELINE_HOURS {
let bucket_start = range
for account in account_ids {
let mut buckets = Vec::with_capacity(ACCOUNT_USAGE_TIMELINE_HOURS as usize);
for index in 0..ACCOUNT_USAGE_TIMELINE_HOURS {
let start = range
.start
.checked_add_signed(TimeDelta::hours(bucket_index))
.checked_add_signed(TimeDelta::hours(index))
.ok_or_else(|| invalid("account request timeline exceeds supported timestamps"))?;
buckets.push(ProviderAccountRequestBucket {
bucket_start,
request_count: account_observed
.remove(&(bucket_index as u64))
bucket_start: start,
request_count: observed
.remove(&(account.clone(), index))
.unwrap_or_default(),
});
}
timelines.insert(account_id.clone(), buckets);
timelines.insert(account.clone(), buckets);
}
Ok(timelines)
}
fn provider_account_model_from_row(
row: &sqlx::postgres::PgRow,
) -> StoreResult<ProviderAccountModelUsageObservation> {
Ok(ProviderAccountModelUsageObservation {
model: get(row, "model")?,
request_count: unsigned(row, "request_count")?,
success_count: unsigned(row, "success_count")?,
input_tokens: optional_unsigned(row, "input_tokens")?,
output_tokens: optional_unsigned(row, "output_tokens")?,
cached_input_tokens: optional_unsigned(row, "cached_input_tokens")?,
cache_write_input_tokens: optional_unsigned(row, "cache_write_input_tokens")?,
reasoning_output_tokens: optional_unsigned(row, "reasoning_output_tokens")?,
image_input_tokens: optional_unsigned(row, "image_input_tokens")?,
image_output_tokens: optional_unsigned(row, "image_output_tokens")?,
image_request_count: unsigned(row, "image_request_count")?,
image_request_failed_count: unsigned(row, "image_request_failed_count")?,
total_tokens: optional_unsigned(row, "total_tokens")?,
cost_coverage: coverage_from_row(row)?,
costs: Vec::new(),
last_used_at: get(row, "last_used_at")?,
})
}
pub(crate) fn provider_account_from_row(
row: &sqlx::postgres::PgRow,
) -> StoreResult<ProviderAccountUsageObservation> {
Ok(ProviderAccountUsageObservation {
account_id: get(row, "id")?,
provider_kind: get(row, "provider_kind")?,
authentication_kind: get(row, "authentication_kind")?,
name: get(row, "name")?,
email: get(row, "email")?,
plan_type: get(row, "plan_type")?,
request_count: unsigned(row, "request_count")?,
success_count: unsigned(row, "success_count")?,
input_tokens: optional_unsigned(row, "input_tokens")?,
output_tokens: optional_unsigned(row, "output_tokens")?,
cached_input_tokens: optional_unsigned(row, "cached_input_tokens")?,
cache_write_input_tokens: optional_unsigned(row, "cache_write_input_tokens")?,
reasoning_output_tokens: optional_unsigned(row, "reasoning_output_tokens")?,
image_input_tokens: optional_unsigned(row, "image_input_tokens")?,
image_output_tokens: optional_unsigned(row, "image_output_tokens")?,
image_request_count: unsigned(row, "image_request_count")?,
image_request_failed_count: unsigned(row, "image_request_failed_count")?,
total_tokens: optional_unsigned(row, "total_tokens")?,
cost_coverage: coverage_from_row(row)?,
costs: Vec::new(),
last_used_at: get(row, "last_used_at")?,
request_buckets: Vec::new(),
models: Vec::new(),
})
}
@@ -1,15 +1,7 @@
//! 时间桶补齐、日历边界与共享费用覆盖映射
//! 时间桶补齐与部署时区的日历边界
use super::super::*;
pub(crate) fn coverage_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<CostCoverage> {
Ok(CostCoverage {
provider_reported_count: unsigned(row, "provider_reported_count")?,
calculated_count: unsigned(row, "calculated_count")?,
unavailable_count: unsigned(row, "unavailable_count")?,
})
}
pub(crate) fn fill_metric_gaps(
range: ObservabilityRange,
granularity: ObservationGranularity,
@@ -69,38 +69,29 @@ left join model_request_observations mr on mr.id = oe.model_request_id
where true";
pub(crate) async fn list_ops_errors(
connection: &mut PgConnection,
connection: &mut ObservabilityConnection,
query: OpsErrorQuery,
) -> StoreResult<OpsErrorPage> {
validate_ops_error_filter(&query.filter)?;
let mut transaction = analytics::begin_snapshot(connection).await?;
let total = count_ops_errors(&mut transaction, query.range, &query.filter).await?;
let offset = observability_page_offset(query.current_page, query.page_size)?;
let mut statement =
QueryBuilder::<Postgres>::new("with candidate_page as materialized (select * from (");
statement.push("select 'model_request'::text as source, mr.id as event_id, mr.completed_at as occurred_at, 'model_request:' || mr.id as stable_sort_id from model_request_observations mr where true");
push_request_error_predicates(&mut statement, query.range, &query.filter);
statement.push(" union all ");
statement.push("select 'ops_event'::text as source, oe.id as event_id, oe.created_at as occurred_at, 'ops_event:' || oe.id as stable_sort_id from ops_events oe left join model_request_observations mr on mr.id=oe.model_request_id where true");
push_ops_event_predicates(&mut statement, query.range, &query.filter);
statement.push(") e order by occurred_at desc, stable_sort_id desc limit ");
statement.push_bind(i64::from(query.page_size.get()));
statement.push(" offset ");
statement.push_bind(offset);
statement.push(") select details.* from candidate_page page cross join lateral (");
statement.push(REQUEST_ERROR_SELECT);
statement.push(" and page.source='model_request' and mr.id=page.event_id union all ");
statement.push(OPS_EVENT_SELECT);
statement.push(" and page.source='ops_event' and oe.id=page.event_id) details order by page.occurred_at desc,page.stable_sort_id desc");
let rows = statement
.build()
.fetch_all(&mut *transaction)
.await
.map_err(|source| postgres_unavailable("list ops errors", source))?;
let items = rows
.iter()
.map(ops_error_from_row)
.collect::<StoreResult<Vec<_>>>()?;
let (total, page) = if query.current_page == 1 {
let total = count_ops_errors(&mut transaction, query.range, &query.filter).await?;
(
total,
(total > 0).then_some(TimePage {
range: query.range,
offset,
total: Some(total),
}),
)
} else {
ops_time_page(&mut transaction, &query, offset).await?
};
let items = match page {
Some(page) => list_ops_error_items_at(&mut transaction, &query, page).await?,
None => Vec::new(),
};
transaction
.commit()
.await
@@ -113,6 +104,95 @@ pub(crate) async fn list_ops_errors(
})
}
async fn list_ops_error_items_at(
connection: &mut PgConnection,
query: &OpsErrorQuery,
page: TimePage,
) -> StoreResult<Vec<OpsErrorRecord>> {
let total =
i64::try_from(page.total.unwrap_or(0)).map_err(|_| invalid("Ops count overflows i64"))?;
let limit = i64::from(query.page_size.get()).min(total - page.offset);
let reverse_offset = total - page.offset - limit;
let ascending = reverse_offset < page.offset;
let offset = if ascending {
reverse_offset
} else {
page.offset
};
let branch_limit = offset + limit;
let order = if ascending { " asc" } else { " desc" };
let mut statement =
QueryBuilder::<Postgres>::new("with candidate_page as materialized (select * from ((");
statement.push("select 'model_request'::text as source, mr.id as event_id, mr.completed_at as occurred_at, 'model_request:' || mr.id as stable_sort_id from model_request_observations mr where true");
push_request_error_predicates(&mut statement, page.range, &query.filter);
statement
.push(" order by mr.completed_at")
.push(order)
.push(", mr.id")
.push(order);
statement.push(" limit ").push_bind(branch_limit);
statement.push(") union all (");
statement.push("select 'ops_event'::text as source, oe.id as event_id, oe.created_at as occurred_at, 'ops_event:' || oe.id as stable_sort_id from ops_events oe left join model_request_observations mr on mr.id=oe.model_request_id where true");
push_ops_event_predicates(&mut statement, page.range, &query.filter);
statement
.push(" order by oe.created_at")
.push(order)
.push(", oe.id")
.push(order);
statement.push(" limit ").push_bind(branch_limit);
statement
.push(")) e order by occurred_at")
.push(order)
.push(", stable_sort_id")
.push(order);
statement
.push(" limit ")
.push_bind(limit)
.push(" offset ")
.push_bind(offset);
statement.push(") select details.* from candidate_page page cross join lateral (");
statement.push(REQUEST_ERROR_SELECT);
statement.push(" and page.source='model_request' and mr.id=page.event_id union all ");
statement.push(OPS_EVENT_SELECT);
statement.push(" and page.source='ops_event' and oe.id=page.event_id) details order by page.occurred_at desc,page.stable_sort_id desc");
let rows = statement
.build()
.fetch_all(connection)
.await
.map_err(|source| postgres_unavailable("list ops errors", source))?;
rows.iter().map(ops_error_from_row).collect()
}
async fn ops_time_page(
connection: &mut PgConnection,
query: &OpsErrorQuery,
offset: i64,
) -> StoreResult<(u64, Option<TimePage>)> {
// 两个来源先各自精确分段,合并分钟计数后再定位同一个全局页码
let mut statement = QueryBuilder::<Postgres>::new(
"select bucket_start, sum(request_count)::bigint as request_count from (
select date_bin(interval '1 minute', mr.completed_at, timestamptz 'epoch') as bucket_start,
count(*)::bigint as request_count
from model_request_observations mr where true",
);
push_request_error_predicates(&mut statement, query.range, &query.filter);
statement.push(" group by bucket_start union all ");
statement.push(
"select date_bin(interval '1 minute', oe.created_at, timestamptz 'epoch') as bucket_start,
count(*)::bigint as request_count
from ops_events oe left join model_request_observations mr on mr.id=oe.model_request_id where true",
);
push_ops_event_predicates(&mut statement, query.range, &query.filter);
statement
.push(" group by bucket_start) counts group by bucket_start order by bucket_start desc");
let rows = statement
.build()
.fetch_all(connection)
.await
.map_err(|source| postgres_unavailable("locate Ops page", source))?;
time_page_from_counts(&rows, query.range, offset, query.page_size)
}
pub(crate) async fn count_ops_errors(
connection: &mut sqlx::PgConnection,
range: ObservabilityRange,
@@ -144,7 +224,6 @@ fn push_request_error_predicates(
// 列表和总数共用此条件,避免流式响应中的错误因 outcome 被漏掉
statement.push(" and mr.error_kind is not null and mr.error_kind <> 'cancelled'");
push_range(statement, "mr.completed_at", range);
push_routing_groups(statement, &filter.routing_group_refs);
for (column, value) in [
("mr.client_api_key_ref", &filter.client_api_key_ref),
("mr.id", &filter.request_id),
@@ -187,7 +266,6 @@ fn push_ops_event_predicates(
filter: &OpsErrorFilter,
) {
push_range(statement, "oe.created_at", range);
push_routing_groups(statement, &filter.routing_group_refs);
for (column, value) in [
("mr.client_api_key_ref", &filter.client_api_key_ref),
("oe.model_request_id", &filter.request_id),
@@ -275,11 +353,3 @@ fn push_prefix_search(statement: &mut QueryBuilder<Postgres>, columns: &[&str],
.push(" escape '\\'");
statement.push(")");
}
fn push_routing_groups(statement: &mut QueryBuilder<Postgres>, groups: &[String]) {
if !groups.is_empty() {
statement.push(" and mr.routing_group_refs && ");
statement.push_bind(groups.to_vec());
statement.push("::text[]");
}
}
@@ -11,12 +11,6 @@ pub(crate) fn push_usage_filter(
filter: &UsageRecordFilter,
alias: &str,
) {
if !filter.routing_group_refs.is_empty() {
query
.push(format!(" and {alias}.routing_group_refs && "))
.push_bind(filter.routing_group_refs.clone())
.push("::text[]");
}
if let Some(value) = &filter.client_api_key_ref {
query.push(format!(" and {alias}.client_api_key_ref = "));
query.push_bind(value.clone());
@@ -86,16 +80,14 @@ pub(crate) fn push_usage_filter(
}
if let Some(value) = &filter.search {
let pattern = literal_prefix_pattern(value);
query.push(format!(" and ({alias}.id like "));
query.push(format!(" and ({alias}.requested_model_id like "));
query.push_bind(pattern.clone());
for column in [
"upstream_model_id",
"client_api_key_ref",
"provider_account_ref",
"provider_account_email_snapshot",
"provider_account_name_snapshot",
"requested_model_id",
"upstream_model_id",
"upstream_request_id",
] {
query.push(format!(" escape '\\' or {alias}.{column} like "));
query.push_bind(pattern.clone());
@@ -174,29 +166,50 @@ pub(crate) const USAGE_RECORD_DETAIL_SELECT: &str =
from model_request_observations mr";
pub(crate) async fn list_usage_records(
connection: &mut PgConnection,
connection: &mut ObservabilityConnection,
query: UsageRecordQuery,
) -> StoreResult<UsageRecordPage> {
query.filter.validate()?;
let mut transaction = analytics::begin_snapshot(connection).await?;
let mut snapshot = AnalyticsSnapshot::new_for_accounts(connection).await?;
// 第一页未填满时,当前快照中的行数已是精确总数,无需再次扫描计数
let (total, items) = if query.current_page == 1 {
let items = list_usage_record_items(&mut *transaction, &query, None).await?;
let items = list_usage_record_items(&mut *snapshot.tx, &query, None).await?;
let total = if items.len() < usize::from(query.page_size.get()) {
items.len() as u64
} else if let Some(counts) = snapshot
.completed_request_counts(query.range, &query.filter)
.await?
{
counts.values().try_fold(0_u64, |sum, count| {
sum.checked_add(*count)
.ok_or_else(|| invalid("usage count overflows u64"))
})?
} else {
count_usage_records(&mut *transaction, query.range, &query.filter).await?
count_usage_records(&mut *snapshot.tx, query.range, &query.filter).await?
};
(total, items)
} else {
let total = count_usage_records(&mut *transaction, query.range, &query.filter).await?;
let items = list_usage_record_items(&mut *transaction, &query, Some(total)).await?;
let (total, page) = if let Some(counts) = snapshot
.completed_request_counts(query.range, &query.filter)
.await?
{
time_page_from_segments(
counts.into_iter().rev(),
query.range,
observability_page_offset(query.current_page, query.page_size)?,
query.page_size,
TimeDelta::minutes(15),
)?
} else {
usage_time_page(&mut snapshot.tx, &query).await?
};
let items = match page {
Some(page) => list_usage_record_items_at(&mut *snapshot.tx, &query, page).await?,
None => Vec::new(),
};
(total, items)
};
transaction
.commit()
.await
.map_err(|error| postgres_unavailable("complete usage page snapshot", error))?;
snapshot.finish().await?;
Ok(UsageRecordPage {
items,
current_page: query.current_page,
@@ -209,12 +222,25 @@ pub(crate) async fn list_usage_record_items<'e, E: sqlx::Executor<'e, Database =
executor: E,
query: &UsageRecordQuery,
total: Option<u64>,
) -> StoreResult<Vec<UsageListRecord>> {
let page = TimePage {
range: query.range,
offset: observability_page_offset(query.current_page, query.page_size)?,
total,
};
list_usage_record_items_at(executor, query, page).await
}
async fn list_usage_record_items_at<'e, E: sqlx::Executor<'e, Database = Postgres>>(
executor: E,
query: &UsageRecordQuery,
page: TimePage,
) -> StoreResult<Vec<UsageListRecord>> {
query.filter.validate()?;
let mut offset = observability_page_offset(query.current_page, query.page_size)?;
let mut offset = page.offset;
let mut limit = i64::from(query.page_size.get());
let mut ascending = false;
if let Some(total) = total {
if let Some(total) = page.total {
let total = i64::try_from(total).map_err(|_| invalid("usage count overflows i64"))?;
if offset >= total {
return Ok(Vec::new());
@@ -231,9 +257,9 @@ pub(crate) async fn list_usage_record_items<'e, E: sqlx::Executor<'e, Database =
"with candidate_page as materialized (select mr.id, mr.started_at from model_request_observations mr",
);
statement.push(" where mr.started_at >= ");
statement.push_bind(query.range.start);
statement.push_bind(page.range.start);
statement.push(" and mr.started_at < ");
statement.push_bind(query.range.end);
statement.push_bind(page.range.end);
push_completed_usage_fact_filter(&mut statement, "mr");
push_usage_filter(&mut statement, &query.filter, "mr");
statement.push(if ascending {
@@ -261,6 +287,93 @@ pub(crate) async fn list_usage_record_items<'e, E: sqlx::Executor<'e, Database =
Ok(items)
}
// 页码保持全局精确排名,实际明细读取只落在覆盖该页的时间区间
pub(crate) struct TimePage {
pub(crate) range: ObservabilityRange,
pub(crate) offset: i64,
pub(crate) total: Option<u64>,
}
pub(crate) fn time_page_from_counts(
rows: &[sqlx::postgres::PgRow],
range: ObservabilityRange,
offset: i64,
page_size: ObservabilityPageSize,
) -> StoreResult<(u64, Option<TimePage>)> {
let counts = rows
.iter()
.map(|row| {
Ok((
get(row, "bucket_start")?,
to_u64(get(row, "request_count")?)?,
))
})
.collect::<StoreResult<Vec<_>>>()?;
time_page_from_segments(counts, range, offset, page_size, TimeDelta::minutes(1))
}
fn time_page_from_segments(
counts: impl IntoIterator<Item = (DateTime<Utc>, u64)>,
range: ObservabilityRange,
offset: i64,
page_size: ObservabilityPageSize,
width: TimeDelta,
) -> StoreResult<(u64, Option<TimePage>)> {
let offset = u64::try_from(offset).map_err(|_| invalid("negative page offset"))?;
let end = offset
.checked_add(u64::from(page_size.get()))
.ok_or_else(|| invalid("page range overflows u64"))?;
let mut total = 0_u64;
let mut page = None;
let mut selected_count = 0_u64;
for (bucket, count) in counts {
let next = total
.checked_add(count)
.ok_or_else(|| invalid("request count overflows u64"))?;
if next > offset && total < end {
let selected = page.get_or_insert_with(|| TimePage {
range: ObservabilityRange {
start: bucket.max(range.start),
end: (bucket + width).min(range.end),
},
offset: (offset - total) as i64,
total: None,
});
selected.range.start = bucket.max(range.start);
selected_count += count;
}
total = next;
}
if let Some(page) = &mut page {
page.total = Some(selected_count);
}
Ok((total, page))
}
async fn usage_time_page(
connection: &mut PgConnection,
query: &UsageRecordQuery,
) -> StoreResult<(u64, Option<TimePage>)> {
let offset = observability_page_offset(query.current_page, query.page_size)?;
let mut statement = QueryBuilder::<Postgres>::new(
"select date_bin(interval '1 minute', mr.started_at, timestamptz 'epoch') as bucket_start,
count(*)::bigint as request_count
from model_request_observations mr where mr.started_at >= ",
);
statement.push_bind(query.range.start);
statement.push(" and mr.started_at < ");
statement.push_bind(query.range.end);
push_completed_usage_fact_filter(&mut statement, "mr");
push_usage_filter(&mut statement, &query.filter, "mr");
statement.push(" group by bucket_start order by bucket_start desc");
let rows = statement
.build()
.fetch_all(connection)
.await
.map_err(|source| postgres_unavailable("locate usage page", source))?;
time_page_from_counts(&rows, query.range, offset, query.page_size)
}
pub(crate) async fn count_usage_records<'e, E: sqlx::Executor<'e, Database = Postgres>>(
executor: E,
range: ObservabilityRange,
@@ -17,12 +17,20 @@ use crate::{StoreBackend, StoreError, StoreResult, postgres_unavailable};
pub struct ObservabilityQueryBudget {
slots: Arc<Semaphore>,
wait_timeout: Duration,
query_timeout: Duration,
}
/// 观测名额与实际 PostgreSQL 连接共同存活,取消后仍覆盖协议收尾和事务回滚
pub struct ObservabilityConnection {
connection: PoolConnection<Postgres>,
permit: Option<OwnedSemaphorePermit>,
query_timeout: Duration,
}
impl ObservabilityConnection {
pub(crate) const fn query_timeout(&self) -> Duration {
self.query_timeout
}
}
impl Deref for ObservabilityConnection {
@@ -53,12 +61,22 @@ impl Drop for ObservabilityConnection {
impl ObservabilityQueryBudget {
/// 预算表示观测类 SQL 最多可同时占用的连接数
pub fn try_new(max_connections: u32, wait_timeout: Duration) -> StoreResult<Self> {
if max_connections == 0 || wait_timeout.is_zero() {
pub fn try_new(
max_connections: u32,
wait_timeout: Duration,
query_timeout: Duration,
) -> StoreResult<Self> {
if max_connections == 0
|| wait_timeout.is_zero()
|| query_timeout.is_zero()
|| query_timeout.as_millis() > i32::MAX as u128
{
return Err(StoreError::InvalidData {
source: None,
entity: "observability query budget",
message: "requires a positive connection count and wait timeout".to_owned(),
message:
"requires positive connection and timeout budgets within PostgreSQL limits"
.to_owned(),
});
}
let permits = usize::try_from(max_connections).map_err(|_| StoreError::InvalidData {
@@ -69,10 +87,16 @@ impl ObservabilityQueryBudget {
Ok(Self {
slots: Arc::new(Semaphore::new(permits)),
wait_timeout,
query_timeout,
})
}
/// 在共享槽位与十秒执行预算内完成整个观测操作
#[must_use]
pub const fn query_timeout(&self) -> Duration {
self.query_timeout
}
/// 在共享槽位与观测总预算内完成排队、连接获取和整个存储操作
///
/// # Errors
///
@@ -87,14 +111,15 @@ impl ObservabilityQueryBudget {
F: FnOnce(ObservabilityConnection) -> Q,
Q: Future<Output = StoreResult<T>>,
{
let permit = self.acquire(operation).await?;
tokio::time::timeout(Duration::from_secs(10), async {
tokio::time::timeout(self.query_timeout, async {
let permit = self.acquire(operation).await?;
let connection = pool.acquire().await.map_err(|source| {
postgres_unavailable("acquire observability connection", source)
})?;
query(ObservabilityConnection {
connection,
permit: Some(permit),
query_timeout: self.query_timeout,
})
.await
})
@@ -102,109 +102,17 @@ impl PgAdminAccountStore {
ObservabilityRange::new(window.range.start, window.range.end)
.map_err(|error| admin_store_error(ENTITY, error))?;
}
let keys = windows
.iter()
.map(|window| window.key.clone())
.collect::<Vec<_>>();
let starts = windows
.iter()
.map(|window| window.range.start)
.collect::<Vec<_>>();
let ends = windows
.iter()
.map(|window| window.range.end)
.collect::<Vec<_>>();
let account_ids = windows
.iter()
.map(|window| window.account_id.clone())
.collect::<Vec<_>>();
let rows = self
.query_budget
self.query_budget
.run(
&self.pool,
"load account usage windows",
|mut connection| async move {
sqlx::query(sqlx::AssertSqlSafe(account_usage_by_windows_sql()))
.bind(account_ids)
.bind(keys)
.bind(starts)
.bind(ends)
.fetch_all(&mut *connection)
crate::postgres::observability::account_usage_windows(&mut connection, windows)
.await
.map_err(|source| {
postgres_unavailable("load provider account quota window usage", source)
})
},
)
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
let mut usage_rows = Vec::with_capacity(windows.len());
let mut costs_by_window = BTreeMap::<(String, String), Vec<AccountCost>>::new();
let mut model_costs = BTreeMap::<(String, String, String), Vec<AccountCost>>::new();
let mut models_by_window = BTreeMap::<(String, String), Vec<AccountModelUsage>>::new();
for row in &rows {
let model_grouping = window_usage_value::<i32>(row, "model_grouping")?;
let currency_grouping = window_usage_value::<i32>(row, "currency_grouping")?;
match (model_grouping, currency_grouping) {
(1, 1) => usage_rows.push(row),
(1, 0) if window_usage_value::<Option<String>>(row, "cost_currency")?.is_some() => {
let (key, cost) = admin_account_usage_window_cost(row)?;
costs_by_window.entry(key).or_default().push(cost);
}
(0, 1) if window_usage_value::<Option<String>>(row, "model")?.is_some() => {
let ((account_id, window_key, _), usage) =
admin_account_usage_window_model(row)?;
models_by_window
.entry((account_id, window_key))
.or_default()
.push(usage);
}
(0, 0)
if window_usage_value::<Option<String>>(row, "model")?.is_some()
&& window_usage_value::<Option<String>>(row, "cost_currency")?
.is_some() =>
{
let (key, cost) = admin_account_usage_window_model_cost(row)?;
model_costs.entry(key).or_default().push(cost);
}
(0 | 1, 0 | 1) => {}
_ => {
return Err(AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
ENTITY,
"account usage window query returned an invalid grouping marker",
));
}
}
}
let mut results = usage_rows
.into_iter()
.map(admin_account_usage_window)
.collect::<AdminStoreResult<Vec<_>>>()?;
for result in &mut results {
result.usage.costs = costs_by_window
.remove(&(result.account_id.clone(), result.key.clone()))
.unwrap_or_default();
result.usage.models = models_by_window
.remove(&(result.account_id.clone(), result.key.clone()))
.unwrap_or_default();
for model in &mut result.usage.models {
model.costs = model_costs
.remove(&(
result.account_id.clone(),
result.key.clone(),
model.model.clone(),
))
.unwrap_or_default();
}
result.usage.models.sort_by(|left, right| {
right
.request_count
.cmp(&left.request_count)
.then_with(|| left.model.cmp(&right.model))
});
}
Ok(results)
.map_err(|error| admin_store_error(ENTITY, error))
}
async fn required_scope(
@@ -376,8 +284,14 @@ impl AccountStore for PgAdminAccountStore {
.filter(|(_, cooldown)| cooldown.is_active(now_system_time))
.map(|(account_id, _)| account_id.clone())
.collect::<Vec<_>>();
let page =
load_admin_account_page(&self.pool, &query, now, active_rate_limited_ids).await?;
let page = load_admin_account_page(
&self.pool,
&self.query_budget,
&query,
now,
active_rate_limited_ids,
)
.await?;
let item_ids = page
.accounts
.iter()
@@ -12,6 +12,7 @@ pub(crate) struct AdminAccountPageRows {
pub(crate) async fn load_admin_account_page(
pool: &PgPool,
budget: &crate::postgres::ObservabilityQueryBudget,
query: &AdminAccountListQuery,
now: DateTime<Utc>,
active_rate_limited_ids: Vec<String>,
@@ -35,35 +36,20 @@ pub(crate) async fn load_admin_account_page(
let page_size = i64::from(query.page_size.get());
let offset = i64::from(query.page - 1).saturating_mul(page_size);
let order_by = admin_account_order(query.sort);
let usage_join = query.sort.is_some_and(|sort| {
let needs_usage = query.sort.is_some_and(|sort| {
matches!(
sort.field,
AdminAccountSortField::Usage | AdminAccountSortField::LastUsedAt
)
});
let completed_usage = completed_usage_fact_predicate("mr");
let usage_join = if usage_join {
format!(
"left join (
select mr.provider_account_ref,
coalesce(sum(coalesce(
mr.total_tokens,
coalesce(mr.input_tokens, 0) + coalesce(mr.output_tokens, 0)
)), 0)::bigint as total_tokens,
max(mr.started_at) as last_used_at
from model_requests mr
where mr.started_at >= $2 - make_interval(days => settings.usage_retention_days::int)
and mr.started_at < $2
and {completed_usage}
group by mr.provider_account_ref
) usage on usage.provider_account_ref = a.id"
)
let usage_join = if needs_usage {
"left join usage_values usage on usage.provider_account_ref = a.id"
} else {
String::new()
""
};
let statement = format!(
"with account_statuses as (
"with usage_values as (select * from unnest($10::text[],$11::bigint[],$12::timestamptz[]) as usage(provider_account_ref,total_tokens,last_used_at)), account_statuses as (
select a.id, a.provider_kind,
case
when not a.enabled then 'disabled'
@@ -150,25 +136,45 @@ pub(crate) async fn load_admin_account_page(
left join outbound_proxies p on p.id = a.outbound_proxy_id
order by page.page_position"
);
// 动态片段只来自上面的封闭排序枚举与固定 usage predicate;所有请求值仍使用 bind
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(active_rate_limited_ids)
.bind(now)
.bind(provider_kind)
.bind(search_prefix)
.bind(status)
.bind(group_mode)
.bind(group_id)
.bind(page_size)
.bind(offset)
.fetch_all(pool)
// 排序事实与本页总数共享同一快照,避免一次查询混入不同源版本
let rows = budget
.run(
pool,
"load admin account page",
|mut connection| async move {
let mut snapshot =
crate::postgres::observability::AnalyticsSnapshot::new_for_accounts(
&mut connection,
)
.await?;
let (usage_ids, usage_tokens, usage_latest) = if needs_usage {
crate::postgres::observability::account_usage_sort_values(&mut snapshot, now)
.await?
} else {
(Vec::new(), Vec::new(), Vec::new())
};
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
.bind(active_rate_limited_ids)
.bind(now)
.bind(provider_kind)
.bind(search_prefix)
.bind(status)
.bind(group_mode)
.bind(group_id)
.bind(page_size)
.bind(offset)
.bind(usage_ids)
.bind(usage_tokens)
.bind(usage_latest)
.fetch_all(&mut *snapshot.tx)
.await
.map_err(|source| postgres_unavailable("load admin account page", source))?;
snapshot.finish().await?;
Ok(rows)
},
)
.await
.map_err(|source| {
admin_store_error(
ENTITY,
postgres_unavailable("load admin account page", source),
)
})?;
.map_err(|error| admin_store_error(ENTITY, error))?;
let metadata = rows.first().ok_or_else(|| {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
@@ -105,164 +105,3 @@ pub(crate) fn admin_account_usage(usage: ProviderAccountUsageObservation) -> Acc
.collect(),
}
}
pub(crate) fn admin_account_usage_window(
row: &sqlx::postgres::PgRow,
) -> AdminStoreResult<AccountUsageWindowResult> {
let provider_reported_count = window_usage_count(row, "provider_reported_count")?;
let calculated_count = window_usage_count(row, "calculated_count")?;
let unavailable_count = window_usage_count(row, "unavailable_count")?;
Ok(AccountUsageWindowResult {
account_id: window_usage_value(row, "account_id")?,
key: window_usage_value(row, "window_key")?,
usage: AccountUsage {
account_id: window_usage_value(row, "account_id")?,
request_count: window_usage_count(row, "request_count")?,
success_count: window_usage_count(row, "success_count")?,
input_tokens: optional_window_usage_count(row, "input_tokens")?,
output_tokens: optional_window_usage_count(row, "output_tokens")?,
cached_input_tokens: optional_window_usage_count(row, "cached_input_tokens")?,
cache_write_input_tokens: optional_window_usage_count(row, "cache_write_input_tokens")?,
reasoning_output_tokens: optional_window_usage_count(row, "reasoning_output_tokens")?,
image_input_tokens: optional_window_usage_count(row, "image_input_tokens")?,
image_output_tokens: optional_window_usage_count(row, "image_output_tokens")?,
image_request_count: window_usage_count(row, "image_request_count")?,
image_request_failed_count: window_usage_count(row, "image_request_failed_count")?,
total_tokens: optional_window_usage_count(row, "total_tokens")?,
cost_coverage: AdminCostCoverage {
provider_reported_count,
calculated_count,
partial_count: 0,
unavailable_count,
not_billable_count: 0,
},
costs: Vec::new(),
last_used_at: window_usage_value(row, "last_used_at")?,
request_buckets: Vec::new(),
models: Vec::new(),
},
})
}
pub(crate) fn admin_account_usage_window_model(
row: &sqlx::postgres::PgRow,
) -> AdminStoreResult<((String, String, String), AccountModelUsage)> {
let account_id: String = window_usage_value(row, "account_id")?;
let window_key: String = window_usage_value(row, "window_key")?;
let model: String = window_usage_value(row, "model")?;
let provider_reported_count = window_usage_count(row, "provider_reported_count")?;
let calculated_count = window_usage_count(row, "calculated_count")?;
let unavailable_count = window_usage_count(row, "unavailable_count")?;
Ok((
(account_id, window_key, model.clone()),
AccountModelUsage {
model,
request_count: window_usage_count(row, "request_count")?,
success_count: window_usage_count(row, "success_count")?,
input_tokens: optional_window_usage_count(row, "input_tokens")?,
output_tokens: optional_window_usage_count(row, "output_tokens")?,
cached_input_tokens: optional_window_usage_count(row, "cached_input_tokens")?,
cache_write_input_tokens: optional_window_usage_count(row, "cache_write_input_tokens")?,
reasoning_output_tokens: optional_window_usage_count(row, "reasoning_output_tokens")?,
image_input_tokens: optional_window_usage_count(row, "image_input_tokens")?,
image_output_tokens: optional_window_usage_count(row, "image_output_tokens")?,
image_request_count: window_usage_count(row, "image_request_count")?,
image_request_failed_count: window_usage_count(row, "image_request_failed_count")?,
total_tokens: optional_window_usage_count(row, "total_tokens")?,
cost_coverage: AdminCostCoverage {
provider_reported_count,
calculated_count,
partial_count: 0,
unavailable_count,
not_billable_count: 0,
},
costs: Vec::new(),
last_used_at: window_usage_value(row, "last_used_at")?,
},
))
}
pub(crate) fn admin_account_usage_window_model_cost(
row: &sqlx::postgres::PgRow,
) -> AdminStoreResult<((String, String, String), AccountCost)> {
let key = (
window_usage_value(row, "account_id")?,
window_usage_value(row, "window_key")?,
window_usage_value(row, "model")?,
);
let amount = window_usage_value::<String>(row, "amount")?;
let amount = AdminDecimalAmount::from_str(&amount).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted account model cost is invalid",
)
})?;
Ok((
key,
AccountCost {
currency: window_usage_value(row, "cost_currency")?,
amount,
},
))
}
pub(crate) fn admin_account_usage_window_cost(
row: &sqlx::postgres::PgRow,
) -> AdminStoreResult<((String, String), AccountCost)> {
let key = (
window_usage_value(row, "account_id")?,
window_usage_value(row, "window_key")?,
);
let amount = window_usage_value::<String>(row, "amount")?;
let amount = AdminDecimalAmount::from_str(&amount).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted account cost is invalid",
)
})?;
Ok((
key,
AccountCost {
currency: window_usage_value(row, "cost_currency")?,
amount,
},
))
}
pub(crate) fn window_usage_count(
row: &sqlx::postgres::PgRow,
column: &'static str,
) -> AdminStoreResult<u64> {
let value = window_usage_value::<i64>(row, column)?;
u64::try_from(value).map_err(|_| invalid_window_usage(column))
}
pub(crate) fn optional_window_usage_count(
row: &sqlx::postgres::PgRow,
column: &'static str,
) -> AdminStoreResult<Option<u64>> {
window_usage_value::<Option<i64>>(row, column)?
.map(|value| u64::try_from(value).map_err(|_| invalid_window_usage(column)))
.transpose()
}
pub(crate) fn window_usage_value<T>(
row: &sqlx::postgres::PgRow,
column: &'static str,
) -> AdminStoreResult<T>
where
for<'row> T: sqlx::Decode<'row, Postgres> + sqlx::Type<Postgres>,
{
row.try_get(column)
.map_err(|_| invalid_window_usage(column))
}
pub(crate) fn invalid_window_usage(column: &'static str) -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
format!("invalid quota window usage column: {column}"),
)
}
@@ -1,6 +1,6 @@
//! 明文 `provider_accounts` 与凭证 revision CAS 的唯一 PostgreSQL owner
use std::{collections::BTreeSet, fmt, str::FromStr};
use std::{collections::BTreeSet, fmt};
use async_trait::async_trait;
use chrono::{DateTime, TimeDelta, Utc};
@@ -9,17 +9,14 @@ use gateway_admin::{
MutationContext, Revision as AdminRevision,
account_groups::AccountGroupRef,
accounts::{
AccountCost, AccountGroupFilter, AccountImportSettings,
AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountPageItem, AccountRecord, AccountSort as AdminAccountSort,
AccountGroupFilter, AccountImportSettings, AccountListQuery as AdminAccountListQuery,
AccountPage, AccountPageItem, AccountRecord, AccountSort as AdminAccountSort,
AccountSortField as AdminAccountSortField, AccountSummary, AccountUpdateResult,
AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, SortDirection as AdminSortDirection,
UpdateAccount,
},
observability::{
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
},
observability::TimeRange,
provider_credentials::{
AuthorizationCommit, AuthorizationCredentialCommit, CredentialDetails,
CredentialImportCommit, CredentialImportResult, CredentialMutationResult,
@@ -54,7 +51,6 @@ use super::{
AdminAuditEvent, ControlPlaneRepository, ObservabilityRange, PgControlPlaneRepository,
PgObservabilityRepository, ProviderAccountUsageObservation, ProviderAccountUsageQuery,
append_admin_audit_event_in_transaction, bump_config_revision_in_transaction,
completed_usage_fact_predicate,
};
mod admin_adapter;
@@ -1,9 +1,6 @@
//! 预测只读采样:先聚合数值、再读取有界的 Provider 文档,不扫描模型排行维度
use gateway_admin::model::quota_forecast_sampling::{
MAX_FORECAST_HISTORY_POINTS, QuotaForecastHistory, QuotaForecastHistoryPoint,
QuotaForecastUsage,
};
use gateway_admin::model::quota_forecast_sampling::QuotaForecastHistory;
use super::*;
@@ -23,124 +20,21 @@ pub(super) async fn load_history(
"forecast range exceeds supported quota windows",
));
}
let rows = budget
budget
.run(
pool,
"load quota forecast history",
|mut connection| async move {
sqlx::query(sqlx::AssertSqlSafe(history_sql()))
.bind(&window.account_id)
.bind(range.start)
.bind(range.end)
.bind(MAX_FORECAST_HISTORY_POINTS as i32)
.fetch_all(&mut *connection)
.await
.map_err(|source| postgres_unavailable("load quota forecast history", source))
let mut snapshot =
super::super::observability::AnalyticsSnapshot::new_for_accounts(
&mut connection,
)
.await?;
let history = snapshot.forecast_history(&window.account_id, range).await?;
snapshot.finish().await?;
Ok(history)
},
)
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
let mut history = QuotaForecastHistory::default();
for row in &rows {
let usage = QuotaForecastUsage {
request_count: window_usage_count(row, "request_count")?,
tokens: window_usage_count(row, "tokens")?,
input_tokens: window_usage_count(row, "input_tokens")?,
output_tokens: window_usage_count(row, "output_tokens")?,
cached_input_tokens: window_usage_count(row, "cached_input_tokens")?,
missing_token_count: window_usage_count(row, "missing_token_count")?,
known_cost_count: window_usage_count(row, "known_cost_count")?,
unavailable_cost_count: window_usage_count(row, "unavailable_cost_count")?,
usd: window_usage_value(row, "usd")?,
excluded_request_count: window_usage_count(row, "excluded_request_count")?,
};
if window_usage_value::<bool>(row, "is_total")? {
history.usage = usage;
history.pending_request_count = window_usage_count(row, "pending_count")?;
} else if let Some(document) = window_usage_value::<
Option<sqlx::types::Json<serde_json::Map<String, serde_json::Value>>>,
>(row, "document")?
{
history.points.push(QuotaForecastHistoryPoint {
started_at: window_usage_value(row, "started_at")?,
completed_at: window_usage_value(row, "completed_at")?,
usage,
provider_observation: ProviderDocument::new(OpaqueProviderData::new(document.0)),
});
}
}
Ok(history)
}
fn history_sql() -> String {
let completed_usage = completed_usage_fact_predicate("mr");
// 一个语句共享 MVCC 快照
// 用 RANGE 帧让相同完成时间的点拥有相同累计值,
// 避免并发请求的任意行顺序制造不同分子;未完成请求只进入待决计数
format!(
"with scoped as materialized (
select mr.id, mr.started_at, mr.completed_at,
mr.completed_at <= $3 and mr.outcome <> 'running' as settled,
coalesce(({completed_usage}), false) as included,
mr.provider_observation_json is not null as has_document,
mr.input_tokens, mr.output_tokens, mr.cached_input_tokens, mr.total_tokens,
mr.cost_amount, mr.cost_currency,
coalesce(mr.cost_source in ('calculated', 'provider_reported')
and mr.cost_amount is not null and mr.cost_currency = 'USD', false) as known_cost
from model_requests mr
where mr.provider_account_ref = $1
and mr.started_at >= $2 and mr.started_at < $3
), facts as materialized (
select *,
(included and settled)::integer as request_count,
case when included and settled then coalesce(total_tokens,
coalesce(input_tokens, 0) + coalesce(output_tokens, 0)) else 0 end as tokens,
case when included and settled then coalesce(input_tokens, 0) else 0 end as inputs,
case when included and settled then coalesce(output_tokens, 0) else 0 end as outputs,
case when included and settled then coalesce(cached_input_tokens, 0) else 0 end as cached,
(included and settled and total_tokens is null
and (input_tokens is null or output_tokens is null))::integer as missing_tokens,
(included and settled and known_cost)::integer as known_costs,
(included and settled and not known_cost)::integer as missing_costs,
case when included and settled and known_cost then cost_amount else 0 end as usd,
(not included and settled)::integer as excluded
from scoped
), cumulative as (
select id, started_at, completed_at, has_document,
least($4 - 1, floor(extract(epoch from (completed_at - $2))
/ greatest(1, extract(epoch from ($3::timestamptz - $2::timestamptz)) / $4))) as bucket,
sum(request_count) over w as request_count,
sum(tokens) over w as tokens,
sum(inputs) over w as input_tokens,
sum(outputs) over w as output_tokens,
sum(cached) over w as cached_input_tokens,
sum(missing_tokens) over w as missing_token_count,
sum(known_costs) over w as known_cost_count,
sum(missing_costs) over w as unavailable_cost_count,
sum(usd) over w as usd,
sum(excluded) over w as excluded_request_count
from facts where settled
window w as (order by completed_at range between unbounded preceding and current row)
), selected as (
select distinct on (bucket) * from cumulative
where has_document order by bucket, completed_at desc, id desc
)
select false as is_total, s.started_at, s.completed_at,
s.request_count::bigint, s.tokens::bigint, s.input_tokens::bigint,
s.output_tokens::bigint, s.cached_input_tokens::bigint, s.missing_token_count::bigint,
s.known_cost_count::bigint, s.unavailable_cost_count::bigint,
s.usd::double precision, s.excluded_request_count::bigint,
0::bigint as pending_count, mr.provider_observation_json as document
from selected s join model_requests mr on mr.id = s.id
union all
select true, $3, $3,
coalesce(sum(request_count), 0)::bigint, coalesce(sum(tokens), 0)::bigint,
coalesce(sum(inputs), 0)::bigint, coalesce(sum(outputs), 0)::bigint,
coalesce(sum(cached), 0)::bigint, coalesce(sum(missing_tokens), 0)::bigint,
coalesce(sum(known_costs), 0)::bigint, coalesce(sum(missing_costs), 0)::bigint,
coalesce(sum(usd), 0)::double precision, coalesce(sum(excluded), 0)::bigint,
count(*) filter (where settled is not true), null::jsonb
from facts
order by completed_at, is_total"
)
.map_err(|error| admin_store_error(ENTITY, error))
}
@@ -9,76 +9,6 @@ pub(crate) const QUOTA_MAX_BYTES: usize = 128 * 1024;
pub(crate) const MAX_ADMIN_IMPORT_BATCH: usize = 200;
pub(crate) const ADMIN_USAGE_CHUNK_SIZE: usize = 200;
/// 四种窗口投影通过一组 `GROUPING SETS` 从同一批匹配请求派生
/// `grouping_* = 1`
/// 表示对应维度未参与该行聚合,调用方据此区分窗口、模型与成本行
pub(crate) fn account_usage_by_windows_sql() -> String {
let completed_usage = completed_usage_fact_predicate("mr");
format!(
"with requested_windows as (
select *
from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[])
as requested(account_id, window_key, window_start, window_end)
), matched as (
select requested.account_id,
requested.window_key,
mr.id as request_id,
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
mr.outcome, mr.input_tokens, mr.output_tokens, mr.cached_input_tokens,
mr.cache_write_input_tokens, mr.reasoning_output_tokens,
mr.image_input_tokens, mr.image_output_tokens,
mr.image_generation_succeeded, mr.total_tokens,
mr.cost_source, mr.cost_amount, mr.cost_currency, mr.started_at
from requested_windows requested
left join model_requests mr
on mr.provider_account_ref = requested.account_id
and mr.started_at >= requested.window_start
and mr.started_at < requested.window_end
and {completed_usage}
)
select account_id,
window_key,
model,
cost_currency,
grouping(model)::integer as model_grouping,
grouping(cost_currency)::integer as currency_grouping,
count(request_id)::bigint as request_count,
count(request_id) filter (where outcome = 'succeeded')::bigint as success_count,
sum(input_tokens)::bigint as input_tokens,
sum(output_tokens)::bigint as output_tokens,
sum(cached_input_tokens)::bigint as cached_input_tokens,
sum(cache_write_input_tokens)::bigint as cache_write_input_tokens,
sum(reasoning_output_tokens)::bigint as reasoning_output_tokens,
sum(image_input_tokens)::bigint as image_input_tokens,
sum(image_output_tokens)::bigint as image_output_tokens,
count(request_id) filter (where image_generation_succeeded is true)::bigint
as image_request_count,
count(request_id) filter (where image_generation_succeeded is false)::bigint
as image_request_failed_count,
coalesce(sum(coalesce(
total_tokens,
coalesce(input_tokens, 0) + coalesce(output_tokens, 0)
)), 0)::bigint as total_tokens,
count(request_id) filter (where cost_source = 'provider_reported')::bigint
as provider_reported_count,
count(request_id) filter (where cost_source = 'calculated')::bigint
as calculated_count,
count(request_id) filter (where cost_source = 'unavailable')::bigint
as unavailable_count,
max(started_at) as last_used_at,
sum(cost_amount)::text as amount
from matched
group by grouping sets (
(account_id, window_key),
(account_id, window_key, cost_currency),
(account_id, window_key, model),
(account_id, window_key, model, cost_currency)
)
order by account_id, window_key, model_grouping desc, currency_grouping desc,
model nulls last, cost_currency nulls last"
)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountAdminScope {
pub provider_kind: String,
@@ -151,7 +151,7 @@ impl RuntimeSettingsRepository for PgRuntimeSettingsRepository {
pub(crate) async fn load_runtime_settings_from_pool(pool: &PgPool) -> StoreResult<RuntimeSettings> {
let row = sqlx::query_as::<_, RuntimeSettingsRow>(
"select provider_request_profiles_json, config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled, codex_privacy_policy_json,
refresh_concurrency, max_concurrent_per_account, request_interval_ms,
refresh_concurrency, max_concurrent_per_account,
rotation_strategy, smart_scheduling_json, model_mappings_json, usage_retention_days, ops_event_retention_days,
audit_retention_days, min_codex_desktop_version,
min_codex_cli_version, updated_at, responses_max_decompressed_body_bytes, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, openai_guardian_reserved_concurrency, openai_account_affinity, max_account_rotations, openai_session_affinity_ttl_hours,
@@ -329,7 +329,7 @@ pub(crate) async fn load_runtime_settings_in_transaction(
) -> StoreResult<RuntimeSettings> {
let row = sqlx::query_as::<_, RuntimeSettingsRow>(
"select provider_request_profiles_json, config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled, codex_privacy_policy_json,
refresh_concurrency, max_concurrent_per_account, request_interval_ms,
refresh_concurrency, max_concurrent_per_account,
rotation_strategy, smart_scheduling_json, model_mappings_json, usage_retention_days, ops_event_retention_days,
audit_retention_days, min_codex_desktop_version,
min_codex_cli_version, updated_at, responses_max_decompressed_body_bytes, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, openai_guardian_reserved_concurrency, openai_account_affinity, max_account_rotations, openai_session_affinity_ttl_hours,
@@ -382,37 +382,36 @@ pub(crate) async fn update_runtime_settings_in_transaction(
refresh_margin_seconds = $1,
refresh_concurrency = $2,
max_concurrent_per_account = $3,
request_interval_ms = $4,
rotation_strategy = $5,
model_mappings_json = $6,
usage_retention_days = $7,
ops_event_retention_days = $8,
audit_retention_days = $9,
min_codex_desktop_version = $10,
min_codex_cli_version = $11,
max_waiting_per_key = $12,
max_waiting_per_account = $13,
concurrency_wait_timeout_seconds = $14,
account_auto_freeze_enabled = $15,
account_auto_freeze_threshold = $16,
account_auto_freeze_window_seconds = $17,
account_auto_freeze_duration_seconds = $18,
account_auto_freeze_probe_enabled = $19,
account_auto_freeze_probe_model = $20,
account_auto_freeze_adaptive_concurrency = $21,
request_location_json = $22,
request_location_enabled = $23,
responses_max_decompressed_body_bytes = $24,
provider_request_profiles_json = (provider_request_profiles_json - $25::text[]) || $26::jsonb,
account_warmup_enabled = $27,
account_warmup_schedule_time = $28,
account_warmup_model = $29,
smart_scheduling_json = $30,
openai_guardian_reserved_concurrency = $31,
openai_account_affinity = $32,
max_account_rotations = $33,
openai_session_affinity_ttl_hours = $34,
codex_privacy_policy_json = $35,
rotation_strategy = $4,
model_mappings_json = $5,
usage_retention_days = $6,
ops_event_retention_days = $7,
audit_retention_days = $8,
min_codex_desktop_version = $9,
min_codex_cli_version = $10,
max_waiting_per_key = $11,
max_waiting_per_account = $12,
concurrency_wait_timeout_seconds = $13,
account_auto_freeze_enabled = $14,
account_auto_freeze_threshold = $15,
account_auto_freeze_window_seconds = $16,
account_auto_freeze_duration_seconds = $17,
account_auto_freeze_probe_enabled = $18,
account_auto_freeze_probe_model = $19,
account_auto_freeze_adaptive_concurrency = $20,
request_location_json = $21,
request_location_enabled = $22,
responses_max_decompressed_body_bytes = $23,
provider_request_profiles_json = (provider_request_profiles_json - $24::text[]) || $25::jsonb,
account_warmup_enabled = $26,
account_warmup_schedule_time = $27,
account_warmup_model = $28,
smart_scheduling_json = $29,
openai_guardian_reserved_concurrency = $30,
openai_account_affinity = $31,
max_account_rotations = $32,
openai_session_affinity_ttl_hours = $33,
codex_privacy_policy_json = $34,
updated_at = now()
where id = 1
returning config_revision",
@@ -420,7 +419,6 @@ pub(crate) async fn update_runtime_settings_in_transaction(
.bind(refresh_margin_seconds)
.bind(i64::from(update.values.refresh_concurrency))
.bind(i64::from(update.values.max_concurrent_per_account))
.bind(i64::try_from(update.values.request_interval_ms).map_err(|_| invalid_numeric())?)
.bind(&update.rotation_strategy)
.bind(sqlx::types::Json(&update.model_mappings))
.bind(i64::from(update.values.usage_retention_days))
@@ -522,7 +520,6 @@ struct RuntimeSettingsRow {
refresh_margin_seconds: i64,
refresh_concurrency: i64,
max_concurrent_per_account: i64,
request_interval_ms: i64,
smart_scheduling_json: sqlx::types::Json<gateway_core::account::SmartSchedulingConfig>,
rotation_strategy: String,
codex_privacy_policy_json:
@@ -581,7 +578,6 @@ fn runtime_settings_from_row(row: RuntimeSettingsRow) -> StoreResult<RuntimeSett
refresh_margin_seconds: to_u64(row.refresh_margin_seconds)?,
refresh_concurrency: to_u32(row.refresh_concurrency)?,
max_concurrent_per_account: to_u32(row.max_concurrent_per_account)?,
request_interval_ms: to_u64(row.request_interval_ms)?,
smart_scheduling: row.smart_scheduling_json.0,
codex_privacy_policy: row.codex_privacy_policy_json.0,
request_location_enabled: row.request_location_enabled,
@@ -213,7 +213,6 @@ struct SnapshotSettingsRow {
pricing_overrides_json: sqlx::types::Json<gateway_core::metering::PricingOverrides>,
config_revision: i64,
max_concurrent_per_account: i64,
request_interval_ms: i64,
smart_scheduling_json: sqlx::types::Json<gateway_core::account::SmartSchedulingConfig>,
rotation_strategy: String,
model_mappings_json: sqlx::types::Json<BTreeMap<String, String>>,
@@ -239,7 +238,7 @@ async fn load_settings(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<(Revision, SettingsValues)> {
let row = sqlx::query_as::<_, SnapshotSettingsRow>(
"select config_revision, max_concurrent_per_account, request_interval_ms, rotation_strategy, smart_scheduling_json, model_mappings_json, min_codex_desktop_version, min_codex_cli_version, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, openai_guardian_reserved_concurrency, openai_account_affinity, max_account_rotations, openai_session_affinity_ttl_hours, request_location_json, request_location_enabled, codex_privacy_policy_json, responses_max_decompressed_body_bytes, provider_request_profiles_json, pricing_overrides_json, pricing_synced_json from runtime_settings where id = 1",
"select config_revision, max_concurrent_per_account, rotation_strategy, smart_scheduling_json, model_mappings_json, min_codex_desktop_version, min_codex_cli_version, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, openai_guardian_reserved_concurrency, openai_account_affinity, max_account_rotations, openai_session_affinity_ttl_hours, request_location_json, request_location_enabled, codex_privacy_policy_json, responses_max_decompressed_body_bytes, provider_request_profiles_json, pricing_overrides_json, pricing_synced_json from runtime_settings where id = 1",
)
.fetch_optional(&mut **transaction)
.await
@@ -253,7 +252,6 @@ async fn load_settings(
revision_from_i64(row.config_revision)?,
SettingsValues::new(
to_u32(row.max_concurrent_per_account)?,
to_u64(row.request_interval_ms)?,
row.rotation_strategy,
row.model_mappings_json.0,
row.min_codex_desktop_version,
@@ -399,7 +399,10 @@ impl ClientAdmissionPort for RedisClientAdmissionRepository {
let request = request.as_str().to_owned();
Box::new(crate::lease_renewal::LeaseRenewal::spawn(
deadline,
Some(cancellation),
Some((
cancellation,
gateway_core::lifecycle::RequestLeaseKind::ClientAdmission,
)),
move |ttl| {
let repository = repository.clone();
let key = key.clone();
@@ -35,7 +35,6 @@ redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now_ms)
local in_flight = redis.call('ZCARD', KEYS[1])
local max_concurrent = tonumber(ARGV[4])
local interval_ms = tonumber(ARGV[5])
local retry_ms = 0
if max_concurrent > 0 and in_flight >= max_concurrent then
@@ -45,11 +44,6 @@ if max_concurrent > 0 and in_flight >= max_concurrent then
end
end
local last_started = tonumber(redis.call('GET', KEYS[3]) or '0')
if last_started > 0 and now_ms - last_started < interval_ms then
retry_ms = math.max(retry_ms, interval_ms - (now_ms - last_started))
end
if retry_ms > 0 then
return {0, '0', '0', tostring(math.max(1, retry_ms))}
end
@@ -69,10 +63,10 @@ local active_ttl = redis.call('PTTL', KEYS[1])
if active_ttl < tonumber(ARGV[3]) then
redis.call('PEXPIRE', KEYS[1], ARGV[3])
end
redis.call('SET', KEYS[3], tostring(now_ms), 'PX', ARGV[6])
redis.call('SET', KEYS[3], tostring(now_ms), 'PX', ARGV[5])
local fence_ttl = redis.call('PTTL', KEYS[2])
if fence_ttl < tonumber(ARGV[6]) then
redis.call('PEXPIRE', KEYS[2], ARGV[6])
if fence_ttl < tonumber(ARGV[5]) then
redis.call('PEXPIRE', KEYS[2], ARGV[5])
end
return {1, tostring(fence), tostring(expires_at), '0'}
"#;
@@ -171,14 +165,13 @@ impl CredentialLeaseRequest {
}
}
/// 对任意明确资源施加并发数和启动间隔约束的计数 lease
/// 对任意明确资源施加并发数约束的计数 lease
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialBoundedLeaseRequest {
pub scope: CredentialLeaseScope,
pub resource_id: String,
pub owner_id: String,
pub max_concurrent: u32,
pub request_interval: Duration,
pub ttl: Duration,
}
@@ -189,7 +182,6 @@ impl CredentialBoundedLeaseRequest {
if self.max_concurrent == 0 && self.scope != CredentialLeaseScope::ProviderAccount {
return Err(invalid("max_concurrent must be positive"));
}
supported_duration(self.request_interval, true, "request interval")?;
supported_duration(self.ttl, false, "lease TTL")?;
Ok(())
}
@@ -260,8 +252,13 @@ impl CredentialLeaseGuard {
entity: "credential lease",
message: "lease has already been released".to_owned(),
})?;
let renewal =
crate::lease_renewal::LeaseRenewal::spawn(deadline, Some(cancellation), move |ttl| {
let renewal = crate::lease_renewal::LeaseRenewal::spawn(
deadline,
Some((
cancellation,
gateway_core::lifecycle::RequestLeaseKind::Credential,
)),
move |ttl| {
let repository = repository.clone();
let mut lease_request = lease_request.clone();
lease_request.ttl = ttl;
@@ -272,7 +269,8 @@ impl CredentialLeaseGuard {
.await
.map(|grant| grant.is_some())
})
});
},
);
Ok(RenewingSchedulingLease {
_renewal: renewal,
_guard: self,
@@ -405,8 +403,8 @@ impl RedisCredentialLeaseRepository {
let fingerprint = resource_fingerprint("credential lease", &request.resource_id)?;
let tag = format!("{{{fingerprint}}}");
let prefix = format!("{}:lease:{}:{tag}", self.namespace, request.scope.as_str());
// 两个容量池独立计数,但账号最小请求间隔仍由同一时间戳原子裁决
let interval_prefix = if request.scope == CredentialLeaseScope::ProviderAccountReserved {
// 两个容量池独立计数,最近调度时间按账号共享,用于选号排序
let signal_prefix = if request.scope == CredentialLeaseScope::ProviderAccountReserved {
format!("{}:lease:account:{tag}", self.namespace)
} else {
prefix.clone()
@@ -414,7 +412,7 @@ impl RedisCredentialLeaseRepository {
Ok([
format!("{prefix}:active"),
format!("{prefix}:fence"),
format!("{interval_prefix}:last-started"),
format!("{signal_prefix}:last-started"),
])
}
@@ -449,7 +447,6 @@ impl RedisCredentialLeaseRepository {
&self,
request: &CredentialLeaseRequest,
max_concurrent: u32,
request_interval: Duration,
) -> StoreResult<LeaseAttempt> {
request.validate()?;
if max_concurrent == 0 && request.scope != CredentialLeaseScope::ProviderAccount {
@@ -459,8 +456,7 @@ impl RedisCredentialLeaseRepository {
let lease_id = Uuid::now_v7().to_string();
let owner = resource_fingerprint("credential lease owner", &request.owner_id)?;
let ttl_millis = duration_millis(request.ttl)?;
let interval_millis = duration_millis(request_interval)?;
let signal_ttl = SIGNAL_TTL_MILLIS.max(ttl_millis).max(interval_millis);
let signal_ttl = SIGNAL_TTL_MILLIS.max(ttl_millis);
let mut connection = self.connection.clone();
let (acquired, fence, expires_at, retry_after): (i64, String, String, String) =
Script::new(ACQUIRE_SCRIPT)
@@ -471,7 +467,6 @@ impl RedisCredentialLeaseRepository {
.arg(&lease_id)
.arg(ttl_millis)
.arg(max_concurrent)
.arg(interval_millis)
.arg(signal_ttl)
.invoke_async(&mut connection)
.await
@@ -629,7 +624,6 @@ impl RedisProviderLeaseCoordinator {
resource_id: request.account_id().as_str().to_owned(),
owner_id: self.owner_id("request"),
max_concurrent: request.max_concurrent().get(),
request_interval: request.request_interval(),
ttl,
})
.await
@@ -661,7 +655,6 @@ impl RedisProviderLeaseCoordinator {
resource_id: OAUTH_REFRESH_CAPACITY_RESOURCE.to_owned(),
owner_id: self.owner_id("refresh-capacity"),
max_concurrent: request.max_concurrent().get(),
request_interval: Duration::ZERO,
ttl: OAUTH_REFRESH_LEASE_TTL,
})
.await
@@ -767,7 +760,7 @@ impl CredentialLeaseRepository for RedisCredentialLeaseRepository {
&self,
request: &CredentialLeaseRequest,
) -> StoreResult<Option<CredentialLeaseGrant>> {
self.acquire_with_limits(request, 1, Duration::ZERO)
self.acquire_with_limits(request, 1)
.await
.map(|attempt| attempt.grant)
}
@@ -831,11 +824,7 @@ impl CredentialLeaseRepository for RedisCredentialLeaseRepository {
request.validate()?;
let lease_request = request.lease_request();
let attempt = self
.acquire_with_limits(
&lease_request,
request.max_concurrent,
request.request_interval,
)
.acquire_with_limits(&lease_request, request.max_concurrent)
.await?;
match attempt.grant {
Some(grant) => Ok(CredentialBoundedLeaseAcquisition::Acquired(

Some files were not shown because too many files have changed in this diff Show More