Optimized healthchecks and config file loading

This commit is contained in:
Ara Sadoyan
2025-08-20 14:03:09 +02:00
parent f8118f9596
commit e304482667
8 changed files with 111 additions and 109 deletions

View File

@@ -28,7 +28,7 @@ struct TaggedAddress {
}
pub async fn start(fp: String, mut toreturn: Sender<Configuration>) {
let config = load_configuration(fp.as_str(), "filepath");
let config = load_configuration(fp.as_str(), "filepath").await;
let headers = DashMap::new();
match config {
Some(config) => {

View File

@@ -15,7 +15,7 @@ pub async fn start(fp: String, mut toreturn: Sender<Configuration>) {
let file_path = fp.as_str();
let parent_dir = Path::new(file_path).parent().unwrap();
let (local_tx, mut local_rx) = tokio::sync::mpsc::channel::<notify::Result<Event>>(1);
let snd = load_configuration(file_path, "filepath");
let snd = load_configuration(file_path, "filepath").await;
match snd {
Some(snd) => {
@@ -53,7 +53,7 @@ pub async fn start(fp: String, mut toreturn: Sender<Configuration>) {
if start.elapsed() > Duration::from_secs(2) {
start = Instant::now();
// info!("Config File changed :=> {:?}", e);
let snd = load_configuration(file_path, "filepath");
let snd = load_configuration(file_path, "filepath").await;
match snd {
Some(snd) => {
toreturn.send(snd).await.unwrap();

View File

@@ -1,7 +1,7 @@
use crate::utils::structs::{InnerMap, UpstreamsDashMap, UpstreamsIdMap};
use crate::utils::tools::*;
use dashmap::DashMap;
use log::{error, info, warn};
use log::{error, warn};
use reqwest::{Client, Version};
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
@@ -9,87 +9,78 @@ use std::time::Duration;
use tokio::time::interval;
use tonic::transport::Endpoint;
#[allow(unused_assignments)]
pub async fn hc2(upslist: Arc<UpstreamsDashMap>, fullist: Arc<UpstreamsDashMap>, idlist: Arc<UpstreamsIdMap>, params: (&str, u64)) {
let mut period = interval(Duration::from_secs(params.1));
let mut first_run = 0;
let client = Client::builder().timeout(Duration::from_secs(2)).danger_accept_invalid_certs(true).build().unwrap();
let client = Client::builder().timeout(Duration::from_secs(params.1)).danger_accept_invalid_certs(true).build().unwrap();
loop {
tokio::select! {
_ = period.tick() => {
let totest : UpstreamsDashMap = DashMap::new();
let fclone : UpstreamsDashMap = clone_dashmap(&fullist);
for val in fclone.iter() {
let host = val.key();
let inner = DashMap::new();
let mut scheme = InnerMap::new();
for path_entry in val.value().iter() {
let path = path_entry.key();
let mut innervec= Vec::new();
for k in path_entry.value().0 .iter().enumerate() {
let mut _link = String::new();
let tls = detect_tls(k.1.address.as_str(), &k.1.port, &client).await;
let mut is_h2 = false;
if tls.1 == Some(Version::HTTP_2) {
is_h2 = true;
}
match tls.0 {
true => _link = format!("https://{}:{}{}", k.1.address, k.1.port, path),
false => _link = format!("http://{}:{}{}", k.1.address, k.1.port, path),
}
scheme = InnerMap {
address: k.1.address.clone(),
port: k.1.port,
is_ssl: tls.0,
is_http2: is_h2,
to_https: k.1.to_https,
rate_limit: k.1.rate_limit,
};
let resp = http_request(_link.as_str(), params.0, "", &client).await;
match resp.0 {
true => {
if resp.1 {
scheme = InnerMap {
address: k.1.address.clone(),
port: k.1.port,
is_ssl: tls.0,
is_http2: is_h2,
to_https: k.1.to_https,
rate_limit: k.1.rate_limit,
};
}
innervec.push(scheme);
}
false => {
warn!("Dead Upstream : {}", _link);
}
}
}
inner.insert(path.clone().to_owned(), (innervec, AtomicUsize::new(0)));
}
totest.insert(host.clone(), inner);
}
if first_run == 1 {
info!("Performing initial hatchecks and upstreams ssl detection");
clone_idmap_into(&totest, &idlist);
info!("Aralez is up and ready to serve requests, the upstreams list is:");
print_upstreams(&totest)
}
first_run+=1;
if ! compare_dashmaps(&totest, &upslist){
clone_dashmap_into(&totest, &upslist);
clone_idmap_into(&totest, &idlist);
}
populate_upstreams(&upslist, &fullist, &idlist, params, &client).await;
}
}
}
}
pub async fn populate_upstreams(upslist: &Arc<UpstreamsDashMap>, fullist: &Arc<UpstreamsDashMap>, idlist: &Arc<UpstreamsIdMap>, params: (&str, u64), client: &Client) {
let totest = build_upstreams(fullist, params.0, client).await;
if !compare_dashmaps(&totest, upslist) {
clone_dashmap_into(&totest, upslist);
clone_idmap_into(&totest, idlist);
}
}
pub async fn initiate_upstreams(fullist: UpstreamsDashMap) -> UpstreamsDashMap {
let client = Client::builder().timeout(Duration::from_secs(2)).danger_accept_invalid_certs(true).build().unwrap();
build_upstreams(&fullist, "HEAD", &client).await
}
async fn build_upstreams(fullist: &UpstreamsDashMap, method: &str, client: &Client) -> UpstreamsDashMap {
let totest: UpstreamsDashMap = DashMap::new();
let fclone = clone_dashmap(fullist);
for val in fclone.iter() {
let host = val.key();
let inner = DashMap::new();
for path_entry in val.value().iter() {
let path = path_entry.key();
let mut innervec = Vec::new();
for (_, upstream) in path_entry.value().0.iter().enumerate() {
let tls = detect_tls(upstream.address.as_str(), &upstream.port, &client).await;
let is_h2 = matches!(tls.1, Some(Version::HTTP_2));
let link = if tls.0 {
format!("https://{}:{}{}", upstream.address, upstream.port, path)
} else {
format!("http://{}:{}{}", upstream.address, upstream.port, path)
};
let mut scheme = InnerMap {
address: upstream.address.clone(),
port: upstream.port,
is_ssl: tls.0,
is_http2: is_h2,
to_https: upstream.to_https,
rate_limit: upstream.rate_limit,
};
let resp = http_request(&link, method, "", &client).await;
if resp.0 {
if resp.1 {
scheme.is_http2 = is_h2; // could be adjusted further
}
innervec.push(scheme);
} else {
warn!("Dead Upstream : {}", link);
}
}
inner.insert(path.clone(), (innervec, AtomicUsize::new(0)));
}
totest.insert(host.clone(), inner);
}
totest
}
async fn http_request(url: &str, method: &str, payload: &str, client: &Client) -> (bool, bool) {
if !["POST", "GET", "HEAD"].contains(&method) {
error!("Method {} not supported. Only GET|POST|HEAD are supported ", method);

View File

@@ -1,11 +1,15 @@
use crate::utils::healthcheck;
use crate::utils::structs::*;
use crate::utils::tools::{clone_dashmap, clone_dashmap_into, print_upstreams};
use dashmap::DashMap;
use log::{error, info, warn};
use std::collections::HashMap;
use std::sync::atomic::AtomicUsize;
// use std::sync::mpsc::{channel, Receiver, Sender};
use std::{env, fs};
// use tokio::sync::oneshot::{Receiver, Sender};
pub fn load_configuration(d: &str, kind: &str) -> Option<Configuration> {
pub async fn load_configuration(d: &str, kind: &str) -> Option<Configuration> {
let yaml_data = match kind {
"filepath" => match fs::read_to_string(d) {
Ok(data) => {
@@ -38,12 +42,12 @@ pub fn load_configuration(d: &str, kind: &str) -> Option<Configuration> {
let mut toreturn = Configuration::default();
populate_headers_and_auth(&mut toreturn, &parsed);
populate_headers_and_auth(&mut toreturn, &parsed).await;
toreturn.typecfg = parsed.provider.clone();
match parsed.provider.as_str() {
"file" => {
populate_file_upstreams(&mut toreturn, &parsed);
populate_file_upstreams(&mut toreturn, &parsed).await;
Some(toreturn)
}
"consul" => {
@@ -62,7 +66,7 @@ pub fn load_configuration(d: &str, kind: &str) -> Option<Configuration> {
}
}
fn populate_headers_and_auth(config: &mut Configuration, parsed: &Config) {
async fn populate_headers_and_auth(config: &mut Configuration, parsed: &Config) {
if let Some(headers) = &parsed.headers {
let mut hl = Vec::new();
for header in headers {
@@ -93,7 +97,8 @@ fn populate_headers_and_auth(config: &mut Configuration, parsed: &Config) {
}
}
fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
async fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
let imtdashmap = UpstreamsDashMap::new();
if let Some(upstreams) = &parsed.upstreams {
for (hostname, host_config) in upstreams {
let path_map = DashMap::new();
@@ -133,11 +138,15 @@ fn populate_file_upstreams(config: &mut Configuration, parsed: &Config) {
path_map.insert(path.clone(), (server_list, AtomicUsize::new(0)));
}
config.headers.insert(hostname.clone(), header_list);
config.upstreams.insert(hostname.clone(), path_map);
imtdashmap.insert(hostname.clone(), path_map);
}
let y = clone_dashmap(&imtdashmap);
let r = healthcheck::initiate_upstreams(y).await;
clone_dashmap_into(&r, &config.upstreams);
println!("Upstream Config:");
print_upstreams(&config.upstreams);
}
}
pub fn parce_main_config(path: &str) -> AppConfig {
let data = fs::read_to_string(path).unwrap();
let reply = DashMap::new();

View File

@@ -101,6 +101,7 @@ pub struct InnerMap {
pub rate_limit: Option<isize>,
}
#[allow(dead_code)]
impl InnerMap {
pub fn new() -> Self {
Self {