feat(key-pool): add LRU sorting via last_used_at column
- Add migration to add last_used_at TIMESTAMPTZ column to provider_keys - Update select_best_key() SQL to sort by last_used_at ASC NULLS FIRST - Update record_key_usage() to set last_used_at = NOW() on each use - Bump SCHEMA_VERSION to 10
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
-- 20260401000001_provider_keys_last_used.sql
|
||||
-- Key Pool LRU: 记录每个 key 最后使用时间
|
||||
ALTER TABLE provider_keys ADD COLUMN IF NOT EXISTS last_used_at TIMESTAMPTZ;
|
||||
@@ -4,7 +4,7 @@ use sqlx::postgres::PgPoolOptions;
|
||||
use sqlx::PgPool;
|
||||
use crate::error::SaasResult;
|
||||
|
||||
const SCHEMA_VERSION: i32 = 9;
|
||||
const SCHEMA_VERSION: i32 = 10;
|
||||
|
||||
/// 初始化数据库
|
||||
pub async fn init_db(database_url: &str) -> SaasResult<PgPool> {
|
||||
|
||||
@@ -51,7 +51,7 @@ pub async fn select_best_key(db: &PgPool, provider_id: &str, enc_key: &[u8; 32])
|
||||
LEFT JOIN key_usage_window uw ON pk.id = uw.key_id AND uw.window_minute = $1
|
||||
WHERE pk.provider_id = $2 AND pk.is_active = TRUE
|
||||
AND (pk.cooldown_until IS NULL OR pk.cooldown_until <= $3)
|
||||
ORDER BY pk.priority ASC"
|
||||
ORDER BY pk.priority ASC, pk.last_used_at ASC NULLS FIRST"
|
||||
).bind(¤t_minute).bind(provider_id).bind(&now).fetch_all(db).await?;
|
||||
|
||||
for (id, key_value, priority, max_rpm, max_tpm, quota_reset_interval, req_count, token_count) in &rows {
|
||||
@@ -165,6 +165,12 @@ pub async fn record_key_usage(
|
||||
.bind(tokens).bind(&chrono::Utc::now().to_rfc3339()).bind(key_id)
|
||||
.execute(db).await?;
|
||||
|
||||
// 更新最后使用时间 (LRU 排序依据)
|
||||
sqlx::query("UPDATE provider_keys SET last_used_at = NOW() WHERE id = $1")
|
||||
.bind(key_id)
|
||||
.execute(db)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user