fix: memory leak replacing long lived spans (#1281)

* fix: replace long lived spans

* docs

* review

* activate self inst on canaries

* test cluster

* use agent_id

* unify span guard

Author: gsanchezgavier

agent-control/src/agent_control/agent_id.rs

@@ -92,6 +92,12 @@ pub(crate) mod tests {
     use crate::agent_control::agent_id::AgentID;
     use crate::agent_control::defaults::AGENT_CONTROL_ID;
 
+    impl Default for AgentID {
+        fn default() -> Self {
+            AgentID::new("default").unwrap()
+        }
+    }
+
     #[test]
     fn agent_control_id() {
         let agent_id = AgentID::new_agent_control_id();

agent-control/src/opamp/operations.rs

@@ -117,7 +117,6 @@ pub fn stop_opamp_client<C: StartedClient>(
             "Stopping OpAMP client for supervised agent type: {}",
             agent_id
         );
-        // TODO We should call disconnect here as this means a graceful shutdown
         client.stop()?;
     }
     Ok(())

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

@@ -1,17 +1,19 @@
 use super::with_start_time::StartTime;
+use crate::agent_control::agent_id::AgentID;
 use crate::agent_type::runtime_config::HealthCheckInterval;
 use crate::event::SubAgentInternalEvent;
 use crate::event::cancellation::CancellationMessage;
 use crate::event::channel::{EventConsumer, EventPublisher};
 
 use crate::k8s;
 use crate::sub_agent::health::with_start_time::HealthWithStartTime;
+use crate::sub_agent::identity::ID_ATTRIBUTE_NAME;
 use crate::sub_agent::supervisor::starter::SupervisorStarterError;
 use crate::utils::thread_context::{NotStartedThreadContext, StartedThreadContext};
 use std::time::{SystemTime, SystemTimeError};
-use tracing::{debug, error};
+use tracing::{debug, error, info_span};
 
-const HEALTH_CHECKER_THREAD_NAME: &str = "health checker";
+const HEALTH_CHECKER_THREAD_NAME: &str = "health_checker";
 
 pub type StatusTime = SystemTime;
 
@@ -212,6 +214,7 @@ pub trait HealthChecker {
 }
 
 pub(crate) fn spawn_health_checker<H>(
+    agent_id: AgentID,
     health_checker: H,
     sub_agent_internal_publisher: EventPublisher<SubAgentInternalEvent>,
     interval: HealthCheckInterval,
@@ -221,6 +224,12 @@ where
     H: HealthChecker + Send + 'static,
 {
     let callback = move |stop_consumer: EventConsumer<CancellationMessage>| loop {
+        let span = info_span!(
+            "health_check",
+            { ID_ATTRIBUTE_NAME } = %agent_id
+        );
+        let _guard = span.enter();
+
         debug!("starting to check health with the configured checker");
 
         let health = health_checker.check_health().unwrap_or_else(|err| {
@@ -378,6 +387,7 @@ pub mod tests {
             });
 
         let started_thread_context = spawn_health_checker(
+            AgentID::default(),
             health_checker,
             health_publisher,
             Duration::from_millis(10).into(), // Give room to publish and consume the events
@@ -441,6 +451,7 @@ pub mod tests {
             });
 
         let started_thread_context = spawn_health_checker(
+            AgentID::default(),
             health_checker,
             health_publisher,
             Duration::from_millis(10).into(), // Give room to publish and consume the events
@@ -499,6 +510,7 @@ pub mod tests {
         let start_time = SystemTime::now();
 
         let started_thread_context = spawn_health_checker(
+            AgentID::default(),
             health_checker,
             health_publisher,
             Duration::from_millis(10).into(), // Give room to publish and consume the events

agent-control/src/sub_agent/identity.rs

@@ -5,6 +5,8 @@ use crate::agent_control::defaults::{
 use crate::agent_type::agent_type_id::AgentTypeID;
 use std::fmt::{Display, Formatter};
 
+pub const ID_ATTRIBUTE_NAME: &str = "id";
+
 // This could be SubAgentIdentity
 #[derive(Clone, Debug, PartialEq)]
 pub struct AgentIdentity {

agent-control/src/sub_agent/k8s/builder.rs

@@ -85,7 +85,7 @@ where
 {
     type NotStartedSubAgent = SubAgent<O::Client, SA, R, H, Y>;
 
-    #[instrument(skip_all, fields(%agent_identity),name = "build_sub_agent")]
+    #[instrument(skip_all, fields(id = %agent_identity.id),name = "build_agent")]
     fn build(
         &self,
         agent_identity: &AgentIdentity,

agent-control/src/sub_agent/k8s/supervisor.rs

@@ -11,7 +11,7 @@ use crate::k8s::labels::Labels;
 use crate::sub_agent::health::health_checker::spawn_health_checker;
 use crate::sub_agent::health::k8s::health_checker::SubAgentHealthChecker;
 use crate::sub_agent::health::with_start_time::StartTime;
-use crate::sub_agent::identity::AgentIdentity;
+use crate::sub_agent::identity::{AgentIdentity, ID_ATTRIBUTE_NAME};
 use crate::sub_agent::supervisor::starter::{SupervisorStarter, SupervisorStarterError};
 use crate::sub_agent::supervisor::stopper::SupervisorStopper;
 use crate::sub_agent::version::k8s::checkers::K8sAgentVersionChecker;
@@ -24,10 +24,10 @@ use k8s_openapi::serde_json;
 use kube::{api::DynamicObject, core::TypeMeta};
 use std::sync::Arc;
 use std::time::Duration;
-use tracing::{debug, error, info, trace, warn};
+use tracing::{debug, error, info, info_span, trace, warn};
 
 const OBJECTS_SUPERVISOR_INTERVAL_SECONDS: u64 = 30;
-const SUPERVISOR_THREAD_NAME: &str = "k8s objects supervisor";
+const SUPERVISOR_THREAD_NAME: &str = "supervisor";
 
 pub struct NotStartedSupervisorK8s {
     agent_identity: AgentIdentity,
@@ -126,7 +126,15 @@ impl NotStartedSupervisorK8s {
     ) -> StartedThreadContext {
         let k8s_client = self.k8s_client.clone();
         let interval = self.interval;
+        let agent_id = self.agent_identity.id.clone();
+
         let callback = move |stop_consumer: EventConsumer<CancellationMessage>| loop {
+            let span = info_span!(
+                "reconcile_resources",
+                { ID_ATTRIBUTE_NAME } = %agent_id
+            );
+            let _guard = span.enter();
+
             // Check and apply k8s objects
             if let Err(err) = Self::apply_resources(resources.iter(), k8s_client.clone()) {
                 warn!(%err, "K8s resources apply failed");
@@ -161,6 +169,7 @@ impl NotStartedSupervisorK8s {
         };
 
         let started_thread_context = spawn_health_checker(
+            self.agent_identity.id.clone(),
             k8s_health_checker,
             sub_agent_internal_publisher,
             health_config.interval,
@@ -182,6 +191,7 @@ impl NotStartedSupervisorK8s {
         )?;
 
         Some(spawn_version_checker(
+            self.agent_identity.id.clone(),
             k8s_version_checker,
             sub_agent_internal_publisher,
             VersionCheckerInterval::default(),

agent-control/src/sub_agent/on_host/builder.rs

@@ -81,7 +81,7 @@ where
 {
     type NotStartedSubAgent = SubAgent<O::Client, SA, R, H, Y>;
 
-    #[instrument(skip_all, fields(%agent_identity),name = "build_sub_agent")]
+    #[instrument(skip_all, fields(id = %agent_identity.id),name = "build_agent")]
     fn build(
         &self,
         agent_identity: &AgentIdentity,

agent-control/src/sub_agent/on_host/supervisor.rs

@@ -12,7 +12,7 @@ use crate::sub_agent::health::health_checker::{
 use crate::sub_agent::health::health_checker::{Healthy, Unhealthy};
 use crate::sub_agent::health::on_host::health_checker::OnHostHealthChecker;
 use crate::sub_agent::health::with_start_time::{HealthWithStartTime, StartTime};
-use crate::sub_agent::identity::AgentIdentity;
+use crate::sub_agent::identity::{AgentIdentity, ID_ATTRIBUTE_NAME};
 use crate::sub_agent::on_host::command::command::CommandError;
 use crate::sub_agent::on_host::command::command_os::CommandOSNotStarted;
 use crate::sub_agent::on_host::command::executable_data::ExecutableData;
@@ -140,6 +140,7 @@ impl NotStartedSupervisorOnHost {
             let health_checker =
                 OnHostHealthChecker::try_new(http_client, health_config.clone(), start_time)?;
             let started_thread_context = spawn_health_checker(
+                self.agent_identity.id.clone(),
                 health_checker,
                 sub_agent_internal_publisher,
                 health_config.interval,
@@ -159,6 +160,7 @@ impl NotStartedSupervisorOnHost {
             OnHostAgentVersionChecker::checked_new(self.agent_identity.agent_type_id.clone())?;
 
         Some(spawn_version_checker(
+            self.agent_identity.id.clone(),
             onhost_version_checker,
             sub_agent_internal_publisher,
             VersionCheckerInterval::default(),
@@ -176,10 +178,16 @@ impl NotStartedSupervisorOnHost {
         _ = wait_for_termination(current_pid.clone(), self.ctx.clone(), shutdown_ctx.clone());
 
         let executable_data_clone = executable_data.clone();
+        let agent_id = self.agent_identity.id.clone();
         // NotStartedThreadContext takes as input a callback that requires a EventConsumer<CancellationMessage>
         // as input. In that specific case it's not used, but we need to pass it to comply with the signature.
         // This should be refactored to work as the other threads used by the supervisor.
         let callback = move |_| loop {
+            let span = info_span!(
+                "start_executable",
+                { ID_ATTRIBUTE_NAME } = %agent_id
+            );
+            let span_guard = span.enter();
             // locks the current_pid to prevent `wait_for_termination` finishing before the process
             // is started and the pid is set.
             // In case starting the process fail the guard will be dropped and `wait_for_termination`
@@ -189,17 +197,11 @@ impl NotStartedSupervisorOnHost {
             // A context cancelled means that the supervisor has been gracefully stopped
             // before the process was started.
             if *Context::get_lock_cvar(&self.ctx).0.lock().unwrap() {
-                debug!(
-                    supervisor = executable_data_clone.bin,
-                    msg = "supervisor stopped before starting the process"
-                );
+                debug!("supervisor stopped before starting the process");
                 break;
             }
 
-            info!(
-                supervisor = executable_data_clone.bin,
-                msg = "starting supervisor process"
-            );
+            info!("starting supervisor process");
 
             shutdown_ctx.reset().unwrap();
             // Signals return exit_code 0, if in the future we need to act on them we can import
@@ -215,7 +217,14 @@ impl NotStartedSupervisorOnHost {
                 HealthWithStartTime::new(init_health.into(), supervisor_start_time).into(),
             );
 
-            let exit_code = start_command(not_started_command, pid_guard)
+            let command_result = start_command(not_started_command, pid_guard, span_guard);
+            let span = info_span!(
+                "stop_executable",
+                { ID_ATTRIBUTE_NAME } = %agent_id
+            );
+            let _span_guard = span.enter();
+
+            let exit_code = command_result
                 .inspect_err(|err| {
                     error!(
                         supervisor = executable_data_clone.bin,
@@ -340,6 +349,7 @@ fn handle_termination(
 fn start_command(
     not_started_command: CommandOSNotStarted,
     mut pid: std::sync::MutexGuard<Option<u32>>,
+    span_guard: tracing::span::Entered<'_>,
 ) -> Result<ExitStatus, CommandError> {
     // run and stream the process
     let started = not_started_command.start()?;
@@ -351,6 +361,8 @@ fn start_command(
     // free the lock so the wait_for_termination can lock it on graceful shutdown
     drop(pid);
 
+    drop(span_guard);
+
     streaming.wait()
 }
 
@@ -361,11 +373,11 @@ fn wait_for_termination(
     ctx: Context<bool>,
     shutdown_ctx: Context<bool>,
 ) -> JoinHandle<()> {
-    let s = info_span!("termination_signal_listener");
+    let span = info_span!("termination_signal");
     spawn_named_thread("OnHost Termination signal listener", move || {
-        let _guards = s.enter();
         let (lck, cvar) = Context::get_lock_cvar(&ctx);
         drop(cvar.wait_while(lck.lock().unwrap(), |finish| !*finish));
+        let _span_guard = span.enter();
 
         // context is unlocked here so locking it again in other thread that is blocking current_pid is safe.
         if let Some(pid) = *current_pid.lock().unwrap() {

agent-control/src/sub_agent/sub_agent.rs

@@ -129,9 +129,9 @@ where
     }
 
     pub fn runtime(self) -> JoinHandle<Result<(), SubAgentError>> {
-        let s = info_span!("agent", id=%self.identity.id);
         spawn_named_thread("Subagent runtime", move || {
-            let _guards = s.enter();
+            let span = info_span!("start_agent", id=%self.identity.id);
+            let _span_guard = span.enter();
 
             let mut supervisor = self.assemble_and_start_supervisor();
             // Stores the current health state for logging purposes.
@@ -142,7 +142,9 @@ where
                 .publish(SubAgentStarted(self.identity.clone(),SystemTime::now()))
                 .inspect_err(|err| error!(error_msg = %err,"Cannot publish sub_agent_event::sub_agent_started"));
 
-            Option::as_ref(&self.maybe_opamp_client).map(|client| client.update_effective_config());
+            self.maybe_opamp_client
+                .as_ref()
+                .map(|client| client.update_effective_config());
 
             // The below two lines are used to create a channel that never receives any message
             // if the sub_agent_opamp_consumer is None. Thus, we avoid erroring if there is no
@@ -166,11 +168,15 @@ where
                 let _ = uptime_reporter.report();
             }
 
+            drop(_span_guard);
+
             // Count the received remote configs during execution
             let mut remote_config_count = 0;
             loop {
                 select! {
                     recv(opamp_receiver.as_ref()) -> opamp_event_res => {
+                        let span = info_span!("process_fleet_event", id=%self.identity.id);
+                        let _span_guard = span.enter();
                         match opamp_event_res {
                             Err(e) => {
                                 debug!(error = %e, select_arm = "sub_agent_opamp_consumer", "Channel closed");
@@ -220,6 +226,8 @@ where
                         }
                     },
                     recv(&self.sub_agent_internal_consumer.as_ref()) -> sub_agent_internal_event_res => {
+                        let span = info_span!("process_event", id=%self.identity.id);
+                        let _span_guard = span.enter();
                         match sub_agent_internal_event_res {
                             Err(e) => {
                                 debug!(error = %e, select_arm = "sub_agent_internal_consumer", "Channel closed");

agent-control/src/sub_agent/version/version_checker.rs

@@ -1,11 +1,13 @@
+use crate::agent_control::agent_id::AgentID;
 use crate::agent_type::version_config::VersionCheckerInterval;
 use crate::event::SubAgentInternalEvent;
 use crate::event::cancellation::CancellationMessage;
 use crate::event::channel::{EventConsumer, EventPublisher};
+use crate::sub_agent::identity::ID_ATTRIBUTE_NAME;
 use crate::utils::thread_context::{NotStartedThreadContext, StartedThreadContext};
-use tracing::{debug, error, info, warn};
+use tracing::{debug, error, info, info_span, warn};
 
-const VERSION_CHECKER_THREAD_NAME: &str = "version checker";
+const VERSION_CHECKER_THREAD_NAME: &str = "version_checker";
 
 pub trait VersionChecker {
     /// Use it to report the agent version for the opamp client
@@ -43,6 +45,7 @@ pub enum VersionCheckError {
 }
 
 pub(crate) fn spawn_version_checker<V>(
+    agent_id: AgentID,
     version_checker: V,
     sub_agent_internal_publisher: EventPublisher<SubAgentInternalEvent>,
     interval: VersionCheckerInterval,
@@ -53,6 +56,12 @@ where
     // Stores if the version was retrieved in last iteration for logging purposes.
     let mut version_retrieved = false;
     let callback = move |stop_consumer: EventConsumer<CancellationMessage>| loop {
+        let span = info_span!(
+            "version_check",
+            { ID_ATTRIBUTE_NAME } = %agent_id
+        );
+        let _guard = span.enter();
+
         debug!("starting to check version with the configured checker");
 
         match version_checker.check_agent_version() {
@@ -99,6 +108,7 @@ pub(crate) fn publish_version_event(
 
 #[cfg(test)]
 pub mod tests {
+    use crate::agent_control::agent_id::AgentID;
     use crate::agent_control::defaults::OPAMP_CHART_VERSION_ATTRIBUTE_KEY;
     use crate::event::SubAgentInternalEvent::AgentVersionInfo;
     use crate::event::channel::pub_sub;
@@ -142,6 +152,7 @@ pub mod tests {
             });
 
         let started_thread_context = spawn_version_checker(
+            AgentID::default(),
             version_checker,
             version_publisher,
             Duration::from_millis(10).into(),