perf(store): batch account lease runtime reads

Read up to 128 accounts per Redis script while preserving input order, pool isolation, shared scheduling timestamps, expiry cleanup, and error propagation. Extend real Redis coverage for mixed lease states and invalid batch data. Co-authored-by: Codex <noreply@openai.com>
2026-11-06 08:06:19 +00:00 · 2026-11-01 23:23:49 +08:00
parent 79df82a4f0
commit c8ffb1a8c8
2 changed files with 154 additions and 49 deletions
@@ -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::<StoreResult<Vec<_>>>()?);
+        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<CredentialRuntimeSignal> {
-        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<Vec<CredentialRuntimeSignal>> {
+        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::<u32>().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::<u32>().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()
    }
 }

@@ -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::<Vec<_>>();
    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::<Vec<_>>();
+    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::<Vec<String>>(&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();
+}