fix(settings): apply refresh margin immediately

Schedule token refreshes at the exact configured lead time and reschedule existing refreshable accounts when that setting changes. Preserve current retry schedules when unrelated settings are saved.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-29 02:34:25 +08:00
co-authored by Codex
parent d5fc778ea7
commit 2e861c03c4
4 changed files with 153 additions and 39 deletions
@@ -667,32 +667,21 @@ impl ProviderRefreshPolicy {
self.concurrency
}
/// 按账号稳定扰动刷新提前量;相同账号与过期时间跨进程得到同一执行时刻。
/// 按配置的精确提前量计算刷新时刻。
pub fn next_attempt_at(
self,
account_id: &ProviderAccountId,
_account_id: &ProviderAccountId,
access_token_expires_at: SystemTime,
observed_at: SystemTime,
) -> Result<SystemTime, ProviderStoreError> {
let remaining = access_token_expires_at
.duration_since(observed_at)
.map_err(|_| invalid_refresh_policy("schedule expired access token"))?;
let base_margin_seconds = self.margin.as_secs();
if base_margin_seconds == 0 {
if self.margin >= remaining {
return Ok(observed_at);
}
let factor = stable_factor(account_id.as_str(), "normal", 850, 1_150);
let jittered_seconds = base_margin_seconds
.saturating_mul(u64::from(factor))
.saturating_add(500)
/ 1_000;
if jittered_seconds >= remaining.as_secs() {
return Ok(observed_at);
}
let lead = jittered_seconds.max(1);
access_token_expires_at
.checked_sub(Duration::from_secs(lead))
.checked_sub(self.margin)
.ok_or_else(|| invalid_refresh_policy("schedule access token refresh"))
}
}
@@ -57,7 +57,7 @@ fn refresh_policy_requires_a_positive_margin() {
}
#[test]
fn refresh_schedule_is_stable_and_never_uses_query_time_margin() {
fn refresh_schedule_should_use_the_exact_configured_margin() {
let account = ProviderAccountId::new("acct_stable").expect("valid account");
let policy = ProviderRefreshPolicy::try_new(
Duration::from_secs(3_600),
@@ -70,13 +70,7 @@ fn refresh_schedule_is_stable_and_never_uses_query_time_margin() {
let first = policy
.next_attempt_at(&account, expires_at, observed_at)
.expect("schedule refresh");
let second = policy
.next_attempt_at(&account, expires_at, observed_at)
.expect("schedule refresh again");
assert_eq!(first, second);
assert!(first > observed_at);
assert!(first < expires_at);
assert_eq!(first, expires_at - Duration::from_secs(3_600));
}
#[test]
@@ -225,6 +225,21 @@ pub(crate) async fn update_runtime_settings_in_transaction(
update: &RuntimeSettingsUpdate,
) -> StoreResult<Revision> {
update.validate()?;
let refresh_margin_seconds =
i64::try_from(update.refresh_margin_seconds).map_err(|_| invalid_numeric())?;
let previous_refresh_margin_seconds = sqlx::query_scalar::<_, i64>(
"select refresh_margin_seconds
from runtime_settings
where id = 1
for update",
)
.fetch_optional(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("lock runtime settings for update"))?
.ok_or_else(|| StoreError::NotFound {
entity: "runtime settings",
id: "1".to_owned(),
})?;
let next = sqlx::query_scalar::<_, i64>(
"update runtime_settings
set config_revision = config_revision + 1,
@@ -243,7 +258,7 @@ pub(crate) async fn update_runtime_settings_in_transaction(
returning config_revision",
)
.bind(update.admin_api_key.as_deref())
.bind(i64::try_from(update.refresh_margin_seconds).map_err(|_| invalid_numeric())?)
.bind(refresh_margin_seconds)
.bind(i64::from(update.refresh_concurrency))
.bind(i64::from(update.max_concurrent_per_account))
.bind(i64::try_from(update.request_interval_ms).map_err(|_| invalid_numeric())?)
@@ -259,9 +274,35 @@ pub(crate) async fn update_runtime_settings_in_transaction(
entity: "runtime settings",
id: "1".to_owned(),
})?;
if previous_refresh_margin_seconds != refresh_margin_seconds {
reschedule_provider_refreshes_in_transaction(transaction, refresh_margin_seconds).await?;
}
Revision::new(u64::try_from(next).map_err(|_| invalid_numeric())?)
}
async fn reschedule_provider_refreshes_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
refresh_margin_seconds: i64,
) -> StoreResult<()> {
sqlx::query(
"update provider_accounts
set next_refresh_at = case
when access_token_expires_at <= now()
or extract(epoch from (access_token_expires_at - now())) <= $1
then now()
else access_token_expires_at - ($1::double precision * interval '1 second')
end,
updated_at = now()
where has_refresh_token
and access_token_expires_at is not null",
)
.bind(refresh_margin_seconds)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("reschedule provider credential refreshes"))?;
Ok(())
}
pub(crate) async fn bump_config_revision_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Revision> {
@@ -1,12 +1,16 @@
use std::collections::BTreeMap;
use gateway_store::postgres::RuntimeSettingsUpdate;
use chrono::{DateTime, TimeDelta, Utc};
use gateway_store::postgres::{
PgRuntimeSettingsRepository, RuntimeSettingsRepository, RuntimeSettingsUpdate,
};
#[test]
fn runtime_settings_keep_account_rotation_global() {
let settings = RuntimeSettingsUpdate {
use super::TestDatabase;
fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate {
RuntimeSettingsUpdate {
admin_api_key: None,
refresh_margin_seconds: 3_600,
refresh_margin_seconds,
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
@@ -18,24 +22,110 @@ fn runtime_settings_keep_account_rotation_global() {
usage_retention_days: 31,
ops_event_retention_days: 30,
audit_retention_days: 90,
};
}
}
#[test]
fn runtime_settings_keep_account_rotation_global() {
let settings = settings_with_margin(3_600);
assert!(settings.validate().is_ok());
}
#[test]
fn runtime_settings_reject_invalid_model_mapping() {
let settings = RuntimeSettingsUpdate {
admin_api_key: None,
refresh_margin_seconds: 3_600,
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
rotation_strategy: "smart".to_owned(),
model_mappings: BTreeMap::from([("".to_owned(), "gpt-5.5".to_owned())]),
usage_retention_days: 31,
ops_event_retention_days: 30,
audit_retention_days: 90,
..settings_with_margin(3_600)
};
assert!(settings.validate().is_err());
}
#[tokio::test]
async fn refresh_margin_change_should_reschedule_existing_refreshable_accounts() {
let Some(database) = TestDatabase::create("refresh_margin_reschedule").await else {
return;
};
let expires_at = timestamp_micros(Utc::now() + TimeDelta::hours(2));
insert_refreshable_account(
&database.pool,
"acct_refresh_margin_changed",
expires_at,
expires_at - TimeDelta::hours(1),
)
.await;
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
repository
.update_runtime_settings(settings_with_margin(1_800))
.await
.expect("update refresh margin");
assert_eq!(
account_next_refresh_at(&database.pool, "acct_refresh_margin_changed").await,
expires_at - TimeDelta::seconds(1_800)
);
database.close().await;
}
#[tokio::test]
async fn unchanged_refresh_margin_should_preserve_existing_retry_schedule() {
let Some(database) = TestDatabase::create("refresh_margin_unchanged").await else {
return;
};
let expires_at = timestamp_micros(Utc::now() + TimeDelta::hours(2));
let retry_at = timestamp_micros(Utc::now() + TimeDelta::minutes(5));
insert_refreshable_account(
&database.pool,
"acct_refresh_margin_unchanged",
expires_at,
retry_at,
)
.await;
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
repository
.update_runtime_settings(settings_with_margin(3_600))
.await
.expect("update unrelated runtime settings");
assert_eq!(
account_next_refresh_at(&database.pool, "acct_refresh_margin_unchanged").await,
retry_at
);
database.close().await;
}
async fn insert_refreshable_account(
pool: &sqlx::PgPool,
account_id: &str,
expires_at: DateTime<Utc>,
next_refresh_at: DateTime<Utc>,
) {
sqlx::query(
"insert into provider_accounts (
id, provider_kind, name, upstream_user_id, authentication_kind,
provider_credentials_json, has_refresh_token, access_token_expires_at,
next_refresh_at, availability, availability_observed_at, created_at, updated_at
) values ($1, 'openai', $1, $1, 'oauth', '{}'::jsonb, true, $2, $3,
'ready', now(), now(), now())",
)
.bind(account_id)
.bind(expires_at)
.bind(next_refresh_at)
.execute(pool)
.await
.expect("insert refreshable account");
}
async fn account_next_refresh_at(pool: &sqlx::PgPool, account_id: &str) -> DateTime<Utc> {
sqlx::query_scalar("select next_refresh_at from provider_accounts where id = $1")
.bind(account_id)
.fetch_one(pool)
.await
.expect("load next refresh time")
}
fn timestamp_micros(value: DateTime<Utc>) -> DateTime<Utc> {
DateTime::from_timestamp_micros(value.timestamp_micros()).expect("valid test timestamp")
}