i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

Rootfix usage counter outbox

Browse Source
This commit is contained in:
fawney19
2026-05-19 00:30:24 +08:00
parent 19c8688eb1
commit 5eed329916
60 changed files with 4910 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/aether-admin/src/observability/monitoring.rs
View File

@@ -727,6 +727,7 @@ pub fn build_admin_monitoring_system_status_payload_response(
active_streams: usize,
path_prefixes: &[&str],
recent_errors: usize,
usage_counter: Value,
) -> Response<Body> {
Json(json!({
"timestamp": timestamp.to_rfc3339(),
@@ -757,6 +758,7 @@ pub fn build_admin_monitoring_system_status_payload_response(
"path_prefixes": path_prefixes,
},
"recent_errors": recent_errors,
"usage_counter": usage_counter,
}))
.into_response()
}

7
crates/aether-admin/src/observability/stats.rs
View File

@@ -667,7 +667,9 @@ pub fn admin_stats_performance_percentiles_empty_response() -> Response<Body> {
Json(json!([])).into_response()
}
pub fn admin_stats_provider_performance_empty_response() -> Response<Body> {
pub fn admin_stats_provider_performance_empty_response(
usage_counter: serde_json::Value,
) -> Response<Body> {
Json(json!({
"summary": {
"request_count": 0,
@@ -678,6 +680,7 @@ pub fn admin_stats_provider_performance_empty_response() -> Response<Body> {
},
"providers": [],
"timeline": [],
"usage_counter": usage_counter,
}))
.into_response()
}
@@ -1093,6 +1096,7 @@ pub fn build_admin_stats_performance_percentiles_response_from_summaries(
pub fn build_admin_stats_provider_performance_response(
performance: &StoredUsageProviderPerformance,
usage_counter: serde_json::Value,
) -> Response<Body> {
let summary = &performance.summary;
let providers = performance
@@ -1158,6 +1162,7 @@ pub fn build_admin_stats_provider_performance_response(
},
"providers": providers,
"timeline": timeline,
"usage_counter": usage_counter,
}))
.into_response()
}

2
crates/aether-admin/src/system.rs
View File

@@ -861,6 +861,7 @@ pub fn build_admin_system_stats_payload(
active_providers: u64,
total_api_keys: u64,
total_requests: u64,
usage_counter: serde_json::Value,
) -> serde_json::Value {
json!({
"users": {
@@ -873,6 +874,7 @@ pub fn build_admin_system_stats_payload(
},
"api_keys": total_api_keys,
"requests": total_requests,
"usage_counter": usage_counter,
})
}

2
crates/aether-data-contracts/src/repository/candidate_selection/types.rs
View File

@@ -108,6 +108,8 @@ fn api_format_matches(left: &str, right: &str) -> bool {
#[async_trait]
pub trait MinimalCandidateSelectionReadRepository: Send + Sync {
fn clear_local_cache(&self) {}
async fn list_for_exact_api_format(
&self,
api_format: &str,

14
crates/aether-data-contracts/src/repository/usage/mod.rs
View File

@@ -2,7 +2,8 @@ mod types;
pub use types::{
parse_usage_body_ref, usage_body_ref, usage_request_metadata_client_family,
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
ApiKeyLastUsedDelta, ManagementTokenCounterDelta, PendingUsageCleanupSummary,
ProviderApiKeyWindowUsageRequest, ProxyNodeCounterDelta, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
@@ -20,10 +21,11 @@ pub use types::{
UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy,
UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupExecutionMode,
UsageCleanupPreviewCounts, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow,
UsageCostSavingsSummaryQuery, UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery,
UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery,
UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery,
UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery,
UsageReadRepository, UsageRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity,
UsageCostSavingsSummaryQuery, UsageCounterFlushSummary, UsageCounterHealthSnapshot,
UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery,
UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy,
UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository,
UsageRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity,
UsageTimeSeriesQuery, UsageWriteRepository,
};

121
crates/aether-data-contracts/src/repository/usage/types.rs
View File

@@ -1518,6 +1518,12 @@ pub trait UsageReadRepository: Send + Sync {
&self,
query: &UsageDailyHeatmapQuery,
) -> Result<Vec<StoredUsageDailySummary>, crate::DataLayerError>;
async fn read_usage_counter_health(
&self,
) -> Result<UsageCounterHealthSnapshot, crate::DataLayerError> {
Ok(UsageCounterHealthSnapshot::default())
}
}
/// Repository write model for a single usage aggregate.
@@ -1690,6 +1696,47 @@ pub trait UsageWriteRepository: Send + Sync {
Ok(PendingUsageCleanupSummary::default())
}
async fn flush_usage_counter_deltas(
&self,
batch_size: usize,
) -> Result<UsageCounterFlushSummary, crate::DataLayerError> {
let _ = batch_size;
Ok(UsageCounterFlushSummary::default())
}
async fn enqueue_proxy_node_counter_delta(
&self,
delta: ProxyNodeCounterDelta,
) -> Result<bool, crate::DataLayerError> {
let _ = delta;
Ok(false)
}
async fn enqueue_management_token_counter_delta(
&self,
delta: ManagementTokenCounterDelta,
) -> Result<bool, crate::DataLayerError> {
let _ = delta;
Ok(false)
}
async fn enqueue_api_key_last_used_delta(
&self,
delta: ApiKeyLastUsedDelta,
) -> Result<bool, crate::DataLayerError> {
let _ = delta;
Ok(false)
}
async fn cleanup_processed_usage_counter_deltas(
&self,
cutoff_unix_secs: u64,
batch_size: usize,
) -> Result<usize, crate::DataLayerError> {
let _ = (cutoff_unix_secs, batch_size);
Ok(0)
}
async fn cleanup_usage(
&self,
window: &UsageCleanupWindow,
@@ -1723,6 +1770,80 @@ pub struct PendingUsageCleanupSummary {
pub recovered: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct UsageCounterFlushSummary {
pub rows_claimed: usize,
pub api_key_targets: usize,
pub provider_api_key_targets: usize,
pub model_targets: usize,
pub provider_monthly_targets: usize,
pub proxy_node_targets: usize,
pub management_token_targets: usize,
pub api_key_last_used_targets: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct UsageCounterHealthSnapshot {
pub pending_rows: u64,
pub processed_rows: u64,
pub oldest_pending_created_at_unix_secs: Option<u64>,
pub latest_processed_at_unix_secs: Option<u64>,
pub pending_by_kind: std::collections::BTreeMap<String, u64>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProxyNodeCounterDelta {
pub node_id: String,
pub total_requests_delta: i64,
pub failed_requests_delta: i64,
pub dns_failures_delta: i64,
pub stream_errors_delta: i64,
}
impl ProxyNodeCounterDelta {
pub fn is_noop(&self) -> bool {
self.node_id.trim().is_empty()
|| (self.total_requests_delta <= 0
&& self.failed_requests_delta <= 0
&& self.dns_failures_delta <= 0
&& self.stream_errors_delta <= 0)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ManagementTokenCounterDelta {
pub token_id: String,
pub usage_count_delta: i64,
pub last_used_at_unix_secs: Option<u64>,
pub last_used_ip: Option<String>,
}
impl ManagementTokenCounterDelta {
pub fn is_noop(&self) -> bool {
self.token_id.trim().is_empty()
|| (self.usage_count_delta <= 0
&& self.last_used_at_unix_secs.is_none()
&& self
.last_used_ip
.as_deref()
.map(str::trim)
.unwrap_or("")
.is_empty())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ApiKeyLastUsedDelta {
pub api_key_id: String,
pub last_used_at_unix_secs: u64,
}
impl ApiKeyLastUsedDelta {
pub fn is_noop(&self) -> bool {
self.api_key_id.trim().is_empty()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct UsageCleanupSummary {
pub body_externalized: usize,

67
crates/aether-data/migrations/postgres/20260518000000_add_usage_counter_deltas.sql Normal file
View File

@@ -0,0 +1,67 @@
CREATE TABLE IF NOT EXISTS public.usage_counter_deltas (
id character varying(36) NOT NULL,
request_id character varying(128) NOT NULL,
kind character varying(64) NOT NULL,
target_id text NOT NULL,
request_count_delta bigint DEFAULT 0 NOT NULL,
total_requests_delta bigint DEFAULT 0 NOT NULL,
success_count_delta bigint DEFAULT 0 NOT NULL,
error_count_delta bigint DEFAULT 0 NOT NULL,
dns_failures_delta bigint DEFAULT 0 NOT NULL,
stream_errors_delta bigint DEFAULT 0 NOT NULL,
total_tokens_delta bigint DEFAULT 0 NOT NULL,
total_cost_usd_delta double precision DEFAULT 0 NOT NULL,
total_response_time_ms_delta bigint DEFAULT 0 NOT NULL,
last_used_at_unix_secs bigint,
last_used_ip text,
candidate_last_used_at_unix_secs bigint,
removed_last_used_at_unix_secs bigint,
usage_created_at_unix_secs bigint,
created_at timestamp with time zone DEFAULT now() NOT NULL,
processed_at timestamp with time zone,
CONSTRAINT usage_counter_deltas_pkey PRIMARY KEY (id),
CONSTRAINT usage_counter_deltas_kind_check CHECK (
kind IN (
'api_key',
'provider_api_key',
'model',
'provider_monthly',
'proxy_node',
'management_token',
'api_key_last_used'
)
)
);
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_unprocessed
ON public.usage_counter_deltas USING btree (created_at, id)
WHERE processed_at IS NULL;
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_processed
ON public.usage_counter_deltas USING btree (processed_at, created_at, id)
WHERE processed_at IS NOT NULL;
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_request_kind
ON public.usage_counter_deltas USING btree (request_id, kind, target_id);
CREATE INDEX IF NOT EXISTS idx_entitlement_usage_entitlement_date
ON public.entitlement_usage_ledgers USING btree (user_entitlement_id, usage_date);
CREATE INDEX IF NOT EXISTS idx_video_tasks_due_poll
ON public.video_tasks USING btree (status, next_poll_at, updated_at)
WHERE next_poll_at IS NOT NULL;
ALTER TABLE IF EXISTS public.api_keys
ALTER COLUMN total_requests TYPE bigint USING COALESCE(total_requests, 0)::bigint;
ALTER TABLE IF EXISTS public.global_models
ALTER COLUMN usage_count TYPE bigint USING COALESCE(usage_count, 0)::bigint;
ALTER TABLE IF EXISTS public.management_tokens
ALTER COLUMN usage_count TYPE bigint USING COALESCE(usage_count, 0)::bigint;
ALTER TABLE IF EXISTS public.provider_api_keys
ALTER COLUMN request_count TYPE bigint USING COALESCE(request_count, 0)::bigint,
ALTER COLUMN success_count TYPE bigint USING COALESCE(success_count, 0)::bigint,
ALTER COLUMN error_count TYPE bigint USING COALESCE(error_count, 0)::bigint,
ALTER COLUMN total_response_time_ms TYPE bigint USING COALESCE(total_response_time_ms, 0)::bigint;

49
crates/aether-data/schema/bootstrap/postgres/001_types_and_tables.sql
View File

@@ -168,7 +168,7 @@ CREATE TABLE IF NOT EXISTS public.api_keys (
name character varying(100),
key_prefix character varying(64),
status character varying(64) DEFAULT 'active'::character varying NOT NULL,
total_requests integer DEFAULT 0,
total_requests bigint DEFAULT 0,
total_tokens bigint DEFAULT '0'::bigint NOT NULL,
total_cost_usd numeric(20,8) DEFAULT '0'::double precision,
is_standalone boolean DEFAULT false NOT NULL,
@@ -288,7 +288,7 @@ CREATE TABLE IF NOT EXISTS public.global_models (
default_tiered_pricing json,
supported_capabilities json,
is_active boolean DEFAULT true NOT NULL,
usage_count integer DEFAULT 0 NOT NULL,
usage_count bigint DEFAULT 0 NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,
config jsonb,
@@ -359,7 +359,7 @@ CREATE TABLE IF NOT EXISTS public.management_tokens (
expires_at timestamp with time zone,
last_used_at timestamp with time zone,
last_used_ip character varying(45),
usage_count integer DEFAULT 0 NOT NULL,
usage_count bigint DEFAULT 0 NOT NULL,
is_active boolean DEFAULT true NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,
@@ -585,10 +585,10 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
last_probe_increase_at timestamp with time zone,
cache_ttl_minutes integer DEFAULT 5 NOT NULL,
max_probe_interval_minutes integer DEFAULT 32 NOT NULL,
request_count integer DEFAULT 0,
success_count integer DEFAULT 0,
error_count integer DEFAULT 0,
total_response_time_ms integer DEFAULT 0,
request_count bigint DEFAULT 0,
success_count bigint DEFAULT 0,
error_count bigint DEFAULT 0,
total_response_time_ms bigint DEFAULT 0,
last_used_at timestamp with time zone,
last_error_at timestamp with time zone,
last_error_msg text,
@@ -856,7 +856,7 @@ CREATE TABLE IF NOT EXISTS public.refund_requests (
CREATE TABLE IF NOT EXISTS public.request_candidates (
id character varying(36) NOT NULL,
request_id character varying(100) NOT NULL,
request_id character varying(128) NOT NULL,
user_id character varying(36),
api_key_id character varying(36),
candidate_index integer NOT NULL,
@@ -1263,6 +1263,39 @@ CREATE TABLE IF NOT EXISTS public.usage (
--
-- Name: usage_counter_deltas; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.usage_counter_deltas (
id character varying(36) NOT NULL,
request_id character varying(128) NOT NULL,
kind character varying(64) NOT NULL,
target_id text NOT NULL,
request_count_delta bigint DEFAULT 0 NOT NULL,
total_requests_delta bigint DEFAULT 0 NOT NULL,
success_count_delta bigint DEFAULT 0 NOT NULL,
error_count_delta bigint DEFAULT 0 NOT NULL,
dns_failures_delta bigint DEFAULT 0 NOT NULL,
stream_errors_delta bigint DEFAULT 0 NOT NULL,
total_tokens_delta bigint DEFAULT 0 NOT NULL,
total_cost_usd_delta double precision DEFAULT 0 NOT NULL,
total_response_time_ms_delta bigint DEFAULT 0 NOT NULL,
last_used_at_unix_secs bigint,
last_used_ip text,
candidate_last_used_at_unix_secs bigint,
removed_last_used_at_unix_secs bigint,
usage_created_at_unix_secs bigint,
created_at timestamp with time zone DEFAULT now() NOT NULL,
processed_at timestamp with time zone,
CONSTRAINT usage_counter_deltas_pkey PRIMARY KEY (id),
CONSTRAINT usage_counter_deltas_kind_check CHECK (
kind IN ('api_key', 'provider_api_key', 'model', 'provider_monthly', 'proxy_node', 'management_token', 'api_key_last_used')
)
);
--
-- Name: user_model_usage_counts; Type: TABLE; Schema: public; Owner: -
--

40
crates/aether-data/schema/bootstrap/postgres/004_indexes.sql
View File

@@ -173,6 +173,14 @@ CREATE INDEX IF NOT EXISTS idx_entitlement_usage_user_date ON public.entitlement
--
-- Name: idx_entitlement_usage_entitlement_date; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_entitlement_usage_entitlement_date ON public.entitlement_usage_ledgers USING btree (user_entitlement_id, usage_date);
--
-- Name: idx_provider_api_keys_provider_active; Type: INDEX; Schema: public; Owner: -
--
@@ -549,6 +557,30 @@ CREATE INDEX IF NOT EXISTS idx_usage_wallet_finalized ON public.usage USING btre
--
-- Name: ix_usage_counter_deltas_request_kind; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_request_kind ON public.usage_counter_deltas USING btree (request_id, kind, target_id);
--
-- Name: ix_usage_counter_deltas_processed; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_processed ON public.usage_counter_deltas USING btree (processed_at, created_at, id) WHERE processed_at IS NOT NULL;
--
-- Name: ix_usage_counter_deltas_unprocessed; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_unprocessed ON public.usage_counter_deltas USING btree (created_at, id) WHERE processed_at IS NULL;
--
-- Name: idx_user_model_usage_model; Type: INDEX; Schema: public; Owner: -
--
@@ -597,6 +629,14 @@ CREATE INDEX IF NOT EXISTS idx_video_tasks_next_poll ON public.video_tasks USING
--
-- Name: idx_video_tasks_due_poll; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_video_tasks_due_poll ON public.video_tasks USING btree (status, next_poll_at, updated_at) WHERE next_poll_at IS NOT NULL;
--
-- Name: idx_video_tasks_request_id; Type: INDEX; Schema: public; Owner: -
--

1
crates/aether-data/schema/generated/mysql/baseline/002_provider_catalog.sql
View File

@@ -286,6 +286,7 @@ CREATE TABLE IF NOT EXISTS video_tasks (
UNIQUE KEY video_tasks_request_id_key (`request_id`),
KEY video_tasks_external_id_idx (`external_task_id`),
KEY video_tasks_next_poll_idx (`next_poll_at`),
KEY video_tasks_due_poll_idx (`status`, `next_poll_at`, `updated_at`),
KEY video_tasks_user_status_idx (`user_id`, `status`),
KEY video_tasks_api_key_id_idx (`api_key_id`),
KEY video_tasks_provider_id_idx (`provider_id`),

3
crates/aether-data/schema/generated/mysql/baseline/005_wallet_billing.sql
View File

@@ -189,7 +189,8 @@ CREATE TABLE IF NOT EXISTS entitlement_usage_ledgers (
`created_at` BIGINT NOT NULL,
PRIMARY KEY (`id`),
UNIQUE KEY uq_entitlement_usage_request (`user_entitlement_id`, `request_id`),
KEY idx_entitlement_usage_user_date (`user_id`, `usage_date`)
KEY idx_entitlement_usage_user_date (`user_id`, `usage_date`),
KEY idx_entitlement_usage_entitlement_date (`user_entitlement_id`, `usage_date`)
);
CREATE TABLE IF NOT EXISTS refund_requests (

27
crates/aether-data/schema/generated/mysql/baseline/006_usage.sql
View File

@@ -111,6 +111,33 @@ CREATE TABLE IF NOT EXISTS `usage` (
KEY usage_wallet_id_idx (`wallet_id`)
);
CREATE TABLE IF NOT EXISTS usage_counter_deltas (
`id` VARCHAR(36) NOT NULL,
`request_id` VARCHAR(128) NOT NULL,
`kind` VARCHAR(64) NOT NULL,
`target_id` TEXT NOT NULL,
`request_count_delta` BIGINT NOT NULL DEFAULT 0,
`total_requests_delta` BIGINT NOT NULL DEFAULT 0,
`success_count_delta` BIGINT NOT NULL DEFAULT 0,
`error_count_delta` BIGINT NOT NULL DEFAULT 0,
`dns_failures_delta` BIGINT NOT NULL DEFAULT 0,
`stream_errors_delta` BIGINT NOT NULL DEFAULT 0,
`total_tokens_delta` BIGINT NOT NULL DEFAULT 0,
`total_cost_usd_delta` DOUBLE NOT NULL DEFAULT 0,
`total_response_time_ms_delta` BIGINT NOT NULL DEFAULT 0,
`last_used_at_unix_secs` BIGINT,
`last_used_ip` TEXT,
`candidate_last_used_at_unix_secs` BIGINT,
`removed_last_used_at_unix_secs` BIGINT,
`usage_created_at_unix_secs` BIGINT,
`created_at` BIGINT NOT NULL,
`processed_at` BIGINT,
PRIMARY KEY (`id`),
KEY ix_usage_counter_deltas_unprocessed (`created_at`, `id`),
KEY ix_usage_counter_deltas_processed (`processed_at`, `created_at`, `id`),
KEY ix_usage_counter_deltas_request_kind (`request_id`, `kind`, `target_id`)
);
CREATE TABLE IF NOT EXISTS usage_settlement_snapshots (
`request_id` VARCHAR(128) NOT NULL,
`billing_status` VARCHAR(64) NOT NULL,

1
crates/aether-data/schema/generated/postgres/baseline/002_provider_catalog.sql
View File

@@ -296,6 +296,7 @@ ALTER TABLE ONLY public.video_tasks ADD CONSTRAINT video_tasks_short_id_key UNIQ
ALTER TABLE ONLY public.video_tasks ADD CONSTRAINT video_tasks_request_id_key UNIQUE (request_id);
CREATE INDEX IF NOT EXISTS video_tasks_external_id_idx ON public.video_tasks USING btree (external_task_id);
CREATE INDEX IF NOT EXISTS video_tasks_next_poll_idx ON public.video_tasks USING btree (next_poll_at);
CREATE INDEX IF NOT EXISTS video_tasks_due_poll_idx ON public.video_tasks USING btree (status, next_poll_at, updated_at);
CREATE INDEX IF NOT EXISTS video_tasks_user_status_idx ON public.video_tasks USING btree (user_id, status);
CREATE INDEX IF NOT EXISTS video_tasks_api_key_id_idx ON public.video_tasks USING btree (api_key_id);
CREATE INDEX IF NOT EXISTS video_tasks_provider_id_idx ON public.video_tasks USING btree (provider_id);

1
crates/aether-data/schema/generated/postgres/baseline/005_wallet_billing.sql
View File

@@ -200,6 +200,7 @@ CREATE TABLE IF NOT EXISTS public.entitlement_usage_ledgers (
ALTER TABLE ONLY public.entitlement_usage_ledgers ADD CONSTRAINT entitlement_usage_ledgers_pkey PRIMARY KEY (id);
ALTER TABLE ONLY public.entitlement_usage_ledgers ADD CONSTRAINT uq_entitlement_usage_request UNIQUE (user_entitlement_id, request_id);
CREATE INDEX IF NOT EXISTS idx_entitlement_usage_user_date ON public.entitlement_usage_ledgers USING btree (user_id, usage_date);
CREATE INDEX IF NOT EXISTS idx_entitlement_usage_entitlement_date ON public.entitlement_usage_ledgers USING btree (user_entitlement_id, usage_date);
CREATE TABLE IF NOT EXISTS public.refund_requests (
id character varying(64) NOT NULL,

28
crates/aether-data/schema/generated/postgres/baseline/006_usage.sql
View File

@@ -112,6 +112,34 @@ CREATE INDEX IF NOT EXISTS usage_request_id_idx ON public.usage USING btree (req
CREATE INDEX IF NOT EXISTS usage_user_id_idx ON public.usage USING btree (user_id);
CREATE INDEX IF NOT EXISTS usage_wallet_id_idx ON public.usage USING btree (wallet_id);
CREATE TABLE IF NOT EXISTS public.usage_counter_deltas (
id character varying(36) NOT NULL,
request_id character varying(128) NOT NULL,
kind character varying(64) NOT NULL,
target_id text NOT NULL,
request_count_delta bigint DEFAULT 0 NOT NULL,
total_requests_delta bigint DEFAULT 0 NOT NULL,
success_count_delta bigint DEFAULT 0 NOT NULL,
error_count_delta bigint DEFAULT 0 NOT NULL,
dns_failures_delta bigint DEFAULT 0 NOT NULL,
stream_errors_delta bigint DEFAULT 0 NOT NULL,
total_tokens_delta bigint DEFAULT 0 NOT NULL,
total_cost_usd_delta double precision DEFAULT 0 NOT NULL,
total_response_time_ms_delta bigint DEFAULT 0 NOT NULL,
last_used_at_unix_secs bigint,
last_used_ip text,
candidate_last_used_at_unix_secs bigint,
removed_last_used_at_unix_secs bigint,
usage_created_at_unix_secs bigint,
created_at timestamp with time zone NOT NULL,
processed_at timestamp with time zone
);
ALTER TABLE ONLY public.usage_counter_deltas ADD CONSTRAINT usage_counter_deltas_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_unprocessed ON public.usage_counter_deltas USING btree (created_at, id);
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_processed ON public.usage_counter_deltas USING btree (processed_at, created_at, id);
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_request_kind ON public.usage_counter_deltas USING btree (request_id, kind, target_id);
CREATE TABLE IF NOT EXISTS public.usage_settlement_snapshots (
request_id character varying(128) NOT NULL,
billing_status character varying(64) NOT NULL,

1
crates/aether-data/schema/generated/sqlite/baseline/002_provider_catalog.sql
View File

@@ -278,6 +278,7 @@ CREATE TABLE IF NOT EXISTS video_tasks (
);
CREATE INDEX IF NOT EXISTS video_tasks_external_id_idx ON video_tasks (external_task_id);
CREATE INDEX IF NOT EXISTS video_tasks_next_poll_idx ON video_tasks (next_poll_at);
CREATE INDEX IF NOT EXISTS video_tasks_due_poll_idx ON video_tasks (status, next_poll_at, updated_at);
CREATE INDEX IF NOT EXISTS video_tasks_user_status_idx ON video_tasks (user_id, status);
CREATE INDEX IF NOT EXISTS video_tasks_api_key_id_idx ON video_tasks (api_key_id);
CREATE INDEX IF NOT EXISTS video_tasks_provider_id_idx ON video_tasks (provider_id);

1
crates/aether-data/schema/generated/sqlite/baseline/005_wallet_billing.sql
View File

@@ -182,6 +182,7 @@ CREATE TABLE IF NOT EXISTS entitlement_usage_ledgers (
UNIQUE (user_entitlement_id, request_id)
);
CREATE INDEX IF NOT EXISTS idx_entitlement_usage_user_date ON entitlement_usage_ledgers (user_id, usage_date);
CREATE INDEX IF NOT EXISTS idx_entitlement_usage_entitlement_date ON entitlement_usage_ledgers (user_entitlement_id, usage_date);
CREATE TABLE IF NOT EXISTS refund_requests (
id TEXT PRIMARY KEY NOT NULL,

26
crates/aether-data/schema/generated/sqlite/baseline/006_usage.sql
View File

@@ -110,6 +110,32 @@ CREATE INDEX IF NOT EXISTS usage_request_id_idx ON "usage" (request_id);
CREATE INDEX IF NOT EXISTS usage_user_id_idx ON "usage" (user_id);
CREATE INDEX IF NOT EXISTS usage_wallet_id_idx ON "usage" (wallet_id);
CREATE TABLE IF NOT EXISTS usage_counter_deltas (
id TEXT PRIMARY KEY NOT NULL,
request_id TEXT NOT NULL,
kind TEXT NOT NULL,
target_id TEXT NOT NULL,
request_count_delta INTEGER NOT NULL DEFAULT 0,
total_requests_delta INTEGER NOT NULL DEFAULT 0,
success_count_delta INTEGER NOT NULL DEFAULT 0,
error_count_delta INTEGER NOT NULL DEFAULT 0,
dns_failures_delta INTEGER NOT NULL DEFAULT 0,
stream_errors_delta INTEGER NOT NULL DEFAULT 0,
total_tokens_delta INTEGER NOT NULL DEFAULT 0,
total_cost_usd_delta REAL NOT NULL DEFAULT 0,
total_response_time_ms_delta INTEGER NOT NULL DEFAULT 0,
last_used_at_unix_secs INTEGER,
last_used_ip TEXT,
candidate_last_used_at_unix_secs INTEGER,
removed_last_used_at_unix_secs INTEGER,
usage_created_at_unix_secs INTEGER,
created_at INTEGER NOT NULL,
processed_at INTEGER
);
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_unprocessed ON usage_counter_deltas (created_at, id);
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_processed ON usage_counter_deltas (processed_at, created_at, id);
CREATE INDEX IF NOT EXISTS ix_usage_counter_deltas_request_kind ON usage_counter_deltas (request_id, kind, target_id);
CREATE TABLE IF NOT EXISTS usage_settlement_snapshots (
request_id TEXT PRIMARY KEY NOT NULL,
billing_status TEXT NOT NULL,

4
crates/aether-data/schema/logical/002_provider_catalog.toml
View File

@@ -1296,6 +1296,10 @@ columns = ["external_task_id"]
name = "video_tasks_next_poll_idx"
columns = ["next_poll_at"]
[[table.video_tasks.indexes]]
name = "video_tasks_due_poll_idx"
columns = ["status", "next_poll_at", "updated_at"]
[[table.video_tasks.indexes]]
name = "video_tasks_user_status_idx"
columns = ["user_id", "status"]

4
crates/aether-data/schema/logical/005_wallet_billing.toml
View File

@@ -770,6 +770,10 @@ columns = ["user_entitlement_id", "request_id"]
name = "idx_entitlement_usage_user_date"
columns = ["user_id", "usage_date"]
[[table.entitlement_usage_ledgers.indexes]]
name = "idx_entitlement_usage_entitlement_date"
columns = ["user_entitlement_id", "usage_date"]
[[table.payment_callbacks.indexes]]
name = "idx_payment_callbacks_order"
columns = ["order_no"]

115
crates/aether-data/schema/logical/006_usage.toml
View File

@@ -554,6 +554,121 @@ columns = ["user_id"]
name = "usage_wallet_id_idx"
columns = ["wallet_id"]
[table.usage_counter_deltas]
domain = "usage"
order = 15
primary_key = ["id"]
[[table.usage_counter_deltas.columns]]
name = "id"
type = "text"
length = 36
[[table.usage_counter_deltas.columns]]
name = "request_id"
type = "text"
length = 128
[[table.usage_counter_deltas.columns]]
name = "kind"
type = "text"
length = 64
[[table.usage_counter_deltas.columns]]
name = "target_id"
type = "text"
[[table.usage_counter_deltas.columns]]
name = "request_count_delta"
type = "int64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "total_requests_delta"
type = "int64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "success_count_delta"
type = "int64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "error_count_delta"
type = "int64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "dns_failures_delta"
type = "int64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "stream_errors_delta"
type = "int64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "total_tokens_delta"
type = "int64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "total_cost_usd_delta"
type = "float64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "total_response_time_ms_delta"
type = "int64"
default = 0
[[table.usage_counter_deltas.columns]]
name = "last_used_at_unix_secs"
type = "unix_seconds"
nullable = true
[[table.usage_counter_deltas.columns]]
name = "last_used_ip"
type = "text"
nullable = true
[[table.usage_counter_deltas.columns]]
name = "candidate_last_used_at_unix_secs"
type = "unix_seconds"
nullable = true
[[table.usage_counter_deltas.columns]]
name = "removed_last_used_at_unix_secs"
type = "unix_seconds"
nullable = true
[[table.usage_counter_deltas.columns]]
name = "usage_created_at_unix_secs"
type = "unix_seconds"
nullable = true
[[table.usage_counter_deltas.columns]]
name = "created_at"
type = "timestamp"
[[table.usage_counter_deltas.columns]]
name = "processed_at"
type = "timestamp"
nullable = true
[[table.usage_counter_deltas.indexes]]
name = "ix_usage_counter_deltas_unprocessed"
columns = ["created_at", "id"]
[[table.usage_counter_deltas.indexes]]
name = "ix_usage_counter_deltas_processed"
columns = ["processed_at", "created_at", "id"]
[[table.usage_counter_deltas.indexes]]
name = "ix_usage_counter_deltas_request_kind"
columns = ["request_id", "kind", "target_id"]
[table.usage_settlement_snapshots]
domain = "usage"
order = 20

2
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260516000000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260518000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

9
crates/aether-data/src/lifecycle/migrate/tests.rs
View File

@@ -306,6 +306,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260512110000,
20260515000000,
20260516000000,
20260518000000,
]
);
}
@@ -377,6 +378,13 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role()
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains(
"ALTER TABLE public.stats_daily_model\n ADD COLUMN IF NOT EXISTS cache_creation_ephemeral_5m_tokens bigint DEFAULT '0'::bigint NOT NULL,"
));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL
.contains("CREATE TABLE IF NOT EXISTS public.usage_counter_deltas"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ix_usage_counter_deltas_unprocessed"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_entitlement_usage_entitlement_date"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_video_tasks_due_poll"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("request_count bigint DEFAULT 0"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL"));
}
#[test]
@@ -1123,6 +1131,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260512110000,
20260515000000,
20260516000000,
20260518000000,
]
);
}

77
crates/aether-data/src/repository/settlement/postgres.rs
View File

@@ -125,6 +125,22 @@ DO UPDATE SET
updated_at = NOW()
"#;
const ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL: &str = r#"
INSERT INTO usage_counter_deltas (
id,
request_id,
kind,
target_id,
total_cost_usd_delta
) VALUES (
$1,
$2,
'provider_monthly',
$3,
$4
)
"#;
#[derive(Debug, Clone)]
pub struct SqlxSettlementRepository {
tx_runner: PostgresTransactionRunner,
@@ -193,6 +209,37 @@ where
Ok(())
}
async fn enqueue_provider_monthly_usage_delta<'e, E>(
executor: E,
request_id: &str,
provider_id: &str,
total_cost_usd_delta: f64,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = sqlx::Postgres>,
{
let request_id = request_id.trim();
let provider_id = provider_id.trim();
if request_id.is_empty() || provider_id.is_empty() || total_cost_usd_delta == 0.0 {
return Ok(());
}
if !total_cost_usd_delta.is_finite() {
return Err(DataLayerError::UnexpectedValue(format!(
"provider monthly usage delta is not finite for {provider_id}"
)));
}
sqlx::query(ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL)
.bind(uuid::Uuid::new_v4().to_string())
.bind(request_id)
.bind(provider_id)
.bind(total_cost_usd_delta)
.execute(executor)
.await
.map_postgres_err()?;
Ok(())
}
#[derive(Debug, Default)]
struct DailyQuotaDebitResult {
debited_usd: f64,
@@ -663,23 +710,13 @@ WHERE id = $1
.as_deref()
.filter(|value| !value.is_empty())
{
let quota_row = sqlx::query(
r#"
UPDATE providers
SET
monthly_used_usd = COALESCE(monthly_used_usd, 0) + $2,
updated_at = NOW()
WHERE id = $1
RETURNING CAST(monthly_used_usd AS DOUBLE PRECISION) AS monthly_used_usd
"#,
enqueue_provider_monthly_usage_delta(
&mut **tx,
&input.request_id,
provider_id,
input.actual_total_cost_usd,
)
.bind(provider_id)
.bind(input.actual_total_cost_usd)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?;
settlement.provider_monthly_used_usd =
quota_row.and_then(|row| row.try_get("monthly_used_usd").ok());
.await?;
}
}
@@ -732,6 +769,14 @@ mod tests {
assert!(!source.contains("UPDATE \"usage\"\nSET\n wallet_id = $2"));
}
#[test]
fn settlement_sql_enqueues_provider_monthly_usage_delta() {
let source = include_str!("postgres.rs");
assert!(super::ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL.contains("usage_counter_deltas"));
assert!(super::ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL.contains("'provider_monthly'"));
assert!(!source.contains("UPDATE providers\nSET\n monthly_used_usd"));
}
#[test]
fn settlement_sql_blocks_standalone_key_owner_wallet_fallback() {
let source = include_str!("postgres.rs");

139
crates/aether-data/src/repository/usage/mod.rs
View File

@@ -363,15 +363,15 @@ mod sqlite;
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::usage::{
usage_request_metadata_client_family, PendingUsageCleanupSummary,
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow,
StoredUsageDashboardProviderCount, StoredUsageDashboardSummary,
StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
usage_request_metadata_client_family, ApiKeyLastUsedDelta, ManagementTokenCounterDelta,
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, ProxyNodeCounterDelta,
StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary,
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow,
StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow,
StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailySummary,
StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount,
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance,
StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary,
StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary,
@@ -380,7 +380,8 @@ pub(crate) use aether_data_contracts::repository::usage::{
UsageAuditListQuery, UsageAuditSummaryQuery, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy,
UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupPreviewCounts,
UsageCleanupSummary, UsageCleanupWindow, UsageCostSavingsSummaryQuery, UsageDailyHeatmapQuery,
UsageCleanupSummary, UsageCleanupWindow, UsageCostSavingsSummaryQuery,
UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageDailyHeatmapQuery,
UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery,
UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy,
UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
@@ -423,7 +424,10 @@ impl ApiKeyUsageDelta {
total_requests: after.total_requests - before.total_requests,
total_tokens: after.total_tokens - before.total_tokens,
total_cost_usd: after.total_cost_usd - before.total_cost_usd,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
candidate_last_used_at_unix_secs: newer_last_used_at(
before.last_used_at_unix_secs,
after.last_used_at_unix_secs,
),
removed_last_used_at_unix_secs: None,
}
}
@@ -530,7 +534,10 @@ impl ProviderApiKeyUsageDelta {
total_tokens: after.total_tokens - before.total_tokens,
total_cost_usd: after.total_cost_usd - before.total_cost_usd,
total_response_time_ms: after.total_response_time_ms - before.total_response_time_ms,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
candidate_last_used_at_unix_secs: newer_last_used_at(
before.last_used_at_unix_secs,
after.last_used_at_unix_secs,
),
removed_last_used_at_unix_secs: None,
usage_created_at_unix_secs: after.usage_created_at_unix_secs,
}
@@ -633,6 +640,9 @@ pub(crate) fn provider_api_key_usage_is_error(
pub(crate) fn provider_api_key_usage_contribution(
usage: &StoredRequestUsageAudit,
) -> Option<ProviderApiKeyUsageContribution> {
if matches!(usage.status.as_str(), "pending" | "streaming") {
return None;
}
let key_id = usage
.provider_api_key_id
.as_deref()
@@ -694,6 +704,9 @@ pub(crate) fn model_usage_contribution(
pub(crate) fn api_key_usage_contribution(
usage: &StoredRequestUsageAudit,
) -> Option<ApiKeyUsageContribution> {
if matches!(usage.status.as_str(), "pending" | "streaming") {
return None;
}
let api_key_id = usage
.api_key_id
.as_deref()
@@ -714,14 +727,23 @@ pub(crate) fn api_key_usage_contribution(
})
}
fn newer_last_used_at(before: Option<u64>, after: Option<u64>) -> Option<u64> {
match (before, after) {
(Some(before), Some(after)) if after > before => Some(after),
(None, Some(after)) => Some(after),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::{
api_key_usage_contribution, incoming_usage_can_recover_terminal_failure,
model_usage_contribution, provider_api_key_usage_contribution,
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure, ModelUsageDelta,
StoredRequestUsageAudit, UpsertUsageRecord,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
ApiKeyUsageDelta, ModelUsageDelta, ProviderApiKeyUsageDelta, StoredRequestUsageAudit,
UpsertUsageRecord,
};
#[test]
@@ -996,6 +1018,95 @@ mod tests {
assert_eq!(contribution.total_tokens, 20);
assert_eq!(contribution.total_cost_usd, 0.25);
assert_eq!(contribution.last_used_at_unix_secs, Some(123));
let mut streaming = usage.clone();
streaming.status = "streaming".to_string();
assert!(api_key_usage_contribution(&streaming).is_none());
let mut pending = usage;
pending.status = "pending".to_string();
assert!(api_key_usage_contribution(&pending).is_none());
}
#[test]
fn provider_api_key_usage_contribution_tracks_terminal_requests_only() {
let usage = StoredRequestUsageAudit::new(
"usage-1".to_string(),
"request-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
None,
None,
"OpenAI".to_string(),
"gpt-5".to_string(),
None,
Some("provider-1".to_string()),
None,
Some("provider-key-1".to_string()),
None,
None,
None,
None,
None,
None,
None,
false,
false,
12,
8,
20,
0.25,
0.25,
Some(200),
None,
None,
Some(120),
None,
"completed".to_string(),
"settled".to_string(),
123,
124,
Some(125),
)
.expect("usage should build");
assert!(provider_api_key_usage_contribution(&usage).is_some());
let mut streaming = usage.clone();
streaming.status = "streaming".to_string();
assert!(provider_api_key_usage_contribution(&streaming).is_none());
let mut pending = usage;
pending.status = "pending".to_string();
assert!(provider_api_key_usage_contribution(&pending).is_none());
}
#[test]
fn usage_delta_between_does_not_emit_duplicate_last_used_candidate() {
let api_key_contribution = super::ApiKeyUsageContribution {
api_key_id: "api-key-1".to_string(),
total_requests: 1,
total_tokens: 20,
total_cost_usd: 0.25,
last_used_at_unix_secs: Some(123),
};
assert!(ApiKeyUsageDelta::between(&api_key_contribution, &api_key_contribution).is_noop());
let provider_contribution = super::ProviderApiKeyUsageContribution {
key_id: "provider-key-1".to_string(),
request_count: 1,
success_count: 1,
error_count: 0,
total_tokens: 20,
total_cost_usd: 0.25,
total_response_time_ms: 120,
last_used_at_unix_secs: Some(123),
usage_created_at_unix_secs: Some(123),
};
assert!(
ProviderApiKeyUsageDelta::between(&provider_contribution, &provider_contribution,)
.is_noop()
);
}
#[test]

1143
crates/aether-data/src/repository/usage/postgres/mod.rs
View File

File diff suppressed because it is too large Load Diff

1
crates/aether-data/src/repository/usage/postgres/queries/apply_api_key_usage_delta_sql.sql
View File

@@ -17,6 +17,7 @@ SET
SELECT MAX(created_at)
FROM "usage"
WHERE api_key_id = $1
AND status NOT IN ('pending', 'streaming')
)
ELSE last_used_at
END

170
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_codex_window_usage_delta_sql.sql
View File

@@ -1,170 +0,0 @@
WITH target_key AS (
SELECT
id,
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
FROM provider_api_keys
WHERE id = $1
AND jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
FOR UPDATE
),
window_items AS (
SELECT
target_key.id,
window_rows.window_item,
window_rows.window_ordinality
FROM target_key
CROSS JOIN LATERAL jsonb_array_elements(target_key.snapshot -> 'quota' -> 'windows')
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
),
parsed_windows AS (
SELECT
window_items.id,
window_items.window_item,
window_items.window_ordinality,
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
CASE
WHEN text_values.reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.reset_at_text) < 19
OR (
length(text_values.reset_at_text) = 19
AND text_values.reset_at_text <= '9223372036854775807'
)
)
THEN text_values.reset_at_text::BIGINT
ELSE NULL
END AS reset_at,
CASE
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
AND (
length(text_values.window_minutes_text) < 19
OR (
length(text_values.window_minutes_text) = 19
AND text_values.window_minutes_text <= '9223372036854775807'
)
)
THEN text_values.window_minutes_text::BIGINT
ELSE CASE lower(BTRIM(COALESCE(window_items.window_item ->> 'code', '')))
WHEN '5h' THEN 300
WHEN 'weekly' THEN 10080
ELSE NULL
END
END AS window_minutes,
CASE
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.usage_reset_at_text) < 19
OR (
length(text_values.usage_reset_at_text) = 19
AND text_values.usage_reset_at_text <= '9223372036854775807'
)
)
THEN text_values.usage_reset_at_text::BIGINT
ELSE NULL
END AS usage_reset_at,
CASE
WHEN text_values.request_count_text ~ '^[0-9]+$'
AND (
length(text_values.request_count_text) < 19
OR (
length(text_values.request_count_text) = 19
AND text_values.request_count_text <= '9223372036854775807'
)
)
THEN text_values.request_count_text::BIGINT
ELSE 0
END AS current_request_count,
CASE
WHEN text_values.total_tokens_text ~ '^[0-9]+$'
AND (
length(text_values.total_tokens_text) < 19
OR (
length(text_values.total_tokens_text) = 19
AND text_values.total_tokens_text <= '9223372036854775807'
)
)
THEN text_values.total_tokens_text::BIGINT
ELSE 0
END AS current_total_tokens,
CASE
WHEN text_values.total_cost_usd_text ~ '^[-+]?[0-9]+([.][0-9]+)?$'
THEN text_values.total_cost_usd_text::DOUBLE PRECISION
ELSE 0
END AS current_total_cost_usd
FROM window_items
CROSS JOIN LATERAL (
SELECT
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'request_count', '')) AS request_count_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_tokens', '')) AS total_tokens_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_cost_usd', '')) AS total_cost_usd_text
) AS text_values
),
window_usage AS (
SELECT
parsed_windows.*,
CASE
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
THEN parsed_windows.window_minutes * 60
ELSE NULL
END AS window_seconds
FROM parsed_windows
),
updated_windows AS (
SELECT
window_usage.id,
jsonb_agg(
CASE
WHEN window_usage.window_code IN ('5h', 'weekly')
AND window_usage.reset_at IS NOT NULL
AND window_usage.window_seconds IS NOT NULL
AND window_usage.reset_at >= window_usage.window_seconds
AND window_usage.reset_at > $2
AND GREATEST(
window_usage.reset_at - window_usage.window_seconds,
COALESCE(window_usage.usage_reset_at, 0)
) <= $2
THEN jsonb_set(
window_usage.window_item,
'{usage}',
jsonb_build_object(
'request_count',
LEAST(
GREATEST(window_usage.current_request_count::NUMERIC + $3::NUMERIC, 0),
9223372036854775807
)::BIGINT,
'total_tokens',
LEAST(
GREATEST(window_usage.current_total_tokens::NUMERIC + $4::NUMERIC, 0),
9223372036854775807
)::BIGINT,
'total_cost_usd',
to_char(
GREATEST(COALESCE(window_usage.current_total_cost_usd, 0) + $5, 0),
'FM999999999999999990.00000000'
)
),
true
)
ELSE window_usage.window_item
END
ORDER BY window_usage.window_ordinality
) AS windows
FROM window_usage
GROUP BY window_usage.id
)
UPDATE provider_api_keys AS keys
SET
status_snapshot = jsonb_set(
target_key.snapshot,
'{quota,windows}',
updated_windows.windows,
true
)::json,
updated_at = NOW()
FROM target_key
JOIN updated_windows ON updated_windows.id = target_key.id
WHERE keys.id = target_key.id

1
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_usage_delta_sql.sql
View File

@@ -20,6 +20,7 @@ SET
SELECT MAX(created_at)
FROM "usage"
WHERE provider_api_key_id = $1
AND status NOT IN ('pending', 'streaming')
)
ELSE last_used_at
END

3
crates/aether-data/src/repository/usage/postgres/queries/rebuild_api_key_usage_stats_sql.sql
View File

@@ -1,7 +1,7 @@
WITH aggregated AS (
SELECT
api_key_id,
COUNT(*)::INTEGER AS total_requests,
COUNT(*)::BIGINT AS total_requests,
COALESCE(SUM(
GREATEST(
COALESCE(
@@ -16,6 +16,7 @@ WITH aggregated AS (
FROM usage_billing_facts AS "usage"
WHERE api_key_id IS NOT NULL
AND BTRIM(api_key_id) <> ''
AND status NOT IN ('pending', 'streaming')
GROUP BY api_key_id
)
UPDATE api_keys

9
crates/aether-data/src/repository/usage/postgres/queries/rebuild_provider_api_key_usage_stats_sql.sql
View File

@@ -1,7 +1,7 @@
WITH aggregated AS (
SELECT
provider_api_key_id,
COUNT(*)::INTEGER AS request_count,
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(
CASE
WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled')
@@ -10,7 +10,7 @@ WITH aggregated AS (
THEN 1
ELSE 0
END
), 0)::INTEGER AS success_count,
), 0)::BIGINT AS success_count,
COALESCE(SUM(
CASE
WHEN status NOT IN ('pending', 'streaming')
@@ -22,7 +22,7 @@ WITH aggregated AS (
THEN 1
ELSE 0
END
), 0)::INTEGER AS error_count,
), 0)::BIGINT AS error_count,
COALESCE(SUM(
GREATEST(
COALESCE(
@@ -42,11 +42,12 @@ WITH aggregated AS (
THEN GREATEST(response_time_ms, 0)
ELSE 0
END
), 0)::INTEGER AS total_response_time_ms,
), 0)::BIGINT AS total_response_time_ms,
MAX(created_at) AS last_used_at
FROM usage_billing_facts AS "usage"
WHERE provider_api_key_id IS NOT NULL
AND BTRIM(provider_api_key_id) <> ''
AND status NOT IN ('pending', 'streaming')
GROUP BY provider_api_key_id
)
UPDATE provider_api_keys

52
crates/aether-data/src/repository/usage/postgres/tests.rs
View File

@@ -234,25 +234,6 @@ fn usage_sql_summarizes_usage_by_provider_api_key_ids_in_database() {
assert!(super::SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
}
#[test]
fn usage_sql_materializes_provider_key_window_usage_in_status_snapshot() {
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("UPDATE provider_api_keys AS keys"));
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("jsonb_set"));
assert!(
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'{quota,windows}'")
);
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'usage'"));
assert!(
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("WHEN '5h' THEN 300")
);
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("WHEN 'weekly' THEN 10080"));
assert!(
!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("usage_billing_facts")
);
}
#[test]
fn usage_sql_rebuilds_provider_key_window_usage_into_status_snapshot() {
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
@@ -280,18 +261,44 @@ fn usage_sql_serializes_request_id_upserts_before_reading_previous_usage() {
.contains("lock_usage_request_id_in_tx(tx, &usage.request_id).await?;"));
}
#[test]
fn usage_sql_moves_shared_counter_updates_behind_outbox() {
let source = include_str!("mod.rs");
assert!(super::INSERT_USAGE_COUNTER_DELTA_SQL.contains("usage_counter_deltas"));
assert!(super::CLAIM_USAGE_COUNTER_DELTAS_SQL.contains("FOR UPDATE SKIP LOCKED"));
assert!(super::MARK_USAGE_COUNTER_DELTAS_PROCESSED_SQL.contains("processed_at = NOW()"));
assert!(super::TRY_LOCK_USAGE_COUNTER_FLUSH_SQL.contains("pg_try_advisory_xact_lock"));
assert!(source.contains("enqueue_api_key_usage_delta_in_tx("));
assert!(source.contains("enqueue_provider_api_key_usage_delta_in_tx("));
assert!(source.contains("enqueue_model_usage_delta_in_tx("));
assert!(source.contains("apply_provider_api_key_main_usage_delta_in_tx("));
assert!(source.contains("USAGE_COUNTER_KIND_PROVIDER_MONTHLY"));
assert!(source.contains("apply_provider_monthly_usage_delta_in_tx("));
}
#[test]
fn usage_sql_exposes_counter_outbox_health() {
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("pending_rows"));
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("oldest_pending_created_at_unix_secs"));
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("latest_processed_at_unix_secs"));
assert!(super::READ_PENDING_USAGE_COUNTER_DELTAS_BY_KIND_SQL.contains("GROUP BY kind"));
}
#[test]
fn usage_sql_rebuild_matches_online_api_key_usage_semantics() {
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::INTEGER"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::BIGINT"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COALESCE("));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("total_tokens,"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("AND BTRIM(api_key_id) <> ''"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL
.contains("AND status NOT IN ('pending', 'streaming')"));
}
#[test]
fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::BIGINT"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("NULLIF(BTRIM(error_message), '') IS NULL"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("COALESCE("));
@@ -300,6 +307,8 @@ fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("AND BTRIM(provider_api_key_id) <> ''"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("AND status NOT IN ('pending', 'streaming')"));
}
#[test]
@@ -503,8 +512,7 @@ fn usage_sql_raw_aggregates_use_canonical_billing_facts() {
.contains("FROM usage_billing_facts AS \"usage\""));
assert!(super::SUMMARIZE_USAGE_TOTALS_BY_USER_IDS_SQL
.contains("FROM usage_billing_facts AS \"usage\""));
assert!(!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("usage_billing_facts AS \"usage\""));
assert!(!source.contains("apply_provider_api_key_codex_window_usage_delta_in_tx"));
}
#[test]

1
crates/aether-testkit/Cargo.toml
View File

@@ -9,6 +9,7 @@ description = "Shared integration test helpers for Aether Rust services"
[dependencies]
async-stream.workspace = true
aether-data.workspace = true
aether-data-contracts.workspace = true
aether-contracts.workspace = true
aether-gateway.workspace = true
aether-http.workspace = true

664
crates/aether-testkit/src/bin/usage_aux_counter_hotspot_baseline.rs Normal file
View File

@@ -0,0 +1,664 @@
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use aether_data::repository::usage::SqlxUsageReadRepository;
use aether_data_contracts::repository::usage::{
ApiKeyLastUsedDelta, ManagementTokenCounterDelta, ProxyNodeCounterDelta,
};
use aether_testkit::{
init_test_runtime_for, prepare_aether_postgres_schema, ManagedPostgresServer,
};
use serde::Serialize;
use sqlx::postgres::PgPoolOptions;
use sqlx::{PgPool, Row};
use tokio::sync::Mutex;
const PROXY_NODE_ID: &str = "proxy-node-hotspot";
const MANAGEMENT_TOKEN_ID: &str = "management-token-hotspot";
const API_KEY_ID: &str = "api-key-last-used-hotspot";
const USER_ID: &str = "usage-aux-hotspot-user";
#[derive(Debug, Clone)]
struct Config {
requests: usize,
concurrency: usize,
max_connections: u32,
flush_batch_size: usize,
flush_interval: Duration,
monitor_interval: Duration,
output_path: Option<PathBuf>,
postgres_url: Option<String>,
}
impl Default for Config {
fn default() -> Self {
Self {
requests: 2_000,
concurrency: 100,
max_connections: 64,
flush_batch_size: 1_000,
flush_interval: Duration::from_millis(100),
monitor_interval: Duration::from_millis(100),
output_path: None,
postgres_url: None,
}
}
}
#[derive(Debug, Serialize)]
struct Report {
suite: &'static str,
config: ReportConfig,
duration_ms: u64,
throughput_rps: u64,
completed_requests: usize,
failed_requests: usize,
p50_ms: u64,
p95_ms: u64,
max_ms: u64,
mean_ms: u64,
flush: FlushReport,
counters: CounterReport,
lock_monitor: LockMonitorReport,
}
#[derive(Debug, Serialize)]
struct ReportConfig {
requests: usize,
concurrency: usize,
max_connections: u32,
flush_batch_size: usize,
flush_interval_ms: u64,
monitor_interval_ms: u64,
managed_postgres: bool,
}
#[derive(Debug, Serialize, Default)]
struct FlushReport {
calls: usize,
rows_claimed: usize,
proxy_node_targets: usize,
management_token_targets: usize,
api_key_last_used_targets: usize,
}
#[derive(Debug, Serialize)]
struct CounterReport {
outbox_pending_rows: i64,
outbox_processed_rows: i64,
proxy_total_requests: i64,
proxy_failed_requests: i64,
proxy_dns_failures: i64,
proxy_stream_errors: i64,
expected_proxy_total_requests: i64,
expected_proxy_failed_requests: i64,
expected_proxy_dns_failures: i64,
expected_proxy_stream_errors: i64,
management_token_usage_count: i64,
expected_management_token_usage_count: i64,
api_key_last_used_at: Option<i64>,
expected_min_api_key_last_used_at: i64,
}
#[derive(Debug, Serialize, Clone, Copy, Default)]
struct LockMonitorReport {
samples: usize,
max_lock_waiters: i64,
max_proxy_node_update_waiters: i64,
max_management_token_update_waiters: i64,
max_api_key_update_waiters: i64,
max_oldest_lock_wait_ms: i64,
}
#[derive(Debug, Clone, Copy, Default)]
struct LockSample {
lock_waiters: i64,
proxy_node_update_waiters: i64,
management_token_update_waiters: i64,
api_key_update_waiters: i64,
oldest_lock_wait_ms: i64,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("usage-aux-counter-hotspot-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let managed_postgres;
let database_url;
let _server;
if let Some(url) = config.postgres_url.as_ref() {
managed_postgres = false;
database_url = url.clone();
_server = None;
} else {
managed_postgres = true;
let server = ManagedPostgresServer::start().await?;
database_url = server.database_url().to_string();
_server = Some(server);
}
prepare_aether_postgres_schema(&database_url).await?;
let pool = PgPoolOptions::new()
.min_connections(1)
.max_connections(config.max_connections)
.acquire_timeout(Duration::from_secs(10))
.connect(&database_url)
.await?;
seed_hotspot_rows(&pool).await?;
let repository = SqlxUsageReadRepository::new(pool.clone());
let stop = Arc::new(AtomicBool::new(false));
let flush_report = Arc::new(Mutex::new(FlushReport::default()));
let lock_report = Arc::new(Mutex::new(LockMonitorReport::default()));
let flush_handle = spawn_flush_loop(
repository.clone(),
Arc::clone(&stop),
Arc::clone(&flush_report),
config.flush_batch_size,
config.flush_interval,
);
let monitor_handle = spawn_lock_monitor(
pool.clone(),
Arc::clone(&stop),
Arc::clone(&lock_report),
config.monitor_interval,
);
let started_at = Instant::now();
let load_result =
run_aux_counter_load(repository.clone(), config.requests, config.concurrency).await;
wait_for_outbox_drain(&pool, &repository, config.flush_batch_size).await?;
let duration_ms = started_at.elapsed().as_millis() as u64;
stop.store(true, Ordering::Release);
flush_handle.await??;
monitor_handle.await??;
let mut flush = flush_report.lock().await;
let final_flush = repository
.flush_usage_counter_deltas(config.flush_batch_size)
.await?;
flush.calls += 1;
flush.rows_claimed += final_flush.rows_claimed;
flush.proxy_node_targets += final_flush.proxy_node_targets;
flush.management_token_targets += final_flush.management_token_targets;
flush.api_key_last_used_targets += final_flush.api_key_last_used_targets;
drop(flush);
let counters = read_counters(&pool, config.requests).await?;
let latencies = load_result.latencies.lock().await.clone();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(latencies);
let completed_requests = load_result.completed.load(Ordering::Acquire);
let throughput_rps = if duration_ms == 0 {
completed_requests as u64
} else {
((completed_requests as u64) * 1_000) / duration_ms.max(1)
};
let report = Report {
suite: "usage_aux_counter_hotspot_baseline",
config: ReportConfig {
requests: config.requests,
concurrency: config.concurrency,
max_connections: config.max_connections,
flush_batch_size: config.flush_batch_size,
flush_interval_ms: config.flush_interval.as_millis() as u64,
monitor_interval_ms: config.monitor_interval.as_millis() as u64,
managed_postgres,
},
duration_ms,
throughput_rps,
completed_requests,
failed_requests: load_result.failed.load(Ordering::Acquire),
p50_ms,
p95_ms,
max_ms,
mean_ms,
flush: Arc::try_unwrap(flush_report)
.unwrap_or_else(|_| panic!("flush report still referenced"))
.into_inner(),
counters,
lock_monitor: Arc::try_unwrap(lock_report)
.unwrap_or_else(|_| panic!("lock report still referenced"))
.into_inner(),
};
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
#[derive(Clone)]
struct LoadResult {
completed: Arc<AtomicUsize>,
failed: Arc<AtomicUsize>,
latencies: Arc<Mutex<Vec<u64>>>,
}
async fn run_aux_counter_load(
repository: SqlxUsageReadRepository,
requests: usize,
concurrency: usize,
) -> LoadResult {
let next = Arc::new(AtomicUsize::new(0));
let completed = Arc::new(AtomicUsize::new(0));
let failed = Arc::new(AtomicUsize::new(0));
let latencies = Arc::new(Mutex::new(Vec::with_capacity(requests)));
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..concurrency {
let repository = repository.clone();
let next = Arc::clone(&next);
let completed = Arc::clone(&completed);
let failed = Arc::clone(&failed);
let latencies = Arc::clone(&latencies);
tasks.spawn(async move {
loop {
let index = next.fetch_add(1, Ordering::AcqRel);
if index >= requests {
break;
}
let started_at = Instant::now();
let result = enqueue_aux_counter_deltas(&repository, index).await;
latencies
.lock()
.await
.push(started_at.elapsed().as_millis() as u64);
completed.fetch_add(1, Ordering::AcqRel);
if result.is_err() {
failed.fetch_add(1, Ordering::AcqRel);
}
}
});
}
while let Some(result) = tasks.join_next().await {
if result.is_err() {
failed.fetch_add(1, Ordering::AcqRel);
}
}
LoadResult {
completed,
failed,
latencies,
}
}
async fn enqueue_aux_counter_deltas(
repository: &SqlxUsageReadRepository,
index: usize,
) -> Result<(), aether_data::DataLayerError> {
repository
.enqueue_proxy_node_counter_delta(proxy_delta_for_index(index))
.await?;
repository
.enqueue_management_token_counter_delta(ManagementTokenCounterDelta {
token_id: MANAGEMENT_TOKEN_ID.to_string(),
usage_count_delta: 1,
last_used_at_unix_secs: Some(now_unix_secs().saturating_add(index as u64)),
last_used_ip: Some(format!("10.0.0.{}", (index % 250) + 1)),
})
.await?;
repository
.enqueue_api_key_last_used_delta(ApiKeyLastUsedDelta {
api_key_id: API_KEY_ID.to_string(),
last_used_at_unix_secs: now_unix_secs().saturating_add(index as u64),
})
.await?;
Ok(())
}
fn proxy_delta_for_index(index: usize) -> ProxyNodeCounterDelta {
ProxyNodeCounterDelta {
node_id: PROXY_NODE_ID.to_string(),
total_requests_delta: 1,
failed_requests_delta: if index.is_multiple_of(10) { 1 } else { 0 },
dns_failures_delta: if index.is_multiple_of(25) { 1 } else { 0 },
stream_errors_delta: if index.is_multiple_of(40) { 1 } else { 0 },
}
}
fn spawn_flush_loop(
repository: SqlxUsageReadRepository,
stop: Arc<AtomicBool>,
report: Arc<Mutex<FlushReport>>,
batch_size: usize,
interval: Duration,
) -> tokio::task::JoinHandle<Result<(), aether_data::DataLayerError>> {
tokio::spawn(async move {
while !stop.load(Ordering::Acquire) {
let summary = repository.flush_usage_counter_deltas(batch_size).await?;
let mut report = report.lock().await;
report.calls += 1;
report.rows_claimed += summary.rows_claimed;
report.proxy_node_targets += summary.proxy_node_targets;
report.management_token_targets += summary.management_token_targets;
report.api_key_last_used_targets += summary.api_key_last_used_targets;
drop(report);
tokio::time::sleep(interval).await;
}
Ok(())
})
}
fn spawn_lock_monitor(
pool: PgPool,
stop: Arc<AtomicBool>,
report: Arc<Mutex<LockMonitorReport>>,
interval: Duration,
) -> tokio::task::JoinHandle<Result<(), sqlx::Error>> {
tokio::spawn(async move {
while !stop.load(Ordering::Acquire) {
let sample = read_lock_sample(&pool).await?;
let mut report = report.lock().await;
report.samples += 1;
report.max_lock_waiters = report.max_lock_waiters.max(sample.lock_waiters);
report.max_proxy_node_update_waiters = report
.max_proxy_node_update_waiters
.max(sample.proxy_node_update_waiters);
report.max_management_token_update_waiters = report
.max_management_token_update_waiters
.max(sample.management_token_update_waiters);
report.max_api_key_update_waiters = report
.max_api_key_update_waiters
.max(sample.api_key_update_waiters);
report.max_oldest_lock_wait_ms = report
.max_oldest_lock_wait_ms
.max(sample.oldest_lock_wait_ms);
drop(report);
tokio::time::sleep(interval).await;
}
Ok(())
})
}
async fn wait_for_outbox_drain(
pool: &PgPool,
repository: &SqlxUsageReadRepository,
batch_size: usize,
) -> Result<(), Box<dyn std::error::Error>> {
let deadline = Instant::now() + Duration::from_secs(30);
loop {
let summary = repository.flush_usage_counter_deltas(batch_size).await?;
let pending: i64 = sqlx::query_scalar(
"SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE processed_at IS NULL",
)
.fetch_one(pool)
.await?;
if pending == 0 && summary.rows_claimed == 0 {
return Ok(());
}
if Instant::now() >= deadline {
return Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
format!("usage counter outbox did not drain; pending={pending}"),
)
.into());
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
async fn seed_hotspot_rows(pool: &PgPool) -> Result<(), sqlx::Error> {
sqlx::query(
r#"
INSERT INTO users (id, username, email_verified)
VALUES ($1, 'usage-aux-hotspot', TRUE)
ON CONFLICT (id) DO NOTHING
"#,
)
.bind(USER_ID)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO api_keys (id, user_id, key_hash, name, is_active)
VALUES ($1, $2, 'usage-aux-hotspot-hash', 'usage aux hotspot', TRUE)
ON CONFLICT (id) DO UPDATE SET
last_used_at = NULL
"#,
)
.bind(API_KEY_ID)
.bind(USER_ID)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO management_tokens (
id, user_id, name, token_hash, token_prefix, usage_count
)
VALUES ($1, $2, 'usage aux hotspot token', 'usage-aux-hotspot-token-hash', 'usageaux', 0)
ON CONFLICT (id) DO UPDATE SET
usage_count = 0,
last_used_at = NULL,
last_used_ip = NULL
"#,
)
.bind(MANAGEMENT_TOKEN_ID)
.bind(USER_ID)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO proxy_nodes (
id, name, ip, port, status, total_requests, failed_requests,
dns_failures, stream_errors
)
VALUES ($1, 'usage aux hotspot proxy', '127.0.0.1', 8080, 'online', 0, 0, 0, 0)
ON CONFLICT (id) DO UPDATE SET
total_requests = 0,
failed_requests = 0,
dns_failures = 0,
stream_errors = 0
"#,
)
.bind(PROXY_NODE_ID)
.execute(pool)
.await?;
sqlx::query(
r#"
DELETE FROM usage_counter_deltas
WHERE target_id IN ($1, $2, $3)
"#,
)
.bind(PROXY_NODE_ID)
.bind(MANAGEMENT_TOKEN_ID)
.bind(API_KEY_ID)
.execute(pool)
.await?;
Ok(())
}
async fn read_lock_sample(pool: &PgPool) -> Result<LockSample, sqlx::Error> {
let row = sqlx::query(
r#"
SELECT
COUNT(*) FILTER (WHERE wait_event_type = 'Lock')::BIGINT AS lock_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE 'UPDATE proxy_nodes%'
)::BIGINT AS proxy_node_update_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE 'UPDATE management_tokens%'
)::BIGINT AS management_token_update_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE 'UPDATE api_keys%'
)::BIGINT AS api_key_update_waiters,
COALESCE(
MAX(EXTRACT(EPOCH FROM (NOW() - query_start)) * 1000)
FILTER (WHERE wait_event_type = 'Lock'),
0
)::BIGINT AS oldest_lock_wait_ms
FROM pg_stat_activity
WHERE datname = current_database()
"#,
)
.fetch_one(pool)
.await?;
Ok(LockSample {
lock_waiters: row.try_get("lock_waiters")?,
proxy_node_update_waiters: row.try_get("proxy_node_update_waiters")?,
management_token_update_waiters: row.try_get("management_token_update_waiters")?,
api_key_update_waiters: row.try_get("api_key_update_waiters")?,
oldest_lock_wait_ms: row.try_get("oldest_lock_wait_ms")?,
})
}
async fn read_counters(pool: &PgPool, requests: usize) -> Result<CounterReport, sqlx::Error> {
let row = sqlx::query(
r#"
SELECT
(SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE processed_at IS NULL) AS outbox_pending_rows,
(SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE processed_at IS NOT NULL) AS outbox_processed_rows,
(SELECT total_requests::BIGINT FROM proxy_nodes WHERE id = $1) AS proxy_total_requests,
(SELECT failed_requests::BIGINT FROM proxy_nodes WHERE id = $1) AS proxy_failed_requests,
(SELECT dns_failures::BIGINT FROM proxy_nodes WHERE id = $1) AS proxy_dns_failures,
(SELECT stream_errors::BIGINT FROM proxy_nodes WHERE id = $1) AS proxy_stream_errors,
(SELECT usage_count::BIGINT FROM management_tokens WHERE id = $2) AS management_token_usage_count,
(
SELECT EXTRACT(EPOCH FROM last_used_at)::BIGINT
FROM api_keys
WHERE id = $3
) AS api_key_last_used_at
"#,
)
.bind(PROXY_NODE_ID)
.bind(MANAGEMENT_TOKEN_ID)
.bind(API_KEY_ID)
.fetch_one(pool)
.await?;
Ok(CounterReport {
outbox_pending_rows: row.try_get("outbox_pending_rows")?,
outbox_processed_rows: row.try_get("outbox_processed_rows")?,
proxy_total_requests: row.try_get("proxy_total_requests")?,
proxy_failed_requests: row.try_get("proxy_failed_requests")?,
proxy_dns_failures: row.try_get("proxy_dns_failures")?,
proxy_stream_errors: row.try_get("proxy_stream_errors")?,
expected_proxy_total_requests: requests as i64,
expected_proxy_failed_requests: count_every(requests, 10),
expected_proxy_dns_failures: count_every(requests, 25),
expected_proxy_stream_errors: count_every(requests, 40),
management_token_usage_count: row.try_get("management_token_usage_count")?,
expected_management_token_usage_count: requests as i64,
api_key_last_used_at: row.try_get("api_key_last_used_at")?,
expected_min_api_key_last_used_at: now_unix_secs().saturating_sub(5) as i64,
})
}
fn count_every(requests: usize, interval: usize) -> i64 {
if requests == 0 {
return 0;
}
((requests - 1) / interval + 1) as i64
}
fn summarize_latencies(mut latencies: Vec<u64>) -> (u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
}
latencies.sort_unstable();
let p50 = percentile(&latencies, 50);
let p95 = percentile(&latencies, 95);
let max = *latencies.last().unwrap_or(&0);
let mean = latencies.iter().sum::<u64>() / latencies.len() as u64;
(p50, p95, max, mean)
}
fn percentile(latencies: &[u64], percentile: usize) -> u64 {
let index = ((latencies.len() - 1) * percentile) / 100;
latencies[index]
}
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
let mut config = Config::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--requests" => config.requests = next_value(&mut iter, "--requests")?.parse()?,
"--concurrency" => {
config.concurrency = next_value(&mut iter, "--concurrency")?.parse()?
}
"--max-connections" => {
config.max_connections = next_value(&mut iter, "--max-connections")?.parse()?
}
"--flush-batch-size" => {
config.flush_batch_size = next_value(&mut iter, "--flush-batch-size")?.parse()?
}
"--flush-interval-ms" => {
config.flush_interval =
Duration::from_millis(next_value(&mut iter, "--flush-interval-ms")?.parse()?)
}
"--monitor-interval-ms" => {
config.monitor_interval =
Duration::from_millis(next_value(&mut iter, "--monitor-interval-ms")?.parse()?)
}
"--postgres-url" => {
config.postgres_url = Some(next_value(&mut iter, "--postgres-url")?)
}
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
if config.requests == 0 || config.concurrency == 0 || config.max_connections == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"requests, concurrency, and max-connections must be positive",
)
.into());
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin usage_aux_counter_hotspot_baseline -- [--requests 2000] [--concurrency 100] [--max-connections 64] [--flush-batch-size 1000] [--flush-interval-ms 100] [--monitor-interval-ms 100] [--postgres-url postgres://...] [--output /tmp/usage_aux_counter_hotspot.json]"
);
}

684
crates/aether-testkit/src/bin/usage_counter_hotspot_baseline.rs Normal file
View File

@@ -0,0 +1,684 @@
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use aether_data::repository::usage::SqlxUsageReadRepository;
use aether_data_contracts::repository::usage::UpsertUsageRecord;
use aether_testkit::{
init_test_runtime_for, prepare_aether_postgres_schema, ManagedPostgresServer,
};
use serde::Serialize;
use serde_json::json;
use sqlx::postgres::PgPoolOptions;
use sqlx::{PgPool, Row};
use tokio::sync::Mutex;
#[derive(Debug, Clone)]
struct Config {
requests: usize,
concurrency: usize,
max_connections: u32,
flush_batch_size: usize,
flush_interval: Duration,
monitor_interval: Duration,
output_path: Option<PathBuf>,
postgres_url: Option<String>,
}
impl Default for Config {
fn default() -> Self {
Self {
requests: 2_000,
concurrency: 100,
max_connections: 64,
flush_batch_size: 1_000,
flush_interval: Duration::from_millis(100),
monitor_interval: Duration::from_millis(100),
output_path: None,
postgres_url: None,
}
}
}
#[derive(Debug, Serialize)]
struct Report {
suite: &'static str,
config: ReportConfig,
duration_ms: u64,
throughput_rps: u64,
completed_requests: usize,
failed_requests: usize,
p50_ms: u64,
p95_ms: u64,
max_ms: u64,
mean_ms: u64,
flush: FlushReport,
counters: CounterReport,
lock_monitor: LockMonitorReport,
}
#[derive(Debug, Serialize)]
struct ReportConfig {
requests: usize,
concurrency: usize,
max_connections: u32,
flush_batch_size: usize,
flush_interval_ms: u64,
monitor_interval_ms: u64,
managed_postgres: bool,
}
#[derive(Debug, Serialize, Default)]
struct FlushReport {
calls: usize,
rows_claimed: usize,
api_key_targets: usize,
provider_api_key_targets: usize,
model_targets: usize,
provider_monthly_targets: usize,
proxy_node_targets: usize,
management_token_targets: usize,
api_key_last_used_targets: usize,
}
#[derive(Debug, Serialize)]
struct CounterReport {
usage_rows: i64,
outbox_pending_rows: i64,
outbox_processed_rows: i64,
api_key_total_requests: i64,
api_key_total_tokens: i64,
provider_key_request_count: i64,
provider_key_success_count: i64,
global_model_usage_count: i64,
}
#[derive(Debug, Serialize, Clone, Copy, Default)]
struct LockMonitorReport {
samples: usize,
max_lock_waiters: i64,
max_api_key_update_waiters: i64,
max_provider_key_update_waiters: i64,
max_global_model_update_waiters: i64,
max_provider_update_waiters: i64,
max_oldest_lock_wait_ms: i64,
}
#[derive(Debug, Clone, Copy, Default)]
struct LockSample {
lock_waiters: i64,
api_key_update_waiters: i64,
provider_key_update_waiters: i64,
global_model_update_waiters: i64,
provider_update_waiters: i64,
oldest_lock_wait_ms: i64,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("usage-counter-hotspot-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let managed_postgres;
let database_url;
let _server;
if let Some(url) = config.postgres_url.as_ref() {
managed_postgres = false;
database_url = url.clone();
_server = None;
} else {
managed_postgres = true;
let server = ManagedPostgresServer::start().await?;
database_url = server.database_url().to_string();
_server = Some(server);
}
prepare_aether_postgres_schema(&database_url).await?;
let pool = PgPoolOptions::new()
.min_connections(1)
.max_connections(config.max_connections)
.acquire_timeout(Duration::from_secs(10))
.connect(&database_url)
.await?;
seed_hotspot_rows(&pool).await?;
let repository = SqlxUsageReadRepository::new(pool.clone());
let stop = Arc::new(AtomicBool::new(false));
let flush_report = Arc::new(Mutex::new(FlushReport::default()));
let lock_report = Arc::new(Mutex::new(LockMonitorReport::default()));
let flush_handle = spawn_flush_loop(
repository.clone(),
Arc::clone(&stop),
Arc::clone(&flush_report),
config.flush_batch_size,
config.flush_interval,
);
let monitor_handle = spawn_lock_monitor(
pool.clone(),
Arc::clone(&stop),
Arc::clone(&lock_report),
config.monitor_interval,
);
let started_at = Instant::now();
let load_result = run_usage_load(repository, config.requests, config.concurrency).await;
wait_for_outbox_drain(
&pool,
&SqlxUsageReadRepository::new(pool.clone()),
config.flush_batch_size,
)
.await?;
let duration_ms = started_at.elapsed().as_millis() as u64;
stop.store(true, Ordering::Release);
flush_handle.await??;
monitor_handle.await??;
let mut flush = flush_report.lock().await;
let final_flush = SqlxUsageReadRepository::new(pool.clone())
.flush_usage_counter_deltas(config.flush_batch_size)
.await?;
flush.calls += 1;
flush.rows_claimed += final_flush.rows_claimed;
flush.api_key_targets += final_flush.api_key_targets;
flush.provider_api_key_targets += final_flush.provider_api_key_targets;
flush.model_targets += final_flush.model_targets;
flush.provider_monthly_targets += final_flush.provider_monthly_targets;
flush.proxy_node_targets += final_flush.proxy_node_targets;
flush.management_token_targets += final_flush.management_token_targets;
flush.api_key_last_used_targets += final_flush.api_key_last_used_targets;
drop(flush);
let counters = read_counters(&pool).await?;
let latencies = load_result.latencies.lock().await.clone();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(latencies);
let completed_requests = load_result.completed.load(Ordering::Acquire);
let throughput_rps = if duration_ms == 0 {
completed_requests as u64
} else {
((completed_requests as u64) * 1_000) / duration_ms.max(1)
};
let report = Report {
suite: "usage_counter_hotspot_baseline",
config: ReportConfig {
requests: config.requests,
concurrency: config.concurrency,
max_connections: config.max_connections,
flush_batch_size: config.flush_batch_size,
flush_interval_ms: config.flush_interval.as_millis() as u64,
monitor_interval_ms: config.monitor_interval.as_millis() as u64,
managed_postgres,
},
duration_ms,
throughput_rps,
completed_requests,
failed_requests: load_result.failed.load(Ordering::Acquire),
p50_ms,
p95_ms,
max_ms,
mean_ms,
flush: Arc::try_unwrap(flush_report)
.unwrap_or_else(|_| panic!("flush report still referenced"))
.into_inner(),
counters,
lock_monitor: Arc::try_unwrap(lock_report)
.unwrap_or_else(|_| panic!("lock report still referenced"))
.into_inner(),
};
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
#[derive(Clone)]
struct LoadResult {
completed: Arc<AtomicUsize>,
failed: Arc<AtomicUsize>,
latencies: Arc<Mutex<Vec<u64>>>,
}
async fn run_usage_load(
repository: SqlxUsageReadRepository,
requests: usize,
concurrency: usize,
) -> LoadResult {
let next = Arc::new(AtomicUsize::new(0));
let completed = Arc::new(AtomicUsize::new(0));
let failed = Arc::new(AtomicUsize::new(0));
let latencies = Arc::new(Mutex::new(Vec::with_capacity(requests)));
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..concurrency {
let repository = repository.clone();
let next = Arc::clone(&next);
let completed = Arc::clone(&completed);
let failed = Arc::clone(&failed);
let latencies = Arc::clone(&latencies);
tasks.spawn(async move {
loop {
let index = next.fetch_add(1, Ordering::AcqRel);
if index >= requests {
break;
}
let started_at = Instant::now();
let result = repository.upsert(usage_record(index)).await;
latencies
.lock()
.await
.push(started_at.elapsed().as_millis() as u64);
completed.fetch_add(1, Ordering::AcqRel);
if result.is_err() {
failed.fetch_add(1, Ordering::AcqRel);
}
}
});
}
while let Some(result) = tasks.join_next().await {
if result.is_err() {
failed.fetch_add(1, Ordering::AcqRel);
}
}
LoadResult {
completed,
failed,
latencies,
}
}
fn spawn_flush_loop(
repository: SqlxUsageReadRepository,
stop: Arc<AtomicBool>,
report: Arc<Mutex<FlushReport>>,
batch_size: usize,
interval: Duration,
) -> tokio::task::JoinHandle<Result<(), aether_data::DataLayerError>> {
tokio::spawn(async move {
while !stop.load(Ordering::Acquire) {
let summary = repository.flush_usage_counter_deltas(batch_size).await?;
let mut report = report.lock().await;
report.calls += 1;
report.rows_claimed += summary.rows_claimed;
report.api_key_targets += summary.api_key_targets;
report.provider_api_key_targets += summary.provider_api_key_targets;
report.model_targets += summary.model_targets;
report.provider_monthly_targets += summary.provider_monthly_targets;
report.proxy_node_targets += summary.proxy_node_targets;
report.management_token_targets += summary.management_token_targets;
report.api_key_last_used_targets += summary.api_key_last_used_targets;
drop(report);
tokio::time::sleep(interval).await;
}
Ok(())
})
}
fn spawn_lock_monitor(
pool: PgPool,
stop: Arc<AtomicBool>,
report: Arc<Mutex<LockMonitorReport>>,
interval: Duration,
) -> tokio::task::JoinHandle<Result<(), sqlx::Error>> {
tokio::spawn(async move {
while !stop.load(Ordering::Acquire) {
let sample = read_lock_sample(&pool).await?;
let mut report = report.lock().await;
report.samples += 1;
report.max_lock_waiters = report.max_lock_waiters.max(sample.lock_waiters);
report.max_api_key_update_waiters = report
.max_api_key_update_waiters
.max(sample.api_key_update_waiters);
report.max_provider_key_update_waiters = report
.max_provider_key_update_waiters
.max(sample.provider_key_update_waiters);
report.max_global_model_update_waiters = report
.max_global_model_update_waiters
.max(sample.global_model_update_waiters);
report.max_provider_update_waiters = report
.max_provider_update_waiters
.max(sample.provider_update_waiters);
report.max_oldest_lock_wait_ms = report
.max_oldest_lock_wait_ms
.max(sample.oldest_lock_wait_ms);
drop(report);
tokio::time::sleep(interval).await;
}
Ok(())
})
}
async fn wait_for_outbox_drain(
pool: &PgPool,
repository: &SqlxUsageReadRepository,
batch_size: usize,
) -> Result<(), Box<dyn std::error::Error>> {
let deadline = Instant::now() + Duration::from_secs(30);
loop {
let summary = repository.flush_usage_counter_deltas(batch_size).await?;
let pending: i64 = sqlx::query_scalar(
"SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE processed_at IS NULL",
)
.fetch_one(pool)
.await?;
if pending == 0 && summary.rows_claimed == 0 {
return Ok(());
}
if Instant::now() >= deadline {
return Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
format!("usage counter outbox did not drain; pending={pending}"),
)
.into());
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
async fn seed_hotspot_rows(pool: &PgPool) -> Result<(), sqlx::Error> {
sqlx::query(
r#"
INSERT INTO users (id, username, email_verified)
VALUES ('user-hotspot', 'usage-hotspot', true)
ON CONFLICT (id) DO NOTHING
"#,
)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO api_keys (id, user_id, key_hash, name, is_active, total_requests, total_tokens, total_cost_usd)
VALUES ('api-key-hotspot', 'user-hotspot', 'hash-hotspot', 'hotspot key', true, 0, 0, 0)
ON CONFLICT (id) DO UPDATE SET
total_requests = 0,
total_tokens = 0,
total_cost_usd = 0,
last_used_at = NULL
"#,
)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO providers (id, name, provider_type, monthly_used_usd)
VALUES ('provider-hotspot', 'Hotspot Provider', 'openai', 0)
ON CONFLICT (id) DO UPDATE SET monthly_used_usd = 0
"#,
)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO provider_api_keys (
id, provider_id, name, total_tokens, total_cost_usd, request_count,
success_count, error_count, total_response_time_ms
)
VALUES ('provider-key-hotspot', 'provider-hotspot', 'Hotspot Provider Key', 0, 0, 0, 0, 0, 0)
ON CONFLICT (id) DO UPDATE SET
total_tokens = 0,
total_cost_usd = 0,
request_count = 0,
success_count = 0,
error_count = 0,
total_response_time_ms = 0,
last_used_at = NULL
"#,
)
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO global_models (id, name, display_name, enabled, is_active, usage_count)
VALUES ('model-hotspot', 'gpt-5', 'gpt-5', true, true, 0)
ON CONFLICT (id) DO UPDATE SET usage_count = 0
"#,
)
.execute(pool)
.await?;
Ok(())
}
fn usage_record(index: usize) -> UpsertUsageRecord {
let now_ms = now_unix_ms().saturating_add(index as u64);
let now_secs = now_ms / 1_000;
UpsertUsageRecord {
request_id: format!("usage-hotspot-{index:08}"),
user_id: Some("user-hotspot".to_string()),
api_key_id: Some("api-key-hotspot".to_string()),
username: None,
api_key_name: None,
provider_name: "openai".to_string(),
model: "gpt-5".to_string(),
target_model: None,
provider_id: Some("provider-hotspot".to_string()),
provider_endpoint_id: None,
provider_api_key_id: Some("provider-key-hotspot".to_string()),
request_type: Some("chat".to_string()),
api_format: Some("openai:chat".to_string()),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
provider_api_family: Some("openai".to_string()),
provider_endpoint_kind: Some("chat".to_string()),
has_format_conversion: Some(false),
is_stream: Some(false),
input_tokens: Some(10),
output_tokens: Some(20),
total_tokens: Some(30),
cache_creation_input_tokens: None,
cache_creation_ephemeral_5m_input_tokens: None,
cache_creation_ephemeral_1h_input_tokens: None,
cache_read_input_tokens: None,
cache_creation_cost_usd: None,
cache_read_cost_usd: None,
output_price_per_1m: None,
total_cost_usd: Some(0.001),
actual_total_cost_usd: Some(0.001),
status_code: Some(200),
error_message: None,
error_category: None,
response_time_ms: Some(100),
first_byte_time_ms: Some(20),
status: "completed".to_string(),
billing_status: "pending".to_string(),
request_headers: None,
request_body: Some(json!({"model": "gpt-5"})),
request_body_ref: None,
request_body_state: None,
provider_request_headers: None,
provider_request_body: None,
provider_request_body_ref: None,
provider_request_body_state: None,
response_headers: None,
response_body: Some(json!({"id": format!("chatcmpl-{index}")})),
response_body_ref: None,
response_body_state: None,
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
client_response_body_state: None,
candidate_id: Some(format!("candidate-{index:08}")),
candidate_index: Some(0),
key_name: None,
planner_kind: Some("hotspot_baseline".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
execution_path: Some("testkit".to_string()),
local_execution_runtime_miss_reason: None,
request_metadata: None,
finalized_at_unix_secs: None,
created_at_unix_ms: Some(now_ms),
updated_at_unix_secs: now_secs,
}
}
async fn read_lock_sample(pool: &PgPool) -> Result<LockSample, sqlx::Error> {
let row = sqlx::query(
r#"
SELECT
COUNT(*) FILTER (WHERE wait_event_type = 'Lock')::BIGINT AS lock_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE 'UPDATE api_keys%'
)::BIGINT AS api_key_update_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE 'UPDATE provider_api_keys%'
)::BIGINT AS provider_key_update_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE 'UPDATE global_models%'
)::BIGINT AS global_model_update_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE 'UPDATE providers%'
)::BIGINT AS provider_update_waiters,
COALESCE(
MAX(EXTRACT(EPOCH FROM (NOW() - query_start)) * 1000)
FILTER (WHERE wait_event_type = 'Lock'),
0
)::BIGINT AS oldest_lock_wait_ms
FROM pg_stat_activity
WHERE datname = current_database()
"#,
)
.fetch_one(pool)
.await?;
Ok(LockSample {
lock_waiters: row.try_get("lock_waiters")?,
api_key_update_waiters: row.try_get("api_key_update_waiters")?,
provider_key_update_waiters: row.try_get("provider_key_update_waiters")?,
global_model_update_waiters: row.try_get("global_model_update_waiters")?,
provider_update_waiters: row.try_get("provider_update_waiters")?,
oldest_lock_wait_ms: row.try_get("oldest_lock_wait_ms")?,
})
}
async fn read_counters(pool: &PgPool) -> Result<CounterReport, sqlx::Error> {
let row = sqlx::query(
r#"
SELECT
(SELECT COUNT(*)::BIGINT FROM usage) AS usage_rows,
(SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE processed_at IS NULL) AS outbox_pending_rows,
(SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE processed_at IS NOT NULL) AS outbox_processed_rows,
(SELECT COALESCE(total_requests, 0)::BIGINT FROM api_keys WHERE id = 'api-key-hotspot') AS api_key_total_requests,
(SELECT COALESCE(total_tokens, 0)::BIGINT FROM api_keys WHERE id = 'api-key-hotspot') AS api_key_total_tokens,
(SELECT COALESCE(request_count, 0)::BIGINT FROM provider_api_keys WHERE id = 'provider-key-hotspot') AS provider_key_request_count,
(SELECT COALESCE(success_count, 0)::BIGINT FROM provider_api_keys WHERE id = 'provider-key-hotspot') AS provider_key_success_count,
(SELECT COALESCE(usage_count, 0)::BIGINT FROM global_models WHERE name = 'gpt-5' LIMIT 1) AS global_model_usage_count
"#,
)
.fetch_one(pool)
.await?;
Ok(CounterReport {
usage_rows: row.try_get("usage_rows")?,
outbox_pending_rows: row.try_get("outbox_pending_rows")?,
outbox_processed_rows: row.try_get("outbox_processed_rows")?,
api_key_total_requests: row.try_get("api_key_total_requests")?,
api_key_total_tokens: row.try_get("api_key_total_tokens")?,
provider_key_request_count: row.try_get("provider_key_request_count")?,
provider_key_success_count: row.try_get("provider_key_success_count")?,
global_model_usage_count: row.try_get("global_model_usage_count")?,
})
}
fn summarize_latencies(mut latencies: Vec<u64>) -> (u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
}
latencies.sort_unstable();
let p50 = percentile(&latencies, 50);
let p95 = percentile(&latencies, 95);
let max = *latencies.last().unwrap_or(&0);
let mean = latencies.iter().sum::<u64>() / latencies.len() as u64;
(p50, p95, max, mean)
}
fn percentile(latencies: &[u64], percentile: usize) -> u64 {
let index = ((latencies.len() - 1) * percentile) / 100;
latencies[index]
}
fn now_unix_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
let mut config = Config::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--requests" => config.requests = next_value(&mut iter, "--requests")?.parse()?,
"--concurrency" => {
config.concurrency = next_value(&mut iter, "--concurrency")?.parse()?
}
"--max-connections" => {
config.max_connections = next_value(&mut iter, "--max-connections")?.parse()?
}
"--flush-batch-size" => {
config.flush_batch_size = next_value(&mut iter, "--flush-batch-size")?.parse()?
}
"--flush-interval-ms" => {
config.flush_interval =
Duration::from_millis(next_value(&mut iter, "--flush-interval-ms")?.parse()?)
}
"--monitor-interval-ms" => {
config.monitor_interval =
Duration::from_millis(next_value(&mut iter, "--monitor-interval-ms")?.parse()?)
}
"--postgres-url" => {
config.postgres_url = Some(next_value(&mut iter, "--postgres-url")?)
}
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
if config.requests == 0 || config.concurrency == 0 || config.max_connections == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"requests, concurrency, and max-connections must be positive",
)
.into());
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin usage_counter_hotspot_baseline -- [--requests 2000] [--concurrency 100] [--max-connections 64] [--flush-batch-size 1000] [--flush-interval-ms 100] [--monitor-interval-ms 100] [--postgres-url postgres://...] [--output /tmp/usage_counter_hotspot.json]"
);
}

673
crates/aether-testkit/src/bin/usage_settlement_hotspot_baseline.rs Normal file
View File

@@ -0,0 +1,673 @@
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use aether_data::repository::settlement::SqlxSettlementRepository;
use aether_data::repository::usage::SqlxUsageReadRepository;
use aether_data_contracts::repository::settlement::{
SettlementWriteRepository, UsageSettlementInput,
};
use aether_testkit::{
init_test_runtime_for, prepare_aether_postgres_schema, ManagedPostgresServer,
};
use serde::Serialize;
use sqlx::postgres::PgPoolOptions;
use sqlx::{PgPool, Row};
use tokio::sync::Mutex;
const PROVIDER_ID: &str = "provider-hotspot";
const REQUEST_PREFIX: &str = "settlement-hotspot";
const COST_PER_REQUEST_USD: f64 = 0.001;
#[derive(Debug, Clone)]
struct Config {
requests: usize,
concurrency: usize,
max_connections: u32,
flush_batch_size: usize,
flush_interval: Duration,
monitor_interval: Duration,
output_path: Option<PathBuf>,
postgres_url: Option<String>,
}
impl Default for Config {
fn default() -> Self {
Self {
requests: 2_000,
concurrency: 100,
max_connections: 64,
flush_batch_size: 1_000,
flush_interval: Duration::from_millis(100),
monitor_interval: Duration::from_millis(100),
output_path: None,
postgres_url: None,
}
}
}
#[derive(Debug, Serialize)]
struct Report {
suite: &'static str,
config: ReportConfig,
duration_ms: u64,
throughput_rps: u64,
completed_requests: usize,
failed_requests: usize,
p50_ms: u64,
p95_ms: u64,
max_ms: u64,
mean_ms: u64,
flush: FlushReport,
counters: CounterReport,
lock_monitor: LockMonitorReport,
}
#[derive(Debug, Serialize)]
struct ReportConfig {
requests: usize,
concurrency: usize,
max_connections: u32,
flush_batch_size: usize,
flush_interval_ms: u64,
monitor_interval_ms: u64,
managed_postgres: bool,
}
#[derive(Debug, Serialize, Default)]
struct FlushReport {
calls: usize,
rows_claimed: usize,
api_key_targets: usize,
provider_api_key_targets: usize,
model_targets: usize,
provider_monthly_targets: usize,
proxy_node_targets: usize,
management_token_targets: usize,
api_key_last_used_targets: usize,
}
#[derive(Debug, Serialize)]
struct CounterReport {
usage_rows: i64,
settled_usage_rows: i64,
settlement_snapshot_rows: i64,
outbox_pending_rows: i64,
outbox_processed_rows: i64,
provider_monthly_outbox_rows: i64,
provider_monthly_used_usd: f64,
expected_provider_monthly_used_usd: f64,
}
#[derive(Debug, Serialize, Clone, Copy, Default)]
struct LockMonitorReport {
samples: usize,
max_lock_waiters: i64,
max_usage_update_waiters: i64,
max_wallet_update_waiters: i64,
max_provider_update_waiters: i64,
max_oldest_lock_wait_ms: i64,
}
#[derive(Debug, Clone, Copy, Default)]
struct LockSample {
lock_waiters: i64,
usage_update_waiters: i64,
wallet_update_waiters: i64,
provider_update_waiters: i64,
oldest_lock_wait_ms: i64,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("usage-settlement-hotspot-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let managed_postgres;
let database_url;
let _server;
if let Some(url) = config.postgres_url.as_ref() {
managed_postgres = false;
database_url = url.clone();
_server = None;
} else {
managed_postgres = true;
let server = ManagedPostgresServer::start().await?;
database_url = server.database_url().to_string();
_server = Some(server);
}
prepare_aether_postgres_schema(&database_url).await?;
let pool = PgPoolOptions::new()
.min_connections(1)
.max_connections(config.max_connections)
.acquire_timeout(Duration::from_secs(10))
.connect(&database_url)
.await?;
seed_settlement_rows(&pool, config.requests).await?;
let usage_repository = SqlxUsageReadRepository::new(pool.clone());
let settlement_repository = SqlxSettlementRepository::new(pool.clone());
let stop = Arc::new(AtomicBool::new(false));
let flush_report = Arc::new(Mutex::new(FlushReport::default()));
let lock_report = Arc::new(Mutex::new(LockMonitorReport::default()));
let flush_handle = spawn_flush_loop(
usage_repository.clone(),
Arc::clone(&stop),
Arc::clone(&flush_report),
config.flush_batch_size,
config.flush_interval,
);
let monitor_handle = spawn_lock_monitor(
pool.clone(),
Arc::clone(&stop),
Arc::clone(&lock_report),
config.monitor_interval,
);
let started_at = Instant::now();
let load_result =
run_settlement_load(settlement_repository, config.requests, config.concurrency).await;
wait_for_outbox_drain(&pool, &usage_repository, config.flush_batch_size).await?;
let duration_ms = started_at.elapsed().as_millis() as u64;
stop.store(true, Ordering::Release);
flush_handle.await??;
monitor_handle.await??;
let mut flush = flush_report.lock().await;
let final_flush = usage_repository
.flush_usage_counter_deltas(config.flush_batch_size)
.await?;
flush.calls += 1;
flush.rows_claimed += final_flush.rows_claimed;
flush.api_key_targets += final_flush.api_key_targets;
flush.provider_api_key_targets += final_flush.provider_api_key_targets;
flush.model_targets += final_flush.model_targets;
flush.provider_monthly_targets += final_flush.provider_monthly_targets;
flush.proxy_node_targets += final_flush.proxy_node_targets;
flush.management_token_targets += final_flush.management_token_targets;
flush.api_key_last_used_targets += final_flush.api_key_last_used_targets;
drop(flush);
let counters = read_counters(&pool, config.requests).await?;
let latencies = load_result.latencies.lock().await.clone();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(latencies);
let completed_requests = load_result.completed.load(Ordering::Acquire);
let throughput_rps = if duration_ms == 0 {
completed_requests as u64
} else {
((completed_requests as u64) * 1_000) / duration_ms.max(1)
};
let report = Report {
suite: "usage_settlement_hotspot_baseline",
config: ReportConfig {
requests: config.requests,
concurrency: config.concurrency,
max_connections: config.max_connections,
flush_batch_size: config.flush_batch_size,
flush_interval_ms: config.flush_interval.as_millis() as u64,
monitor_interval_ms: config.monitor_interval.as_millis() as u64,
managed_postgres,
},
duration_ms,
throughput_rps,
completed_requests,
failed_requests: load_result.failed.load(Ordering::Acquire),
p50_ms,
p95_ms,
max_ms,
mean_ms,
flush: Arc::try_unwrap(flush_report)
.unwrap_or_else(|_| panic!("flush report still referenced"))
.into_inner(),
counters,
lock_monitor: Arc::try_unwrap(lock_report)
.unwrap_or_else(|_| panic!("lock report still referenced"))
.into_inner(),
};
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
#[derive(Clone)]
struct LoadResult {
completed: Arc<AtomicUsize>,
failed: Arc<AtomicUsize>,
latencies: Arc<Mutex<Vec<u64>>>,
}
async fn run_settlement_load(
repository: SqlxSettlementRepository,
requests: usize,
concurrency: usize,
) -> LoadResult {
let next = Arc::new(AtomicUsize::new(0));
let completed = Arc::new(AtomicUsize::new(0));
let failed = Arc::new(AtomicUsize::new(0));
let latencies = Arc::new(Mutex::new(Vec::with_capacity(requests)));
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..concurrency {
let repository = repository.clone();
let next = Arc::clone(&next);
let completed = Arc::clone(&completed);
let failed = Arc::clone(&failed);
let latencies = Arc::clone(&latencies);
tasks.spawn(async move {
loop {
let index = next.fetch_add(1, Ordering::AcqRel);
if index >= requests {
break;
}
let started_at = Instant::now();
let result = repository.settle_usage(settlement_input(index)).await;
latencies
.lock()
.await
.push(started_at.elapsed().as_millis() as u64);
completed.fetch_add(1, Ordering::AcqRel);
match result {
Ok(Some(settlement)) if settlement.billing_status == "settled" => {}
_ => {
failed.fetch_add(1, Ordering::AcqRel);
}
}
}
});
}
while let Some(result) = tasks.join_next().await {
if result.is_err() {
failed.fetch_add(1, Ordering::AcqRel);
}
}
LoadResult {
completed,
failed,
latencies,
}
}
fn spawn_flush_loop(
repository: SqlxUsageReadRepository,
stop: Arc<AtomicBool>,
report: Arc<Mutex<FlushReport>>,
batch_size: usize,
interval: Duration,
) -> tokio::task::JoinHandle<Result<(), aether_data::DataLayerError>> {
tokio::spawn(async move {
while !stop.load(Ordering::Acquire) {
let summary = repository.flush_usage_counter_deltas(batch_size).await?;
let mut report = report.lock().await;
report.calls += 1;
report.rows_claimed += summary.rows_claimed;
report.api_key_targets += summary.api_key_targets;
report.provider_api_key_targets += summary.provider_api_key_targets;
report.model_targets += summary.model_targets;
report.provider_monthly_targets += summary.provider_monthly_targets;
report.proxy_node_targets += summary.proxy_node_targets;
report.management_token_targets += summary.management_token_targets;
report.api_key_last_used_targets += summary.api_key_last_used_targets;
drop(report);
tokio::time::sleep(interval).await;
}
Ok(())
})
}
fn spawn_lock_monitor(
pool: PgPool,
stop: Arc<AtomicBool>,
report: Arc<Mutex<LockMonitorReport>>,
interval: Duration,
) -> tokio::task::JoinHandle<Result<(), sqlx::Error>> {
tokio::spawn(async move {
while !stop.load(Ordering::Acquire) {
let sample = read_lock_sample(&pool).await?;
let mut report = report.lock().await;
report.samples += 1;
report.max_lock_waiters = report.max_lock_waiters.max(sample.lock_waiters);
report.max_usage_update_waiters = report
.max_usage_update_waiters
.max(sample.usage_update_waiters);
report.max_wallet_update_waiters = report
.max_wallet_update_waiters
.max(sample.wallet_update_waiters);
report.max_provider_update_waiters = report
.max_provider_update_waiters
.max(sample.provider_update_waiters);
report.max_oldest_lock_wait_ms = report
.max_oldest_lock_wait_ms
.max(sample.oldest_lock_wait_ms);
drop(report);
tokio::time::sleep(interval).await;
}
Ok(())
})
}
async fn wait_for_outbox_drain(
pool: &PgPool,
repository: &SqlxUsageReadRepository,
batch_size: usize,
) -> Result<(), Box<dyn std::error::Error>> {
let deadline = Instant::now() + Duration::from_secs(30);
loop {
let summary = repository.flush_usage_counter_deltas(batch_size).await?;
let pending: i64 = sqlx::query_scalar(
"SELECT COUNT(*)::BIGINT FROM usage_counter_deltas WHERE processed_at IS NULL",
)
.fetch_one(pool)
.await?;
if pending == 0 && summary.rows_claimed == 0 {
return Ok(());
}
if Instant::now() >= deadline {
return Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
format!("usage counter outbox did not drain; pending={pending}"),
)
.into());
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
async fn seed_settlement_rows(pool: &PgPool, requests: usize) -> Result<(), sqlx::Error> {
sqlx::query(
r#"
INSERT INTO providers (id, name, provider_type, monthly_used_usd)
VALUES ($1, 'Hotspot Provider', 'openai', 0)
ON CONFLICT (id) DO UPDATE SET monthly_used_usd = 0
"#,
)
.bind(PROVIDER_ID)
.execute(pool)
.await?;
sqlx::query(
r#"
DELETE FROM usage_counter_deltas
WHERE request_id LIKE $1
"#,
)
.bind(format!("{REQUEST_PREFIX}-%"))
.execute(pool)
.await?;
sqlx::query(
r#"
DELETE FROM usage_settlement_snapshots
WHERE request_id LIKE $1
"#,
)
.bind(format!("{REQUEST_PREFIX}-%"))
.execute(pool)
.await?;
sqlx::query(
r#"
DELETE FROM "usage"
WHERE request_id LIKE $1
"#,
)
.bind(format!("{REQUEST_PREFIX}-%"))
.execute(pool)
.await?;
sqlx::query(
r#"
INSERT INTO "usage" (
id,
request_id,
provider_name,
model,
provider_id,
status,
billing_status,
total_cost_usd,
actual_total_cost_usd,
created_at_unix_ms,
updated_at_unix_secs
)
SELECT
'settlement-usage-' || LPAD(gs::TEXT, 8, '0'),
$2 || '-' || LPAD(gs::TEXT, 8, '0'),
'Hotspot Provider',
'gpt-5',
$3,
'completed',
'pending',
0,
$4,
$5 + gs,
$6
FROM generate_series(0, $1::INTEGER - 1) AS gs
"#,
)
.bind(i32::try_from(requests).unwrap_or(i32::MAX))
.bind(REQUEST_PREFIX)
.bind(PROVIDER_ID)
.bind(COST_PER_REQUEST_USD)
.bind(now_unix_ms() as i64)
.bind(now_unix_secs() as i64)
.execute(pool)
.await?;
Ok(())
}
fn settlement_input(index: usize) -> UsageSettlementInput {
UsageSettlementInput {
request_id: format!("{REQUEST_PREFIX}-{index:08}"),
user_id: None,
api_key_id: None,
api_key_is_standalone: false,
provider_id: Some(PROVIDER_ID.to_string()),
status: "completed".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 0.0,
actual_total_cost_usd: COST_PER_REQUEST_USD,
finalized_at_unix_secs: Some(now_unix_secs().saturating_add(index as u64)),
}
}
async fn read_lock_sample(pool: &PgPool) -> Result<LockSample, sqlx::Error> {
let row = sqlx::query(
r#"
SELECT
COUNT(*) FILTER (WHERE wait_event_type = 'Lock')::BIGINT AS lock_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE '%UPDATE "usage"%'
)::BIGINT AS usage_update_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE '%UPDATE wallets%'
)::BIGINT AS wallet_update_waiters,
COUNT(*) FILTER (
WHERE wait_event_type = 'Lock' AND query LIKE '%UPDATE providers%'
)::BIGINT AS provider_update_waiters,
COALESCE(
MAX(EXTRACT(EPOCH FROM (NOW() - query_start)) * 1000)
FILTER (WHERE wait_event_type = 'Lock'),
0
)::BIGINT AS oldest_lock_wait_ms
FROM pg_stat_activity
WHERE datname = current_database()
"#,
)
.fetch_one(pool)
.await?;
Ok(LockSample {
lock_waiters: row.try_get("lock_waiters")?,
usage_update_waiters: row.try_get("usage_update_waiters")?,
wallet_update_waiters: row.try_get("wallet_update_waiters")?,
provider_update_waiters: row.try_get("provider_update_waiters")?,
oldest_lock_wait_ms: row.try_get("oldest_lock_wait_ms")?,
})
}
async fn read_counters(pool: &PgPool, requests: usize) -> Result<CounterReport, sqlx::Error> {
let row = sqlx::query(
r#"
SELECT
(SELECT COUNT(*)::BIGINT FROM "usage" WHERE request_id LIKE $1) AS usage_rows,
(
SELECT COUNT(*)::BIGINT
FROM "usage"
WHERE request_id LIKE $1 AND billing_status = 'settled'
) AS settled_usage_rows,
(
SELECT COUNT(*)::BIGINT
FROM usage_settlement_snapshots
WHERE request_id LIKE $1
) AS settlement_snapshot_rows,
(
SELECT COUNT(*)::BIGINT
FROM usage_counter_deltas
WHERE request_id LIKE $1 AND processed_at IS NULL
) AS outbox_pending_rows,
(
SELECT COUNT(*)::BIGINT
FROM usage_counter_deltas
WHERE request_id LIKE $1 AND processed_at IS NOT NULL
) AS outbox_processed_rows,
(
SELECT COUNT(*)::BIGINT
FROM usage_counter_deltas
WHERE request_id LIKE $1 AND kind = 'provider_monthly'
) AS provider_monthly_outbox_rows,
(
SELECT CAST(monthly_used_usd AS DOUBLE PRECISION)
FROM providers
WHERE id = $2
) AS provider_monthly_used_usd
"#,
)
.bind(format!("{REQUEST_PREFIX}-%"))
.bind(PROVIDER_ID)
.fetch_one(pool)
.await?;
Ok(CounterReport {
usage_rows: row.try_get("usage_rows")?,
settled_usage_rows: row.try_get("settled_usage_rows")?,
settlement_snapshot_rows: row.try_get("settlement_snapshot_rows")?,
outbox_pending_rows: row.try_get("outbox_pending_rows")?,
outbox_processed_rows: row.try_get("outbox_processed_rows")?,
provider_monthly_outbox_rows: row.try_get("provider_monthly_outbox_rows")?,
provider_monthly_used_usd: row.try_get("provider_monthly_used_usd")?,
expected_provider_monthly_used_usd: (requests as f64) * COST_PER_REQUEST_USD,
})
}
fn summarize_latencies(mut latencies: Vec<u64>) -> (u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
}
latencies.sort_unstable();
let p50 = percentile(&latencies, 50);
let p95 = percentile(&latencies, 95);
let max = *latencies.last().unwrap_or(&0);
let mean = latencies.iter().sum::<u64>() / latencies.len() as u64;
(p50, p95, max, mean)
}
fn percentile(latencies: &[u64], percentile: usize) -> u64 {
let index = ((latencies.len() - 1) * percentile) / 100;
latencies[index]
}
fn now_unix_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
let mut config = Config::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--requests" => config.requests = next_value(&mut iter, "--requests")?.parse()?,
"--concurrency" => {
config.concurrency = next_value(&mut iter, "--concurrency")?.parse()?
}
"--max-connections" => {
config.max_connections = next_value(&mut iter, "--max-connections")?.parse()?
}
"--flush-batch-size" => {
config.flush_batch_size = next_value(&mut iter, "--flush-batch-size")?.parse()?
}
"--flush-interval-ms" => {
config.flush_interval =
Duration::from_millis(next_value(&mut iter, "--flush-interval-ms")?.parse()?)
}
"--monitor-interval-ms" => {
config.monitor_interval =
Duration::from_millis(next_value(&mut iter, "--monitor-interval-ms")?.parse()?)
}
"--postgres-url" => {
config.postgres_url = Some(next_value(&mut iter, "--postgres-url")?)
}
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
if config.requests == 0 || config.concurrency == 0 || config.max_connections == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"requests, concurrency, and max-connections must be positive",
)
.into());
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin usage_settlement_hotspot_baseline -- [--requests 2000] [--concurrency 100] [--max-connections 64] [--flush-batch-size 1000] [--flush-interval-ms 100] [--monitor-interval-ms 100] [--postgres-url postgres://...] [--output /tmp/usage_settlement_hotspot.json]"
);
}