Fix #17 - Add TCP connection timeout

Cargo.toml

@@ -5,7 +5,7 @@ authors = ["ufoscout <ufoscout@gmail.com>"]
 edition = "2018"
 
 [dependencies]
-port_check = "0.1"
+port_check = "0.1.3"
 
 [dev-dependencies]
 atomic-counter = "1.0"

README.md

@@ -74,7 +74,8 @@ This is discussed further [here](https://stackoverflow.com/questions/30063907/us
 # Additional configuration options
 The behaviour of the wait utility can be configured with the following environment variables:
 - *WAIT_HOSTS*: comma separated list of pairs host:port for which you want to wait.
-- *WAIT_HOSTS_TIMEOUT*: max number of seconds to wait for the hosts to be available before failure. The default is 30 seconds.
+- *WAIT_HOSTS_TIMEOUT*: max number of seconds to wait for all the hosts to be available before failure. The default is 30 seconds.
+- *WAIT_HOST_CONNECT_TIMEOUT*: The timeout of a single TCP connection to a remote host before attempting a new connection. The default is 5 seconds.
 - *WAIT_BEFORE_HOSTS*: number of seconds to wait (sleep) before start checking for the hosts availability
 - *WAIT_AFTER_HOSTS*: number of seconds to wait (sleep) once all the hosts are available
 - *WAIT_SLEEP_INTERVAL*: number of seconds to sleep between retries. The default is 1 second.

src/lib.rs

@@ -1,9 +1,12 @@
+use std::time::Duration;
+
 pub mod env_reader;
 pub mod sleeper;
 
 pub struct Config {
     pub hosts: String,
-    pub timeout: u64,
+    pub global_timeout: u64,
+    pub tcp_connection_timeout: u64,
     pub wait_before: u64,
     pub wait_after: u64,
     pub wait_sleep_interval: u64,
@@ -19,7 +22,8 @@ pub fn wait(
     println!("{}", LINE_SEPARATOR);
     println!("docker-compose-wait - starting with configuration:");
     println!(" - Hosts to be waiting for: [{}]", config.hosts);
-    println!(" - Timeout before failure: {} seconds ", config.timeout);
+    println!(" - Timeout before failure: {} seconds ", config.global_timeout);
+    println!(" - TCP connection timeout before retry: {} seconds ", config.tcp_connection_timeout);
     println!(
         " - Sleeping time before checking for hosts availability: {} seconds",
         config.wait_before
@@ -43,12 +47,12 @@ pub fn wait(
         sleep.reset();
         for host in config.hosts.trim().split(',') {
             println!("Checking availability of {}", host);
-            while !port_check::is_port_reachable(&host.trim().to_string()) {
+            while !port_check::is_port_reachable_with_timeout(&host.trim().to_string().parse().expect("The host IP should be valid"), Duration::from_secs(config.tcp_connection_timeout)) {
                 println!("Host {} not yet available...", host);
-                if sleep.elapsed(config.timeout) {
+                if sleep.elapsed(config.global_timeout) {
                     println!(
                         "Timeout! After {} seconds some hosts are still not reachable",
-                        config.timeout
+                        config.global_timeout
                     );
                     on_timeout();
                     return;
@@ -76,10 +80,14 @@ pub fn wait(
 pub fn config_from_env() -> Config {
     Config {
         hosts: crate::env_reader::env_var(&"WAIT_HOSTS".to_string(), "".to_string()),
-        timeout: to_int(
+        global_timeout: to_int(
             &crate::env_reader::env_var(&"WAIT_HOSTS_TIMEOUT".to_string(), "".to_string()),
             30,
         ),
+        tcp_connection_timeout: to_int(
+            &crate::env_reader::env_var(&"WAIT_HOST_CONNECT_TIMEOUT".to_string(), "".to_string()),
+            5,
+        ),
         wait_before: to_int(
             &crate::env_reader::env_var(&"WAIT_BEFORE_HOSTS".to_string(), "".to_string()),
             0,
@@ -139,23 +147,25 @@ mod test {
     }
 
     #[test]
-    fn default_timeout_should_be_30() {
+    fn config_should_use_default_values() {
         let _guard = TEST_MUTEX.lock().unwrap();
-        set_env("", "", "10o", "10", "");
+        set_env("", "", "10o", "10", "", "abc");
         let config = config_from_env();
         assert_eq!("".to_string(), config.hosts);
-        assert_eq!(30, config.timeout);
+        assert_eq!(30, config.global_timeout);
+        assert_eq!(5, config.tcp_connection_timeout);
         assert_eq!(0, config.wait_before);
         assert_eq!(10, config.wait_after);
     }
 
     #[test]
     fn should_get_config_values_from_env() {
         let _guard = TEST_MUTEX.lock().unwrap();
-        set_env("localhost:1234", "20", "2", "3", "4");
+        set_env("localhost:1234", "20", "2", "3", "4", "23");
         let config = config_from_env();
         assert_eq!("localhost:1234".to_string(), config.hosts);
-        assert_eq!(20, config.timeout);
+        assert_eq!(20, config.global_timeout);
+        assert_eq!(23, config.tcp_connection_timeout);
         assert_eq!(2, config.wait_before);
         assert_eq!(3, config.wait_after);
         assert_eq!(4, config.wait_sleep_interval);
@@ -164,19 +174,21 @@ mod test {
     #[test]
     fn should_get_default_config_values() {
         let _guard = TEST_MUTEX.lock().unwrap();
-        set_env("localhost:1234", "", "", "", "");
+        set_env("localhost:1234", "", "", "", "", "");
         let config = config_from_env();
         assert_eq!("localhost:1234".to_string(), config.hosts);
-        assert_eq!(30, config.timeout);
+        assert_eq!(30, config.global_timeout);
+        assert_eq!(5, config.tcp_connection_timeout);
         assert_eq!(0, config.wait_before);
         assert_eq!(0, config.wait_after);
         assert_eq!(1, config.wait_sleep_interval);
     }
 
-    fn set_env(hosts: &str, timeout: &str, before: &str, after: &str, sleep: &str) {
+    fn set_env(hosts: &str, timeout: &str, before: &str, after: &str, sleep: &str, tcp_timeout: &str) {
         env::set_var("WAIT_BEFORE_HOSTS", before.to_string());
         env::set_var("WAIT_AFTER_HOSTS", after.to_string());
         env::set_var("WAIT_HOSTS_TIMEOUT", timeout.to_string());
+        env::set_var("WAIT_HOST_CONNECT_TIMEOUT", tcp_timeout.to_string());
         env::set_var("WAIT_HOSTS", hosts.to_string());
         env::set_var("WAIT_SLEEP_INTERVAL", sleep.to_string());
     }

tests/integration_test.rs

@@ -11,7 +11,7 @@ fn should_wait_5_seconds_before() {
     let mut sleeper = MillisSleeper::default();
     wait::wait(
         &mut sleeper,
-        &new_config("", 1, wait_for, 0, 1),
+        &new_config("", 1, wait_for, 0, 1, 1),
         &mut on_timeout,
     );
     assert!(millis_elapsed(start) >= wait_for)
@@ -24,7 +24,7 @@ fn should_wait_10_seconds_after() {
     let mut sleeper = MillisSleeper::default();
     wait::wait(
         &mut sleeper,
-        &new_config("", 1, 0, wait_for, 1),
+        &new_config("", 1, 0, wait_for, 1, 1),
         &mut on_timeout,
     );
     assert!(millis_elapsed(start) >= wait_for)
@@ -37,7 +37,7 @@ fn should_wait_before_and_after() {
     let mut sleeper = MillisSleeper::default();
     wait::wait(
         &mut sleeper,
-        &new_config("", 1, wait_for, wait_for, 1),
+        &new_config("", 1, wait_for, wait_for, 1, 1),
         &mut on_timeout,
     );
     assert!(millis_elapsed(start) >= (wait_for + wait_for))
@@ -47,16 +47,26 @@ fn should_wait_before_and_after() {
 fn should_execute_without_wait() {
     let start = Instant::now();
     let mut sleeper = MillisSleeper::default();
-    wait::wait(&mut sleeper, &new_config("", 1, 0, 0, 1), &mut on_timeout);
+    wait::wait(&mut sleeper, &new_config("", 1, 0, 0, 1, 1), &mut on_timeout);
     assert!(millis_elapsed(start) <= 5)
 }
 
+#[test]
+fn should_sleep_the_specified_time_between_checks() {
+    let start = Instant::now();
+    let mut sleeper = MillisSleeper::default();
+    wait::wait(&mut sleeper, &new_config("198.19.255.255:1", 2_000, 0, 0, 10, 1), &mut on_timeout);
+    let elapsed = millis_elapsed(start);
+    assert!(elapsed >= 2010);
+    assert!(elapsed < 3000);
+}
+
 #[test]
 fn should_exit_on_timeout() {
     let timeout = 25;
     let wait_before = 30;
     let wait_after = 300;
-    let hosts = "localhost:".to_string() + &free_port().to_string();
+    let hosts = format!("127.0.0.1:{}", free_port());
     let start = Instant::now();
     let mut sleeper = MillisSleeper::default();
 
@@ -68,7 +78,7 @@ fn should_exit_on_timeout() {
 
     wait::wait(
         &mut sleeper,
-        &new_config(&hosts, timeout, wait_before, wait_after, 1),
+        &new_config(&hosts, timeout, wait_before, wait_after, 1, 1),
         &mut fun,
     );
 
@@ -101,7 +111,7 @@ fn should_identify_the_open_port() {
     thread::sleep(time::Duration::from_millis(250));
     wait::wait(
         &mut sleeper,
-        &new_config(&hosts, timeout, wait_before, wait_after, 1),
+        &new_config(&hosts, timeout, wait_before, wait_after, 1, 1),
         &mut fun,
     );
 
@@ -137,7 +147,7 @@ fn should_wait_multiple_hosts() {
     thread::sleep(time::Duration::from_millis(250));
     wait::wait(
         &mut sleeper,
-        &new_config(&hosts, timeout, wait_before, wait_after, 1),
+        &new_config(&hosts, timeout, wait_before, wait_after, 1, 1),
         &mut fun,
     );
 
@@ -170,7 +180,7 @@ fn should_fail_if_not_all_hosts_are_available() {
     thread::sleep(time::Duration::from_millis(250));
     wait::wait(
         &mut sleeper,
-        &new_config(&hosts, timeout, wait_before, wait_after, 1),
+        &new_config(&hosts, timeout, wait_before, wait_after, 1, 1),
         &mut fun,
     );
 
@@ -182,10 +192,11 @@ fn should_fail_if_not_all_hosts_are_available() {
 
 fn on_timeout() {}
 
-fn new_config(hosts: &str, timeout: u64, before: u64, after: u64, sleep: u64) -> wait::Config {
+fn new_config(hosts: &str, timeout: u64, before: u64, after: u64, sleep: u64, tcp_connection_timeout: u64) -> wait::Config {
     wait::Config {
         hosts: hosts.to_string(),
-        timeout: timeout,
+        global_timeout: timeout,
+        tcp_connection_timeout,
         wait_before: before,
         wait_after: after,
         wait_sleep_interval: sleep,
@@ -216,7 +227,7 @@ fn free_port() -> u16 {
 }
 
 fn millis_elapsed(start: Instant) -> u64 {
-    let elapsed = start.elapsed().subsec_nanos() / 1000000;
+    let elapsed = start.elapsed().as_millis();
     println!("Millis elapsed {}", elapsed);
     elapsed as u64
 }