fix: Add stop signal to status server (#1126)

gsanchezgavier · web-flow · commit c635b1dd8a57 · 2025-03-17T16:10:11.000+01:00
diff --git a/agent-control/src/agent_control/http_server/async_bridge.rs b/agent-control/src/agent_control/http_server/async_bridge.rs
@@ -13,6 +13,7 @@ pub fn run_async_sync_bridge(
     async_suba_publisher: UnboundedSender<SubAgentEvent>,
     agent_control_consumer: EventConsumer<AgentControlEvent>,
     sub_agent_consumer: EventConsumer<SubAgentEvent>,
+    stop_rx: EventConsumer<()>,
 ) -> JoinHandle<()> {
     spawn_named_thread("Async-Sync bridge", move || loop {
         select! {
@@ -36,24 +37,28 @@ pub fn run_async_sync_bridge(
                 }
             },
             recv(&sub_agent_consumer.as_ref()) -> suba_event_res => {
-                    match suba_event_res {
-                        Ok(sub_agent_event) => {
-                            let _ = async_suba_publisher.send(sub_agent_event).inspect_err(|err| {
-                                error!(
-                                    error_msg = %err,
-                                    "cannot forward agent control event"
-                                );
-                            });
-                        }
-                        Err(err) => {
-                            debug!(
+                match suba_event_res {
+                    Ok(sub_agent_event) => {
+                        let _ = async_suba_publisher.send(sub_agent_event).inspect_err(|err| {
+                            error!(
                                 error_msg = %err,
-                                "status server bridge channel closed"
+                                "cannot forward agent control event"
                             );
-                            break;
-                        }
+                        });
+                    }
+                    Err(err) => {
+                        debug!(
+                            error_msg = %err,
+                            "status server bridge channel closed"
+                        );
+                        break;
                     }
                 }
+            },
+            recv(&stop_rx.as_ref()) -> _ => {
+                debug!("status server bridge stopping");
+                break;
+            }
         }
     })
 }
diff --git a/agent-control/src/agent_control/http_server/config.rs b/agent-control/src/agent_control/http_server/config.rs
@@ -63,6 +63,12 @@ mod tests {
         server_config: ServerConfig,
     }
 
+    impl From<u16> for Port {
+        fn from(value: u16) -> Self {
+            Port(value)
+        }
+    }
+
     #[test]
     fn test_deserialize_default() {
         struct Test {
diff --git a/agent-control/src/agent_control/http_server/runner.rs b/agent-control/src/agent_control/http_server/runner.rs
@@ -1,20 +1,20 @@
 use crate::agent_control::config::OpAMPClientConfig;
 use crate::agent_control::http_server::async_bridge::run_async_sync_bridge;
 use crate::agent_control::http_server::config::ServerConfig;
+use crate::event::cancellation::CancellationMessage;
 use crate::event::channel::EventConsumer;
 use crate::event::{AgentControlEvent, SubAgentEvent};
-use crate::utils::threads::spawn_named_thread;
+use crate::utils::thread_context::{NotStartedThreadContext, StartedThreadContext};
 use crossbeam::select;
 use std::sync::Arc;
-use std::thread::JoinHandle;
 use tokio::runtime::Runtime;
 use tokio::sync::mpsc;
 use tracing::{debug, error, info};
 
 /// Runner will be responsible for spawning the OS Thread for the HTTP Server
 /// and owning the JoinHandle. It controls the server stop implementing drop
 pub struct Runner {
-    join_handle: Option<JoinHandle<()>>,
+    thread_context: Option<StartedThreadContext>,
 }
 
 impl Runner {
@@ -30,21 +30,29 @@ impl Runner {
         sub_agent_consumer: EventConsumer<SubAgentEvent>,
         maybe_opamp_client_config: Option<OpAMPClientConfig>,
     ) -> Self {
-        let join_handle = if config.enabled {
-            Self::spawn_server(
-                config,
-                runtime,
-                agent_control_consumer,
-                sub_agent_consumer,
-                maybe_opamp_client_config,
-            )
+        let thread_context = if config.enabled {
+            let callback = move |stop_consumer: EventConsumer<CancellationMessage>| {
+                Self::spawn_server(
+                    config,
+                    runtime,
+                    agent_control_consumer,
+                    sub_agent_consumer,
+                    maybe_opamp_client_config,
+                    stop_consumer,
+                )
+            };
+            NotStartedThreadContext::new("Http server", callback).start()
         } else {
+            let callback = move |stop_consumer: EventConsumer<CancellationMessage>| {
+                Self::noop_consumer_loop(agent_control_consumer, sub_agent_consumer, stop_consumer)
+            };
             // Spawn a thread with a no-action consumer to drain the channel and
             // avoid memory leaks
-            Self::spawn_noop_consumer(agent_control_consumer, sub_agent_consumer)
+            NotStartedThreadContext::new("No-action consumer", callback).start()
         };
+
         Runner {
-            join_handle: Some(join_handle),
+            thread_context: Some(thread_context),
         }
     }
 
@@ -54,47 +62,49 @@ impl Runner {
         agent_control_consumer: EventConsumer<AgentControlEvent>,
         sub_agent_consumer: EventConsumer<SubAgentEvent>,
         maybe_opamp_client_config: Option<OpAMPClientConfig>,
-    ) -> JoinHandle<()> {
-        spawn_named_thread("Http server", move || {
-            // Create 2 unbounded channel to send the Agent Control and Sub Agent Sync events
-            // to the Async Status Server
-            let (async_agent_control_event_publisher, async_agent_control_event_consumer) =
-                mpsc::unbounded_channel::<AgentControlEvent>();
-            let (async_sub_agent_event_publisher, async_sub_agent_event_consumer) =
-                mpsc::unbounded_channel::<SubAgentEvent>();
-            // Run an OS Thread that listens to sync channel and forwards the events
-            // to an async channel
-            let bridge_join_handle = run_async_sync_bridge(
-                async_agent_control_event_publisher,
-                async_sub_agent_event_publisher,
-                agent_control_consumer,
-                sub_agent_consumer,
-            );
-
-            // Run the async status server
-            let _ = runtime
-                .block_on(
-                    crate::agent_control::http_server::server::run_status_server(
-                        config.clone(),
-                        async_agent_control_event_consumer,
-                        async_sub_agent_event_consumer,
-                        maybe_opamp_client_config,
-                    ),
-                )
-                .inspect_err(|err| {
-                    error!(error_msg = %err, "error running status server");
-                });
+        stop_rx: EventConsumer<CancellationMessage>,
+    ) {
+        // Create 2 unbounded channel to send the Agent Control and Sub Agent Sync events
+        // to the Async Status Server
+        let (async_agent_control_event_publisher, async_agent_control_event_consumer) =
+            mpsc::unbounded_channel::<AgentControlEvent>();
+        let (async_sub_agent_event_publisher, async_sub_agent_event_consumer) =
+            mpsc::unbounded_channel::<SubAgentEvent>();
 
-            // Wait until the bridge is closed
-            bridge_join_handle.join().unwrap();
-        })
+        // Run an OS Thread that listens to sync channel and forwards the events
+        // to an async channel
+        let bridge_join_handle = run_async_sync_bridge(
+            async_agent_control_event_publisher,
+            async_sub_agent_event_publisher,
+            agent_control_consumer,
+            sub_agent_consumer,
+            stop_rx,
+        );
+
+        // Run the async status server
+        let _ = runtime
+            .block_on(
+                crate::agent_control::http_server::server::run_status_server(
+                    config.clone(),
+                    async_agent_control_event_consumer,
+                    async_sub_agent_event_consumer,
+                    maybe_opamp_client_config,
+                ),
+            )
+            .inspect_err(|err| {
+                error!(error_msg = %err, "error running status server");
+            });
+
+        // Wait until the bridge is closed
+        bridge_join_handle.join().unwrap();
     }
 
-    fn spawn_noop_consumer(
+    fn noop_consumer_loop(
         agent_control_consumer: EventConsumer<AgentControlEvent>,
         sub_agent_consumer: EventConsumer<SubAgentEvent>,
-    ) -> JoinHandle<()> {
-        spawn_named_thread("No-action consumer", move || loop {
+        stop_rx: EventConsumer<CancellationMessage>,
+    ) {
+        loop {
             select! {
                 recv(agent_control_consumer.as_ref()) -> agent_control_consumer_res => {
                     match agent_control_consumer_res {
@@ -119,19 +129,145 @@ impl Runner {
                             break;
                         }
                     }
-                }
+                },
+                recv(stop_rx.as_ref()) -> _ => {
+                    debug!("http server event drain processor stopped");
+                    break;
+                },
             }
-        })
+        }
     }
 }
 
 impl Drop for Runner {
     fn drop(&mut self) {
-        if let Some(join_handle) = self.join_handle.take() {
-            info!("waiting for status server to stop gracefully...");
-            join_handle
-                .join()
-                .expect("error waiting for server join handle")
-        }
+        info!("waiting for status server to stop gracefully...");
+
+        let Some(thread_context) = self.thread_context.take() else {
+            return;
+        };
+
+        let _ = thread_context
+            .stop()
+            .inspect(|_| debug!("status server runner thread stopped"))
+            .inspect_err(|error_msg| {
+                error!(
+                    %error_msg,
+                    "Error stopping Status Server"
+                )
+            });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+    use std::thread::sleep;
+    use std::time::Duration;
+
+    use tracing_test::internal::logs_with_scope_contain;
+    use tracing_test::traced_test;
+
+    use crate::agent_control::http_server::config::ServerConfig;
+    use crate::event::channel::pub_sub;
+    use crate::event::AgentControlEvent;
+
+    use super::Runner;
+
+    #[test]
+    #[traced_test]
+    fn test_noop_consumer_stops_gracefully_when_dropped() {
+        let runtime = Arc::new(
+            tokio::runtime::Builder::new_multi_thread()
+                .enable_all()
+                .build()
+                .unwrap(),
+        );
+        let (_agent_control_publisher, agent_control_consumer) = pub_sub::<AgentControlEvent>();
+        let (_sub_agent_publisher, sub_agent_consumer) = pub_sub();
+        let _http_server_runner = Runner::start(
+            ServerConfig::default(),
+            runtime,
+            agent_control_consumer,
+            sub_agent_consumer,
+            None,
+        );
+        drop(_http_server_runner);
+
+        // wait for logs to be flushed
+        sleep(Duration::from_millis(100));
+        assert!(logs_with_scope_contain(
+            "newrelic_agent_control::agent_control::http_server::runner",
+            "http server event drain processor stopped",
+        ));
+    }
+    #[test]
+    #[traced_test]
+    fn test_server_stops_gracefully_when_dropped() {
+        let runtime = Arc::new(
+            tokio::runtime::Builder::new_multi_thread()
+                .enable_all()
+                .build()
+                .unwrap(),
+        );
+        let (_agent_control_publisher, agent_control_consumer) = pub_sub::<AgentControlEvent>();
+        let (_sub_agent_publisher, sub_agent_consumer) = pub_sub();
+        let _http_server_runner = Runner::start(
+            ServerConfig {
+                enabled: true,
+                port: 0.into(),
+                ..Default::default()
+            },
+            runtime,
+            agent_control_consumer,
+            sub_agent_consumer,
+            None,
+        );
+        // server warm up
+        sleep(Duration::from_millis(100));
+
+        drop(_http_server_runner);
+
+        // wait for logs to be flushed
+        sleep(Duration::from_millis(100));
+        assert!(logs_with_scope_contain(
+            "newrelic_agent_control::agent_control::http_server::server",
+            "status server gracefully stopped",
+        ));
+    }
+    #[test]
+    #[traced_test]
+    fn test_server_stops_gracefully_when_external_channels_close() {
+        let runtime = Arc::new(
+            tokio::runtime::Builder::new_multi_thread()
+                .enable_all()
+                .build()
+                .unwrap(),
+        );
+        let (_agent_control_publisher, agent_control_consumer) = pub_sub::<AgentControlEvent>();
+        let (_sub_agent_publisher, sub_agent_consumer) = pub_sub();
+        let _http_server_runner = Runner::start(
+            ServerConfig {
+                enabled: true,
+                port: 0.into(),
+                ..Default::default()
+            },
+            runtime,
+            agent_control_consumer,
+            sub_agent_consumer,
+            None,
+        );
+        // server warm up
+        sleep(Duration::from_millis(100));
+
+        drop(_agent_control_publisher);
+        drop(_sub_agent_publisher);
+
+        // wait for logs to be flushed
+        sleep(Duration::from_millis(100));
+        assert!(logs_with_scope_contain(
+            "newrelic_agent_control::agent_control::http_server::server",
+            "status server gracefully stopped",
+        ));
     }
 }
diff --git a/agent-control/src/agent_control/http_server/server.rs b/agent-control/src/agent_control/http_server/server.rs
@@ -70,6 +70,8 @@ pub async fn run_status_server(
     debug!("waiting for status server join handle");
     server_join_handle.await?;
 
+    debug!("status server gracefully stopped");
+
     Ok(())
 }
 
diff --git a/agent-control/src/agent_control/http_server/status_updater.rs b/agent-control/src/agent_control/http_server/status_updater.rs

Original file line number	Diff line number	Diff line change
`@@ -70,6 +70,8 @@ pub async fn run_status_server(`
`70`	`70`	`debug!("waiting for status server join handle");`
`71`	`71`	`server_join_handle.await?;`
`72`	`72`
	`73`	`+ debug!("status server gracefully stopped");`
	`74`	`+`
`73`	`75`	`Ok(())`
`74`	`76`	`}`
`75`	`77`