mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-09 00:05:47 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0fc25bfce9 | ||
|
|
eed8c93ae2 | ||
|
|
d865a00768 | ||
|
|
a0671fb70b | ||
|
|
fce543d4ce | ||
|
|
219888d5c2 | ||
|
|
5b80340135 | ||
|
|
b2dab66a80 | ||
|
|
f0cf8e7198 | ||
|
|
a9a86071bf | ||
|
|
27cef052c9 | ||
|
|
ef40088303 |
Generated
+2
-2
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
|
||||
+4
-4
@@ -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 {
|
||||
|
||||
@@ -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 [(¤t, "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(¤t, &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 = [¤t, &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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(¤cy)
|
||||
.fetch_one(&mut *self.tx)
|
||||
.await
|
||||
.map_err(|error| postgres_unavailable("repair account cost coverage", error))?;
|
||||
if !exists {
|
||||
state.costs.remove(¤cy);
|
||||
}
|
||||
}
|
||||
}
|
||||
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, ¤t).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
Reference in New Issue
Block a user