Skip to content

perf: reduce homepage database hot paths and chunk large key queries #287

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mozi1924 wants to merge 4 commits into
base: main
Choose a base branch
from
Open

perf: reduce homepage database hot paths and chunk large key queries #287

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0da87fc
fix: push homepage hot-path filtering into storage queries
mozi1924 May 28, 2026
845e4da
fix: chunk large key-list storage queries to avoid SQLite limits
mozi1924 May 28, 2026
1d3e9e5
fix: address dashboard hot path review feedback
mozi1924 May 29, 2026
4118c05
fix: optimize member dashboard request log filters
mozi1924 May 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions crates/core/migrations/057_account_manager.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ CREATE TABLE IF NOT EXISTS app_wallet_ledger_entries (

CREATE INDEX IF NOT EXISTS idx_app_wallet_ledger_wallet_created ON app_wallet_ledger_entries(wallet_id, created_at DESC);
CREATE INDEX IF NOT EXISTS idx_app_wallet_ledger_api_key ON app_wallet_ledger_entries(api_key_id);
CREATE INDEX IF NOT EXISTS idx_app_wallet_ledger_request_log_kind ON app_wallet_ledger_entries(request_log_id, entry_kind);

CREATE TABLE IF NOT EXISTS api_key_owners (
key_id TEXT PRIMARY KEY REFERENCES api_keys(id) ON DELETE CASCADE,
Expand Down
28 changes: 28 additions & 0 deletions crates/core/src/storage/account_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -475,6 +475,34 @@ impl Storage {
Ok(out)
}

/// 函数 `list_api_key_ids_for_user`
///
///
/// 时间: 2026-05-28
///
/// # 参数
/// - self: 参数 self
/// - user_id: 参数 user_id
///
/// # 返回
/// 返回函数执行结果
pub fn list_api_key_ids_for_user(&self, user_id: &str) -> Result<Vec<String>> {
Copy link
Copy Markdown
Collaborator

@KilimiaoSix KilimiaoSix May 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里新增了按 owner_kind + owner_user_id 查询 key id 的索引路径，但服务层实际使用的 crate::list_api_key_ids_for_user 仍然调用 storage.list_api_key_owners() 拉全量 owner 后在 Rust 里过滤。因此成员 Dashboard、startup snapshot、API Key 页面仍会走旧热点路径，PR 的主要优化没有真正生效。建议把服务层 helper 改为调用这个新的 storage 方法，并补一个覆盖非 admin 路径的测试。

let normalized_user_id = user_id.trim();
if normalized_user_id.is_empty() {
return Ok(Vec::new());
}

// 中文注释：这里直接按 owner_user_id 走索引取 key_id，避免把整张 api_key_owners 表拉回 Rust 再过滤。
let mut stmt = self.conn.prepare(
"SELECT key_id
FROM api_key_owners
WHERE owner_kind = 'user' AND owner_user_id = ?1
ORDER BY key_id ASC",
)?;
let rows = stmt.query_map([normalized_user_id], |row| row.get(0))?;
rows.collect()
}

pub fn api_key_owner_count(&self) -> Result<i64> {
self.conn
.query_row("SELECT COUNT(*) FROM api_key_owners", [], |row| row.get(0))
Expand Down
289 changes: 164 additions & 125 deletions crates/core/src/storage/api_key_quota_limits.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,125 +1,164 @@
use std::collections::HashMap;

use rusqlite::{OptionalExtension, Result};

use super::{now_ts, Storage};

impl Storage {
pub fn upsert_api_key_quota_limit(
&self,
key_id: &str,
quota_limit_tokens: Option<i64>,
) -> Result<()> {
let normalized = quota_limit_tokens.filter(|value| *value > 0);
let Some(limit) = normalized else {
self.conn.execute(
"DELETE FROM api_key_quota_limits WHERE key_id = ?1",
[key_id],
)?;
return Ok(());
};

let now = now_ts();
self.conn.execute(
"INSERT INTO api_key_quota_limits (
key_id, quota_limit_tokens, created_at, updated_at
) VALUES (?1, ?2, ?3, ?3)
ON CONFLICT(key_id) DO UPDATE SET
quota_limit_tokens = excluded.quota_limit_tokens,
updated_at = excluded.updated_at",
(key_id, limit, now),
)?;
Ok(())
}

pub fn find_api_key_quota_limit(&self, key_id: &str) -> Result<Option<i64>> {
self.conn
.query_row(
"SELECT quota_limit_tokens
FROM api_key_quota_limits
WHERE key_id = ?1
LIMIT 1",
[key_id],
|row| row.get(0),
)
.optional()
}

pub fn list_api_key_quota_limits(&self) -> Result<HashMap<String, i64>> {
let mut stmt = self.conn.prepare(
"SELECT key_id, quota_limit_tokens
FROM api_key_quota_limits
WHERE quota_limit_tokens > 0",
)?;
let mut rows = stmt.query([])?;
let mut out = HashMap::new();
while let Some(row) = rows.next()? {
out.insert(row.get(0)?, row.get(1)?);
}
Ok(out)
}

pub fn api_key_total_token_usage(&self, key_id: &str) -> Result<i64> {
let mut stmt = self.conn.prepare(
"WITH all_stats AS (
SELECT
key_id,
input_tokens,
cached_input_tokens,
output_tokens,
total_tokens
FROM request_token_stats
UNION ALL
SELECT
NULLIF(key_id, '') AS key_id,
input_tokens,
cached_input_tokens,
output_tokens,
total_tokens
FROM request_token_stat_rollups
)
SELECT
IFNULL(
SUM(
CASE
WHEN total_tokens IS NOT NULL THEN
CASE WHEN total_tokens > 0 THEN total_tokens ELSE 0 END
ELSE
CASE
WHEN IFNULL(input_tokens, 0) - IFNULL(cached_input_tokens, 0) + IFNULL(output_tokens, 0) > 0
THEN IFNULL(input_tokens, 0) - IFNULL(cached_input_tokens, 0) + IFNULL(output_tokens, 0)
ELSE 0
END
END
),
0
) AS total_tokens
FROM all_stats
WHERE key_id = ?1",
)?;
let mut rows = stmt.query([key_id])?;
if let Some(row) = rows.next()? {
let total: i64 = row.get(0)?;
return Ok(total.max(0));
}
Ok(0)
}

pub(super) fn ensure_api_key_quota_limits_table(&self) -> Result<()> {
self.conn.execute(
"CREATE TABLE IF NOT EXISTS api_key_quota_limits (
key_id TEXT PRIMARY KEY REFERENCES api_keys(id) ON DELETE CASCADE,
quota_limit_tokens INTEGER NOT NULL,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
)",
[],
)?;
self.conn.execute(
"CREATE INDEX IF NOT EXISTS idx_api_key_quota_limits_updated_at
ON api_key_quota_limits(updated_at DESC)",
[],
)?;
Ok(())
}
}
use std::collections::HashMap;

use rusqlite::{params_from_iter, OptionalExtension, Result};

use super::key_id_filters::{key_id_in_clause, normalize_key_ids, SQLITE_IN_CLAUSE_BATCH_SIZE};
use super::{now_ts, Storage};

impl Storage {
pub fn upsert_api_key_quota_limit(
&self,
key_id: &str,
quota_limit_tokens: Option<i64>,
) -> Result<()> {
let normalized = quota_limit_tokens.filter(|value| *value > 0);
let Some(limit) = normalized else {
self.conn.execute(
"DELETE FROM api_key_quota_limits WHERE key_id = ?1",
[key_id],
)?;
return Ok(());
};

let now = now_ts();
self.conn.execute(
"INSERT INTO api_key_quota_limits (
key_id, quota_limit_tokens, created_at, updated_at
) VALUES (?1, ?2, ?3, ?3)
ON CONFLICT(key_id) DO UPDATE SET
quota_limit_tokens = excluded.quota_limit_tokens,
updated_at = excluded.updated_at",
(key_id, limit, now),
)?;
Ok(())
}

pub fn find_api_key_quota_limit(&self, key_id: &str) -> Result<Option<i64>> {
self.conn
.query_row(
"SELECT quota_limit_tokens
FROM api_key_quota_limits
WHERE key_id = ?1
LIMIT 1",
[key_id],
|row| row.get(0),
)
.optional()
}

pub fn list_api_key_quota_limits(&self) -> Result<HashMap<String, i64>> {
let mut stmt = self.conn.prepare(
"SELECT key_id, quota_limit_tokens
FROM api_key_quota_limits
WHERE quota_limit_tokens > 0",
)?;
let mut rows = stmt.query([])?;
let mut out = HashMap::new();
while let Some(row) = rows.next()? {
out.insert(row.get(0)?, row.get(1)?);
}
Ok(out)
}

pub fn list_api_key_quota_limits_for_ids(
&self,
key_ids: &[String],
) -> Result<HashMap<String, i64>> {
let key_ids = normalize_key_ids(key_ids);
if key_ids.is_empty() {
return Ok(HashMap::new());
}

let mut out = HashMap::new();
for chunk in key_ids.chunks(SQLITE_IN_CLAUSE_BATCH_SIZE) {
out.extend(list_api_key_quota_limits_for_ids_chunk(self, chunk)?);
}
Ok(out)
}

pub fn api_key_total_token_usage(&self, key_id: &str) -> Result<i64> {
let mut stmt = self.conn.prepare(
"WITH all_stats AS (
SELECT
key_id,
input_tokens,
cached_input_tokens,
output_tokens,
total_tokens
FROM request_token_stats
UNION ALL
SELECT
NULLIF(key_id, '') AS key_id,
input_tokens,
cached_input_tokens,
output_tokens,
total_tokens
FROM request_token_stat_rollups
)
SELECT
IFNULL(
SUM(
CASE
WHEN total_tokens IS NOT NULL THEN
CASE WHEN total_tokens > 0 THEN total_tokens ELSE 0 END
ELSE
CASE
WHEN IFNULL(input_tokens, 0) - IFNULL(cached_input_tokens, 0) + IFNULL(output_tokens, 0) > 0
THEN IFNULL(input_tokens, 0) - IFNULL(cached_input_tokens, 0) + IFNULL(output_tokens, 0)
ELSE 0
END
END
),
0
) AS total_tokens
FROM all_stats
WHERE key_id = ?1",
)?;
let mut rows = stmt.query([key_id])?;
if let Some(row) = rows.next()? {
let total: i64 = row.get(0)?;
return Ok(total.max(0));
}
Ok(0)
}

pub(super) fn ensure_api_key_quota_limits_table(&self) -> Result<()> {
self.conn.execute(
"CREATE TABLE IF NOT EXISTS api_key_quota_limits (
key_id TEXT PRIMARY KEY REFERENCES api_keys(id) ON DELETE CASCADE,
quota_limit_tokens INTEGER NOT NULL,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
)",
[],
)?;
self.conn.execute(
"CREATE INDEX IF NOT EXISTS idx_api_key_quota_limits_updated_at
ON api_key_quota_limits(updated_at DESC)",
[],
)?;
Ok(())
}
}

fn list_api_key_quota_limits_for_ids_chunk(
storage: &Storage,
key_ids: &[String],
) -> Result<HashMap<String, i64>> {
let Some((clause, params)) = key_id_in_clause("key_id", key_ids) else {
return Ok(HashMap::new());
};
let sql = format!(
"SELECT key_id, quota_limit_tokens
FROM api_key_quota_limits
WHERE quota_limit_tokens > 0
AND {clause}"
);
let mut stmt = storage.conn.prepare(&sql)?;
let mut rows = stmt.query(params_from_iter(params.iter()))?;
let mut out = HashMap::new();
while let Some(row) = rows.next()? {
out.insert(row.get(0)?, row.get(1)?);
}
Ok(out)
}
Loading