mirror of
https://github.com/sadoyan/aralez.git
synced 2026-06-28 02:12:22 +08:00
resolved possible divide by zero issue from #46
This commit is contained in:
@@ -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};
|
||||
|
||||
@@ -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<Vec<(String, Arc<str>)>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait GetHost {
|
||||
fn find_sticky_backend(&self, servers: &[Arc<InnerMap>], backend_id: Option<&str>) -> Option<Arc<InnerMap>>;
|
||||
fn pick_backend(&self, servers: &[Arc<InnerMap>], index: &AtomicUsize, backend_id: Option<&str>) -> Option<Arc<InnerMap>>;
|
||||
fn get_host(&self, peer: &str, path: &str, backend_id: Option<&str>) -> Option<Arc<InnerMap>>;
|
||||
|
||||
fn get_header(&self, peer: &str, path: &str) -> Option<GetHostsReturHeaders>;
|
||||
// fn get_upstreams(&self) -> Arc<UpstreamsDashMap>;
|
||||
}
|
||||
#[async_trait]
|
||||
impl GetHost for LB {
|
||||
fn find_sticky_backend(&self, servers: &[Arc<InnerMap>], backend_id: Option<&str>) -> Option<Arc<InnerMap>> {
|
||||
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<InnerMap>], index: &AtomicUsize, backend_id: Option<&str>) -> Option<Arc<InnerMap>> {
|
||||
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<Arc<InnerMap>> {
|
||||
let host_entry = self.ump_upst.get(peer)?;
|
||||
let mut end = path.len();
|
||||
@@ -26,17 +39,10 @@ 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;
|
||||
} else {
|
||||
@@ -45,19 +51,8 @@ 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
|
||||
|
||||
Reference in New Issue
Block a user