feat(k8s): add installation health check

Author: sigilioso

agent-control/src/cli/errors.rs

@@ -12,6 +12,9 @@ pub enum CliError {
 
     #[error("failed to apply resource: {0}")]
     ApplyResource(String),
+
+    #[error("installation check failure: {0}")]
+    InstallationCheck(String),
 }
 
 impl CliError {
@@ -26,7 +29,7 @@ impl CliError {
         match self {
             CliError::K8sClient(_) => ExitCode::from(69),
             CliError::Tracing(_) => ExitCode::from(70),
-            CliError::ApplyResource(_) => ExitCode::from(1),
+            CliError::ApplyResource(_) | CliError::InstallationCheck(_) => ExitCode::from(1),
         }
     }
 }

agent-control/src/cli/install_agent_control.rs

@@ -2,16 +2,21 @@ use crate::agent_control::config::{helmrelease_v2_type_meta, helmrepository_type
 use crate::cli::errors::CliError;
 use crate::cli::utils::parse_key_value_pairs;
 use crate::k8s::annotations::Annotations;
+#[cfg_attr(test, mockall_double::double)]
 use crate::k8s::client::SyncK8sClient;
 use crate::k8s::labels::Labels;
+use crate::sub_agent::health::health_checker::HealthChecker;
+use crate::sub_agent::health::k8s::health_checker::SubAgentHealthChecker;
+use crate::sub_agent::health::with_start_time::StartTime;
 use crate::sub_agent::identity::AgentIdentity;
 use clap::Parser;
 use kube::{
     Resource,
     api::{DynamicObject, ObjectMeta},
-    core::Duration,
 };
 use std::sync::Arc;
+use std::thread::sleep;
+use std::time::Duration;
 use std::{collections::BTreeMap, str::FromStr};
 use tracing::{debug, info};
 
@@ -20,6 +25,10 @@ const REPOSITORY_URL: &str = "https://helm-charts.newrelic.com";
 const FIVE_MINUTES: &str = "5m";
 const AC_DEPLOYMENT_CHART_NAME: &str = "agent-control-deployment";
 
+const INSTALLATION_CHECK_DEFAULT_INITIAL_DELAY: Duration = Duration::from_secs(10);
+const INSTALLATION_CHECK_DEFAULT_MAX_RETRIES: i32 = 10;
+const INSTALLATION_CHECK_DEFAULT_RETRY_INTERVAL: Duration = Duration::from_secs(3);
+
 #[derive(Debug, Parser)]
 pub struct AgentControlInstallData {
     /// Release name
@@ -50,6 +59,10 @@ pub struct AgentControlInstallData {
     /// [k8s labels]: https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
     #[arg(long)]
     pub extra_labels: Option<String>,
+
+    /// Skip the installation check if set
+    #[arg(long)]
+    pub skip_installation_check: bool,
 }
 
 pub fn install_agent_control(
@@ -58,6 +71,7 @@ pub fn install_agent_control(
 ) -> Result<(), CliError> {
     info!("Installing agent control");
 
+    let skip_check = data.skip_installation_check;
     let dynamic_objects = Vec::<DynamicObject>::from(data);
 
     let k8s_client = k8s_client(namespace.clone())?;
@@ -66,12 +80,16 @@ pub fn install_agent_control(
     // For example, what happens if the user applies a remote configuration with a lower version
     // that includes a breaking change?
     info!("Applying agent control resources");
-    for object in dynamic_objects {
-        apply_resource(&k8s_client, &object)?;
+    for object in dynamic_objects.iter() {
+        apply_resource(&k8s_client, object)?;
     }
     info!("Agent control resources applied successfully");
 
-    info!("Agent control installed successfully");
+    if !skip_check {
+        info!("Checking Agent control installation");
+        check_installation(k8s_client, dynamic_objects)?;
+        info!("Agent control installed successfully");
+    }
 
     Ok(())
 }
@@ -143,7 +161,7 @@ fn helm_repository(
         data: serde_json::json!({
             "spec": {
                 "url": REPOSITORY_URL,
-                "interval": Duration::from_str(FIVE_MINUTES).expect("Hardcoded value should be correct"),
+                "interval": kube::core::Duration::from_str(FIVE_MINUTES).expect("Hardcoded value should be correct"),
             }
         }),
     }
@@ -155,8 +173,10 @@ fn helm_release(
     labels: BTreeMap<String, String>,
     annotations: BTreeMap<String, String>,
 ) -> DynamicObject {
-    let interval = Duration::from_str(FIVE_MINUTES).expect("Hardcoded value should be correct");
-    let timeout = Duration::from_str(FIVE_MINUTES).expect("Hardcoded value should be correct");
+    let interval =
+        kube::core::Duration::from_str(FIVE_MINUTES).expect("Hardcoded value should be correct");
+    let timeout =
+        kube::core::Duration::from_str(FIVE_MINUTES).expect("Hardcoded value should be correct");
     let mut data = serde_json::json!({
         "spec": {
             "interval": interval,
@@ -206,6 +226,39 @@ fn secrets_to_json(secrets: BTreeMap<String, String>) -> serde_json::Value {
     serde_json::json!(items)
 }
 
+fn check_installation(
+    k8s_client: SyncK8sClient,
+    objects: Vec<DynamicObject>,
+) -> Result<(), CliError> {
+    let health_checker =
+        SubAgentHealthChecker::try_new(Arc::new(k8s_client), Arc::new(objects), StartTime::now())
+            .map_err(|err| {
+                CliError::InstallationCheck(format!("could not build health-checker: {err}"))
+            })?
+            .ok_or_else(|| {
+                CliError::InstallationCheck("no resources to check health were found".to_string())
+            })?;
+
+    // An initial delay is needed because the api-server can take a while to actually apply the changes and we could
+    // perform the health check to previous objects which could lead to false positives.
+    sleep(INSTALLATION_CHECK_DEFAULT_INITIAL_DELAY);
+    let format_err = |err| {
+        format!(
+            "installation check failed after {INSTALLATION_CHECK_DEFAULT_MAX_RETRIES} attempts: {err}"
+        )
+    };
+    let health = health_checker
+        .check_health_with_retry(
+            INSTALLATION_CHECK_DEFAULT_MAX_RETRIES,
+            INSTALLATION_CHECK_DEFAULT_RETRY_INTERVAL,
+        )
+        .map_err(|err| CliError::InstallationCheck(format_err(err.to_string())))?;
+    if let Some(err) = health.last_error() {
+        return Err(CliError::InstallationCheck(format_err(err)));
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -219,6 +272,7 @@ mod tests {
             chart_version: VERSION.to_string(),
             secrets: None,
             extra_labels: None,
+            skip_installation_check: false,
         }
     }
 

agent-control/src/sub_agent/health/health_checker.rs

@@ -10,7 +10,8 @@ use crate::sub_agent::health::with_start_time::HealthWithStartTime;
 use crate::sub_agent::identity::ID_ATTRIBUTE_NAME;
 use crate::sub_agent::supervisor::starter::SupervisorStarterError;
 use crate::utils::thread_context::{NotStartedThreadContext, StartedThreadContext};
-use std::time::{SystemTime, SystemTimeError};
+use std::thread::sleep;
+use std::time::{Duration, SystemTime, SystemTimeError};
 use tracing::{debug, error, info_span};
 
 const HEALTH_CHECKER_THREAD_NAME: &str = "health_checker";
@@ -211,6 +212,36 @@ pub trait HealthChecker {
     /// See OpAMP's [spec](https://github.com/open-telemetry/opamp-spec/blob/main/specification.md#componenthealthstatus)
     /// for more details.
     fn check_health(&self) -> Result<HealthWithStartTime, HealthCheckerError>;
+
+    /// Checks health and perform retries if the result is unhealthy or there was an error executing health check.
+    /// The retries are performed as specified by provided limit and retry_interval.
+    fn check_health_with_retry(
+        &self,
+        limit: i32,
+        retry_interval: Duration,
+    ) -> Result<HealthWithStartTime, HealthCheckerError> {
+        let mut last_health = Err(HealthCheckerError::Generic("initial value".into()));
+        for attempt in 1..=limit {
+            debug!("Checking health with retries {attempt}/{limit}");
+            last_health = self.check_health();
+            match last_health.as_ref() {
+                Ok(health) => {
+                    if health.is_healthy() {
+                        debug!("Health check result was healthy");
+                        return last_health;
+                    }
+                    if let Some(err) = health.last_error() {
+                        debug!("Health check result was unhealthy: {err}");
+                    }
+                }
+                Err(err) => {
+                    debug!("Failure to check health: {err}");
+                }
+            }
+            sleep(retry_interval);
+        }
+        last_health
+    }
 }
 
 pub(crate) fn spawn_health_checker<H>(
@@ -292,6 +323,7 @@ pub mod tests {
     use crate::event::channel::pub_sub;
 
     use super::*;
+    use assert_matches::assert_matches;
     use mockall::{Sequence, mock};
 
     impl Default for Healthy {
@@ -358,6 +390,95 @@ pub mod tests {
         }
     }
 
+    #[test]
+    fn test_health_check_with_retry_success_on_first_attempt() {
+        let health_checker = MockHealthCheck::new_healthy();
+
+        let result = health_checker.check_health_with_retry(3, Duration::from_millis(10));
+
+        assert_matches!(result, Ok(health) => {
+            assert!(health.is_healthy());
+        });
+    }
+
+    #[test]
+    fn test_health_check_with_retry_success_after_retries() {
+        let mut health_checker = MockHealthCheck::new();
+        let mut seq = Sequence::new();
+
+        health_checker
+            .expect_check_health()
+            .once()
+            .in_sequence(&mut seq)
+            .returning(|| {
+                Err(HealthCheckerError::Generic(
+                    "error on first attempt".to_string(),
+                ))
+            });
+        health_checker
+            .expect_check_health()
+            .once()
+            .in_sequence(&mut seq)
+            .returning(|| {
+                Ok(HealthWithStartTime::from_unhealthy(
+                    Unhealthy::new(
+                        "Unhealthy".to_string(),
+                        "Unhealthy on second attempt".to_string(),
+                    ),
+                    UNIX_EPOCH,
+                ))
+            });
+        health_checker
+            .expect_check_health()
+            .once()
+            .in_sequence(&mut seq)
+            .returning(|| {
+                Ok(HealthWithStartTime::from_healthy(
+                    Healthy::default(),
+                    UNIX_EPOCH,
+                ))
+            });
+
+        let result = health_checker.check_health_with_retry(3, Duration::from_millis(10));
+
+        assert_matches!(result, Ok(health) => {
+            assert!(health.is_healthy());
+        });
+    }
+
+    #[test]
+    fn test_health_check_with_retry_failure_after_all_attempts() {
+        let mut health_checker = MockHealthCheck::new();
+        health_checker
+            .expect_check_health()
+            .times(3)
+            .returning(|| Err(HealthCheckerError::Generic("persistent error".to_string())));
+
+        let result = health_checker.check_health_with_retry(3, Duration::from_millis(10));
+
+        assert_matches!(result, Err(HealthCheckerError::Generic(s)) => {
+            assert_eq!(s, "persistent error".to_string());
+        });
+    }
+
+    #[test]
+    fn test_health_check_with_retry_unhealthy_result() {
+        let mut health_checker = MockHealthCheck::new();
+        health_checker.expect_check_health().times(3).returning(|| {
+            Ok(HealthWithStartTime::from_unhealthy(
+                Unhealthy::new("Unhealthy".to_string(), "persistent unhealthy".to_string()),
+                UNIX_EPOCH,
+            ))
+        });
+
+        let result = health_checker.check_health_with_retry(3, Duration::from_millis(10));
+
+        assert_matches!(result, Ok(health) => {
+            assert!(!health.is_healthy());
+            assert_eq!(health.last_error().unwrap(), "persistent unhealthy".to_string());
+        });
+    }
+
     #[test]
     fn test_spawn_health_checker() {
         let (health_publisher, health_consumer) = pub_sub();

agent-control/tests/k8s/agent_control_cli/dynamic_objects.rs

@@ -27,6 +27,7 @@ fn k8s_cli_install_agent_control_creates_resources() {
         .arg("chart=podinfo, env=testing, app=ac");
     cmd.arg("--secrets")
         .arg("secret1=default.yaml,secret2=values.yaml,secret3=fixed.yaml");
+    cmd.arg("--skip-installation-check"); // Skipping checks because we are merely checking that the resources are created.
     cmd.assert().success();
 
     let k8s_client = SyncK8sClient::try_new(runtime.clone(), namespace.clone()).unwrap();