temp remove command line arg parser

This commit is contained in:
Ara Sadoyan
2025-03-20 18:33:21 +01:00
parent 7501e367c6
commit 5524fd5011
9 changed files with 430 additions and 562 deletions

View File

@@ -2,38 +2,59 @@ use crate::utils::discovery::{APIUpstreamProvider, Discovery, FromFileProvider};
use crate::utils::tools::*;
use crate::utils::*;
use async_trait::async_trait;
use dashmap::DashMap;
use futures::channel::mpsc;
use futures::StreamExt;
use log::{info, warn};
use log::{debug, error, info, warn};
use pingora::prelude::*;
use pingora_core::prelude::HttpPeer;
use pingora_core::server::ShutdownWatch;
use pingora_core::services::background::BackgroundService;
use pingora_http::{RequestHeader, ResponseHeader};
use pingora_proxy::{ProxyHttp, Session};
use std::ops::Deref;
use std::sync::atomic::Ordering;
use std::sync::Arc;
pub struct LB {
pub ump_upst: Arc<UpstreamsDashMap>,
pub ump_full: Arc<UpstreamsDashMap>,
pub config: Arc<DashMap<String, String>>,
pub local: Arc<(String, u16)>,
}
#[async_trait]
impl BackgroundService for LB {
async fn start(&self, mut shutdown: ShutdownWatch) {
println!("Starting example background service");
let (tux, mut rux) = mpsc::channel::<UpstreamsDashMap>(0);
let file_load2 = FromFileProvider {
path: "etc/upstreams.yaml".to_string(),
};
info!("Starting background service");
let (tx, mut rx) = mpsc::channel::<UpstreamsDashMap>(0);
let api_load = APIUpstreamProvider;
let from_file = self.config.get("upstreams_conf");
match from_file {
Some(from_file) => {
let file_load = FromFileProvider { path: from_file.to_string() };
let tx_file = tx.clone();
let _ = tokio::spawn(async move { file_load.start(tx_file).await });
}
None => {
error!("Can't read config file");
}
}
let config_address = self.config.get("config_address");
match config_address {
Some(config_address) => {
let api_load = APIUpstreamProvider {
address: config_address.to_string(),
};
let tx_api = tx.clone();
let _ = tokio::spawn(async move { api_load.start(tx_api).await });
}
None => {
error!("Can't read config file");
}
}
let tux_file = tux.clone();
let tux_api = tux.clone();
let _ = tokio::spawn(async move { file_load2.start(tux_file).await });
let _ = tokio::spawn(async move { api_load.start(tux_api).await });
let uu = self.ump_upst.clone();
let ff = self.ump_full.clone();
let _ = tokio::spawn(async move { healthcheck::hc2(uu, ff).await });
@@ -43,7 +64,7 @@ impl BackgroundService for LB {
_ = shutdown.changed() => {
break;
}
val = rux.next() => {
val = rx.next() => {
match val {
Some(ss) => {
let foo = compare_dashmaps(&*self.ump_full, &ss);
@@ -67,26 +88,31 @@ pub trait GetHost {
}
#[async_trait]
impl GetHost for LB {
async fn get_host(&self, peer: &str, path: &str, upgrade: bool) -> Option<(String, u16, bool, String)> {
let mut _proto = "";
if upgrade {
_proto = "wsoc";
} else {
_proto = "http"
}
let host_entry = self.ump_upst.get(peer).unwrap();
let x = if let Some(entry) = host_entry.get(path) {
let (servers, index) = entry.value();
if servers.is_empty() {
return None;
async fn get_host(&self, peer: &str, path: &str, _upgrade: bool) -> Option<(String, u16, bool, String)> {
let _proto = "";
// if upgrade {
// _proto = "wsoc";
// } else {
// _proto = "http"
// }
let host_entry = self.ump_upst.get(peer);
match host_entry {
Some(host_entry) => {
let upstream = if let Some(entry) = host_entry.get(path) {
let (servers, index) = entry.value();
if servers.is_empty() {
return None;
}
let idx = index.fetch_add(1, Ordering::Relaxed) % servers.len();
// info!("Host: {}, Path: {}, Peer: {:?}, len: {}, idx: {}", peer, path, servers[idx], servers.len(), idx);
Some(servers[idx].clone())
} else {
None
};
upstream
}
let idx = index.fetch_add(1, Ordering::Relaxed) % servers.len();
println!("{} {:?} => len: {}, idx: {}", peer, servers[idx], servers.len(), idx);
Some(servers[idx].clone())
} else {
None
};
x
None => None,
}
}
}
@@ -95,31 +121,26 @@ impl ProxyHttp for LB {
type CTX = ();
fn new_ctx(&self) -> Self::CTX {}
async fn upstream_peer(&self, session: &mut Session, _ctx: &mut Self::CTX) -> Result<Box<HttpPeer>> {
// let before = Instant::now();
let host_name = session.req_header().headers.get("host");
match host_name {
Some(host) => {
let header_host = host.to_str().unwrap().split(':').collect::<Vec<&str>>();
let ddr = self.get_host(header_host[0], session.req_header().uri.path(), session.is_upgrade_req());
match ddr.await {
Some((host, port, ssl, _proto)) => {
let peer = Box::new(HttpPeer::new((host, port), ssl, String::new()));
// info!("{:?}, Time => {:.2?}", session.request_summary(), before.elapsed());
Ok(peer)
}
None => {
warn!("Returning default list => {:?}, {:?}", host_name, session.req_header().uri);
let peer = Box::new(HttpPeer::new(("127.0.0.1", 3000), false, String::new()));
// info!("{:?}, Time => {:.2?}", session.request_summary(), before.elapsed());
warn!("Upstream not found. Host: {:?}, Path: {}", host, session.req_header().uri);
let peer = Box::new(HttpPeer::new(self.local.deref(), false, String::new()));
Ok(peer)
}
}
}
None => {
warn!("Returning default list => {:?}, {:?}", host_name, session.req_header().uri);
let peer = Box::new(HttpPeer::new(("127.0.0.1", 3000), false, String::new()));
// info!("{:?}, Time => {:.2?}", session.request_summary(), before.elapsed());
warn!("Upstream not found. Host: {:?}, Path: {}", host_name, session.req_header().uri);
let peer = Box::new(HttpPeer::new(self.local.deref(), false, String::new()));
Ok(peer)
}
}
@@ -159,15 +180,32 @@ impl ProxyHttp for LB {
}
Ok(())
}
async fn response_filter(&self, _session: &mut Session, _upstream_response: &mut ResponseHeader, _ctx: &mut Self::CTX) -> pingora_core::Result<()>
async fn response_filter(&self, _session: &mut Session, _upstream_response: &mut ResponseHeader, _ctx: &mut Self::CTX) -> Result<()>
where
Self::CTX: Send + Sync,
{
// -------------------------------------------------------------------------------------------- //
// let header_host = _session.req_header().headers.get("host");
// match header_host {
// Some(header_host) => {
// let data = self.ump_full.get(header_host.to_str().unwrap()).unwrap();
// let detail = data.get(_session.req_header().uri.path());
// let thelast = detail.unwrap().value().0[0].clone().3;
// println!(" Host ==> {}", thelast);
// }
// None => {}
// }
// _upstream_response.insert_header("X-DoSome-About", "Yaaaaaaaaaaaaaaa").unwrap();
// -------------------------------------------------------------------------------------------- //
_upstream_response.insert_header("X-Proxied-From", "Fooooooooooooooo").unwrap();
Ok(())
}
// async fn logging(&self, session: &mut Session, _e: Option<&pingora::Error>, ctx: &mut Self::CTX) {
// let response_code = session.response_written().map_or(0, |resp| resp.status.as_u16());
// info!("{}, response code: {response_code}", self.request_summary(session, ctx));
// }
async fn logging(&self, session: &mut Session, _e: Option<&pingora::Error>, ctx: &mut Self::CTX) {
let response_code = session.response_written().map_or(0, |resp| resp.status.as_u16());
debug!("{}, response code: {response_code}", self.request_summary(session, ctx));
}
}

View File

@@ -1,47 +1,74 @@
use crate::utils::tools::*;
use crate::web::proxyhttp::LB;
// use clap::{arg, Parser};
use dashmap::DashMap;
use pingora_core::prelude::{background_service, Opt};
use pingora_core::server::Server;
use std::env;
use std::sync::Arc;
// #[derive(Parser, Debug)]
// #[command(version, about, long_about = None)]
// struct Args {
// #[arg(short, long)]
// address: String,
// #[arg(short, long)]
// port: String,
// }
pub fn run() {
env_logger::init();
let parameters = Some(Opt::parse_args()).unwrap();
let file = parameters.conf.clone().unwrap();
let maincfg = crate::utils::parceyaml::parce_main_config(file.as_str());
// let mut server = Server::new(None).unwrap();
let mut server = Server::new(Some(Opt::parse_args())).unwrap();
let mut local_conf: (String, u16) = ("0.0.0.0".to_string(), 0);
if let Some((ip, port_str)) = maincfg.get("config_address").unwrap().split_once(':') {
if let Ok(port) = port_str.parse::<u16>() {
local_conf = (ip.to_string(), port);
}
}
let mut server = Server::new(parameters).unwrap();
server.bootstrap();
let uf: UpstreamsDashMap = DashMap::new();
let ff: UpstreamsDashMap = DashMap::new();
let uf_config = Arc::new(uf);
let ff_config = Arc::new(ff);
let cfg = Arc::new(maincfg);
let local = Arc::new(local_conf);
let lb = LB {
ump_upst: uf_config.clone(),
ump_full: ff_config.clone(),
config: cfg.clone(),
local: local.clone(),
};
let bg = LB {
ump_upst: uf_config.clone(),
ump_full: ff_config.clone(),
config: cfg.clone(),
local: local.clone(),
};
// env_logger::Env::new();
// env_logger::init();
let log_level = cfg.get("log_level").unwrap();
match log_level.as_str() {
"info" => env::set_var("RUST_LOG", "info"),
"error" => env::set_var("RUST_LOG", "error"),
"warn" => env::set_var("RUST_LOG", "warn"),
"debug" => env::set_var("RUST_LOG", "debug"),
"trace" => env::set_var("RUST_LOG", "trace"),
"off" => env::set_var("RUST_LOG", "off"),
_ => {
println!("Error reading log level, defaulting to: INFO");
env::set_var("RUST_LOG", "info")
}
}
env_logger::builder()
// .format_timestamp(None)
// .format_module_path(false)
// .format_source_path(false)
// .format_target(false)
.init();
let bg_srvc = background_service("bgsrvc", bg);
let mut proxy = pingora_proxy::http_proxy_service(&server.configuration, lb);
let bindaddress = cfg.get("proxy_address_http").unwrap();
// let args = Args::parse();
// let addr = format!("{}:{}", args.address, args.port);
proxy.add_tcp("0.0.0.0:6193");
proxy.add_tcp(bindaddress.as_str());
server.add_service(proxy);
server.add_service(bg_srvc);

View File

@@ -6,14 +6,11 @@ use axum::routing::{delete, get, head, post, put};
use axum::Router;
use futures::channel::mpsc::Sender;
use futures::SinkExt;
use log::info;
use tokio::net::TcpListener;
// struct UpstreamData {
// servers: UpstreamsDashMap,
// }
#[allow(unused_mut)]
pub async fn run_server(mut toreturn: Sender<UpstreamsDashMap>) {
pub async fn run_server(bindaddress: String, mut toreturn: Sender<UpstreamsDashMap>) {
let mut tr = toreturn.clone();
let app = Router::new()
.route("/{*wildcard}", get(getconfig))
@@ -25,13 +22,22 @@ pub async fn run_server(mut toreturn: Sender<UpstreamsDashMap>) {
"/conf",
post(|up: String| async move {
let serverlist = crate::utils::parceyaml::load_yaml_to_dashmap(up.as_str(), "content");
let _ = tr.send(serverlist).await.unwrap();
Response::builder().status(StatusCode::CREATED).body(Body::from("Config, conf file, updated!\n")).unwrap()
match serverlist {
Some(serverlist) => {
let _ = tr.send(serverlist).await.unwrap();
Response::builder().status(StatusCode::CREATED).body(Body::from("Config, conf file, updated!\n")).unwrap()
}
None => Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(Body::from("Failed to parce config file!\n"))
.unwrap(),
}
})
.with_state("state"),
);
let listener = TcpListener::bind("0.0.0.0:3000").await.unwrap();
println!("Axum API server running on port 3000");
let listener = TcpListener::bind(bindaddress.clone()).await.unwrap();
info!("Starting the API server on: {}", bindaddress);
axum::serve(listener, app).await.unwrap();
}