fix(openai): recheck cross-account capacity while queued

2026-11-06 08:06:19 +00:00 · 2026-11-06 15:02:16 +08:00
parent 6de1beb1df
commit 4e4cd96478
4 changed files with 137 additions and 28 deletions
@@ -11,7 +11,7 @@ use futures_timer::Delay;

 const MAX_TOTAL_WAITING: usize = 1_024;
 const CAPACITY_RECHECK_INTERVAL: Duration = Duration::from_millis(100);
-const BINDING_RECHECK_INTERVAL: Duration = Duration::from_millis(500);
+const NON_HEAD_RECHECK_INTERVAL: Duration = Duration::from_millis(500);

 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
 pub struct ConcurrencyQueuePolicy {
@@ -235,11 +235,11 @@ impl<K: Eq + Hash> WaitTicket<K> {
        self.turn().await
    }

-    async fn retry_or_recheck_binding(&self) -> Result<(), QueueRejection> {
+    async fn retry_or_recheck(&self) -> Result<(), QueueRejection> {
        if !self.is_head() {
-            // 非队首由位置释放唤醒，低频兜底只用于检查外部绑定迁移
+            // 非队首由位置释放唤醒，低频重查覆盖外部调度状态变化
            let turn = self.turn().fuse();
-            let recheck = Delay::new(BINDING_RECHECK_INTERVAL).fuse();
+            let recheck = Delay::new(NON_HEAD_RECHECK_INTERVAL).fuse();
            pin_mut!(turn, recheck);
            return select_biased! {
                result = turn => result,
@@ -326,14 +326,14 @@ impl<'a, K: Clone + Eq + Hash> CapacityWait<'a, K> {
        self
    }

-    /// 外部绑定可在队列等待期间变化，非队首只需低频检查绑定并由 can_try 复核资格
+    /// 允许非队首低频复核外部调度状态，当前队列准入仍由 can_try 裁决
    #[must_use]
    pub const fn with_periodic_recheck(mut self, enabled: bool) -> Self {
        self.periodic_recheck = enabled;
        self
    }

-    /// 当前等待位置，供调用方在非队首时只检查外部绑定而避免重复加载候选池
+    /// 当前等待位置，供调用方按请求归属限定重查范围
    #[must_use]
    pub fn queued_key(&self) -> Option<&K> {
        self.ticket.as_ref().map(WaitTicket::key)
@@ -388,7 +388,7 @@ impl<'a, K: Clone + Eq + Hash> CapacityWait<'a, K> {
        }
        if let Some(ticket) = &self.ticket {
            if self.periodic_recheck {
-                ticket.retry_or_recheck_binding().await?;
+                ticket.retry_or_recheck().await?;
            } else {
                ticket.retry().await?;
            }
@@ -543,22 +543,6 @@ impl CodexCredentialSelector {
        } else {
            ProviderConcurrencyPool::Shared
        };
-        let mut waiting = CapacityWait::new(
-            if prioritized {
-                &self.reserved_waiting
-            } else {
-                &self.waiting
-            },
-            queue_policy,
-            request.attempt.deadline().at(),
-            request.attempt.concurrency_wait_budget(),
-        )
-        .with_periodic_recheck(binding_key.is_some())
-        .with_priority(if prioritized {
-            WaitPriority::High
-        } else {
-            WaitPriority::Normal
-        });
        let continuation_account = match request.attempt.continuation_attempt() {
            ContinuationAttempt::Native => request
                .attempt
@@ -581,6 +565,22 @@ impl CodexCredentialSelector {
            return Err(CredentialSelectionError::NoEligibleCredential);
        }
        let pinned_account = required_account.or_else(|| continuation_account.clone());
+        let mut waiting = CapacityWait::new(
+            if prioritized {
+                &self.reserved_waiting
+            } else {
+                &self.waiting
+            },
+            queue_policy,
+            request.attempt.deadline().at(),
+            request.attempt.concurrency_wait_budget(),
+        )
+        .with_periodic_recheck(binding_key.is_some() || pinned_account.is_none())
+        .with_priority(if prioritized {
+            WaitPriority::High
+        } else {
+            WaitPriority::Normal
+        });
        let projection_key = match (model, request.attempt.account_scope()) {
            (CodexSelectionModel::Requested(model), Some(scope))
                if !diagnostic
@@ -623,13 +623,18 @@ impl CodexCredentialSelector {
                return Err(CredentialSelectionError::ContinuationOwnerChanged);
            }
            let force_full = std::mem::take(&mut refresh_candidates);
-            // 等待期间只复核当前账号，绑定迁移、资格变化或取得队首空位后才展开候选池
+            // 只有固定归属的等待者可缩小重查范围，可分流请求还需发现其他账号释放的容量
            let queued_account = waiting
                .queued_key()
-                .filter(|_| {
+                .filter(|&queued| {
                    !force_full
                        && queued_binding == binding
                        && request.attempt.request_policy_context().is_none()
+                        && (pinned_account.as_ref() == Some(queued)
+                            || (follow_only || affinity_mode == AccountAffinity::Strict)
+                                && binding
+                                    .as_ref()
+                                    .is_some_and(|binding| binding.account_id() == queued))
                })
                .cloned();
            queued_binding = binding.clone();
@@ -304,7 +304,7 @@ async fn projection_fill_merges_equal_keys_and_releases_cancelled_owners() {
 }

 #[tokio::test]
-async fn waiting_requests_recheck_only_their_account_until_capacity_opens() {
+async fn strict_bound_waiters_recheck_only_their_account_until_capacity_opens() {
    let store = Arc::new(MemoryAccountStore::default());
    create_account(&store, "acct_first", "test-first");
    create_account(&store, "acct_second", "test-second");
@@ -357,3 +357,106 @@ async fn waiting_requests_recheck_only_their_account_until_capacity_opens() {
        assert_eq!(lease.capacity_snapshot().unwrap().total_slots(), 4);
    }
 }
+
+#[tokio::test]
+async fn movable_waiters_use_other_accounts_even_before_queue_promotion() {
+    for (mode, has_session, bind_session) in [
+        (AccountAffinity::Relaxed, false, false),
+        (AccountAffinity::Strict, true, false),
+        (AccountAffinity::Preferred, true, true),
+    ] {
+        for has_waiter_ahead in [false, true] {
+            let store = Arc::new(MemoryAccountStore::default());
+            create_account(&store, "acct_first", "test-first");
+            create_account(&store, "acct_second", "test-second");
+            let first = ProviderAccountId::new("acct_first").unwrap();
+            let second = ProviderAccountId::new("acct_second").unwrap();
+            let leases = Arc::new(TestLeaseCoordinator::default());
+            let feedback = Arc::new(AccountFeedbackStats::default());
+            for _ in 0..4 {
+                feedback.report(
+                    &ProviderKind::new("openai").unwrap(),
+                    &second,
+                    AccountAttemptFeedback::Failed {
+                        first_output_ms: None,
+                    },
+                );
+            }
+            let selector = selector_with_runtime(
+                &store,
+                leases.clone(),
+                Arc::new(MemorySessionAffinity::default()),
+                feedback,
+                Arc::new(MemoryCooldownPort::new()),
+            );
+            // 健康分保证请求排入 A，已有队首时也不因 B 的队列较短而提前排入 B
+            let policy = account_policy()
+                .with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy {
+                    max_waiting: 10,
+                    timeout: Duration::from_secs(5),
+                })
+                .with_smart_scheduling(
+                    gateway_core::account::SmartSchedulingConfig::new([
+                        0.0, 0.0, 1.0, 0.0, 0.0, 0.2,
+                    ])
+                    .unwrap(),
+                );
+            let head_attempt =
+                projection_attempt("key_codex_contract", contract_account_scope(), policy);
+            let head_key = ProviderSessionAffinityKey::try_new("fixed-head").unwrap();
+            assert_eq!(
+                select(&selector, &head_attempt, &head_key, "gpt-5.4")
+                    .await
+                    .unwrap()
+                    .account_id(),
+                &first,
+            );
+            let attempt = projection_attempt(
+                "key_codex_contract",
+                contract_account_scope(),
+                policy.with_openai_account_affinity(mode),
+            );
+            let key = ProviderSessionAffinityKey::try_new("movable-waiter").unwrap();
+            if bind_session {
+                assert_eq!(
+                    select(&selector, &attempt, &key, "gpt-5.4")
+                        .await
+                        .unwrap()
+                        .account_id(),
+                    &first,
+                );
+            }
+            for id in [&first, &second] {
+                leases.signals.lock().unwrap().insert(id.clone(), signal(2));
+            }
+            let mut head = Box::pin(select(&selector, &head_attempt, &head_key, "gpt-5.4"));
+            if has_waiter_ahead {
+                assert!(head.as_mut().now_or_never().is_none());
+            }
+            let request = SelectCodexCredential {
+                upstream_model: "gpt-5.4",
+                request_url: &Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap(),
+                attempt: &attempt,
+                session_affinity_key: has_session.then_some(&key),
+            };
+            let mut pending = Box::pin(selector.select(&request));
+            assert!(pending.as_mut().now_or_never().is_none());
+            // A 及其队首始终占位，只释放 B，验证请求能在晋升或超时前使用新空位
+            leases.signals.lock().unwrap().remove(&second);
+            let selected = tokio::time::timeout(Duration::from_secs(2), pending)
+                .await
+                .unwrap_or_else(|_| {
+                    panic!(
+                        "missed free account: mode={mode:?}, has_waiter_ahead={has_waiter_ahead}"
+                    )
+                })
+                .unwrap();
+            assert_eq!(selected.account_id(), &second);
+            let capacity = selected.capacity_snapshot().unwrap();
+            assert_eq!((capacity.used_slots(), capacity.total_slots()), (3, 4));
+            if has_waiter_ahead {
+                assert!(head.as_mut().now_or_never().is_none());
+            }
+        }
+    }
+}
@@ -838,8 +838,9 @@ Core 的 `concurrency` 拥有中立的有界等待位置、优先级与同级 FI
 账号并发明细由 Redis 租约关联的最小入口事实投影，不读取请求历史。事实与租约原子创建，按 fencing 校验续期和释放，异常退出后随 TTL 回收；管理端只对普通池做有界分页读取

 新请求不能越过同级或更高优先级的等待者，取消或完成等待通过 Drop 回收位置并立即唤醒后继
-队首每 100ms 重查容量，带绑定的非队首每 500ms 检查迁移，普通非队首等待晋升通知
-OpenAI 内置调度在绑定稳定时只重读排队目标的资格、容量、目录和冷却，迁移、资格变化或发现空位后展开完整候选池
+队首每 100ms 重查容量，需要复核会话绑定或跨账号容量的非队首每 500ms 重查，其余非队首等待晋升通知
+OpenAI 内置调度仅在排队目标与固定归属或严格会话绑定一致时，单独重读该账号的资格、容量、目录和冷却
+绑定迁移、资格变化或发现空位后恢复完整候选评估，可跨账号调度的等待者按重查周期评估整个候选池
 完整容量观测和发送前准入发生在展开之后，策略插件仍使用完整候选合同，等待重查保持在同一次未发送准备阶段，不增加 attempt

 Key 的 RPM 在成功准入时才计数，金额限制在入队前及成功准入后检查。