Add ability to specify server timeout amount in ms from config file

mozilla · Aug 12, 2024 · 4788cba · 4788cba
1 parent 9958e7c
commit 4788cba
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 8 deletions.
diff --git a/src/commands.rs b/src/commands.rs
@@ -611,7 +611,7 @@ pub fn run_command(cmd: Command) -> Result<i32> {
     // Config isn't required for all commands, but if it's broken then we should flag
     // it early and loudly.
     let config = &Config::load()?;
-    let startup_timeout = config.server_startup_timeout;
+    let startup_timeout = config.server_timing.startup_timeout;
 
     match cmd {
         Command::ShowStats(fmt, advanced) => {

diff --git a/src/config.rs b/src/config.rs
@@ -555,9 +555,18 @@ impl Default for DistConfig {
 pub struct FileConfig {
     pub cache: CacheConfigs,
     pub dist: DistConfig,
+    // pub timing: ServerTimingConfig,
     pub server_startup_timeout_ms: Option<u64>,
+    pub server_shutdown_timeout_ms: Option<u64>,
+    pub port: Option<u16>,
 }
 
+// #[derive(Debug, Default, Serialize, Deserialize, Eq, PartialEq)]
+// pub struct ServerTimingConfig {
+//     pub server_startup_timeout_ms: Option<u64>,
+//     pub server_shutdown_timeout_ms: Option<u64>,
+// }
+
 // If the file doesn't exist or we can't read it, log the issue and proceed. If the
 // config exists but doesn't parse then something is wrong - return an error.
 pub fn try_read_config_file<T: DeserializeOwned>(path: &Path) -> Result<Option<T>> {
@@ -945,7 +954,14 @@ pub struct Config {
     pub cache: Option<CacheType>,
     pub fallback_cache: DiskCacheConfig,
     pub dist: DistConfig,
-    pub server_startup_timeout: Option<std::time::Duration>,
+    pub server_timing: ServerTiming,
+    pub port: Option<u16>,
+}
+
+#[derive(Debug, Default, PartialEq, Eq)]
+pub struct ServerTiming {
+    pub startup_timeout: Option<std::time::Duration>,
+    pub shutdown_timeout: Option<std::time::Duration>,
 }
 
 impl Config {
@@ -967,11 +983,20 @@ impl Config {
             cache,
             dist,
             server_startup_timeout_ms,
+            server_shutdown_timeout_ms,
+            port,
         } = file_conf;
         conf_caches.merge(cache);
 
         let server_startup_timeout =
             server_startup_timeout_ms.map(std::time::Duration::from_millis);
+        let server_shutdown_timeout =
+            server_shutdown_timeout_ms.map(std::time::Duration::from_millis);
+        let server_timing = ServerTiming {
+            startup_timeout: server_startup_timeout,
+            shutdown_timeout: server_shutdown_timeout,
+        };
+
 
         let EnvConfig { cache } = env_conf;
         conf_caches.merge(cache);
@@ -981,7 +1006,8 @@ impl Config {
             cache: caches,
             fallback_cache,
             dist,
-            server_startup_timeout,
+            server_timing,
+            port,
         }
     }
 }
@@ -1281,6 +1307,8 @@ fn config_overrides() {
         },
         dist: Default::default(),
         server_startup_timeout_ms: None,
+        server_shutdown_timeout_ms: None,
+        port: None,
     };
 
     assert_eq!(
@@ -1302,7 +1330,8 @@ fn config_overrides() {
                 rw_mode: CacheModeConfig::ReadWrite,
             },
             dist: Default::default(),
-            server_startup_timeout: None,
+            server_timing: Default::default(),
+            port: None,
         }
     );
 }
@@ -1577,7 +1606,26 @@ no_credentials = true
                 toolchain_cache_size: 5368709120,
                 rewrite_includes_only: false,
             },
-            server_startup_timeout_ms: Some(10000),
+            server_startup_timeout_ms: Some(10_000),
+            server_shutdown_timeout_ms: None,
+            port: None,
+        }
+    )
+}
+
+#[test]
+fn test_port_config() {
+    // just set up a config file with just port, then have it read it in, and ensure port is defined in the struct
+    const CONFIG_STR: &str = "port = 8080";
+    let file_config: FileConfig = toml::from_str(CONFIG_STR).expect("Is valid toml.");
+    assert_eq!(
+        file_config,
+        FileConfig {
+            cache: Default::default(),
+            dist: Default::default(),
+            server_startup_timeout_ms: None,
+            server_shutdown_timeout_ms: None,
+            port: Some(8080),
         }
     )
 }
diff --git a/src/server.rs b/src/server.rs
@@ -658,11 +658,12 @@ impl<C: CommandCreatorSync> SccacheServer<C> {
             }
         })?;
 
-        const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
+        let config = Config::load().unwrap_or_default();
+        let shutdown_timeout: Duration = config.server_timing.shutdown_timeout.unwrap_or(Duration::from_secs(10));
         info!(
             "moving into the shutdown phase now, waiting at most {} seconds \
              for all client requests to complete",
-            SHUTDOWN_TIMEOUT.as_secs()
+             shutdown_timeout.as_secs()
         );
 
         // Once our server has shut down either due to inactivity or a manual
@@ -672,7 +673,7 @@ impl<C: CommandCreatorSync> SccacheServer<C> {
         //
         // Note that we cap the amount of time this can take, however, as we
         // don't want to wait *too* long.
-        runtime.block_on(async { time::timeout(SHUTDOWN_TIMEOUT, wait).await })?;
+        runtime.block_on(async { time::timeout(shutdown_timeout, wait).await })?;
 
         info!("ok, fully shutting down now");
 

diff --git a/tests/harness/mod.rs b/tests/harness/mod.rs
@@ -165,6 +165,8 @@ pub fn sccache_client_cfg(
             rewrite_includes_only: false, // TODO
         },
         server_startup_timeout_ms: None,
+        server_shutdown_timeout_ms: None,
+        port: None,
     }
 }
 

diff --git a/tests/oauth.rs b/tests/oauth.rs
@@ -60,6 +60,8 @@ fn config_with_dist_auth(
             rewrite_includes_only: true,
         },
         server_startup_timeout_ms: None,
+        server_shutdown_timeout_ms: None,
+        port: None,
     }
 }