diff --git a/src/utils/tools.rs b/src/utils/tools.rs index b401abc..e59de1d 100644 --- a/src/utils/tools.rs +++ b/src/utils/tools.rs @@ -1,6 +1,6 @@ use crate::tls::load; use crate::tls::load::CertificateConfig; -use crate::utils::structs::{Extraparams, InnerMap, InnerMapForJson, UpstreamSnapshotForJson, UpstreamsDashMap, UpstreamsIdMap}; +use crate::utils::structs::{Extraparams, InnerMapForJson, UpstreamSnapshotForJson, UpstreamsDashMap, UpstreamsIdMap}; use dashmap::DashMap; use log::{error, info}; use notify::{event::ModifyKind, Config, EventKind, RecommendedWatcher, RecursiveMode, Watcher}; diff --git a/src/web/gethosts.rs b/src/web/gethosts.rs index 452e0a4..e95d90b 100644 --- a/src/web/gethosts.rs +++ b/src/web/gethosts.rs @@ -1,7 +1,6 @@ use crate::utils::structs::InnerMap; use crate::web::proxyhttp::LB; -use async_trait::async_trait; -use std::sync::atomic::Ordering; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; #[derive(Debug, Clone)] @@ -10,15 +9,29 @@ pub struct GetHostsReturHeaders { pub server_headers: Option)>>, } -#[async_trait] pub trait GetHost { + fn find_sticky_backend(&self, servers: &[Arc], backend_id: Option<&str>) -> Option>; + fn pick_backend(&self, servers: &[Arc], index: &AtomicUsize, backend_id: Option<&str>) -> Option>; fn get_host(&self, peer: &str, path: &str, backend_id: Option<&str>) -> Option>; - fn get_header(&self, peer: &str, path: &str) -> Option; - // fn get_upstreams(&self) -> Arc; } -#[async_trait] impl GetHost for LB { + fn find_sticky_backend(&self, servers: &[Arc], backend_id: Option<&str>) -> Option> { + let b = backend_id?; + let bb = self.ump_byid.get(b)?; + let target = bb.value(); + servers.iter().any(|s| s.address == target.address && s.port == target.port).then(|| target.clone()) + } + fn pick_backend(&self, servers: &[Arc], index: &AtomicUsize, backend_id: Option<&str>) -> Option> { + if servers.is_empty() { + return None; + } + if let Some(target) = self.find_sticky_backend(servers, backend_id) { + return Some(target); + } + let idx = index.fetch_add(1, Ordering::Relaxed) % servers.len(); + Some(servers[idx].clone()) + } fn get_host(&self, peer: &str, path: &str, backend_id: Option<&str>) -> Option> { let host_entry = self.ump_upst.get(peer)?; let mut end = path.len(); @@ -26,16 +39,9 @@ impl GetHost for LB { let slice = &path[..end]; if let Some(entry) = host_entry.get(slice) { let (servers, index) = entry.value(); - if let Some(b) = backend_id { - if let Some(bb) = self.ump_byid.get(b) { - let target = bb.value(); - if servers.iter().any(|s| s.address == target.address && s.port == target.port) { - return Some(target.clone()); - } - } + if let Some(backend) = self.pick_backend(servers, index, backend_id) { + return Some(backend); } - let idx = index.fetch_add(1, Ordering::Relaxed) % servers.len(); - return Some(servers[idx].clone()); } if let Some(pos) = slice.rfind('/') { end = pos; @@ -45,21 +51,10 @@ impl GetHost for LB { } if let Some(entry) = host_entry.get("/") { let (servers, index) = entry.value(); - if !servers.is_empty() { - - if let Some(b) = backend_id { - if let Some(bb) = self.ump_byid.get(b) { - let target = bb.value(); - if servers.iter().any(|s| s.address == target.address && s.port == target.port) { - return Some(target.clone()); - } - } - } - - let idx = index.fetch_add(1, Ordering::Relaxed) % servers.len(); - return Some(servers[idx].clone()); - } + if let Some(backend) = self.pick_backend(servers, index, backend_id) { + return Some(backend); } + } None }