mirror of
https://github.com/sadoyan/aralez.git
synced 2026-04-30 23:08:40 +08:00
Optimized healthchecks and config file loading
This commit is contained in:
@@ -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) => {
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -101,6 +101,7 @@ pub struct InnerMap {
|
||||
pub rate_limit: Option<isize>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
impl InnerMap {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
|
||||
Reference in New Issue
Block a user