From 970672a392cf8ce8e3bc82ae01dbdb4826b3847e Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Tue, 25 Jul 2023 16:06:56 -0500
Subject: [PATCH] Make client timeout configurable

---
 README.md         | 4 ++++
 src/config.rs     | 9 +++++++++
 src/data/state.rs | 1 +
 src/main.rs       | 6 +++++-
 src/requests.rs   | 7 ++++---
 5 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 320e2ce..e1d7733 100644
--- a/README.md
+++ b/README.md
@@ -105,6 +105,7 @@ LOCAL_DOMAINS=masto.asonix.dog
 LOCAL_BLURB="<p>Welcome to my cool relay where I have cool relay things happening. I hope you enjoy your stay!</p>"
 PROMETHEUS_ADDR=0.0.0.0
 PROMETHEUS_PORT=9000
+CLIENT_TIMEOUT=10
 CLIENT_POOL_SIZE=20
 DELIVER_CONCURRENCY=8
 ```
@@ -156,6 +157,9 @@ Optional - description for the relay
 Optional - Address to bind to for serving the prometheus scrape endpoint
 ##### `PROMETHEUS_PORT`
 Optional - Port to bind to for serving the prometheus scrape endpoint
+##### `CLIENT_TIMEOUT`
+Optional - How long the relay will hold open a connection (in seconds) to a remote server during
+fetches and deliveries. This defaults to 10
 ##### `CLIENT_POOL_SIZE`
 Optional - How many connections the relay should maintain per thread. This value will be multiplied
 by the number of cores available to the relay. This defaults to 20, so a 4-core machine will have a
diff --git a/src/config.rs b/src/config.rs
index b7c9a2d..414de76 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -46,6 +46,7 @@ pub(crate) struct ParsedConfig {
     prometheus_addr: Option<IpAddr>,
     prometheus_port: Option<u16>,
     deliver_concurrency: u64,
+    client_timeout: u64,
     client_pool_size: usize,
 }
 
@@ -71,6 +72,7 @@ pub struct Config {
     local_blurb: Option<String>,
     prometheus_config: Option<PrometheusConfig>,
     deliver_concurrency: u64,
+    client_timeout: u64,
     client_pool_size: usize,
 }
 
@@ -140,6 +142,7 @@ impl std::fmt::Debug for Config {
             .field("local_blurb", &self.local_blurb)
             .field("prometheus_config", &self.prometheus_config)
             .field("deliver_concurrency", &self.deliver_concurrency)
+            .field("client_timeout", &self.client_timeout)
             .field("client_pool_size", &self.client_pool_size)
             .finish()
     }
@@ -171,6 +174,7 @@ impl Config {
             .set_default("prometheus_addr", None as Option<&str>)?
             .set_default("prometheus_port", None as Option<u16>)?
             .set_default("deliver_concurrency", 8u64)?
+            .set_default("client_timeout", 10u64)?
             .set_default("client_pool_size", 20u64)?
             .add_source(Environment::default())
             .build()?;
@@ -244,10 +248,15 @@ impl Config {
             local_blurb: config.local_blurb,
             prometheus_config,
             deliver_concurrency: config.deliver_concurrency,
+            client_timeout: config.client_timeout,
             client_pool_size: config.client_pool_size,
         })
     }
 
+    pub(crate) fn client_timeout(&self) -> u64 {
+        self.client_timeout
+    }
+
     pub(crate) fn deliver_concurrency(&self) -> u64 {
         self.deliver_concurrency
     }
diff --git a/src/data/state.rs b/src/data/state.rs
index d6c484b..e673668 100644
--- a/src/data/state.rs
+++ b/src/data/state.rs
@@ -48,6 +48,7 @@ impl State {
             self.breakers.clone(),
             self.last_online.clone(),
             config.client_pool_size(),
+            config.client_timeout(),
         )
     }
 
diff --git a/src/main.rs b/src/main.rs
index 6ce2f22..e837c4d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -148,7 +148,11 @@ fn client_main(config: Config, args: Args) -> JoinHandle<Result<(), anyhow::Erro
 }
 
 async fn do_client_main(config: Config, args: Args) -> Result<(), anyhow::Error> {
-    let client = requests::build_client(&config.user_agent(), config.client_pool_size());
+    let client = requests::build_client(
+        &config.user_agent(),
+        config.client_pool_size(),
+        config.client_timeout(),
+    );
 
     if !args.blocks().is_empty() || !args.allowed().is_empty() {
         if args.undo() {
diff --git a/src/requests.rs b/src/requests.rs
index a1fc8ef..f7d2169 100644
--- a/src/requests.rs
+++ b/src/requests.rs
@@ -166,7 +166,7 @@ thread_local! {
     static CLIENT: std::cell::OnceCell<Client> = std::cell::OnceCell::new();
 }
 
-pub(crate) fn build_client(user_agent: &str, pool_size: usize) -> Client {
+pub(crate) fn build_client(user_agent: &str, pool_size: usize, timeout_seconds: u64) -> Client {
     CLIENT.with(|client| {
         client
             .get_or_init(|| {
@@ -176,7 +176,7 @@ pub(crate) fn build_client(user_agent: &str, pool_size: usize) -> Client {
                     .connector(connector)
                     .wrap(Tracing)
                     .add_default_header(("User-Agent", user_agent.to_string()))
-                    .timeout(Duration::from_secs(15))
+                    .timeout(Duration::from_secs(timeout_seconds))
                     .finish()
             })
             .clone()
@@ -191,10 +191,11 @@ impl Requests {
         breakers: Breakers,
         last_online: Arc<LastOnline>,
         pool_size: usize,
+        timeout_seconds: u64,
     ) -> Self {
         Requests {
             pool_size,
-            client: build_client(&user_agent, pool_size),
+            client: build_client(&user_agent, pool_size, timeout_seconds),
             key_id,
             user_agent,
             private_key,