From d3acce5aa63a068be95119993a69e6030557c228 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Wed, 20 Aug 2025 16:07:08 +0200 Subject: [PATCH 01/11] fix(tracer-flare): fix design of TracerFlareManager --- datadog-tracer-flare/src/error.rs | 5 -- datadog-tracer-flare/src/lib.rs | 106 ++++++++---------------------- datadog-tracer-flare/src/zip.rs | 32 ++++----- 3 files changed, 42 insertions(+), 101 deletions(-) diff --git a/datadog-tracer-flare/src/error.rs b/datadog-tracer-flare/src/error.rs index cdc6f9a1cc..b2ef620940 100644 --- a/datadog-tracer-flare/src/error.rs +++ b/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}"), } diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index 907e8f5b6c..2c7b961752 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/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, /// As a featured option so we can use the component with no Listener #[cfg(feature = "listener")] pub listener: Option, } -#[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 { - 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), @@ -346,6 +300,17 @@ pub async fn run_remote_config_listener( 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())); + } + } + } } } Err(e) => { @@ -359,7 +324,7 @@ pub async fn run_remote_config_listener( #[cfg(test)] mod tests { use crate::{ - check_remote_config_file, FlareError, LogLevel, ReturnAction, State, TracerFlareManager, + check_remote_config_file, FlareError, LogLevel, ReturnAction, TracerFlareManager, }; use datadog_remote_config::{ config::{ @@ -410,17 +375,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,22 +401,13 @@ 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!(result.unwrap(), ReturnAction::Send); assert_eq!( - tracer_flare.state, - State::Sending { - agent_task: task, - log_level: "debug".to_string() - } + tracer_flare.agent_task, + Some(task) ); } diff --git a/datadog-tracer-flare/src/zip.rs b/datadog-tracer-flare/src/zip.rs index 04eede6b4e..3c9cf9127f 100644 --- a/datadog-tracer-flare/src/zip.rs +++ b/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,23 +229,18 @@ 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), - _ => { - return Err(FlareError::SendError( - "Trying to send the flare without AGENT_TASK received".to_string(), - )) - } +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(), + )) }; 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 +287,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 +307,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 +341,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 +350,16 @@ async fn send(zip: File, tracer_flare: &mut TracerFlareManager) -> Result<(), Fl /// ``` pub async fn zip_and_send( files: Vec, + 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 } From d43829d6d15895b0a411d8397ff828d9a2a1bb13 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Wed, 20 Aug 2025 16:14:27 +0200 Subject: [PATCH 02/11] chore(tracer-flare): fix format --- datadog-tracer-flare/src/lib.rs | 16 ++++++---------- datadog-tracer-flare/src/zip.rs | 14 ++++++++++---- 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index 2c7b961752..15d3967865 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/datadog-tracer-flare/src/lib.rs @@ -299,11 +299,12 @@ pub async fn run_remote_config_listener( if action != Ok(ReturnAction::None) { return action; } - } - else if let Change::Remove(file) = change { + } else if let Change::Remove(file) = change { match file.contents().as_ref() { Ok(data) => match data { - RemoteConfigData::TracerFlareConfig(_) => return Ok(ReturnAction::Unset), + RemoteConfigData::TracerFlareConfig(_) => { + return Ok(ReturnAction::Unset) + } _ => continue, }, Err(e) => { @@ -323,9 +324,7 @@ pub async fn run_remote_config_listener( #[cfg(test)] mod tests { - use crate::{ - check_remote_config_file, FlareError, LogLevel, ReturnAction, TracerFlareManager, - }; + use crate::{check_remote_config_file, FlareError, LogLevel, ReturnAction, TracerFlareManager}; use datadog_remote_config::{ config::{ agent_config::{AgentConfig, AgentConfigFile}, @@ -405,10 +404,7 @@ mod tests { let result = check_remote_config_file(file, &mut tracer_flare); assert!(result.is_ok()); assert_eq!(result.unwrap(), ReturnAction::Send); - assert_eq!( - tracer_flare.agent_task, - Some(task) - ); + assert_eq!(tracer_flare.agent_task, Some(task)); } #[test] diff --git a/datadog-tracer-flare/src/zip.rs b/datadog-tracer-flare/src/zip.rs index 3c9cf9127f..867135edde 100644 --- a/datadog-tracer-flare/src/zip.rs +++ b/datadog-tracer-flare/src/zip.rs @@ -229,12 +229,18 @@ 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, log_level: String, tracer_flare: &mut TracerFlareManager) -> Result<(), FlareError> { +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(), - )) + _ => { + return Err(FlareError::SendError( + "Trying to send the flare without AGENT_TASK received".to_string(), + )) + } }; let payload = generate_payload( From eaf105a5003a68acb21f04690834ccace6af16b6 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Mon, 29 Sep 2025 15:14:32 +0200 Subject: [PATCH 03/11] fix(tracer-flare): fix design and handling of RemoteConfigFile --- datadog-tracer-flare/src/lib.rs | 60 +++++++++++++++++++++++++++------ datadog-tracer-flare/src/zip.rs | 22 ++++++------ 2 files changed, 62 insertions(+), 20 deletions(-) diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index 15d3967865..795f31acee 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/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,10 +29,35 @@ 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 agent_task: Option, + pub state: ReturnAction, /// As a featured option so we can use the component with no Listener #[cfg(feature = "listener")] pub listener: Option, @@ -42,6 +69,7 @@ impl Default for TracerFlareManager { agent_url: hyper::Uri::default().to_string(), language: "rust".to_string(), agent_task: None, + state: ReturnAction::None, #[cfg(feature = "listener")] listener: None, } @@ -141,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, @@ -153,7 +181,7 @@ 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. Set(LogLevel), @@ -207,10 +235,18 @@ 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 { let log_level = log_level.as_str().try_into()?; - return Ok(ReturnAction::Set(log_level)); + match tracer_flare.state { + ReturnAction::Set(level) => { + return Ok(ReturnAction::Set(max(level, log_level))) + } + _ => return Ok(ReturnAction::Set(log_level)), + } } } } @@ -220,13 +256,13 @@ pub fn check_remote_config_file( 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 @@ -295,15 +331,18 @@ 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(_) => { - return Ok(ReturnAction::Unset) + if tracer_flare.state == ReturnAction::None { + tracer_flare.state = ReturnAction::Unset; + } } _ => continue, }, @@ -319,7 +358,7 @@ pub async fn run_remote_config_listener( } } - Ok(ReturnAction::None) + Ok(tracer_flare.state) } #[cfg(test)] @@ -402,6 +441,7 @@ mod tests { .unwrap(); 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::Send); assert_eq!(tracer_flare.agent_task, Some(task)); diff --git a/datadog-tracer-flare/src/zip.rs b/datadog-tracer-flare/src/zip.rs index 867135edde..118871797c 100644 --- a/datadog-tracer-flare/src/zip.rs +++ b/datadog-tracer-flare/src/zip.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +use datadog_remote_config::config::agent_task::AgentTaskFile; use ddcommon::{hyper_migration, Endpoint}; use hyper::{body::Bytes, Method}; use std::{ @@ -232,17 +233,9 @@ fn generate_payload( async fn send( zip: File, log_level: String, + agent_task: AgentTaskFile, 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(), - )) - } - }; - let payload = generate_payload( zip, &tracer_flare.language, @@ -359,11 +352,20 @@ pub async fn zip_and_send( log_level: String, tracer_flare: &mut TracerFlareManager, ) -> Result<(), FlareError> { + let agent_task = match &tracer_flare.agent_task { + Some(agent_task) => agent_task.to_owned(), + _ => { + return Err(FlareError::SendError( + "Trying to send the flare without AGENT_TASK received".to_string(), + )) + } + }; + let zip = zip_files(files)?; // APMSP-2118 - TODO: Implement obfuscation of sensitive data - let response = send(zip, log_level, tracer_flare).await; + let response = send(zip, log_level, agent_task, tracer_flare).await; tracer_flare.agent_task = None; From fdca20a1034e52b0dc85076b9a49e665fd553e99 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Mon, 29 Sep 2025 15:20:55 +0200 Subject: [PATCH 04/11] fix(tracer-flare): fix zip function --- datadog-tracer-flare/src/zip.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/datadog-tracer-flare/src/zip.rs b/datadog-tracer-flare/src/zip.rs index 118871797c..9c733b0e36 100644 --- a/datadog-tracer-flare/src/zip.rs +++ b/datadog-tracer-flare/src/zip.rs @@ -368,6 +368,7 @@ pub async fn zip_and_send( let response = send(zip, log_level, agent_task, tracer_flare).await; tracer_flare.agent_task = None; + tracer_flare.state = crate::ReturnAction::None; response } From 95905163feeccb3b26a014c823ff2fb0318d3dc4 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Wed, 1 Oct 2025 16:23:50 +0200 Subject: [PATCH 05/11] fix(tracer-flare): fix design to return vec of action to perform --- datadog-tracer-flare/src/lib.rs | 70 +++++++++++---------------------- datadog-tracer-flare/src/zip.rs | 46 ++++++++++++++-------- 2 files changed, 53 insertions(+), 63 deletions(-) diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index 795f31acee..25264e139e 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/datadog-tracer-flare/src/lib.rs @@ -9,8 +9,6 @@ pub mod error; pub mod zip; -use std::cmp::max; - use datadog_remote_config::{ config::agent_task::AgentTaskFile, file_storage::RawFile, RemoteConfigData, }; @@ -56,8 +54,6 @@ use crate::error::FlareError; pub struct TracerFlareManager { pub agent_url: String, pub language: String, - pub agent_task: Option, - pub state: ReturnAction, /// As a featured option so we can use the component with no Listener #[cfg(feature = "listener")] pub listener: Option, @@ -68,8 +64,6 @@ impl Default for TracerFlareManager { TracerFlareManager { agent_url: hyper::Uri::default().to_string(), language: "rust".to_string(), - agent_task: None, - state: ReturnAction::None, #[cfg(feature = "listener")] listener: None, } @@ -169,7 +163,7 @@ impl TracerFlareManager { } /// Enum that hold the different log level possible -#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] +#[derive(Debug, PartialEq)] pub enum LogLevel { Trace, Debug, @@ -181,14 +175,14 @@ pub enum LogLevel { } /// Enum that hold the different returned action to do after listening -#[derive(Debug, PartialEq, Clone, Copy)] +#[derive(Debug, PartialEq)] pub enum ReturnAction { /// If AGENT_CONFIG received with the right properties. Set(LogLevel), /// If AGENT_CONFIG is removed. Unset, /// If AGENT_TASK received with the right properties. - Send, + Send(AgentTaskFile), /// If anything else received. None, } @@ -229,47 +223,38 @@ pub type Listener = SingleChangesFetcher Result { let config = file.contents(); 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 { let log_level = log_level.as_str().try_into()?; - match tracer_flare.state { - ReturnAction::Set(level) => { - return Ok(ReturnAction::Set(max(level, log_level))) - } - _ => return Ok(ReturnAction::Set(log_level)), - } + return Ok(ReturnAction::Set(log_level)); } } } RemoteConfigData::TracerFlareTask(agent_task) => { if agent_task.task_type.eq("tracer_flare") { - tracer_flare.agent_task = Some(agent_task.to_owned()); - return Ok(ReturnAction::Send); + return Ok(ReturnAction::Send(agent_task.to_owned())); } } - _ => return Ok(tracer_flare.state), + _ => return Ok(ReturnAction::None), }, Err(e) => { return Err(FlareError::ParsingError(e.to_string())); } } - Ok(tracer_flare.state) + Ok(ReturnAction::None) } /// Function that listens to RemoteConfig on the agent using the TracerFlareManager instance /// /// This function uses the listener contained within the TracerFlareManager to fetch /// RemoteConfig changes from the agent and processes them to determine the -/// appropriate action to take. +/// appropriate actions to take. /// /// # Arguments /// @@ -279,7 +264,7 @@ pub fn check_remote_config_file( /// /// # Returns /// -/// * `Ok(ReturnAction)` - If successful. +/// * `Ok(Vec)` - If successful. /// * `FlareError(msg)` - If something fail. /// /// # Examples @@ -317,7 +302,7 @@ pub fn check_remote_config_file( #[cfg(feature = "listener")] pub async fn run_remote_config_listener( tracer_flare: &mut TracerFlareManager, -) -> Result { +) -> Result, FlareError> { let listener = match &mut tracer_flare.listener { Some(listener) => listener, None => { @@ -326,23 +311,22 @@ pub async fn run_remote_config_listener( )) } }; + let mut actions = Vec::new(); match listener.fetch_changes().await { Ok(changes) => { - println!("Got {} changes.", changes.len()); for change in changes { if let Change::Add(file) = change { - match check_remote_config_file(file, tracer_flare) { - Ok(ReturnAction::Send) => return Ok(ReturnAction::Send), - Ok(action) => tracer_flare.state = action, + match check_remote_config_file(file) { + //, tracer_flare) { + Ok(action) => actions.push(action), Err(err) => return Err(err), } } else if let Change::Remove(file) = change { match file.contents().as_ref() { Ok(data) => match data { + // add action RemoteConfigData::TracerFlareConfig(_) => { - if tracer_flare.state == ReturnAction::None { - tracer_flare.state = ReturnAction::Unset; - } + actions.push(ReturnAction::Unset); } _ => continue, }, @@ -358,12 +342,12 @@ pub async fn run_remote_config_listener( } } - Ok(tracer_flare.state) + Ok(actions) } #[cfg(test)] mod tests { - use crate::{check_remote_config_file, FlareError, LogLevel, ReturnAction, TracerFlareManager}; + use crate::{check_remote_config_file, FlareError, LogLevel, ReturnAction}; use datadog_remote_config::{ config::{ agent_config::{AgentConfig, AgentConfigFile}, @@ -410,8 +394,7 @@ mod tests { let file = storage .store(1, path.clone(), serde_json::to_vec(&config).unwrap()) .unwrap(); - let mut tracer_flare = TracerFlareManager::default(); - let result = check_remote_config_file(file, &mut tracer_flare); + let result = check_remote_config_file(file); assert!(result.is_ok()); assert_eq!(result.unwrap(), ReturnAction::Set(LogLevel::Info)); } @@ -439,12 +422,9 @@ mod tests { let file = storage .store(1, path.clone(), serde_json::to_vec(&task).unwrap()) .unwrap(); - let mut tracer_flare = TracerFlareManager::default(); - let result = check_remote_config_file(file, &mut tracer_flare); - println!("Res : {:?}", result); + let result = check_remote_config_file(file); assert!(result.is_ok()); - assert_eq!(result.unwrap(), ReturnAction::Send); - assert_eq!(tracer_flare.agent_task, Some(task)); + assert_eq!(result.unwrap(), ReturnAction::Send(task)); } #[test] @@ -465,8 +445,7 @@ mod tests { let file = storage .store(1, path.clone(), serde_json::to_vec(&config).unwrap()) .unwrap(); - let mut tracer_flare = TracerFlareManager::default(); - let result = check_remote_config_file(file, &mut tracer_flare); + let result = check_remote_config_file(file); assert!(result.is_ok()); assert_eq!(result.unwrap(), ReturnAction::None); } @@ -484,8 +463,7 @@ mod tests { let file = storage .store(1, path.clone(), b"invalid json".to_vec()) .unwrap(); - let mut tracer_flare = TracerFlareManager::default(); - let result = check_remote_config_file(file, &mut tracer_flare); + let result = check_remote_config_file(file); assert!(result.is_err()); assert!(matches!(result.unwrap_err(), FlareError::ParsingError(_))); } diff --git a/datadog-tracer-flare/src/zip.rs b/datadog-tracer-flare/src/zip.rs index 9c733b0e36..40feca1eaa 100644 --- a/datadog-tracer-flare/src/zip.rs +++ b/datadog-tracer-flare/src/zip.rs @@ -17,7 +17,7 @@ use tempfile::tempfile; use walkdir::WalkDir; use zip::{write::FileOptions, ZipWriter}; -use crate::{error::FlareError, TracerFlareManager}; +use crate::{error::FlareError, ReturnAction, TracerFlareManager}; /// Adds a single file to the zip archive with the specified options and relative path fn add_file_to_zip( @@ -55,7 +55,6 @@ fn add_file_to_zip( /// /// * `files` - A vector of strings representing the paths of files and directories to include in /// the zip archive. -/// * `temp_file` - A temporary file where the zip will be created. /// /// # Returns /// @@ -215,6 +214,7 @@ fn generate_payload( /// /// * `zip` - A file handle to the zip archive to be sent /// * `log_level` - Log level of the tracer +/// * `agent_task` - Agent /// * `tracer_flare` - TracerFlareManager instance containing the agent configuration /// /// # Returns @@ -226,7 +226,6 @@ fn generate_payload( /// /// This function will return an error if: /// - The zip file cannot be read into memory -/// - No agent task was received by the tracer_flare /// - The agent URL is invalid /// - The HTTP request fails after retries /// - The agent returns a non-success HTTP status code @@ -234,7 +233,7 @@ async fn send( zip: File, log_level: String, agent_task: AgentTaskFile, - tracer_flare: &mut TracerFlareManager, + tracer_flare: &TracerFlareManager, ) -> Result<(), FlareError> { let payload = generate_payload( zip, @@ -308,7 +307,8 @@ async fn send( /// 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). +/// * `send_action` - ReturnAction to perform by the tracer flare. Must be a Send action or the +/// function will return an Error. /// /// # Returns /// @@ -328,19 +328,35 @@ async fn send( /// ```rust no_run /// use datadog_tracer_flare::zip::zip_and_send; /// use datadog_tracer_flare::TracerFlareManager; +/// use datadog_tracer_flare::ReturnAction; +/// use datadog_remote_config::config::agent_task::AgentTaskFile; +/// use datadog_remote_config::config::agent_task::AgentTask; +/// use std::num::NonZeroU64; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { -/// let mut tracer_flare = TracerFlareManager::default(); +/// let tracer_flare = TracerFlareManager::default(); /// /// // ... listen to remote config and receiving an agent task ... /// +/// // simulate the receiving of a Send action +/// let task = AgentTaskFile { +/// args: AgentTask { +/// case_id: NonZeroU64::new(123).unwrap(), +/// hostname: "test-host".to_string(), +/// user_handle: "test@example.com".to_string(), +/// }, +/// task_type: "tracer_flare".to_string(), +/// uuid: "test-uuid".to_string(), +/// }; +/// let send_action = ReturnAction::Send(task); +/// /// let files = vec![ /// "/path/to/logs".to_string(), /// "/path/to/config.txt".to_string(), /// ]; /// -/// match zip_and_send(files, "debug".to_string(), &mut tracer_flare).await { +/// match zip_and_send(files, "debug".to_string(), &tracer_flare, send_action).await { /// Ok(_) => println!("Flare sent successfully"), /// Err(e) => eprintln!("Failed to send flare: {}", e), /// } @@ -350,13 +366,14 @@ async fn send( pub async fn zip_and_send( files: Vec, log_level: String, - tracer_flare: &mut TracerFlareManager, + tracer_flare: &TracerFlareManager, + send_action: ReturnAction, ) -> Result<(), FlareError> { - let agent_task = match &tracer_flare.agent_task { - Some(agent_task) => agent_task.to_owned(), + let agent_task = match send_action { + ReturnAction::Send(agent_task) => agent_task, _ => { return Err(FlareError::SendError( - "Trying to send the flare without AGENT_TASK received".to_string(), + "Trying to send the flare with a non Send Action".to_string(), )) } }; @@ -365,12 +382,7 @@ pub async fn zip_and_send( // APMSP-2118 - TODO: Implement obfuscation of sensitive data - let response = send(zip, log_level, agent_task, tracer_flare).await; - - tracer_flare.agent_task = None; - tracer_flare.state = crate::ReturnAction::None; - - response + send(zip, log_level, agent_task, tracer_flare).await } #[cfg(test)] From a6cd28d2431bbc91229583ccabed00dc2541e706 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Wed, 1 Oct 2025 17:19:27 +0200 Subject: [PATCH 06/11] fix: clippy --- datadog-tracer-flare/src/zip.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog-tracer-flare/src/zip.rs b/datadog-tracer-flare/src/zip.rs index 40feca1eaa..d01ec44e9c 100644 --- a/datadog-tracer-flare/src/zip.rs +++ b/datadog-tracer-flare/src/zip.rs @@ -308,7 +308,7 @@ async fn send( /// * `log_level` - Log level of the tracer. /// * `tracer_flare` - TracerFlareManager instance containing the agent configuration and task data. /// * `send_action` - ReturnAction to perform by the tracer flare. Must be a Send action or the -/// function will return an Error. +/// function will return an Error. /// /// # Returns /// From e42d3d3b9ead5d8b0dead4713b6451c9984a8900 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Wed, 1 Oct 2025 17:31:59 +0200 Subject: [PATCH 07/11] fix: remove comment --- datadog-tracer-flare/src/lib.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index 25264e139e..d3ee4741a5 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/datadog-tracer-flare/src/lib.rs @@ -223,7 +223,6 @@ pub type Listener = SingleChangesFetcher Result { let config = file.contents(); match config.as_ref() { @@ -317,14 +316,12 @@ pub async fn run_remote_config_listener( for change in changes { if let Change::Add(file) = change { match check_remote_config_file(file) { - //, tracer_flare) { Ok(action) => actions.push(action), Err(err) => return Err(err), } } else if let Change::Remove(file) = change { match file.contents().as_ref() { Ok(data) => match data { - // add action RemoteConfigData::TracerFlareConfig(_) => { actions.push(ReturnAction::Unset); } From cf8cf294a403837ceaac68e4b46fd8d44fce7d31 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Wed, 1 Oct 2025 17:34:47 +0200 Subject: [PATCH 08/11] fix: fmt --- datadog-tracer-flare/src/lib.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index d3ee4741a5..3af221975a 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/datadog-tracer-flare/src/lib.rs @@ -221,9 +221,7 @@ pub type Listener = SingleChangesFetcher Result { +pub fn check_remote_config_file(file: RemoteConfigFile) -> Result { let config = file.contents(); match config.as_ref() { Ok(data) => match data { From 8cc455af5a097dee11152983c3447191cee7b91c Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Tue, 7 Oct 2025 22:39:08 +0200 Subject: [PATCH 09/11] fix(tracer-flare): update design to handle action by priority --- datadog-tracer-flare/src/lib.rs | 106 ++++++++++++++++++++++++++------ datadog-tracer-flare/src/zip.rs | 16 ++--- 2 files changed, 95 insertions(+), 27 deletions(-) diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index 3af221975a..a7a75db4ad 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/datadog-tracer-flare/src/lib.rs @@ -54,6 +54,7 @@ use crate::error::FlareError; pub struct TracerFlareManager { pub agent_url: String, pub language: String, + pub collecting: bool, /// As a featured option so we can use the component with no Listener #[cfg(feature = "listener")] pub listener: Option, @@ -64,6 +65,7 @@ impl Default for TracerFlareManager { TracerFlareManager { agent_url: hyper::Uri::default().to_string(), language: "rust".to_string(), + collecting: false, #[cfg(feature = "listener")] listener: None, } @@ -71,7 +73,7 @@ impl Default for TracerFlareManager { } impl TracerFlareManager { - /// Function that creates a new TracerFlareManager instance with basic configuration. + /// Creates a new TracerFlareManager instance with basic configuration. /// /// # Arguments /// @@ -91,7 +93,7 @@ impl TracerFlareManager { } } - /// Function that creates a new TracerFlareManager instance and initializes its RemoteConfig + /// Creates a new TracerFlareManager instance and initializes its RemoteConfig /// listener with the provided configuration parameters. /// /// # Arguments @@ -162,8 +164,8 @@ impl TracerFlareManager { } } -/// Enum that hold the different log level possible -#[derive(Debug, PartialEq)] +/// Enum that holds the different log levels possible +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)] pub enum LogLevel { Trace, Debug, @@ -174,19 +176,50 @@ pub enum LogLevel { Off, } -/// Enum that hold the different returned action to do after listening -#[derive(Debug, PartialEq)] +/// Enum that holds the different return actions to perform after listening +#[derive(Debug, PartialEq, Clone)] pub enum ReturnAction { + /// If AGENT_TASK received with the right properties. + Send(AgentTaskFile), /// If AGENT_CONFIG received with the right properties. Set(LogLevel), /// If AGENT_CONFIG is removed. Unset, - /// If AGENT_TASK received with the right properties. - Send(AgentTaskFile), /// If anything else received. None, } +impl ReturnAction { + /// A priority is used to know which action to handle when receiving multiple RemoteConfigFile + /// at the same time. Here is the specific order implemented : + /// 1. Add an AGENT_TASK : `Send(agent_task)` + /// 2. Add an AGENT_CONFIG : `Set(log_level)` + /// 3. Remove an AGENT_CONFIG : `Unset` + /// 4. Anything else : `None` + fn priority(self, other: Self) -> Self { + match &self { + ReturnAction::Send(_) => return self, + ReturnAction::Set(self_level) => match &other { + ReturnAction::Send(_) => return other, + ReturnAction::Set(other_level) => { + if self_level <= other_level { + return self; + } + return other; + } + _ => return self, + }, + ReturnAction::Unset => { + if other == ReturnAction::None { + return self; + } + return other; + } + _ => return other, + } + } +} + impl TryFrom<&str> for LogLevel { type Error = FlareError; @@ -209,19 +242,17 @@ pub type RemoteConfigFile = std::sync::Arc>>; /// Check the `RemoteConfigFile` and return the action that tracer flare needs -/// to perform. This function also updates the `TracerFlareManager` state based on the -/// received configuration. +/// to perform. /// /// # Arguments /// /// * `file` - RemoteConfigFile received by the Listener. -/// * `tracer_flare` - TracerFlareManager object to update with the received configuration. /// /// # Returns /// /// * `Ok(ReturnAction)` - If successful. /// * `FlareError(msg)` - If something fail. -pub fn check_remote_config_file(file: RemoteConfigFile) -> Result { +fn check_remote_config_file(file: RemoteConfigFile) -> Result { let config = file.contents(); match config.as_ref() { Ok(data) => match data { @@ -247,11 +278,40 @@ pub fn check_remote_config_file(file: RemoteConfigFile) -> Result Result { + let action = check_remote_config_file(file); + if let Ok(ReturnAction::Set(_)) = action { + if tracer_flare.collecting { + return Ok(ReturnAction::None); + } + tracer_flare.collecting = true; + } else if let Ok(ReturnAction::Send(_)) = action { + tracer_flare.collecting = false; + } + return action; +} + /// Function that listens to RemoteConfig on the agent using the TracerFlareManager instance /// /// This function uses the listener contained within the TracerFlareManager to fetch /// RemoteConfig changes from the agent and processes them to determine the -/// appropriate actions to take. +/// appropriate action to take. /// /// # Arguments /// @@ -261,7 +321,7 @@ pub fn check_remote_config_file(file: RemoteConfigFile) -> Result)` - If successful. +/// * `Ok(ReturnAction)` - If successful. /// * `FlareError(msg)` - If something fail. /// /// # Examples @@ -299,7 +359,7 @@ pub fn check_remote_config_file(file: RemoteConfigFile) -> Result Result, FlareError> { +) -> Result { let listener = match &mut tracer_flare.listener { Some(listener) => listener, None => { @@ -308,20 +368,22 @@ pub async fn run_remote_config_listener( )) } }; - let mut actions = Vec::new(); + let mut state = ReturnAction::None; match listener.fetch_changes().await { Ok(changes) => { for change in changes { if let Change::Add(file) = change { match check_remote_config_file(file) { - Ok(action) => actions.push(action), + Ok(action) => state = ReturnAction::priority(action, state), Err(err) => return Err(err), } } else if let Change::Remove(file) = change { match file.contents().as_ref() { Ok(data) => match data { RemoteConfigData::TracerFlareConfig(_) => { - actions.push(ReturnAction::Unset); + if state == ReturnAction::None { + state = ReturnAction::Unset; + } } _ => continue, }, @@ -337,7 +399,13 @@ pub async fn run_remote_config_listener( } } - Ok(actions) + if let ReturnAction::Set(_) = state { + tracer_flare.collecting = true; + } else if let ReturnAction::Send(_) = state { + tracer_flare.collecting = false; + } + + Ok(state) } #[cfg(test)] diff --git a/datadog-tracer-flare/src/zip.rs b/datadog-tracer-flare/src/zip.rs index d01ec44e9c..2330736c06 100644 --- a/datadog-tracer-flare/src/zip.rs +++ b/datadog-tracer-flare/src/zip.rs @@ -111,7 +111,7 @@ fn zip_files(files: Vec) -> Result { add_file_to_zip(&mut zip, &path, None, &options)?; } else { return Err(FlareError::ZipError(format!( - "Invalid or inexisting file: {}", + "Invalid or non-existent file: {}", path.to_string_lossy() ))); } @@ -327,19 +327,17 @@ async fn send( /// /// ```rust no_run /// use datadog_tracer_flare::zip::zip_and_send; -/// use datadog_tracer_flare::TracerFlareManager; -/// use datadog_tracer_flare::ReturnAction; -/// use datadog_remote_config::config::agent_task::AgentTaskFile; -/// use datadog_remote_config::config::agent_task::AgentTask; +/// use datadog_tracer_flare::{TracerFlareManager, ReturnAction}; +/// use datadog_remote_config::config::agent_task::{AgentTaskFile, AgentTask}; /// use std::num::NonZeroU64; /// /// #[tokio::main] /// async fn main() -> Result<(), Box> { /// let tracer_flare = TracerFlareManager::default(); /// -/// // ... listen to remote config and receiving an agent task ... +/// // ... listen to remote config and receive an agent task ... /// -/// // simulate the receiving of a Send action +/// // Simulate receiving a Send action from remote config /// let task = AgentTaskFile { /// args: AgentTask { /// case_id: NonZeroU64::new(123).unwrap(), @@ -382,7 +380,9 @@ pub async fn zip_and_send( // APMSP-2118 - TODO: Implement obfuscation of sensitive data - send(zip, log_level, agent_task, tracer_flare).await + let response = send(zip, log_level, agent_task, tracer_flare).await; + + response } #[cfg(test)] From d1b051f0719c7bf9380ce827c40c968ef24ba8d0 Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Tue, 7 Oct 2025 22:59:25 +0200 Subject: [PATCH 10/11] fix: clippy --- datadog-tracer-flare/src/lib.rs | 14 +++++++------- datadog-tracer-flare/src/zip.rs | 4 +--- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index a7a75db4ad..b1d7e53d07 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/datadog-tracer-flare/src/lib.rs @@ -198,24 +198,24 @@ impl ReturnAction { /// 4. Anything else : `None` fn priority(self, other: Self) -> Self { match &self { - ReturnAction::Send(_) => return self, + ReturnAction::Send(_) => self, ReturnAction::Set(self_level) => match &other { - ReturnAction::Send(_) => return other, + ReturnAction::Send(_) => other, ReturnAction::Set(other_level) => { if self_level <= other_level { return self; } - return other; + other } - _ => return self, + _ => self, }, ReturnAction::Unset => { if other == ReturnAction::None { return self; } - return other; + other } - _ => return other, + _ => other, } } } @@ -304,7 +304,7 @@ pub fn handle_remote_config_file( } else if let Ok(ReturnAction::Send(_)) = action { tracer_flare.collecting = false; } - return action; + action } /// Function that listens to RemoteConfig on the agent using the TracerFlareManager instance diff --git a/datadog-tracer-flare/src/zip.rs b/datadog-tracer-flare/src/zip.rs index 2330736c06..35af39f77b 100644 --- a/datadog-tracer-flare/src/zip.rs +++ b/datadog-tracer-flare/src/zip.rs @@ -380,9 +380,7 @@ pub async fn zip_and_send( // APMSP-2118 - TODO: Implement obfuscation of sensitive data - let response = send(zip, log_level, agent_task, tracer_flare).await; - - response + send(zip, log_level, agent_task, tracer_flare).await } #[cfg(test)] From 0cf60477388fb7ba66b45ceb625a1046a5015b0f Mon Sep 17 00:00:00 2001 From: Anais Raison Date: Mon, 13 Oct 2025 15:02:33 +0200 Subject: [PATCH 11/11] fix: apply review --- datadog-tracer-flare/src/lib.rs | 301 +++++++++++++++++++++++--------- 1 file changed, 218 insertions(+), 83 deletions(-) diff --git a/datadog-tracer-flare/src/lib.rs b/datadog-tracer-flare/src/lib.rs index b1d7e53d07..c3dae955e5 100644 --- a/datadog-tracer-flare/src/lib.rs +++ b/datadog-tracer-flare/src/lib.rs @@ -27,30 +27,29 @@ use { use crate::error::FlareError; -/// A manager for handling tracer flare functionality and remote configuration. +/// Manager for tracer flare functionality with optional remote configuration support. /// /// 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 +/// - **No listener mode**: Stores agent URL and language configuration for flare operations +/// - **Listener 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) +/// # Fields /// -/// 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 +/// - `agent_url`: The agent endpoint URL for flare transmission +/// - `language`: The tracer language identifier +/// - `collecting`: Current collection state (true when actively collecting) +/// - `listener`: Optional remote config listener (requires "listener" feature) /// -/// 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 +/// # Typical usage flow +/// +/// 1. Create manager with [`new`](Self::new) for usage without listener or +/// [`new_with_listener`](Self::new_with_listener) for usage with listener +/// 2. Call [`run_remote_config_listener`] periodically to fetch and process remote config changes +/// 3. Handle returned [`ReturnAction`]: `Send(agent_task)`, `Set(log_level)`, `Unset`, or `None` +/// 4. Use the `collecting` field to track current flare collection state pub struct TracerFlareManager { pub agent_url: String, pub language: String, @@ -162,9 +161,40 @@ impl TracerFlareManager { Ok(tracer_flare) } + + /// Handle the `RemoteConfigFile` and return the action that tracer flare needs + /// to perform. This function also updates the `TracerFlareManager` state based on the + /// received configuration. + /// + /// # Arguments + /// + /// * `file` - RemoteConfigFile received by the Listener. + /// * `tracer_flare` - TracerFlareManager object to update with the received configuration. + /// + /// # Returns + /// + /// * `Ok(ReturnAction)` - If successful. + /// * `FlareError(msg)` - If something fail. + pub fn handle_remote_config_file( + &mut self, + file: RemoteConfigFile, + ) -> Result { + let action = file.try_into(); + if let Ok(ReturnAction::Set(_)) = action { + if self.collecting { + return Ok(ReturnAction::None); + } + self.collecting = true; + } else if Ok(ReturnAction::None) != action { + // If action is Send, Unset or an error, we need to stop collecting + self.collecting = false; + } + action + } } /// Enum that holds the different log levels possible +/// Do not change the order of the variants because we rely on Ord #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)] pub enum LogLevel { Trace, @@ -180,10 +210,16 @@ pub enum LogLevel { #[derive(Debug, PartialEq, Clone)] pub enum ReturnAction { /// If AGENT_TASK received with the right properties. + /// + /// Trigger to collect the flare and send it to the agent. Send(AgentTaskFile), /// If AGENT_CONFIG received with the right properties. + /// + /// Trigger to set the log level of the tracer. Set(LogLevel), /// If AGENT_CONFIG is removed. + /// + /// Trigger to and unset the log level. Unset, /// If anything else received. None, @@ -241,70 +277,45 @@ pub type RemoteConfigFile = std::sync::Arc>>; -/// Check the `RemoteConfigFile` and return the action that tracer flare needs -/// to perform. -/// -/// # Arguments -/// -/// * `file` - RemoteConfigFile received by the Listener. -/// -/// # Returns -/// -/// * `Ok(ReturnAction)` - If successful. -/// * `FlareError(msg)` - If something fail. -fn check_remote_config_file(file: RemoteConfigFile) -> Result { - let config = file.contents(); - match config.as_ref() { - Ok(data) => match data { - RemoteConfigData::TracerFlareConfig(agent_config) => { - if agent_config.name.starts_with("flare-log-level.") { - if let Some(log_level) = &agent_config.config.log_level { - let log_level = log_level.as_str().try_into()?; - return Ok(ReturnAction::Set(log_level)); +impl TryFrom for ReturnAction { + type Error = FlareError; + + /// Check the `RemoteConfigFile` and return the action that tracer flare needs + /// to perform. + /// + /// # Arguments + /// + /// * `file` - RemoteConfigFile received by the Listener. + /// + /// # Returns + /// + /// * `Ok(ReturnAction)` - If successful. + /// * `FlareError(msg)` - If something fail. + fn try_from(file: RemoteConfigFile) -> Result { + let config = file.contents(); + match config.as_ref() { + Ok(data) => match data { + RemoteConfigData::TracerFlareConfig(agent_config) => { + if agent_config.name.starts_with("flare-log-level.") { + if let Some(log_level) = &agent_config.config.log_level { + let log_level = log_level.as_str().try_into()?; + return Ok(ReturnAction::Set(log_level)); + } } } - } - RemoteConfigData::TracerFlareTask(agent_task) => { - if agent_task.task_type.eq("tracer_flare") { - return Ok(ReturnAction::Send(agent_task.to_owned())); + RemoteConfigData::TracerFlareTask(agent_task) => { + if agent_task.task_type.eq("tracer_flare") { + return Ok(ReturnAction::Send(agent_task.to_owned())); + } } + _ => return Ok(ReturnAction::None), + }, + Err(e) => { + return Err(FlareError::ParsingError(e.to_string())); } - _ => return Ok(ReturnAction::None), - }, - Err(e) => { - return Err(FlareError::ParsingError(e.to_string())); } + Ok(ReturnAction::None) } - Ok(ReturnAction::None) -} - -/// Handle the `RemoteConfigFile` and return the action that tracer flare needs -/// to perform. This function also updates the `TracerFlareManager` state based on the -/// received configuration. -/// -/// # Arguments -/// -/// * `file` - RemoteConfigFile received by the Listener. -/// * `tracer_flare` - TracerFlareManager object to update with the received configuration. -/// -/// # Returns -/// -/// * `Ok(ReturnAction)` - If successful. -/// * `FlareError(msg)` - If something fail. -pub fn handle_remote_config_file( - file: RemoteConfigFile, - tracer_flare: &mut TracerFlareManager, -) -> Result { - let action = check_remote_config_file(file); - if let Ok(ReturnAction::Set(_)) = action { - if tracer_flare.collecting { - return Ok(ReturnAction::None); - } - tracer_flare.collecting = true; - } else if let Ok(ReturnAction::Send(_)) = action { - tracer_flare.collecting = false; - } - action } /// Function that listens to RemoteConfig on the agent using the TracerFlareManager instance @@ -373,7 +384,7 @@ pub async fn run_remote_config_listener( Ok(changes) => { for change in changes { if let Change::Add(file) = change { - match check_remote_config_file(file) { + match file.try_into() { Ok(action) => state = ReturnAction::priority(action, state), Err(err) => return Err(err), } @@ -410,7 +421,7 @@ pub async fn run_remote_config_listener( #[cfg(test)] mod tests { - use crate::{check_remote_config_file, FlareError, LogLevel, ReturnAction}; + use crate::{FlareError, LogLevel, ReturnAction}; use datadog_remote_config::{ config::{ agent_config::{AgentConfig, AgentConfigFile}, @@ -438,7 +449,73 @@ mod tests { } #[test] - fn test_check_remote_config_file_with_valid_log_level() { + fn test_log_level_ordering() { + // Test that the ordering is maintained as expected (Trace < Debug < Info < Warn < Error < + // Critical < Off) + assert!(LogLevel::Trace < LogLevel::Debug); + assert!(LogLevel::Debug < LogLevel::Info); + assert!(LogLevel::Info < LogLevel::Warn); + assert!(LogLevel::Warn < LogLevel::Error); + assert!(LogLevel::Error < LogLevel::Critical); + assert!(LogLevel::Critical < LogLevel::Off); + } + + #[test] + fn test_priority_in_return_action() { + // Test that when two Set actions are compared, the one with lower log level wins + let send_action = ReturnAction::Send(AgentTaskFile { + args: AgentTask { + case_id: NonZeroU64::new(123).unwrap(), + hostname: "test-host".to_string(), + user_handle: "test@example.com".to_string(), + }, + task_type: "tracer_flare".to_string(), + uuid: "test_uuid".to_string(), + }); + let trace_action = ReturnAction::Set(LogLevel::Trace); + let off_action = ReturnAction::Set(LogLevel::Off); + let unset_action = ReturnAction::Unset; + let none_action = ReturnAction::None; + + // Lower log levels should have priority (trace < debug < info < ... < off) + assert_eq!( + send_action.clone().priority(trace_action.clone()), + send_action + ); + assert_eq!( + trace_action.clone().priority(off_action.clone()), + trace_action + ); + assert_eq!( + off_action.clone().priority(unset_action.clone()), + off_action + ); + assert_eq!( + unset_action.clone().priority(none_action.clone()), + unset_action + ); + + // Test reverse order + assert_eq!( + trace_action.clone().priority(send_action.clone()), + send_action + ); + assert_eq!( + off_action.clone().priority(trace_action.clone()), + trace_action + ); + assert_eq!( + unset_action.clone().priority(off_action.clone()), + off_action + ); + assert_eq!( + none_action.clone().priority(unset_action.clone()), + unset_action + ); + } + + #[test] + fn test_remote_config_with_valid_log_level() { let storage = ParsedFileStorage::default(); let path = Arc::new(RemoteConfigPath { product: RemoteConfigProduct::AgentConfig, @@ -457,13 +534,13 @@ mod tests { let file = storage .store(1, path.clone(), serde_json::to_vec(&config).unwrap()) .unwrap(); - let result = check_remote_config_file(file); + let result = ReturnAction::try_from(file); assert!(result.is_ok()); assert_eq!(result.unwrap(), ReturnAction::Set(LogLevel::Info)); } #[test] - fn test_check_remote_config_file_with_send_task() { + fn test_remote_config_with_send_task() { let storage = ParsedFileStorage::default(); let path = Arc::new(RemoteConfigPath { product: RemoteConfigProduct::AgentTask, @@ -485,13 +562,13 @@ mod tests { let file = storage .store(1, path.clone(), serde_json::to_vec(&task).unwrap()) .unwrap(); - let result = check_remote_config_file(file); + let result = ReturnAction::try_from(file); assert!(result.is_ok()); assert_eq!(result.unwrap(), ReturnAction::Send(task)); } #[test] - fn test_check_remote_config_file_with_invalid_config() { + fn test_remote_config_with_invalid_config() { let storage = ParsedFileStorage::default(); let path = Arc::new(RemoteConfigPath { product: RemoteConfigProduct::AgentConfig, @@ -508,11 +585,69 @@ mod tests { let file = storage .store(1, path.clone(), serde_json::to_vec(&config).unwrap()) .unwrap(); - let result = check_remote_config_file(file); + let result = ReturnAction::try_from(file); assert!(result.is_ok()); assert_eq!(result.unwrap(), ReturnAction::None); } + #[test] + fn test_handle_remote_config_file() { + use crate::TracerFlareManager; + let mut tracer_flare = TracerFlareManager::new("http://localhost:8126", "rust"); + let storage = ParsedFileStorage::default(); + + let agent_config_file = storage + .store( + 1, + Arc::new(RemoteConfigPath { + product: RemoteConfigProduct::AgentConfig, + config_id: "test".to_string(), + name: "flare-log-level.test".to_string(), + source: RemoteConfigSource::Datadog(1), + }), + serde_json::to_vec(&AgentConfigFile { + name: "flare-log-level.test".to_string(), + config: AgentConfig { + log_level: Some("info".to_string()), + }, + }) + .unwrap(), + ) + .unwrap(); + + // First AGENT_CONFIG + assert!(!tracer_flare.collecting); + let result = tracer_flare + .handle_remote_config_file(agent_config_file.clone()) + .unwrap(); + assert_eq!(result, ReturnAction::Set(LogLevel::Info)); + assert!(tracer_flare.collecting); + + // Second AGENT_CONFIG + let result = tracer_flare + .handle_remote_config_file(agent_config_file) + .unwrap(); + assert_eq!(result, ReturnAction::None); + assert!(tracer_flare.collecting); + + // Non-None actions stop collecting + let error_file = storage + .store( + 2, + Arc::new(RemoteConfigPath { + product: RemoteConfigProduct::AgentConfig, + config_id: "error".to_string(), + name: "error".to_string(), + source: RemoteConfigSource::Datadog(1), + }), + b"invalid".to_vec(), + ) + .unwrap(); + + let _ = tracer_flare.handle_remote_config_file(error_file); + assert!(!tracer_flare.collecting); + } + #[test] fn test_check_remote_config_file_with_parsing_error() { let storage = ParsedFileStorage::default(); @@ -526,7 +661,7 @@ mod tests { let file = storage .store(1, path.clone(), b"invalid json".to_vec()) .unwrap(); - let result = check_remote_config_file(file); + let result = ReturnAction::try_from(file); assert!(result.is_err()); assert!(matches!(result.unwrap_err(), FlareError::ParsingError(_))); }