#[macro_use] extern crate anyhow; #[macro_use] extern crate log; #[macro_use] extern crate serde_derive; pub mod api; pub mod newrelic_logger; pub mod speak_proxy; use anyhow::Result as AnyhowResult; use crate::newrelic_logger::NewRelicLogger; use crate::speak_proxy::speak_proxy_with_retry; use hyper::server::conn::AddrStream; use hyper::service::{service_fn, make_service_fn}; use hyper::{Body, Request, Response, Server}; use hyper::{StatusCode, Method}; use rand::seq::IteratorRandom; use std::collections::HashMap; use std::convert::Infallible; use std::env; use std::fmt::Debug; use std::fmt::Display; use std::net::IpAddr; use std::str::FromStr; use std::sync::Arc; const ENV_DEFAULT_ROUTE : &'static str = "DEFAULT_ROUTE"; const ENV_MAX_RETRIES : &'static str = "MAX_RETRIES"; const ENV_NEWRELIC_API_KEY : &'static str = "NEWRELIC_API_KEY"; const ENV_PROXY_CONFIG_FILE : &'static str = "PROXY_CONFIG_FILE"; const ENV_RETRY_WAIT_MS : &'static str = "RETRY_WAIT_MS"; const ENV_RUST_LOG : &'static str = "RUST_LOG"; const ENV_SERVICE_PORT : &'static str = "SERVICE_PORT"; const DEFAULT_DEFAULT_ROUTE : &'static str = "http://127.0.0.1:12345"; const DEFAULT_MAX_RETRIES : u8 = 3; const DEFAULT_NEWRELIC_API_KEY : &'static str = ""; const DEFAULT_PROXY_CONFIG_FILE : &'static str = "proxy_configs.toml"; const DEFAULT_RETRY_WAIT_MS : u64 = 400; const DEFAULT_RUST_LOG : &'static str = "tokio_reactor=warn,hyper=info,debug"; const DEFAULT_SERVICE_PORT : u16 = 5555; #[derive(Deserialize, Debug, Clone)] struct ProxyConfigs { pub backends: Vec, } #[derive(Deserialize, Debug, Clone)] struct ProxyConfig { pub voice : String, pub ip: String, pub port: String, } impl ProxyConfigs { pub fn load_from_file(filename: &str) -> AnyhowResult { let contents = std::fs::read_to_string(filename)?; let configs = toml::from_str(&contents)?; Ok(configs) } } fn get_env_string(env_name: &str, default: &str) -> String { match env::var(env_name).as_ref().ok() { Some(s) => s.to_string(), None => { warn!("Env var '{}' not supplied. Using default '{}'.", env_name, default); default.to_string() }, } } // the trait `std::error::Error` is not implemented for `::Err` fn get_env_num(env_name: &str, default: T) -> AnyhowResult where T: FromStr + Display, T::Err: Debug { match env::var(env_name).as_ref().ok() { None => { warn!("Env var '{}' not supplied. Using default '{}'.", env_name, default); Ok(default) }, Some(val) => { match val.parse::() { Err(e) => bail!("Can't parse value '{:?}'. Error: {:?}", val, e), Ok(val) => Ok(val), } }, } } pub struct Router { pub routes: HashMap, pub default_route: String, } impl Router { fn get_random_host(&self) -> String { let mut rng = rand::thread_rng(); self.routes.values() .choose(&mut rng) .map(|url| url.to_string()) .unwrap_or(self.default_route.clone()) } fn get_speaker_host(&self, voice: &str) -> String { self.routes.get(voice) .map(|url| url.to_string()) .unwrap_or(self.default_route.clone()) } } // fn debug_request(req: Request) -> Result, Infallible> { // let body_str = format!("{:?}", req); // Ok(Response::new(Body::from(body_str))) // } fn health_check_response(_req: Request) -> Result, Infallible> { Ok(Response::new(Body::from("healthy"))) } pub async fn handle( client_ip: IpAddr, req: Request, router: Arc, newrelic_logger: Arc, server_params: ServerParams) -> Result, Infallible> { match (req.method(), req.uri().path()) { (&Method::POST, "/speak") => { let nr_transaction = newrelic_logger.web_transaction("/speak"); let result = speak_proxy_with_retry( client_ip, req, router, nr_transaction, server_params.max_retries, server_params.retry_wait_ms, &server_params.server_hostname, "/speak").await; match result { Ok(response) => Ok(response), Err(error) => Ok(error.to_response()), } }, (&Method::POST, "/speak_spectrogram") => { let nr_transaction = newrelic_logger.web_transaction("/speak_spectrogram"); let result = speak_proxy_with_retry( client_ip, req, router, nr_transaction, server_params.max_retries, server_params.retry_wait_ms, &server_params.server_hostname, "/speak_spectrogram").await; match result { Ok(response) => Ok(response), Err(error) => Ok(error.to_response()), } }, (&Method::GET, "/proxy_health") => { let _droppable_transaction = newrelic_logger.web_transaction("/proxy_health"); health_check_response(req) }, _ => { let _droppable_transaction = newrelic_logger .web_transaction("unmatched endpoint"); let forward = router.get_random_host(); info!("Forwarding to `{}` random host: {}", &forward, req.uri()); match hyper_reverse_proxy::call(client_ip, &forward, req).await { Ok(response) => { Ok(response) }, Err(_) => { Ok(Response::builder() .status(StatusCode::INTERNAL_SERVER_ERROR) .body(Body::from("unknown proxy error")) .unwrap()) } } }, } } #[derive(Debug, Clone)] pub struct ServerParams { pub max_retries: u8, pub retry_wait_ms: u64, pub server_hostname: String, } #[tokio::main] async fn main() { if env::var(ENV_RUST_LOG) .as_ref() .ok() .is_none() { println!("Setting default logging level to \"{}\", override with env var {}.", DEFAULT_RUST_LOG, ENV_RUST_LOG); std::env::set_var(ENV_RUST_LOG, DEFAULT_RUST_LOG); } env_logger::init(); let default_route = get_env_string(ENV_DEFAULT_ROUTE, DEFAULT_DEFAULT_ROUTE); let service_port = get_env_num::(ENV_SERVICE_PORT, DEFAULT_SERVICE_PORT) .expect("should have service port env var"); let newrelic_api_key = get_env_string(ENV_NEWRELIC_API_KEY, DEFAULT_NEWRELIC_API_KEY); info!("Default route: {}", default_route); let proxy_configs_file = get_env_string(ENV_PROXY_CONFIG_FILE, DEFAULT_PROXY_CONFIG_FILE); info!("Proxy config file: {}", proxy_configs_file); let max_retries = get_env_num::(ENV_MAX_RETRIES, DEFAULT_MAX_RETRIES) .expect("should have max retries env var"); info!("Max retries: {}", max_retries); let retry_wait_ms = get_env_num::(ENV_RETRY_WAIT_MS, DEFAULT_RETRY_WAIT_MS) .expect("should have retry wait ms env var"); info!("Retry wait: {}", retry_wait_ms); let proxy_configs = ProxyConfigs::load_from_file(&proxy_configs_file) .expect("should have configs"); info!("Proxy configs: {:?}", proxy_configs); let server_hostname = hostname::get() .ok() .and_then(|h| h.into_string().ok()) .unwrap_or("proxy-unknown".to_string()); let server_params = ServerParams { max_retries, retry_wait_ms, server_hostname, }; let mut route_map = HashMap::new(); for backend in proxy_configs.backends.iter() { let host = format!("http://{}:{}", backend.ip, backend.port); route_map.insert(backend.voice.clone(), host); } let router = Arc::new(Router { routes: route_map, default_route: default_route, }); let newrelic_logger = if newrelic_api_key.is_empty() { NewRelicLogger::null_instance() } else { NewRelicLogger::try_new_or_null("voder-proxy", &newrelic_api_key) }; let newrelic_logger = Arc::new(newrelic_logger); let make_service = make_service_fn(move |conn: &AddrStream| { let remote_addr = conn.remote_addr().ip(); let router2 = router.clone(); let newrelic_logger2 = newrelic_logger.clone(); let server_params = server_params.clone(); async move { let router3 = router2.clone(); let newrelic_logger3 = newrelic_logger2.clone(); let server_params = server_params.clone(); Ok::<_, Infallible>(service_fn(move |req| { let router4 = router3.clone(); let newrelic_logger4 = newrelic_logger3.clone(); handle(remote_addr, req, router4, newrelic_logger4, server_params.clone()) })) } }); let addr = ([0, 0, 0, 0], service_port).into(); info!("Running server on {:?}", addr); let server = Server::bind(&addr).serve(make_service); if let Err(e) = server.await { eprintln!("server error: {}", e); } }