diff --git a/backend/crates/gateway-store/src/redis/credential_leases/mod.rs b/backend/crates/gateway-store/src/redis/credential_leases/mod.rs index 83d028817..e1387a492 100644 --- a/backend/crates/gateway-store/src/redis/credential_leases/mod.rs +++ b/backend/crates/gateway-store/src/redis/credential_leases/mod.rs @@ -7,7 +7,6 @@ use std::time::Duration; use async_trait::async_trait; use chrono::{DateTime, Utc}; -use futures::future::join_all; use gateway_core::account::{AccountRuntimeSignals, ProviderAccountId}; use gateway_core::lifecycle::REQUEST_LEASE_TTL; use gateway_core::policy::ClientApiKeyId; @@ -27,6 +26,8 @@ use super::{MAX_REDIS_EXACT_INTEGER, namespace, resource_fingerprint}; mod details; const SIGNAL_TTL_MILLIS: u64 = 24 * 60 * 60 * 1_000; +// 限制单次脚本处理的账号数，避免大池查询长时间占用 Redis +const RUNTIME_SIGNAL_BATCH_SIZE: usize = 128; const OAUTH_REFRESH_LEASE_TTL: Duration = Duration::from_secs(5 * 60); const OAUTH_REFRESH_CAPACITY_RESOURCE: &str = "oauth-refresh-global"; @@ -109,10 +110,14 @@ return removed const SIGNAL_SCRIPT: &str = r#" local clock = redis.call('TIME') local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000) -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now_ms) -local in_flight = redis.call('ZCARD', KEYS[1]) -local last_started = redis.call('GET', KEYS[2]) or '0' -return {tostring(in_flight), tostring(last_started)} +local signals = {} +for i = 1, #KEYS, 2 do + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now_ms) + local in_flight = redis.call('ZCARD', KEYS[i]) + local last_started = redis.call('GET', KEYS[i + 1]) or '0' + signals[#signals + 1] = {tostring(in_flight), tostring(last_started)} +end +return signals "#; const ADVANCE_SCHEDULING_CURSOR_SCRIPT: &str = r#" @@ -528,45 +533,56 @@ impl RedisCredentialLeaseRepository { require_nonempty("credential runtime signal", "resource_id", resource_id)?; } let mut signals = Vec::with_capacity(resource_ids.len()); - for ids in resource_ids.chunks(super::ACCOUNT_STATE_READ_CONCURRENCY) { - let batch = join_all(ids.iter().cloned().map(|id| self.load_signal(id, pool))).await; - signals.extend(batch.into_iter().collect::>>()?); + for ids in resource_ids.chunks(RUNTIME_SIGNAL_BATCH_SIZE) { + signals.extend(self.load_signals(ids, pool).await?); } Ok(signals) } - async fn load_signal( + async fn load_signals( &self, - resource_id: String, + resource_ids: &[String], pool: ProviderConcurrencyPool, - ) -> StoreResult { - let request = CredentialLeaseRequest { - scope: CredentialLeaseScope::scheduling(pool), - resource_id: resource_id.clone(), - owner_id: "signal-reader".to_owned(), - ttl: Duration::from_secs(1), - }; - let keys = self.keys(&request)?; + ) -> StoreResult> { + let script = Script::new(SIGNAL_SCRIPT); + let mut invocation = script.prepare_invoke(); + for resource_id in resource_ids { + let request = CredentialLeaseRequest { + scope: CredentialLeaseScope::scheduling(pool), + resource_id: resource_id.clone(), + owner_id: "signal-reader".to_owned(), + ttl: Duration::from_secs(1), + }; + let keys = self.keys(&request)?; + invocation.key(&keys[0]).key(&keys[2]); + } let mut connection = self.connection.clone(); - let (in_flight, last_started): (String, String) = Script::new(SIGNAL_SCRIPT) - .key(&keys[0]) - .key(&keys[2]) + let signals: Vec<(String, String)> = invocation .invoke_async(&mut connection) .await - .map_err(|source| redis_unavailable("load credential runtime signal", source))?; - let in_flight = in_flight.parse::().map_err(|source| { - invalid("Redis returned an invalid in-flight count").with_source(source) - })?; - let last_started_at = if last_started == "0" { - None - } else { - Some(timestamp(&last_started)?) - }; - Ok(CredentialRuntimeSignal { - resource_id, - in_flight, - last_started_at, - }) + .map_err(|source| redis_unavailable("load credential runtime signals", source))?; + if signals.len() != resource_ids.len() { + return Err(invalid("Redis returned an invalid runtime signal count")); + } + resource_ids + .iter() + .zip(signals) + .map(|(resource_id, (in_flight, last_started))| { + let in_flight = in_flight.parse::().map_err(|source| { + invalid("Redis returned an invalid in-flight count").with_source(source) + })?; + let last_started_at = if last_started == "0" { + None + } else { + Some(timestamp(&last_started)?) + }; + Ok(CredentialRuntimeSignal { + resource_id: resource_id.clone(), + in_flight, + last_started_at, + }) + }) + .collect() } } diff --git a/backend/crates/gateway-store/tests/redis/credential_leases/mod.rs b/backend/crates/gateway-store/tests/redis/credential_leases/mod.rs index abf91bffd..7735236b5 100644 --- a/backend/crates/gateway-store/tests/redis/credential_leases/mod.rs +++ b/backend/crates/gateway-store/tests/redis/credential_leases/mod.rs @@ -2,6 +2,7 @@ mod details; +use std::collections::BTreeMap; use std::time::Duration; use gateway_store::redis::{ @@ -741,33 +742,121 @@ async fn runtime_signals_preserve_order_and_counts_across_batches() { let Some((repository, _, _)) = repository().await else { return; }; - let ids = (0..1000) + let mut ids = (0..3000) .map(|index| format!("acct_batch_{index}")) .collect::>(); let mut guards = Vec::new(); - for index in [0, 127, 128, 500, 999] { - let request = scheduling_request(&ids[index], "batch-test", 2); - guards.push(acquired( + let mut expected = BTreeMap::new(); + for (index, count) in [(0, 1), (127, 2), (128, 1), (500, 2), (2999, 1)] { + for _ in 0..count { + let request = scheduling_request(&ids[index], "batch-test", 2); + let guard = acquired( + repository + .try_acquire_bounded_lease(&request) + .await + .unwrap(), + ); + let started = guard.grant().unwrap().expires_at - chrono::Duration::seconds(60); + expected.insert(ids[index].clone(), (count, Some(started))); + guards.push(guard); + } + } + // 保留池不计入普通池负载，但两个池共用最近调度时间 + for index in [257, 500] { + let request = CredentialBoundedLeaseRequest { + scope: CredentialLeaseScope::ProviderAccountReserved, + ..scheduling_request(&ids[index], "reserved-batch-test", 1) + }; + let guard = acquired( repository .try_acquire_bounded_lease(&request) .await .unwrap(), - )); + ); + let started = guard.grant().unwrap().expires_at - chrono::Duration::seconds(60); + expected.entry(ids[index].clone()).or_insert((0, None)).1 = Some(started); + guards.push(guard); } - let started = std::time::Instant::now(); + let expired_request = CredentialBoundedLeaseRequest { + ttl: Duration::from_millis(25), + ..scheduling_request(&ids[256], "expired-batch-test", 1) + }; + let expired = acquired( + repository + .try_acquire_bounded_lease(&expired_request) + .await + .unwrap(), + ); + guards.push(expired); + // 长租约保留同一个 ZSET，确保读取会清除已过期成员而非只依赖整个 key 到期 + let live_request = scheduling_request(&ids[256], "live-batch-test", 2); + let live = acquired( + repository + .try_acquire_bounded_lease(&live_request) + .await + .unwrap(), + ); + let started = live.grant().unwrap().expires_at - chrono::Duration::seconds(60); + expected.insert(ids[256].clone(), (1, Some(started))); + guards.push(live); + tokio::time::sleep(Duration::from_millis(50)).await; + + // 调用方的顺序与重复项不因批处理而被排序或去重 + ids.extend([ids[127].clone(), ids[257].clone()]); + ids.reverse(); let signals = repository.credential_runtime_signals(&ids).await.unwrap(); assert_eq!(signals.len(), ids.len()); - for (index, (signal, id)) in signals.iter().zip(&ids).enumerate() { + for (signal, id) in signals.iter().zip(&ids) { assert_eq!(&signal.resource_id, id); - let acquired = [0, 127, 128, 500, 999].contains(&index); - assert_eq!(signal.in_flight, u32::from(acquired)); - assert_eq!(signal.last_started_at.is_some(), acquired); + let (count, started) = expected.get(id).copied().unwrap_or((0, None)); + assert_eq!(signal.in_flight, count, "{id}"); + assert_eq!(signal.last_started_at, started, "{id}"); } - eprintln!( - "1000 Redis runtime signals: elapsed={:?}", - started.elapsed() - ); for guard in guards { guard.release().await.unwrap(); } } + +#[tokio::test] +async fn runtime_signals_reject_invalid_data_instead_of_returning_a_partial_pool() { + let Some((repository, mut connection, namespace)) = repository().await else { + return; + }; + assert!( + repository + .credential_runtime_signals(&[]) + .await + .unwrap() + .is_empty() + ); + let ids = (0..300) + .map(|index| format!("acct_invalid_signal_{index}")) + .collect::>(); + let request = scheduling_request(&ids[299], "invalid-signal-test", 1); + let guard = acquired( + repository + .try_acquire_bounded_lease(&request) + .await + .unwrap(), + ); + let keys = redis::cmd("KEYS") + .arg(format!("{namespace}:*:last-started")) + .query_async::>(&mut connection) + .await + .unwrap(); + assert_eq!(keys.len(), 1); + redis::cmd("SET") + .arg(&keys[0]) + .arg("invalid-timestamp") + .query_async::<()>(&mut connection) + .await + .unwrap(); + assert!(matches!( + repository + .credential_runtime_signals(&ids) + .await + .unwrap_err(), + gateway_store::StoreError::InvalidData { .. } + )); + guard.release().await.unwrap(); +}