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

@@ -9,6 +9,8 @@
 pub mod error;
 pub mod zip;
 
+use std::cmp::max;
+
 use datadog_remote_config::{
     config::agent_task::AgentTaskFile, file_storage::RawFile, RemoteConfigData,
 };
@@ -27,33 +29,47 @@ use {
 
 use crate::error::FlareError;
 
+/// A manager for handling tracer flare functionality and remote configuration.
+///
+/// The TracerFlareManager serves as the central coordinator for tracer flare operations,
+/// managing the lifecycle of flare collection and transmission. It operates in two modes:
+///
+/// - **Basic mode**: Stores agent URL and language configuration for flare operations
+/// - **Remote config mode**: Listens to remote configuration updates to automatically trigger flare
+///   collection and transmission
+///
+/// The manager maintains:
+/// - Agent connection details (URL, language)
+/// - Current agent task information (when received via remote config)
+/// - Remote configuration listener (when enabled)
+///
+/// Key responsibilities:
+/// - Parsing remote configuration files to determine flare actions (Set log level, Send flare)
+/// - Storing agent task metadata needed for flare transmission
+/// - Coordinating with the zip module to package and send flares to the agent
+///
+/// Typical usage flow:
+/// 1. Create manager with basic config or remote config listener
+/// 2. Listen for remote config changes that trigger flare actions
+/// 3. When a flare is requested, use the stored agent task to send the flare
+/// 4. Manager state is reset after flare transmission
 pub struct TracerFlareManager {
     pub agent_url: String,
     pub language: String,
-    pub state: State,
+    pub agent_task: Option<AgentTaskFile>,
+    pub state: ReturnAction,
     /// 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,
+            state: ReturnAction::None,
             #[cfg(feature = "listener")]
             listener: None,
         }
@@ -153,7 +169,7 @@ impl TracerFlareManager {
 }
 
 /// Enum that hold the different log level possible
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
 pub enum LogLevel {
     Trace,
     Debug,
@@ -165,12 +181,14 @@ pub enum LogLevel {
 }
 
 /// Enum that hold the different returned action to do after listening
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Clone, Copy)]
 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 +197,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),
@@ -217,62 +235,34 @@ pub fn check_remote_config_file(
     match config.as_ref() {
         Ok(data) => match data {
             RemoteConfigData::TracerFlareConfig(agent_config) => {
+                if let ReturnAction::Send = tracer_flare.state {
+                    return Ok(tracer_flare.state);
+                }
                 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));
+                        match tracer_flare.state {
+                            ReturnAction::Set(level) => {
+                                return Ok(ReturnAction::Set(max(level, 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),
+            _ => return Ok(tracer_flare.state),
         },
         Err(e) => {
             return Err(FlareError::ParsingError(e.to_string()));
         }
     }
-    Ok(ReturnAction::None)
+    Ok(tracer_flare.state)
 }
 
 /// Function that listens to RemoteConfig on the agent using the TracerFlareManager instance
@@ -341,9 +331,24 @@ pub async fn run_remote_config_listener(
             println!("Got {} changes.", changes.len());
             for change in changes {
                 if let Change::Add(file) = change {
-                    let action = check_remote_config_file(file, tracer_flare);
-                    if action != Ok(ReturnAction::None) {
-                        return action;
+                    match check_remote_config_file(file, tracer_flare) {
+                        Ok(ReturnAction::Send) => return Ok(ReturnAction::Send),
+                        Ok(action) => tracer_flare.state = action,
+                        Err(err) => return Err(err),
+                    }
+                } else if let Change::Remove(file) = change {
+                    match file.contents().as_ref() {
+                        Ok(data) => match data {
+                            RemoteConfigData::TracerFlareConfig(_) => {
+                                if tracer_flare.state == ReturnAction::None {
+                                    tracer_flare.state = ReturnAction::Unset;
+                                }
+                            }
+                            _ => continue,
+                        },
+                        Err(e) => {
+                            return Err(FlareError::ParsingError(e.to_string()));
+                        }
                     }
                 }
             }
@@ -353,14 +358,12 @@ pub async fn run_remote_config_listener(
         }
     }
 
-    Ok(ReturnAction::None)
+    Ok(tracer_flare.state)
 }
 
 #[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 +413,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 +439,12 @@ 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);
+        println!("Res : {:?}", result);
         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]