diff --git a/core/src/libraries/helpers/healthcheck.rs b/core/src/libraries/helpers/healthcheck.rs index 6544f1d7..74f60571 100644 --- a/core/src/libraries/helpers/healthcheck.rs +++ b/core/src/libraries/helpers/healthcheck.rs @@ -4,6 +4,7 @@ use hyper::{body, Client, Uri}; use log::{debug, trace}; +use redis::{aio::ConnectionLike, AsyncCommands, RedisResult}; use std::time::Duration; use tokio::time::sleep; use tokio::time::timeout; @@ -51,3 +52,36 @@ pub async fn wait_for(url: &str, timeout_duration: Duration) -> Result( + key: &str, + timeout_duration: Duration, + con: &mut C, +) -> Result<(), ()> { + let check_interval = Duration::from_millis(250); + let mut remaining_duration = timeout_duration; + + debug!("Awaiting existence of redis key {}", key); + + loop { + let result: RedisResult = con.exists(key).await; + + if let Ok(exists) = result { + if exists { + return Ok(()); + } else { + trace!("Expected redis key does not exist yet"); + } + } else { + trace!("Unable to check existence of redis key! {:?}", result); + } + + if remaining_duration.as_secs() == 0 { + debug!("Timeout while waiting for redis key {}", key); + return Err(()); + } + + sleep(check_interval).await; + remaining_duration -= check_interval; + } +} diff --git a/core/src/libraries/helpers/mod.rs b/core/src/libraries/helpers/mod.rs index c88e85f2..5a42e88b 100644 --- a/core/src/libraries/helpers/mod.rs +++ b/core/src/libraries/helpers/mod.rs @@ -13,7 +13,7 @@ pub mod lua; pub use backoff::Backoff; pub use capabilities::*; -pub use healthcheck::wait_for; +pub use healthcheck::{wait_for, wait_for_key}; pub use timeout::Timeout; /// Splits the input string into two parts at the first occurence of the separator diff --git a/core/src/services/manager/tasks/create_session.rs b/core/src/services/manager/tasks/create_session.rs index 63eb9fb4..21cdcb0b 100644 --- a/core/src/services/manager/tasks/create_session.rs +++ b/core/src/services/manager/tasks/create_session.rs @@ -1,6 +1,6 @@ use super::super::{context::SessionCreationContext, RequestError, SessionReplyValue}; use crate::libraries::helpers::{ - keys, parse_browser_string, wait_for, CapabilitiesRequest, Timeout, + keys, parse_browser_string, wait_for, wait_for_key, CapabilitiesRequest, Timeout, }; use crate::libraries::lifecycle::logging::{LogCode, SessionLogger}; use crate::libraries::metrics::MetricsEntry; @@ -256,7 +256,7 @@ mod subtasks { pub async fn await_healthcheck( con: &mut C, session_id: &str, - ) -> Result { + ) -> Result<(), RequestError> { let (host, port): (String, String) = con .hget(keys::session::upstream(session_id), &["host", "port"]) .map_err(RequestError::RedisError) @@ -264,10 +264,27 @@ mod subtasks { let url = format!("http://{}:{}/status", host, port); let timeout = Timeout::NodeStartup.get(con).await as u64; - - wait_for(&url, Duration::from_secs(timeout)) + let timeout_duration = Duration::from_secs(timeout); + let healthcheck_start = Instant::now(); + + // Wait for the node heartbeat in redis first to save some HTTP calls + wait_for_key( + &keys::session::heartbeat::node(session_id), + timeout_duration, + con, + ) + .map_err(|_| RequestError::HealthCheckTimeout) + .await?; + + // Spend the remaining timeout duration HTTP polling the webdrivers status endpoint + let remaining_duration = + timeout_duration - Instant::now().duration_since(healthcheck_start); + + wait_for(&url, remaining_duration) .map_err(|_| RequestError::HealthCheckTimeout) - .await + .await?; + + Ok(()) } mod helpers {