"tests/git@developer.sourcefind.cn:renzhc/diffusers_dcu.git" did not exist on "52f1378e64367c655461d4b6d5b2e19585b92791"
Unverified Commit c3faf2d6 authored by luna's avatar luna Committed by GitHub
Browse files

[router] select first healthy worker on proxied get requests (#10827)

parent 9209b209
...@@ -103,10 +103,11 @@ impl Router { ...@@ -103,10 +103,11 @@ impl Router {
fn select_first_worker(&self) -> Result<String, String> { fn select_first_worker(&self) -> Result<String, String> {
let workers = self.worker_registry.get_all(); let workers = self.worker_registry.get_all();
if workers.is_empty() { let healthy_workers: Vec<_> = workers.iter().filter(|w| w.is_healthy()).collect();
if healthy_workers.is_empty() {
Err("No workers are available".to_string()) Err("No workers are available".to_string())
} else { } else {
Ok(workers[0].url().to_string()) Ok(healthy_workers[0].url().to_string())
} }
} }
...@@ -1119,6 +1120,13 @@ mod tests { ...@@ -1119,6 +1120,13 @@ mod tests {
} }
} }
fn create_test_unhealthy_router() -> Router {
let router = create_test_regular_router();
let workers = router.worker_registry.get_all();
workers[0].set_healthy(false);
router
}
#[test] #[test]
fn test_router_get_worker_urls_regular() { fn test_router_get_worker_urls_regular() {
let router = create_test_regular_router(); let router = create_test_regular_router();
...@@ -1140,4 +1148,16 @@ mod tests { ...@@ -1140,4 +1148,16 @@ mod tests {
// DashMap doesn't guarantee order, so just check we get one of the workers // DashMap doesn't guarantee order, so just check we get one of the workers
assert!(url == "http://worker1:8080" || url == "http://worker2:8080"); assert!(url == "http://worker1:8080" || url == "http://worker2:8080");
} }
#[test]
fn test_select_first_worker_with_unhealthy_worker() {
let router = create_test_unhealthy_router();
let result = router.select_first_worker();
assert!(result.is_ok());
let url = result.unwrap();
let worker = router.worker_registry.get_by_url(&url).unwrap();
assert!(worker.is_healthy());
}
} }
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment