codex_helper_core/
lb.rs

1use std::collections::HashMap;
2use std::collections::HashSet;
3use std::sync::{Arc, Mutex};
4
5use crate::config::{ServiceConfig, UpstreamConfig};
6use crate::runtime_identity::ProviderEndpointKey;
7use tracing::info;
8
9pub const FAILURE_THRESHOLD: u32 = 3;
10pub const COOLDOWN_SECS: u64 = 30;
11
12#[derive(Debug, Clone, Copy)]
13pub struct CooldownBackoff {
14    pub factor: u64,
15    pub max_secs: u64,
16}
17
18impl CooldownBackoff {
19    pub(crate) fn effective_cooldown_secs(&self, base_secs: u64, penalty_streak: u32) -> u64 {
20        if base_secs == 0 {
21            return 0;
22        }
23        if self.factor <= 1 {
24            return base_secs;
25        }
26        let cap = if self.max_secs == 0 {
27            base_secs
28        } else {
29            self.max_secs.max(base_secs)
30        };
31
32        let mut secs = base_secs;
33        for _ in 0..penalty_streak.min(64) {
34            secs = secs.saturating_mul(self.factor);
35            if secs >= cap {
36                return cap;
37            }
38        }
39        secs.min(cap)
40    }
41}
42
43#[derive(Debug, Clone, Default)]
44pub struct LbState {
45    pub failure_counts: Vec<u32>,
46    pub cooldown_until: Vec<Option<std::time::Instant>>,
47    pub usage_exhausted: Vec<bool>,
48    pub last_good_index: Option<usize>,
49    pub penalty_streak: Vec<u32>,
50    pub(crate) upstream_signature: Vec<String>,
51}
52
53impl LbState {
54    pub(crate) fn ensure_layout(&mut self, service_name: &str, upstreams: &[UpstreamConfig]) {
55        let signature = upstreams
56            .iter()
57            .enumerate()
58            .map(|(idx, upstream)| upstream_signature_key(service_name, idx, upstream))
59            .collect::<Vec<_>>();
60        let legacy_signature = upstreams
61            .iter()
62            .map(|upstream| upstream.base_url.clone())
63            .collect::<Vec<_>>();
64
65        if has_duplicate_signatures(&signature) {
66            self.reset_for_layout(signature);
67            return;
68        }
69
70        let len = upstreams.len();
71        if self.upstream_signature == signature
72            && self.failure_counts.len() == len
73            && self.cooldown_until.len() == len
74            && self.usage_exhausted.len() == len
75            && self.penalty_streak.len() == len
76        {
77            return;
78        }
79
80        self.migrate_layout(signature, legacy_signature);
81    }
82
83    fn reset_for_layout(&mut self, signature: Vec<String>) {
84        let len = signature.len();
85        self.failure_counts = vec![0; len];
86        self.cooldown_until = vec![None; len];
87        self.usage_exhausted = vec![false; len];
88        self.penalty_streak = vec![0; len];
89        // upstream 布局变化时，原来的粘性索引不再可信，直接清空。
90        self.last_good_index = None;
91        self.upstream_signature = signature;
92    }
93
94    fn migrate_layout(&mut self, signature: Vec<String>, legacy_signature: Vec<String>) {
95        if self.upstream_signature.is_empty() {
96            self.reset_for_layout(signature);
97            return;
98        }
99
100        let old_signature = std::mem::take(&mut self.upstream_signature);
101        if has_duplicate_signatures(&old_signature) {
102            self.reset_for_layout(signature);
103            return;
104        }
105
106        let old_index_by_signature = old_signature
107            .iter()
108            .enumerate()
109            .map(|(idx, key)| (key.clone(), idx))
110            .collect::<std::collections::HashMap<_, _>>();
111        let legacy_fallback_enabled = !has_duplicate_signatures(&legacy_signature);
112
113        let old_failure_counts = std::mem::take(&mut self.failure_counts);
114        let old_cooldown_until = std::mem::take(&mut self.cooldown_until);
115        let old_usage_exhausted = std::mem::take(&mut self.usage_exhausted);
116        let old_penalty_streak = std::mem::take(&mut self.penalty_streak);
117        let old_last_good_index = self.last_good_index.take();
118
119        let len = signature.len();
120        self.failure_counts = vec![0; len];
121        self.cooldown_until = vec![None; len];
122        self.usage_exhausted = vec![false; len];
123        self.penalty_streak = vec![0; len];
124
125        for (new_idx, key) in signature.iter().enumerate() {
126            let old_idx = old_index_by_signature.get(key).copied().or_else(|| {
127                legacy_fallback_enabled
128                    .then(|| legacy_signature.get(new_idx))
129                    .flatten()
130                    .and_then(|legacy_key| old_index_by_signature.get(legacy_key).copied())
131            });
132            let Some(old_idx) = old_idx else {
133                continue;
134            };
135            self.failure_counts[new_idx] = old_failure_counts.get(old_idx).copied().unwrap_or(0);
136            self.cooldown_until[new_idx] = old_cooldown_until.get(old_idx).and_then(|until| *until);
137            self.usage_exhausted[new_idx] =
138                old_usage_exhausted.get(old_idx).copied().unwrap_or(false);
139            self.penalty_streak[new_idx] = old_penalty_streak.get(old_idx).copied().unwrap_or(0);
140        }
141
142        self.last_good_index = old_last_good_index.and_then(|old_idx| {
143            old_signature.get(old_idx).and_then(|key| {
144                signature
145                    .iter()
146                    .position(|new_key| new_key == key)
147                    .or_else(|| {
148                        legacy_fallback_enabled
149                            .then(|| {
150                                legacy_signature
151                                    .iter()
152                                    .position(|legacy_key| legacy_key == key)
153                            })
154                            .flatten()
155                    })
156            })
157        });
158        self.upstream_signature = signature;
159    }
160}
161
162fn has_duplicate_signatures(values: &[String]) -> bool {
163    let mut seen = HashSet::new();
164    values.iter().any(|value| !seen.insert(value))
165}
166
167fn upstream_signature_key(
168    service_name: &str,
169    upstream_index: usize,
170    upstream: &UpstreamConfig,
171) -> String {
172    let provider_id = upstream
173        .tags
174        .get("provider_id")
175        .cloned()
176        .unwrap_or_else(|| format!("{service_name}#{upstream_index}"));
177    let endpoint_id = upstream
178        .tags
179        .get("endpoint_id")
180        .cloned()
181        .unwrap_or_else(|| upstream_index.to_string());
182    let provider_endpoint = ProviderEndpointKey::new(service_name, provider_id, endpoint_id);
183    format!("{}|{}", provider_endpoint.stable_key(), upstream.base_url)
184}
185
186/// Upstream selection result
187#[derive(Debug, Clone)]
188pub struct SelectedUpstream {
189    pub station_name: String,
190    pub index: usize,
191    pub upstream: UpstreamConfig,
192}
193
194/// 简单的负载选择器，当前仅按权重随机，未来可扩展为按 usage / 失败次数等切换。
195#[derive(Clone)]
196pub struct LoadBalancer {
197    pub service: Arc<ServiceConfig>,
198    pub states: Arc<Mutex<HashMap<String, LbState>>>,
199}
200
201impl LoadBalancer {
202    pub fn new(service: Arc<ServiceConfig>, states: Arc<Mutex<HashMap<String, LbState>>>) -> Self {
203        Self { service, states }
204    }
205
206    #[cfg(test)]
207    pub fn select_upstream(&self) -> Option<SelectedUpstream> {
208        self.select_upstream_avoiding(&HashSet::new())
209    }
210
211    pub fn select_upstream_avoiding(&self, avoid: &HashSet<usize>) -> Option<SelectedUpstream> {
212        self.select_upstream_avoiding_inner(avoid, false)
213    }
214
215    pub fn select_upstream_avoiding_strict(
216        &self,
217        avoid: &HashSet<usize>,
218    ) -> Option<SelectedUpstream> {
219        self.select_upstream_avoiding_inner(avoid, true)
220    }
221
222    fn select_upstream_avoiding_inner(
223        &self,
224        avoid: &HashSet<usize>,
225        strict: bool,
226    ) -> Option<SelectedUpstream> {
227        if self.service.upstreams.is_empty() {
228            return None;
229        }
230
231        let mut map = match self.states.lock() {
232            Ok(m) => m,
233            Err(e) => e.into_inner(),
234        };
235        let entry = map.entry(self.service.name.clone()).or_default();
236        entry.ensure_layout(self.service.name.as_str(), &self.service.upstreams);
237
238        let now = std::time::Instant::now();
239
240        // 更新冷却状态：如果冷却期已过，重置失败计数和冷却时间。
241        for idx in 0..self.service.upstreams.len() {
242            if let Some(until) = entry.cooldown_until.get(idx).and_then(|v| *v)
243                && now >= until
244            {
245                entry.failure_counts[idx] = 0;
246                if let Some(slot) = entry.cooldown_until.get_mut(idx) {
247                    *slot = None;
248                }
249            }
250        }
251
252        // 优先使用最近一次“成功”的 upstream，实现粘性路由：
253        // 一旦已经切换到可用线路，就尽量保持在该线路上，而不是每次都从头熔断。
254        if let Some(idx) = entry.last_good_index
255            && idx < self.service.upstreams.len()
256            && entry.failure_counts[idx] < FAILURE_THRESHOLD
257            && !entry.usage_exhausted.get(idx).copied().unwrap_or(false)
258            && !avoid.contains(&idx)
259        {
260            let upstream = self.service.upstreams[idx].clone();
261            return Some(SelectedUpstream {
262                station_name: self.service.name.clone(),
263                index: idx,
264                upstream,
265            });
266        }
267
268        // 第一轮：按顺序选择第一个「未熔断 + 未标记用量用尽」的 upstream。
269        if let Some(idx) = self
270            .service
271            .upstreams
272            .iter()
273            .enumerate()
274            .find_map(|(idx, _)| {
275                if avoid.contains(&idx) {
276                    return None;
277                }
278                if entry.failure_counts[idx] >= FAILURE_THRESHOLD {
279                    return None;
280                }
281                if entry.usage_exhausted.get(idx).copied().unwrap_or(false) {
282                    return None;
283                }
284                Some(idx)
285            })
286        {
287            let upstream = self.service.upstreams[idx].clone();
288            return Some(SelectedUpstream {
289                station_name: self.service.name.clone(),
290                index: idx,
291                upstream,
292            });
293        }
294
295        // 第二轮：非 strict 模式允许熔断兜底，但不能打已确认 usage_exhausted 的上游。
296        if let Some(idx) = self
297            .service
298            .upstreams
299            .iter()
300            .enumerate()
301            .find_map(|(idx, _)| {
302                if avoid.contains(&idx) {
303                    return None;
304                }
305                if entry.usage_exhausted.get(idx).copied().unwrap_or(false) {
306                    return None;
307                }
308                if strict && entry.failure_counts[idx] >= FAILURE_THRESHOLD {
309                    return None;
310                }
311                Some(idx)
312            })
313        {
314            let upstream = self.service.upstreams[idx].clone();
315            return Some(SelectedUpstream {
316                station_name: self.service.name.clone(),
317                index: idx,
318                upstream,
319            });
320        }
321
322        None
323    }
324
325    pub fn penalize_with_backoff(
326        &self,
327        index: usize,
328        cooldown_secs: u64,
329        reason: &str,
330        backoff: CooldownBackoff,
331    ) {
332        let mut map = match self.states.lock() {
333            Ok(m) => m,
334            Err(_) => return,
335        };
336        let entry = map
337            .entry(self.service.name.clone())
338            .or_insert_with(LbState::default);
339        entry.ensure_layout(self.service.name.as_str(), &self.service.upstreams);
340        if index >= entry.failure_counts.len() {
341            return;
342        }
343
344        let streak = entry.penalty_streak.get(index).copied().unwrap_or(0);
345        let effective_secs = backoff.effective_cooldown_secs(cooldown_secs, streak);
346
347        entry.failure_counts[index] = FAILURE_THRESHOLD;
348        if let Some(slot) = entry.cooldown_until.get_mut(index) {
349            *slot =
350                Some(std::time::Instant::now() + std::time::Duration::from_secs(effective_secs));
351        }
352        if let Some(slot) = entry.penalty_streak.get_mut(index) {
353            *slot = streak.saturating_add(1);
354        }
355        if entry.last_good_index == Some(index) {
356            entry.last_good_index = None;
357        }
358        info!(
359            "lb: upstream '{}' index {} penalized for {}s (reason: {})",
360            self.service.name, index, effective_secs, reason
361        );
362    }
363
364    pub fn record_result_with_backoff(
365        &self,
366        index: usize,
367        success: bool,
368        failure_threshold_cooldown_secs: u64,
369        backoff: CooldownBackoff,
370    ) {
371        let mut map = match self.states.lock() {
372            Ok(m) => m,
373            Err(_) => return,
374        };
375        let entry = map
376            .entry(self.service.name.clone())
377            .or_insert_with(LbState::default);
378        entry.ensure_layout(self.service.name.as_str(), &self.service.upstreams);
379        if index >= entry.failure_counts.len() {
380            return;
381        }
382        if success {
383            entry.failure_counts[index] = 0;
384            if let Some(slot) = entry.cooldown_until.get_mut(index) {
385                *slot = None;
386            }
387            if let Some(slot) = entry.penalty_streak.get_mut(index) {
388                *slot = 0;
389            }
390            // 成功请求会将该 upstream 记为“最近可用线路”，后续优先继续使用。
391            entry.last_good_index = Some(index);
392        } else {
393            entry.failure_counts[index] = entry.failure_counts[index].saturating_add(1);
394            if entry.failure_counts[index] >= FAILURE_THRESHOLD
395                && let Some(slot) = entry.cooldown_until.get_mut(index)
396            {
397                let base_secs = if failure_threshold_cooldown_secs == 0 {
398                    COOLDOWN_SECS
399                } else {
400                    failure_threshold_cooldown_secs
401                };
402                let streak = entry.penalty_streak.get(index).copied().unwrap_or(0);
403                let effective_secs = backoff.effective_cooldown_secs(base_secs, streak);
404                let now = std::time::Instant::now();
405                let new_until = now + std::time::Duration::from_secs(effective_secs);
406                let should_update = match *slot {
407                    Some(existing) => new_until > existing,
408                    None => true,
409                };
410                if should_update {
411                    *slot = Some(new_until);
412                }
413                if let Some(slot) = entry.penalty_streak.get_mut(index) {
414                    *slot = streak.saturating_add(1);
415                }
416                info!(
417                    "lb: upstream '{}' index {} reached failure threshold {} (count = {}), entering cooldown for {}s",
418                    self.service.name,
419                    index,
420                    FAILURE_THRESHOLD,
421                    entry.failure_counts[index],
422                    effective_secs
423                );
424                // 触发熔断时，如当前 last_good_index 指向该线路，则清空，允许后续选择其他线路。
425                if entry.last_good_index == Some(index) {
426                    entry.last_good_index = None;
427                }
428            }
429        }
430    }
431}
432
433#[cfg(test)]
434mod tests {
435    use super::*;
436    use crate::config::{ServiceConfig, UpstreamAuth, UpstreamConfig};
437
438    fn make_service(name: &str, urls: &[&str]) -> ServiceConfig {
439        ServiceConfig {
440            name: name.to_string(),
441            alias: None,
442            enabled: true,
443            level: 1,
444            upstreams: urls
445                .iter()
446                .map(|u| UpstreamConfig {
447                    base_url: u.to_string(),
448                    auth: UpstreamAuth {
449                        auth_token: Some("sk-test".to_string()),
450                        auth_token_env: None,
451                        api_key: None,
452                        api_key_env: None,
453                    },
454                    tags: HashMap::new(),
455                    supported_models: HashMap::new(),
456                    model_mapping: HashMap::new(),
457                })
458                .collect(),
459        }
460    }
461
462    fn make_provider_endpoint_service(
463        name: &str,
464        upstreams: &[(&str, &str, &str)],
465    ) -> ServiceConfig {
466        ServiceConfig {
467            name: name.to_string(),
468            alias: None,
469            enabled: true,
470            level: 1,
471            upstreams: upstreams
472                .iter()
473                .map(|(base_url, provider_id, endpoint_id)| UpstreamConfig {
474                    base_url: (*base_url).to_string(),
475                    auth: UpstreamAuth {
476                        auth_token: Some("sk-test".to_string()),
477                        auth_token_env: None,
478                        api_key: None,
479                        api_key_env: None,
480                    },
481                    tags: HashMap::from([
482                        ("provider_id".to_string(), (*provider_id).to_string()),
483                        ("endpoint_id".to_string(), (*endpoint_id).to_string()),
484                    ]),
485                    supported_models: HashMap::new(),
486                    model_mapping: HashMap::new(),
487                })
488                .collect(),
489        }
490    }
491
492    #[test]
493    fn lb_prefers_non_exhausted_upstream_when_available() {
494        let service = make_service(
495            "codex-main",
496            &["https://primary.example", "https://backup.example"],
497        );
498        let states = Arc::new(Mutex::new(HashMap::new()));
499        let lb = LoadBalancer::new(Arc::new(service), states.clone());
500
501        // 初次选择应选第一个 upstream（index 0）。
502        let first = lb.select_upstream().expect("should select an upstream");
503        assert_eq!(first.index, 0);
504
505        // 标记 index 0 为 usage_exhausted，index 1 为可用。
506        {
507            let mut guard = states.lock().unwrap();
508            let entry = guard
509                .entry("codex-main".to_string())
510                .or_insert_with(LbState::default);
511            entry.ensure_layout(lb.service.name.as_str(), &lb.service.upstreams);
512            entry.usage_exhausted[0] = true;
513            entry.usage_exhausted[1] = false;
514        }
515
516        // 此时应优先选择未 exhausted 的 index 1。
517        let second = lb.select_upstream().expect("should select backup upstream");
518        assert_eq!(second.index, 1);
519    }
520
521    #[test]
522    fn lb_does_not_select_when_all_usage_exhausted() {
523        let service = make_service(
524            "codex-main",
525            &["https://primary.example", "https://backup.example"],
526        );
527        let states = Arc::new(Mutex::new(HashMap::new()));
528        let lb = LoadBalancer::new(Arc::new(service), states.clone());
529
530        let _ = lb.select_upstream();
531
532        {
533            let mut guard = states.lock().unwrap();
534            let entry = guard
535                .entry("codex-main".to_string())
536                .or_insert_with(LbState::default);
537            entry.ensure_layout(lb.service.name.as_str(), &lb.service.upstreams);
538            entry.usage_exhausted[0] = true;
539            entry.usage_exhausted[1] = true;
540        }
541
542        assert!(
543            lb.select_upstream().is_none(),
544            "trusted usage exhaustion must not keep hammering a depleted upstream"
545        );
546    }
547
548    #[test]
549    fn lb_strict_mode_does_not_select_when_all_usage_exhausted() {
550        let service = make_service(
551            "codex-main",
552            &["https://primary.example", "https://backup.example"],
553        );
554        let states = Arc::new(Mutex::new(HashMap::new()));
555        let lb = LoadBalancer::new(Arc::new(service), states.clone());
556
557        {
558            let mut guard = states.lock().unwrap();
559            let entry = guard
560                .entry("codex-main".to_string())
561                .or_insert_with(LbState::default);
562            entry.ensure_layout(lb.service.name.as_str(), &lb.service.upstreams);
563            entry.usage_exhausted[0] = true;
564            entry.usage_exhausted[1] = true;
565        }
566
567        assert!(
568            lb.select_upstream_avoiding_strict(&HashSet::new())
569                .is_none(),
570            "strict mode must not ignore trusted usage exhaustion"
571        );
572    }
573
574    #[test]
575    fn lb_resets_state_when_upstream_layout_changes() {
576        let states = Arc::new(Mutex::new(HashMap::new()));
577        let initial = LoadBalancer::new(
578            Arc::new(make_service(
579                "codex-main",
580                &["https://primary.example", "https://backup.example"],
581            )),
582            states.clone(),
583        );
584        initial.record_result_with_backoff(
585            0,
586            false,
587            COOLDOWN_SECS,
588            CooldownBackoff {
589                factor: 1,
590                max_secs: 0,
591            },
592        );
593
594        {
595            let guard = states.lock().unwrap();
596            let entry = guard.get("codex-main").expect("state exists");
597            assert_eq!(entry.failure_counts, vec![1, 0]);
598        }
599
600        let reordered = LoadBalancer::new(
601            Arc::new(make_service(
602                "codex-main",
603                &["https://backup.example", "https://primary.example"],
604            )),
605            states.clone(),
606        );
607        let selected = reordered
608            .select_upstream()
609            .expect("should select an upstream");
610        assert_eq!(selected.index, 0);
611
612        let guard = states.lock().unwrap();
613        let entry = guard.get("codex-main").expect("state exists");
614        assert_eq!(entry.failure_counts, vec![0, 0]);
615        assert_eq!(entry.last_good_index, None);
616    }
617
618    #[test]
619    fn lb_migrates_state_when_provider_endpoint_order_changes() {
620        let states = Arc::new(Mutex::new(HashMap::new()));
621        let initial = LoadBalancer::new(
622            Arc::new(make_provider_endpoint_service(
623                "routing",
624                &[
625                    ("https://primary.example", "primary", "default"),
626                    ("https://backup.example", "backup", "default"),
627                ],
628            )),
629            states.clone(),
630        );
631
632        {
633            let mut guard = states.lock().unwrap();
634            let entry = guard
635                .entry("routing".to_string())
636                .or_insert_with(LbState::default);
637            entry.ensure_layout(initial.service.name.as_str(), &initial.service.upstreams);
638            entry.failure_counts[0] = 2;
639            entry.cooldown_until[0] =
640                Some(std::time::Instant::now() + std::time::Duration::from_secs(30));
641            entry.penalty_streak[0] = 3;
642            entry.usage_exhausted[1] = true;
643            entry.last_good_index = Some(1);
644        }
645
646        let reordered = LoadBalancer::new(
647            Arc::new(make_provider_endpoint_service(
648                "routing",
649                &[
650                    ("https://backup.example", "backup", "default"),
651                    ("https://primary.example", "primary", "default"),
652                ],
653            )),
654            states.clone(),
655        );
656        let selected = reordered
657            .select_upstream()
658            .expect("should select a migrated non-exhausted upstream");
659        assert_eq!(selected.index, 1);
660
661        let guard = states.lock().unwrap();
662        let entry = guard.get("routing").expect("state exists");
663        assert_eq!(entry.failure_counts, vec![0, 2]);
664        assert_eq!(entry.usage_exhausted, vec![true, false]);
665        assert_eq!(entry.penalty_streak, vec![0, 3]);
666        assert!(entry.cooldown_until[0].is_none());
667        assert!(entry.cooldown_until[1].is_some());
668        assert_eq!(entry.last_good_index, Some(0));
669    }
670
671    #[test]
672    fn lb_migrates_legacy_base_url_signature_when_endpoint_identity_is_unambiguous() {
673        let states = Arc::new(Mutex::new(HashMap::new()));
674        let primary_url = "https://primary.example";
675        let backup_url = "https://backup.example";
676
677        {
678            let mut guard = states.lock().unwrap();
679            guard.insert(
680                "routing".to_string(),
681                LbState {
682                    failure_counts: vec![FAILURE_THRESHOLD, 0],
683                    cooldown_until: vec![None, None],
684                    usage_exhausted: vec![false, true],
685                    last_good_index: Some(1),
686                    penalty_streak: vec![2, 0],
687                    upstream_signature: vec![primary_url.to_string(), backup_url.to_string()],
688                },
689            );
690        }
691
692        let reordered = LoadBalancer::new(
693            Arc::new(make_provider_endpoint_service(
694                "routing",
695                &[
696                    (backup_url, "backup", "default"),
697                    (primary_url, "primary", "default"),
698                ],
699            )),
700            states.clone(),
701        );
702        {
703            let mut guard = states.lock().unwrap();
704            let entry = guard.get_mut("routing").expect("state exists");
705            entry.ensure_layout(
706                reordered.service.name.as_str(),
707                &reordered.service.upstreams,
708            );
709        }
710
711        let guard = states.lock().unwrap();
712        let entry = guard.get("routing").expect("state exists");
713        assert_eq!(entry.failure_counts, vec![0, FAILURE_THRESHOLD]);
714        assert_eq!(entry.usage_exhausted, vec![true, false]);
715        assert_eq!(entry.penalty_streak, vec![0, 2]);
716        assert_eq!(entry.last_good_index, Some(0));
717    }
718
719    #[test]
720    fn lb_replaces_state_when_provider_endpoint_base_url_changes() {
721        let states = Arc::new(Mutex::new(HashMap::new()));
722        let initial = LoadBalancer::new(
723            Arc::new(make_provider_endpoint_service(
724                "routing",
725                &[("https://old.example", "input", "default")],
726            )),
727            states.clone(),
728        );
729
730        {
731            let mut guard = states.lock().unwrap();
732            let entry = guard
733                .entry("routing".to_string())
734                .or_insert_with(LbState::default);
735            entry.ensure_layout(initial.service.name.as_str(), &initial.service.upstreams);
736            entry.failure_counts[0] = FAILURE_THRESHOLD;
737            entry.cooldown_until[0] =
738                Some(std::time::Instant::now() + std::time::Duration::from_secs(30));
739            entry.usage_exhausted[0] = true;
740            entry.penalty_streak[0] = 2;
741            entry.last_good_index = Some(0);
742        }
743
744        let updated = LoadBalancer::new(
745            Arc::new(make_provider_endpoint_service(
746                "routing",
747                &[("https://new.example", "input", "default")],
748            )),
749            states.clone(),
750        );
751        let selected = updated
752            .select_upstream()
753            .expect("new endpoint URL should be selectable after state replacement");
754        assert_eq!(selected.index, 0);
755
756        let guard = states.lock().unwrap();
757        let entry = guard.get("routing").expect("state exists");
758        assert_eq!(entry.failure_counts, vec![0]);
759        assert_eq!(entry.cooldown_until, vec![None]);
760        assert_eq!(entry.usage_exhausted, vec![false]);
761        assert_eq!(entry.penalty_streak, vec![0]);
762        assert_eq!(entry.last_good_index, None);
763    }
764
765    #[test]
766    fn lb_avoids_upstreams_past_failure_threshold() {
767        let service = make_service(
768            "codex-main",
769            &["https://primary.example", "https://backup.example"],
770        );
771        let states = Arc::new(Mutex::new(HashMap::new()));
772        let lb = LoadBalancer::new(Arc::new(service), states.clone());
773
774        let disabled_backoff = CooldownBackoff {
775            factor: 1,
776            max_secs: 0,
777        };
778
779        // 对 primary 连续记录 FAILURE_THRESHOLD 次失败。
780        for _ in 0..FAILURE_THRESHOLD {
781            lb.record_result_with_backoff(0, false, COOLDOWN_SECS, disabled_backoff);
782        }
783
784        // 此时应选择 backup（index 1），因为 index 0 已达到失败阈值。
785        let selected = lb
786            .select_upstream()
787            .expect("should select backup after failures");
788        assert_eq!(selected.index, 1);
789    }
790
791    #[test]
792    fn lb_cooldown_expiry_restores_upstream_selection() {
793        let service = make_service(
794            "codex-main",
795            &["https://primary.example", "https://backup.example"],
796        );
797        let states = Arc::new(Mutex::new(HashMap::new()));
798        let lb = LoadBalancer::new(Arc::new(service), states.clone());
799
800        let disabled_backoff = CooldownBackoff {
801            factor: 1,
802            max_secs: 0,
803        };
804
805        for _ in 0..FAILURE_THRESHOLD {
806            lb.record_result_with_backoff(0, false, 2, disabled_backoff);
807        }
808
809        {
810            let guard = states.lock().unwrap();
811            let entry = guard.get("codex-main").expect("lb state exists");
812            assert_eq!(entry.failure_counts[0], FAILURE_THRESHOLD);
813            assert!(entry.cooldown_until[0].is_some());
814        }
815
816        let during_cooldown = lb
817            .select_upstream()
818            .expect("should select backup while primary cools down");
819        assert_eq!(during_cooldown.index, 1);
820
821        {
822            let mut guard = states.lock().unwrap();
823            let entry = guard.get_mut("codex-main").expect("lb state exists");
824            entry.cooldown_until[0] =
825                Some(std::time::Instant::now() - std::time::Duration::from_secs(1));
826        }
827
828        let recovered = lb
829            .select_upstream()
830            .expect("should select primary after cooldown expiry");
831        assert_eq!(recovered.index, 0);
832
833        {
834            let guard = states.lock().unwrap();
835            let entry = guard.get("codex-main").expect("lb state exists");
836            assert_eq!(entry.failure_counts[0], 0);
837            assert!(entry.cooldown_until[0].is_none());
838        }
839    }
840
841    #[test]
842    fn lb_threshold_cooldown_backoff_grows_and_success_resets_streak() {
843        let service = make_service(
844            "codex-main",
845            &["https://primary.example", "https://backup.example"],
846        );
847        let states = Arc::new(Mutex::new(HashMap::new()));
848        let lb = LoadBalancer::new(Arc::new(service), states.clone());
849
850        let backoff = CooldownBackoff {
851            factor: 2,
852            max_secs: 10,
853        };
854
855        for _ in 0..FAILURE_THRESHOLD {
856            lb.record_result_with_backoff(0, false, 2, backoff);
857        }
858
859        let first_remaining_secs = {
860            let guard = states.lock().unwrap();
861            let entry = guard.get("codex-main").expect("lb state exists");
862            assert_eq!(entry.penalty_streak[0], 1);
863            entry.cooldown_until[0]
864                .map(|until| {
865                    until
866                        .saturating_duration_since(std::time::Instant::now())
867                        .as_secs()
868                })
869                .expect("first cooldown exists")
870        };
871        assert!(first_remaining_secs <= 2);
872
873        {
874            let mut guard = states.lock().unwrap();
875            let entry = guard.get_mut("codex-main").expect("lb state exists");
876            entry.cooldown_until[0] =
877                Some(std::time::Instant::now() - std::time::Duration::from_secs(1));
878        }
879        let _ = lb.select_upstream();
880
881        for _ in 0..FAILURE_THRESHOLD {
882            lb.record_result_with_backoff(0, false, 2, backoff);
883        }
884
885        let second_remaining_secs = {
886            let guard = states.lock().unwrap();
887            let entry = guard.get("codex-main").expect("lb state exists");
888            assert_eq!(entry.penalty_streak[0], 2);
889            entry.cooldown_until[0]
890                .map(|until| {
891                    until
892                        .saturating_duration_since(std::time::Instant::now())
893                        .as_secs()
894                })
895                .expect("second cooldown exists")
896        };
897        assert!(second_remaining_secs <= 4);
898        assert!(second_remaining_secs >= first_remaining_secs);
899
900        lb.record_result_with_backoff(0, true, 2, backoff);
901
902        {
903            let guard = states.lock().unwrap();
904            let entry = guard.get("codex-main").expect("lb state exists");
905            assert_eq!(entry.failure_counts[0], 0);
906            assert!(entry.cooldown_until[0].is_none());
907            assert_eq!(entry.penalty_streak[0], 0);
908            assert_eq!(entry.last_good_index, Some(0));
909        }
910    }
911}
codex_helper_core/lb.rs

codex_helper_core/
lb.rs
