Skip to content

Commit

Permalink
feat: expose timeouts in CLI args for connecting to Edge and/or upstream
Browse files Browse the repository at this point in the history
  • Loading branch information
@sighphyre
sighphyre committed Sep 13, 2023
1 parent 022b361 commit ac2aebf
Show file tree
Hide file tree
Showing 6 changed files with 108 additions and 12 deletions.
2 changes: 2 additions & 0 deletions server/src/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,6 +142,8 @@ async fn build_edge(args: &EdgeArgs) -> EdgeResult<EdgeInfo> {
args.skip_ssl_verification,
args.client_identity.clone(),
args.upstream_certificate_file.clone(),
Duration::seconds(args.upstream_request_timeout),
Duration::seconds(args.upstream_socket_timeout),
)
})
.map(|c| c.with_custom_client_headers(args.custom_client_headers.clone()))
Expand Down
12 changes: 12 additions & 0 deletions server/src/cli.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,14 @@ pub struct EdgeArgs {
/// Service account token. Used to create client tokens if receiving a frontend token we don't have data for
#[clap(long, global = true, env)]
pub service_account_token: Option<String>,

/// Timeout for requests to the upstream server
#[clap(long, env, default_value_t = 5)]
pub upstream_request_timeout: i64,

/// Socket timeout for requests to upstream
#[clap(long, env, default_value_t = 5)]
pub upstream_socket_timeout: i64,
}

pub fn string_to_header_tuple(s: &str) -> Result<(String, String), String> {
Expand Down Expand Up @@ -223,6 +231,10 @@ pub struct CliArgs {
/// Because returning all toggles regardless of their state is a potential security vulnerability, these endpoints can be disabled
#[clap(long, env, default_value_t = false, global = true)]
pub disable_all_endpoint: bool,

/// Timeout for requests to Edge
#[clap(long, env, default_value_t = 5)]
pub edge_request_timeout: u64,
}

#[derive(Args, Debug, Clone)]
Expand Down
16 changes: 16 additions & 0 deletions server/src/http/feature_refresher.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -403,6 +403,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let features_cache = Arc::new(DashMap::default());
let engines_cache = Arc::new(DashMap::default());
Expand Down Expand Up @@ -432,6 +434,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let features_cache = Arc::new(DashMap::default());
let engines_cache = Arc::new(DashMap::default());
Expand Down Expand Up @@ -465,6 +469,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let features_cache = Arc::new(DashMap::default());
let engines_cache = Arc::new(DashMap::default());
Expand Down Expand Up @@ -505,6 +511,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let features_cache = Arc::new(DashMap::default());
let engines_cache = Arc::new(DashMap::default());
Expand Down Expand Up @@ -554,6 +562,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let features_cache = Arc::new(DashMap::default());
let engines_cache = Arc::new(DashMap::default());
Expand Down Expand Up @@ -607,6 +617,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let features_cache = Arc::new(DashMap::default());
let engines_cache = Arc::new(DashMap::default());
Expand Down Expand Up @@ -645,6 +657,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let features_cache = Arc::new(DashMap::default());
let engines_cache = Arc::new(DashMap::default());
Expand Down Expand Up @@ -678,6 +692,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let features_cache = Arc::new(DashMap::default());
let engines_cache = Arc::new(DashMap::default());
Expand Down
78 changes: 67 additions & 11 deletions server/src/http/unleash_client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@ use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::str::FromStr;
use std::time::Duration;

use actix_web::http::header::EntityTag;
use chrono::Duration;
use chrono::Utc;
use lazy_static::lazy_static;
use prometheus::{register_histogram_vec, register_int_gauge_vec, HistogramVec, IntGaugeVec, Opts};
Expand Down Expand Up @@ -126,6 +126,8 @@ fn new_reqwest_client(
skip_ssl_verification: bool,
client_identity: Option<ClientIdentity>,
upstream_certificate_file: Option<PathBuf>,
connect_timeout: Duration,
socket_timeout: Duration,
) -> EdgeResult<Client> {
build_identity(client_identity)
.and_then(|builder| {
Expand Down Expand Up @@ -153,7 +155,8 @@ fn new_reqwest_client(
.user_agent(format!("unleash-edge-{}", crate::types::build::PKG_VERSION))
.default_headers(header_map)
.danger_accept_invalid_certs(skip_ssl_verification)
.timeout(Duration::from_secs(5))
.timeout(socket_timeout.to_std().unwrap())
.connect_timeout(connect_timeout.to_std().unwrap())
.build()
.map_err(|e| EdgeError::ClientBuildError(format!("{e:?}")))
})
Expand All @@ -170,6 +173,8 @@ impl UnleashClient {
skip_ssl_verification: bool,
client_identity: Option<ClientIdentity>,
upstream_certificate_file: Option<PathBuf>,
connect_timeout: Duration,
socket_timeout: Duration,
) -> Self {
Self {
urls: UnleashUrls::from_base_url(server_url),
Expand All @@ -178,6 +183,8 @@ impl UnleashClient {
skip_ssl_verification,
client_identity,
upstream_certificate_file,
connect_timeout,
socket_timeout,
)
.unwrap(),
custom_headers: Default::default(),
Expand All @@ -190,6 +197,8 @@ impl UnleashClient {
client_identity: Option<ClientIdentity>,
upstream_certificate_file: Option<PathBuf>,
service_account_token: String,
connect_timeout: Duration,
socket_timeout: Duration,
) -> Self {
Self {
urls: UnleashUrls::from_base_url(server_url),
Expand All @@ -198,6 +207,8 @@ impl UnleashClient {
skip_ssl_verification,
client_identity,
upstream_certificate_file,
connect_timeout,
socket_timeout,
)
.unwrap(),
custom_headers: Default::default(),
Expand All @@ -212,7 +223,15 @@ impl UnleashClient {
let instance_id = instance_id_opt.unwrap_or_else(|| Ulid::new().to_string());
Ok(Self {
urls: UnleashUrls::from_str(server_url)?,
backing_client: new_reqwest_client(instance_id, false, None, None).unwrap(),
backing_client: new_reqwest_client(
instance_id,
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
)
.unwrap(),
custom_headers: Default::default(),
service_account_token: Default::default(),
})
Expand All @@ -224,7 +243,15 @@ impl UnleashClient {

Ok(Self {
urls: UnleashUrls::from_str(server_url)?,
backing_client: new_reqwest_client(Ulid::new().to_string(), false, None, None).unwrap(),
backing_client: new_reqwest_client(
Ulid::new().to_string(),
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
)
.unwrap(),
custom_headers: Default::default(),
service_account_token: Some(sa_token.into()),
})
Expand All @@ -236,7 +263,15 @@ impl UnleashClient {

Ok(Self {
urls: UnleashUrls::from_str(server_url)?,
backing_client: new_reqwest_client(Ulid::new().to_string(), true, None, None).unwrap(),
backing_client: new_reqwest_client(
Ulid::new().to_string(),
true,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
)
.unwrap(),
custom_headers: Default::default(),
service_account_token: Default::default(),
})
Expand Down Expand Up @@ -487,7 +522,7 @@ impl UnleashClient {
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use std::{str::FromStr, time::Duration};
use std::str::FromStr;

use actix_http::{body::MessageBody, HttpService, TlsAcceptorConfig};
use actix_http_test::{test_server, TestServer};
Expand All @@ -499,7 +534,7 @@ mod tests {
http::header::EntityTag,
web, App, HttpResponse,
};
use chrono::Utc;
use chrono::{Duration, Utc};
use unleash_types::client_features::{ClientFeature, ClientFeatures};

use rand::distributions::Alphanumeric;
Expand Down Expand Up @@ -630,7 +665,7 @@ mod tests {
};
let server_config = tls::config(tls_options).unwrap();
let tls_acceptor_config =
TlsAcceptorConfig::default().handshake_timeout(Duration::from_secs(5));
TlsAcceptorConfig::default().handshake_timeout(std::time::Duration::from_secs(5));
HttpService::new(map_config(
App::new()
.wrap(Etag)
Expand Down Expand Up @@ -855,7 +890,14 @@ mod tests {
pkcs12_identity_file: Some(PathBuf::from(pfx)),
pkcs12_passphrase: Some(passphrase.into()),
};
let client = new_reqwest_client("test_pkcs12".into(), false, Some(identity), None);
let client = new_reqwest_client(
"test_pkcs12".into(),
false,
Some(identity),
None,
Duration::seconds(5),
Duration::seconds(5),
);
assert!(client.is_ok());
}

Expand All @@ -869,7 +911,14 @@ mod tests {
pkcs12_identity_file: Some(PathBuf::from(pfx)),
pkcs12_passphrase: Some(passphrase.into()),
};
let client = new_reqwest_client("test_pkcs12".into(), false, Some(identity), None);
let client = new_reqwest_client(
"test_pkcs12".into(),
false,
Some(identity),
None,
Duration::seconds(5),
Duration::seconds(5),
);
assert!(client.is_err());
}

Expand All @@ -883,7 +932,14 @@ mod tests {
pkcs12_identity_file: None,
pkcs12_passphrase: None,
};
let client = new_reqwest_client("test_pkcs8".into(), false, Some(identity), None);
let client = new_reqwest_client(
"test_pkcs8".into(),
false,
Some(identity),
None,
Duration::seconds(5),
Duration::seconds(5),
);
assert!(client.is_ok());
}

Expand Down
6 changes: 5 additions & 1 deletion server/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ async fn main() -> Result<(), anyhow::Error> {
let schedule_args = args.clone();
let mode_arg = args.clone().mode;
let http_args = args.clone().http;
let request_timeout = args.edge_request_timeout;
let trust_proxy = args.clone().trust_proxy;
let base_path = http_args.base_path.clone();
let (metrics_handler, request_metrics) = prom_metrics::instantiate(None);
Expand Down Expand Up @@ -130,7 +131,10 @@ async fn main() -> Result<(), anyhow::Error> {
} else {
server.bind(http_args.http_server_tuple())
};
let server = server?.workers(http_args.workers).shutdown_timeout(5);
let server = server?
.workers(http_args.workers)
.shutdown_timeout(5)
.client_request_timeout(std::time::Duration::from_secs(request_timeout));

match schedule_args.mode {
cli::EdgeMode::Edge(edge) => {
Expand Down
6 changes: 6 additions & 0 deletions server/src/middleware/client_token_from_frontend_token.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,8 @@ mod tests {
None,
None,
upstream_sa.token.to_string(),
Duration::seconds(5),
Duration::seconds(5),
);
let arced_client = Arc::new(unleash_client);
let local_features_cache: Arc<DashMap<String, ClientFeatures>> =
Expand Down Expand Up @@ -207,6 +209,8 @@ mod tests {
None,
None,
upstream_sa.to_string(),
Duration::seconds(5),
Duration::seconds(5),
);
let arced_client = Arc::new(unleash_client);
let local_features_cache: Arc<DashMap<String, ClientFeatures>> =
Expand Down Expand Up @@ -267,6 +271,8 @@ mod tests {
false,
None,
None,
Duration::seconds(5),
Duration::seconds(5),
);
let local_features_cache: Arc<DashMap<String, ClientFeatures>> =
Arc::new(DashMap::default());
Expand Down

0 comments on commit ac2aebf

Please sign in to comment.