chore: simplify agent-control command (#1617)

Author: sigilioso

agent-control/src/agent_control/run.rs

@@ -191,7 +191,7 @@ impl AgentControlRunner {
 
 #[cfg(debug_assertions)]
 /// Set path override if local_dir, remote_dir, and logs_dir flags are set
-pub fn set_debug_dirs(base_paths: BasePaths, cli: &crate::flags::Flags) -> BasePaths {
+pub fn set_debug_dirs(base_paths: BasePaths, cli: &crate::command::Command) -> BasePaths {
     let mut base_paths = base_paths;
 
     if let Some(ref local_path) = cli.local_dir {

agent-control/src/bin/main_k8s.rs

@@ -7,9 +7,9 @@
 use newrelic_agent_control::agent_control::run::{
     AgentControlRunConfig, AgentControlRunner, Environment,
 };
+use newrelic_agent_control::command::Command;
 use newrelic_agent_control::event::ApplicationEvent;
 use newrelic_agent_control::event::channel::{EventPublisher, pub_sub};
-use newrelic_agent_control::flags::{Command, Flags};
 use newrelic_agent_control::http::tls::install_rustls_default_crypto_provider;
 use newrelic_agent_control::instrumentation::tracing::TracingGuardBox;
 use std::error::Error;
@@ -19,35 +19,7 @@ use tracing::{error, info, trace};
 const AGENT_CONTROL_MODE: Environment = Environment::K8s;
 
 fn main() -> ExitCode {
-    let Ok(command) = Flags::init(AGENT_CONTROL_MODE)
-        .inspect_err(|init_err| println!("Error parsing Flags: {init_err}"))
-    else {
-        return ExitCode::FAILURE;
-    };
-
-    let (agent_control_config, tracer) = match command {
-        // Agent Control command call instructs normal operation. Continue with required data.
-        Command::InitAgentControl(agent_control_init_config, tracer) => {
-            (agent_control_init_config, tracer)
-        }
-
-        // Agent Control command call was a "one-shot" operation. Exit successfully after performing.
-        Command::OneShot(op) => {
-            op.run_one_shot(AGENT_CONTROL_MODE);
-            return ExitCode::SUCCESS;
-        }
-    };
-
-    match _main(*agent_control_config, tracer) {
-        Err(e) => {
-            error!("The agent control main process exited with an error: {e}");
-            ExitCode::FAILURE
-        }
-        Ok(()) => {
-            info!("The agent control main process exited successfully");
-            ExitCode::SUCCESS
-        }
-    }
+    Command::run(AGENT_CONTROL_MODE, _main)
 }
 
 /// This is the actual main function.

agent-control/src/bin/main_onhost.rs

@@ -9,9 +9,9 @@ use newrelic_agent_control::agent_control::pid_cache::PIDCache;
 use newrelic_agent_control::agent_control::run::{
     AgentControlRunConfig, AgentControlRunner, Environment,
 };
+use newrelic_agent_control::command::Command;
 use newrelic_agent_control::event::ApplicationEvent;
 use newrelic_agent_control::event::channel::{EventPublisher, pub_sub};
-use newrelic_agent_control::flags::{Command, Flags};
 use newrelic_agent_control::http::tls::install_rustls_default_crypto_provider;
 use newrelic_agent_control::instrumentation::tracing::TracingGuardBox;
 use std::error::Error;
@@ -21,35 +21,7 @@ use tracing::{error, info, trace};
 const AGENT_CONTROL_MODE: Environment = Environment::OnHost;
 
 fn main() -> ExitCode {
-    let Ok(command) = Flags::init(AGENT_CONTROL_MODE)
-        .inspect_err(|init_err| println!("Error parsing Flags: {init_err}"))
-    else {
-        return ExitCode::FAILURE;
-    };
-
-    let (agent_control_config, tracer) = match command {
-        // Agent Control command call instructs normal operation. Continue with required data.
-        Command::InitAgentControl(agent_control_init_config, tracer) => {
-            (agent_control_init_config, tracer)
-        }
-
-        // Agent Control command call was a "one-shot" operation. Exit successfully after performing.
-        Command::OneShot(op) => {
-            op.run_one_shot(AGENT_CONTROL_MODE);
-            return ExitCode::SUCCESS;
-        }
-    };
-
-    match _main(*agent_control_config, tracer) {
-        Err(e) => {
-            error!("The agent control main process exited with an error: {e}");
-            ExitCode::FAILURE
-        }
-        Ok(()) => {
-            info!("The agent control main process exited successfully");
-            ExitCode::SUCCESS
-        }
-    }
+    Command::run(AGENT_CONTROL_MODE, _main)
 }
 
 /// This is the actual main function.

agent-control/src/flags.rs agent-control/src/command.rsagent-control/src/flags.rs renamed to agent-control/src/command.rs

@@ -19,40 +19,32 @@ use crate::{
     utils::binary_metadata::binary_metadata,
 };
 use clap::Parser;
+use std::error::Error;
+use std::process::ExitCode;
 use std::sync::Arc;
-use thiserror::Error;
-use tracing::info;
+use tracing::{error, info};
 
 /// All possible errors that can happen while running the initialization.
-#[derive(Debug, Error)]
+#[derive(Debug, thiserror::Error)]
 pub enum InitError {
     /// Could not initialize tracer
-    #[error("could not initialize tracer: `{0}`")]
+    #[error("could not initialize tracer: {0}")]
     TracerError(#[from] TracingError),
     /// K8s config is missing
-    #[error("k8s config missing while running on k8s ")]
+    #[error("k8s config missing while running on k8s")]
     K8sConfig(),
     /// The config could not be read
-    #[error("could not read Agent Control config from `{0}`: `{1}`")]
+    #[error("could not read Agent Control config from {0}: {1}")]
     LoaderError(String, String),
     /// The configuration is invalid
-    #[error("invalid configuration: `{0}`")]
+    #[error("invalid configuration: {0}")]
     InvalidConfig(String),
 }
 
-/// What action was requested from the initialization?
-pub enum Command {
-    /// Normal operation requested. Get the required config and continue.
-    InitAgentControl(Box<AgentControlRunConfig>, Vec<TracingGuardBox>),
-    /// Do a "one-shot" operation and exit successfully.
-    /// In the future, many different operations could be added here.
-    OneShot(OneShotCommand),
-}
-
 /// Command line arguments for Agent Control, as parsed by [`clap`].
 #[derive(Parser, Debug)]
 #[command(author, about, long_about = None)] // Read from `Cargo.toml`
-pub struct Flags {
+pub struct Command {
     #[arg(long)]
     print_debug_info: bool,
 
@@ -75,9 +67,22 @@ pub struct Flags {
     pub logs_dir: Option<std::path::PathBuf>,
 }
 
-impl Flags {
-    /// Parses command line arguments and decides how the application runs.
-    pub fn init(mode: Environment) -> Result<Command, InitError> {
+impl Command {
+    /// Checks if the flag to show the version was set
+    fn print_version(&self) -> bool {
+        self.version
+    }
+
+    /// Checks if the flag to show debug information was set
+    fn print_debug_info(&self) -> bool {
+        self.print_debug_info
+    }
+
+    /// Runs the provided main function or shows the binary information according to flags
+    pub fn run<F: Fn(AgentControlRunConfig, Vec<TracingGuardBox>) -> Result<(), Box<dyn Error>>>(
+        mode: Environment,
+        main_fn: F,
+    ) -> ExitCode {
         // Get command line args
         let flags = Self::parse();
 
@@ -87,14 +92,44 @@ impl Flags {
         #[cfg(debug_assertions)]
         let base_paths = set_debug_dirs(base_paths, &flags);
 
-        // If the version flag is set, print the version and exit
+        // Handle flags requiring different execution mode
         if flags.print_version() {
-            return Ok(Command::OneShot(OneShotCommand::PrintVersion));
+            println!("{}", binary_metadata(mode));
+            return ExitCode::SUCCESS;
         }
         if flags.print_debug_info() {
-            return Ok(Command::OneShot(OneShotCommand::PrintDebugInfo(flags)));
+            println!("Printing debug info");
+            println!("Agent Control Mode: {mode:?}");
+            println!("FLAGS: {flags:#?}");
+            return ExitCode::SUCCESS;
+        }
+
+        let Ok((run_config, tracer)) =
+            Self::init_agent_control(mode, base_paths).inspect_err(|err| {
+                // Using print because the tracer might have failed to start
+                println!("Error on Agent Control initialization: {err}");
+            })
+        else {
+            return ExitCode::FAILURE;
+        };
+
+        match main_fn(run_config, tracer) {
+            Ok(_) => {
+                info!("The agent control main process exited successfully");
+                ExitCode::SUCCESS
+            }
+            Err(err) => {
+                error!("The agent control main process exited with an error: {err}");
+                ExitCode::FAILURE
+            }
         }
+    }
 
+    /// Builds the Agent Control configuration required to execute the application.
+    fn init_agent_control(
+        mode: Environment,
+        base_paths: BasePaths,
+    ) -> Result<(AgentControlRunConfig, Vec<TracingGuardBox>), InitError> {
         let agent_control_repository =
             ConfigRepositoryFile::new(base_paths.local_dir.clone(), base_paths.remote_dir.clone());
 
@@ -148,39 +183,6 @@ impl Flags {
             },
             agent_type_var_constraints,
         };
-
-        Ok(Command::InitAgentControl(Box::new(run_config), tracer))
-    }
-
-    fn print_version(&self) -> bool {
-        self.version
-    }
-
-    fn print_debug_info(&self) -> bool {
-        self.print_debug_info
-    }
-}
-
-/// One-shot operations that can be performed by the agent-control
-pub enum OneShotCommand {
-    /// Print the version of the agent-control and exits
-    PrintVersion,
-    /// Print debug information and exits
-    PrintDebugInfo(Flags),
-}
-
-impl OneShotCommand {
-    /// Runs the one-shot operation
-    pub fn run_one_shot(&self, env: Environment) {
-        match self {
-            OneShotCommand::PrintVersion => {
-                println!("{}", binary_metadata(env));
-            }
-            OneShotCommand::PrintDebugInfo(flags) => {
-                println!("Printing debug info");
-                println!("Agent Control Mode: {env:?}");
-                println!("FLAGS: {flags:#?}");
-            }
-        }
+        Ok((run_config, tracer))
     }
 }

agent-control/src/lib.rs

@@ -6,10 +6,10 @@
 pub mod agent_control;
 pub mod agent_type;
 pub mod cli;
+pub mod command;
 pub mod config_migrate;
 pub mod context;
 pub mod event;
-pub mod flags;
 pub mod health;
 pub mod http;
 pub mod instrumentation;

agent-control/tests/on_host/cli.rs

@@ -255,13 +255,12 @@ fn runs_with_no_config() -> Result<(), Box<dyn std::error::Error>> {
     //   - (\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}) matches the timestamp format.
     // Any character match ".*" is used as the raw logging output contains the raw colors unicode
     // values: \u{1b}[2m2024\u{1b}[0m \u{1b}[32m INFO\u{1b}[0m \u{1b}[2mnewrelic_agent_control\u{1b}[0m\u{1b}[2m:\u{1b}[0m Creating the global context
-    cmd.assert().failure().stdout(
-        predicate::str::is_match(format!(
-            ".*could not read Agent Control config from `{}`.*",
+    cmd.assert()
+        .failure()
+        .stdout(predicate::str::contains(format!(
+            "could not read Agent Control config from {}",
             dir.path().to_string_lossy()
-        ))
-        .unwrap(),
-    );
+        )));
 
     // Env cleanup
     unsafe { env::remove_var(env_var_name) };