Refector redirector to able close eBPF cleanly.

proxy_agent/src/common/error.rs

@@ -134,6 +134,9 @@ pub enum BpfErrorType {
     #[error("Failed to attach program '{0}' with error: {1}")]
     AttachBpfProgram(String, String),
 
+    #[error("Failed to close BPF object: {0}")]
+    CloseBpfObject(String),
+
     #[error("Failed to convert program to '{0}' with error: {1}")]
     ConvertBpfProgram(String, String),
 

proxy_agent/src/redirector.rs

@@ -36,7 +36,7 @@
 //! let status = agent_status_shared_state.get_status(AgentStatusModule::Redirector).await;
 //!
 //! // Close the redirector to offload the eBPF program
-//! redirector::close(redirector_shared_state.clone(), agent_status_shared_state.clone()).await;
+//! redirector::close(redirector_shared_state.clone(), agent_status_shared_state.clone()).await.unwrap();
 //! ```
 
 #[cfg(windows)]
@@ -68,10 +68,72 @@ use std::net::Ipv4Addr;
 use std::path::PathBuf;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
+use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
 
 const ALERT_ONLY_EVENT_LOG_INTERVAL: Duration = Duration::from_secs(60 * 15); // 15 minutes
 
+/// Runtime structure for managing event handling tasks associated with a BPF object.
+#[derive(Default)]
+pub(crate) struct EventTaskRuntime {
+    event_cancellation_token: Option<CancellationToken>,
+    /// Event handling tasks associated with this BPF object.
+    /// 1. The event-processing task created during start_internal.
+    /// 2. The platform ring-buffer reader task created during subscribe_alert_only.
+    event_tasks: Vec<JoinHandle<()>>,
+}
+
+impl EventTaskRuntime {
+    fn set_cancellation_token(&mut self, cancellation_token: CancellationToken) {
+        self.event_cancellation_token = Some(cancellation_token);
+    }
+
+    fn add_task(&mut self, task: JoinHandle<()>) {
+        self.event_tasks.push(task);
+    }
+
+    /// Shuts down all event handling tasks associated with this BPF object.
+    /// Cancels the event cancellation token and awaits the completion of all tasks.
+    /// Returns the first error encountered, if any.
+    pub(crate) async fn shutdown(mut self) -> Result<()> {
+        if let Some(cancellation_token) = self.event_cancellation_token.take() {
+            cancellation_token.cancel();
+        }
+
+        let mut first_error = None;
+        for task in self.event_tasks.drain(..) {
+            if let Err(e) = task.await {
+                let error = Error::Bpf(BpfErrorType::CloseBpfObject(format!(
+                    "eBPF event task failed while stopping: {e}"
+                )));
+                logger::write_error(error.to_string());
+                if first_error.is_none() {
+                    first_error = Some(error);
+                }
+            }
+        }
+
+        if let Some(error) = first_error {
+            return Err(error);
+        }
+
+        Ok(())
+    }
+}
+
+/// Implements the `Drop` trait for `EventTaskRuntime`.
+/// Ensures that all event handling tasks are aborted and the cancellation token is triggered when the runtime is dropped.
+impl Drop for EventTaskRuntime {
+    fn drop(&mut self) {
+        if let Some(cancellation_token) = self.event_cancellation_token.take() {
+            cancellation_token.cancel();
+        }
+        for task in self.event_tasks.drain(..) {
+            task.abort();
+        }
+    }
+}
+
 #[cfg(not(windows))]
 pub use linux::BpfObject;
 #[cfg(windows)]
@@ -190,20 +252,30 @@ impl AlertOnlyEventRateLimiter {
 pub struct Redirector {
     local_port: u16,
     shared_state: SharedState,
+    cancellation_token: CancellationToken,
 }
 
 impl Redirector {
-    pub fn new(local_port: u16, shared_state: &SharedState) -> Self {
+    pub fn new(
+        local_port: u16,
+        shared_state: &SharedState,
+        cancellation_token: CancellationToken,
+    ) -> Self {
         Redirector {
             local_port,
             shared_state: shared_state.clone(),
+            cancellation_token,
         }
     }
 
     const MAX_RETRIES: usize = 5;
     const RETRY_INTERVAL_MS: u64 = 10;
 
     pub async fn start(&self) {
+        if self.cancellation_token.is_cancelled() {
+            return;
+        }
+
         let message = "eBPF redirector is starting";
         if let Err(e) = self
             .shared_state
@@ -217,8 +289,30 @@ impl Redirector {
         }
 
         let level = match self.start_impl().await {
-            Ok(_) => LoggerLevel::Info,
-            Err(_) => LoggerLevel::Error,
+            Ok(_) if !self.cancellation_token.is_cancelled() => LoggerLevel::Info,
+            Ok(_) => {
+                if let Err(e) = close(
+                    self.shared_state.get_redirector_shared_state(),
+                    self.shared_state.get_agent_status_shared_state(),
+                )
+                .await
+                {
+                    logger::write_error(format!(
+                        "Failed to close redirector after startup cancellation: {e}"
+                    ));
+                    LoggerLevel::Error
+                } else {
+                    LoggerLevel::Info
+                }
+            }
+            Err(_) => {
+                let _ = self
+                    .shared_state
+                    .get_agent_status_shared_state()
+                    .set_module_state(ModuleState::STOPPED, AgentStatusModule::Redirector)
+                    .await;
+                LoggerLevel::Error
+            }
         };
         event_logger::write_event(
             level,
@@ -231,6 +325,9 @@ impl Redirector {
 
     async fn start_impl(&self) -> Result<()> {
         for _ in 0..Self::MAX_RETRIES {
+            if self.cancellation_token.is_cancelled() {
+                return Err(Error::Bpf(BpfErrorType::FailedToStartRedirector));
+            }
             match self.start_internal().await {
                 Ok(_) => return Ok(()),
                 Err(e) => {
@@ -268,14 +365,21 @@ impl Redirector {
         self.attach_bpf_prog(&mut bpf_object)?;
         logger::write_information("Success attached bpf prog.".to_string());
 
-        match bpf_object.subscribe_alert_only(self.shared_state.get_cancellation_token().clone()) {
+        let event_cancellation_token = self.cancellation_token.child_token();
+        bpf_object
+            .event_runtime_mut()
+            .set_cancellation_token(event_cancellation_token.clone());
+        match bpf_object.subscribe_alert_only(event_cancellation_token.clone()) {
             Ok(receiver) => {
                 // Handle the received alert-only events here, spawn a task to process them
-                tokio::spawn(process_alert_only_events(
+                let event_processor_task = tokio::spawn(process_alert_only_events(
                     receiver,
                     self.shared_state.get_proxy_server_shared_state(),
-                    self.shared_state.get_cancellation_token(),
+                    event_cancellation_token,
                 ));
+                bpf_object
+                    .event_runtime_mut()
+                    .add_task(event_processor_task);
             }
             Err(e) => {
                 logger::write_error(format!("Failed to subscribe to alert-only events: {e}"));
@@ -539,17 +643,21 @@ pub async fn update_local_ip_bind_monitor_only(
 pub async fn close(
     redirector_shared_state: RedirectorSharedState,
     agent_status_shared_state: AgentStatusSharedState,
-) {
-    let _ = agent_status_shared_state
-        .set_module_state(ModuleState::STOPPED, AgentStatusModule::Redirector)
-        .await;
-
-    // reset ebpf object
+) -> Result<()> {
     #[cfg(windows)]
     {
-        windows::close_bpf_object(redirector_shared_state.clone()).await;
+        windows::close_bpf_object(redirector_shared_state).await?;
     }
-    let _ = redirector_shared_state.clear_bpf_object().await;
+
+    #[cfg(not(windows))]
+    {
+        linux::close_bpf_object(redirector_shared_state).await?;
+    }
+
+    agent_status_shared_state
+        .set_module_state(ModuleState::STOPPED, AgentStatusModule::Redirector)
+        .await?;
+    Ok(())
 }
 
 #[cfg(not(windows))]