fix(tracer-flare): fix design of TracerFlareManager

datadog-tracer-flare/src/error.rs

@@ -10,8 +10,6 @@ pub enum FlareError {
     ListeningError(String),
     /// Parsing of config failed.
     ParsingError(String),
-    /// Send the flare was asking without being prepared.
-    RemoteConfigError(String),
     /// Sending the flare failed.
     SendError(String),
     /// Creating the zipped flare failed.
@@ -23,9 +21,6 @@ impl std::fmt::Display for FlareError {
         match self {
             FlareError::ListeningError(msg) => write!(f, "Listening failed with: {msg}"),
             FlareError::ParsingError(msg) => write!(f, "Parsing failed with: {msg}"),
-            FlareError::RemoteConfigError(msg) => {
-                write!(f, "RemoteConfig file processed in a wrong order: {msg}")
-            }
             FlareError::SendError(msg) => write!(f, "Sending the flare failed with: {msg}"),
             FlareError::ZipError(msg) => write!(f, "Creating the zip failed with: {msg}"),
         }

datadog-tracer-flare/src/lib.rs

@@ -30,30 +30,18 @@ use crate::error::FlareError;
 pub struct TracerFlareManager {
     pub agent_url: String,
     pub language: String,
-    pub state: State,
+    pub agent_task: Option<AgentTaskFile>,
     /// As a featured option so we can use the component with no Listener
     #[cfg(feature = "listener")]
     pub listener: Option<Listener>,
 }
 
-#[derive(Debug, PartialEq)]
-pub enum State {
-    Idle,
-    Collecting {
-        log_level: String,
-    },
-    Sending {
-        agent_task: AgentTaskFile,
-        log_level: String,
-    },
-}
-
 impl Default for TracerFlareManager {
     fn default() -> Self {
         TracerFlareManager {
             agent_url: hyper::Uri::default().to_string(),
             language: "rust".to_string(),
-            state: State::Idle,
+            agent_task: None,
             #[cfg(feature = "listener")]
             listener: None,
         }
@@ -168,9 +156,11 @@ pub enum LogLevel {
 #[derive(Debug, PartialEq)]
 pub enum ReturnAction {
     /// If AGENT_CONFIG received with the right properties.
-    Start(LogLevel),
+    Set(LogLevel),
+    /// If AGENT_CONFIG is removed.
+    Unset,
     /// If AGENT_TASK received with the right properties.
-    Stop,
+    Send,
     /// If anything else received.
     None,
 }
@@ -179,7 +169,7 @@ impl TryFrom<&str> for LogLevel {
     type Error = FlareError;
 
     fn try_from(level: &str) -> Result<Self, FlareError> {
-        match level {
+        match level.to_lowercase().as_str() {
             "trace" => Ok(LogLevel::Trace),
             "debug" => Ok(LogLevel::Debug),
             "info" => Ok(LogLevel::Info),
@@ -219,51 +209,15 @@ pub fn check_remote_config_file(
             RemoteConfigData::TracerFlareConfig(agent_config) => {
                 if agent_config.name.starts_with("flare-log-level.") {
                     if let Some(log_level) = &agent_config.config.log_level {
-                        if let State::Collecting { log_level: _ } = tracer_flare.state {
-                            return Err(FlareError::RemoteConfigError(
-                                "Cannot start a flare while one is already running".to_string(),
-                            ));
-                        }
-                        if let State::Sending {
-                            agent_task: _,
-                            log_level: _,
-                        } = tracer_flare.state
-                        {
-                            return Err(FlareError::RemoteConfigError(
-                                "Cannot start a flare while one is waiting to be sent".to_string(),
-                            ));
-                        }
-                        // Idle state
-                        tracer_flare.state = State::Collecting {
-                            log_level: log_level.to_string(),
-                        };
                         let log_level = log_level.as_str().try_into()?;
-                        return Ok(ReturnAction::Start(log_level));
+                        return Ok(ReturnAction::Set(log_level));
                     }
                 }
             }
             RemoteConfigData::TracerFlareTask(agent_task) => {
                 if agent_task.task_type.eq("tracer_flare") {
-                    if let State::Collecting { log_level } = &tracer_flare.state {
-                        tracer_flare.state = State::Sending {
-                            agent_task: agent_task.clone(),
-                            log_level: log_level.to_string(),
-                        };
-                        return Ok(ReturnAction::Stop);
-                    }
-                    if let State::Sending {
-                        agent_task: _,
-                        log_level: _,
-                    } = tracer_flare.state
-                    {
-                        return Err(FlareError::RemoteConfigError(
-                            "Cannot stop a flare that it is already waiting to be sent".to_string(),
-                        ));
-                    }
-                    // Idle state
-                    return Err(FlareError::RemoteConfigError(
-                        "Cannot stop an inexisting flare".to_string(),
-                    ));
+                    tracer_flare.agent_task = Some(agent_task.to_owned());
+                    return Ok(ReturnAction::Send);
                 }
             }
             _ => return Ok(ReturnAction::None),
@@ -345,6 +299,18 @@ pub async fn run_remote_config_listener(
                     if action != Ok(ReturnAction::None) {
                         return action;
                     }
+                } else if let Change::Remove(file) = change {
+                    match file.contents().as_ref() {
+                        Ok(data) => match data {
+                            RemoteConfigData::TracerFlareConfig(_) => {
+                                return Ok(ReturnAction::Unset)
+                            }
+                            _ => continue,
+                        },
+                        Err(e) => {
+                            return Err(FlareError::ParsingError(e.to_string()));
+                        }
+                    }
                 }
             }
         }
@@ -358,9 +324,7 @@ pub async fn run_remote_config_listener(
 
 #[cfg(test)]
 mod tests {
-    use crate::{
-        check_remote_config_file, FlareError, LogLevel, ReturnAction, State, TracerFlareManager,
-    };
+    use crate::{check_remote_config_file, FlareError, LogLevel, ReturnAction, TracerFlareManager};
     use datadog_remote_config::{
         config::{
             agent_config::{AgentConfig, AgentConfigFile},
@@ -410,17 +374,11 @@ mod tests {
         let mut tracer_flare = TracerFlareManager::default();
         let result = check_remote_config_file(file, &mut tracer_flare);
         assert!(result.is_ok());
-        assert_eq!(result.unwrap(), ReturnAction::Start(LogLevel::Info));
-        assert_eq!(
-            tracer_flare.state,
-            State::Collecting {
-                log_level: "info".to_string()
-            }
-        );
+        assert_eq!(result.unwrap(), ReturnAction::Set(LogLevel::Info));
     }
 
     #[test]
-    fn test_check_remote_config_file_with_stop_task() {
+    fn test_check_remote_config_file_with_send_task() {
         let storage = ParsedFileStorage::default();
         let path = Arc::new(RemoteConfigPath {
             product: RemoteConfigProduct::AgentTask,
@@ -442,23 +400,11 @@ mod tests {
         let file = storage
             .store(1, path.clone(), serde_json::to_vec(&task).unwrap())
             .unwrap();
-        let mut tracer_flare = TracerFlareManager {
-            // Emulate the start action
-            state: State::Collecting {
-                log_level: "debug".to_string(),
-            },
-            ..Default::default()
-        };
+        let mut tracer_flare = TracerFlareManager::default();
         let result = check_remote_config_file(file, &mut tracer_flare);
         assert!(result.is_ok());
-        assert_eq!(result.unwrap(), ReturnAction::Stop);
-        assert_eq!(
-            tracer_flare.state,
-            State::Sending {
-                agent_task: task,
-                log_level: "debug".to_string()
-            }
-        );
+        assert_eq!(result.unwrap(), ReturnAction::Send);
+        assert_eq!(tracer_flare.agent_task, Some(task));
     }
 
     #[test]

datadog-tracer-flare/src/zip.rs

@@ -16,7 +16,7 @@ use tempfile::tempfile;
 use walkdir::WalkDir;
 use zip::{write::FileOptions, ZipWriter};
 
-use crate::{error::FlareError, State, TracerFlareManager};
+use crate::{error::FlareError, TracerFlareManager};
 
 /// Adds a single file to the zip archive with the specified options and relative path
 fn add_file_to_zip(
@@ -213,6 +213,7 @@ fn generate_payload(
 /// # Arguments
 ///
 /// * `zip` - A file handle to the zip archive to be sent
+/// * `log_level` - Log level of the tracer
 /// * `tracer_flare` - TracerFlareManager instance containing the agent configuration
 ///
 /// # Returns
@@ -228,12 +229,13 @@ fn generate_payload(
 /// - The agent URL is invalid
 /// - The HTTP request fails after retries
 /// - The agent returns a non-success HTTP status code
-async fn send(zip: File, tracer_flare: &mut TracerFlareManager) -> Result<(), FlareError> {
-    let (agent_task, log_level) = match &tracer_flare.state {
-        State::Sending {
-            agent_task,
-            log_level,
-        } => (agent_task, log_level),
+async fn send(
+    zip: File,
+    log_level: String,
+    tracer_flare: &mut TracerFlareManager,
+) -> Result<(), FlareError> {
+    let agent_task = match &tracer_flare.agent_task {
+        Some(agent_task) => agent_task,
         _ => {
             return Err(FlareError::SendError(
                 "Trying to send the flare without AGENT_TASK received".to_string(),
@@ -244,7 +246,7 @@ async fn send(zip: File, tracer_flare: &mut TracerFlareManager) -> Result<(), Fl
     let payload = generate_payload(
         zip,
         &tracer_flare.language,
-        log_level,
+        &log_level,
         &agent_task.args.case_id,
         &agent_task.args.hostname,
         &agent_task.args.user_handle,
@@ -291,10 +293,8 @@ async fn send(zip: File, tracer_flare: &mut TracerFlareManager) -> Result<(), Fl
                 let response = hyper_migration::into_response(body);
                 let status = response.status();
                 if status.is_success() {
-                    // Should we return something specific ?
                     Ok(())
                 } else {
-                    // Maybe just put a warning log message ?
                     Err(FlareError::SendError(format!(
                         "Agent returned non-success status for flare send: HTTP {status}"
                     )))
@@ -313,6 +313,7 @@ async fn send(zip: File, tracer_flare: &mut TracerFlareManager) -> Result<(), Fl
 ///
 /// * `files` - A vector of strings representing the paths of files and directories to include in
 ///   the zip archive.
+/// * `log_level` - Log level of the tracer.
 /// * `tracer_flare` - TracerFlareManager instance containing the agent configuration and task data.
 ///   The state will be reset after sending (agent_task set to None, running set to false).
 ///
@@ -346,7 +347,7 @@ async fn send(zip: File, tracer_flare: &mut TracerFlareManager) -> Result<(), Fl
 ///         "/path/to/config.txt".to_string(),
 ///     ];
 ///
-///     match zip_and_send(files, &mut tracer_flare).await {
+///     match zip_and_send(files, "debug".to_string(), &mut tracer_flare).await {
 ///         Ok(_) => println!("Flare sent successfully"),
 ///         Err(e) => eprintln!("Failed to send flare: {}", e),
 ///     }
@@ -355,15 +356,16 @@ async fn send(zip: File, tracer_flare: &mut TracerFlareManager) -> Result<(), Fl
 /// ```
 pub async fn zip_and_send(
     files: Vec<String>,
+    log_level: String,
     tracer_flare: &mut TracerFlareManager,
 ) -> Result<(), FlareError> {
     let zip = zip_files(files)?;
 
     // APMSP-2118 - TODO: Implement obfuscation of sensitive data
 
-    let response = send(zip, tracer_flare).await;
+    let response = send(zip, log_level, tracer_flare).await;
 
-    tracer_flare.state = State::Idle;
+    tracer_flare.agent_task = None;
 
     response
 }