[#158]: added program handlers function in the common crate. Remove duplicated code in metrics module

Author: LorenzoTettamanti

core/common/Cargo.toml

@@ -19,3 +19,4 @@ aya = "0.13.1"
 
 [features]
 map-handlers = []
+program-handlers = []

core/common/src/lib.rs

@@ -1,4 +1,7 @@
 pub mod constants;
 pub mod formatters;
 pub mod logger;
+#[cfg(feature = "map-handlers")]
 pub mod map_handlers;
+#[cfg(feature = "program-handlers")]
+pub mod program_handlers;

core/common/src/map_handlers.rs

@@ -17,10 +17,6 @@ use tracing::{error, info};
 //
 // this function init the bpfs maps used in the main program
 //
-//  index 0: events_map
-//  index 1: veth_map
-//  index 2: blocklist map
-//  index 3: tcp_registry map
 //
 
 #[cfg(feature = "map-handlers")]

core/common/src/program_handlers.rs

@@ -0,0 +1,42 @@
+use aya::{Ebpf, programs::KProbe};
+use std::convert::TryInto;
+use std::sync::{Arc, Mutex};
+use tracing::{error, info};
+
+#[cfg(feature = "program-handlers")]
+pub fn load_program(
+    bpf: Arc<Mutex<Ebpf>>,
+    program_name: &str,
+    actual_program: &str,
+) -> Result<(), anyhow::Error> {
+    let mut bpf_new = bpf.lock().expect("Cannot get value from lock");
+
+    // Load and attach the eBPF programs
+    let program: &mut KProbe = bpf_new
+        .program_mut(program_name)
+        .ok_or_else(|| anyhow::anyhow!("Program {} not found", program_name))?
+        .try_into()
+        .map_err(|e| anyhow::anyhow!("Failed to convert program: {:?}", e))?;
+
+    program
+        .load()
+        .map_err(|e| anyhow::anyhow!("Cannot load program: {}. Error: {}", &program_name, e))?;
+
+    match program.attach(actual_program, 0) {
+        Ok(_) => info!("{} program attached successfully", actual_program),
+        Err(e) => {
+            error!("Error attaching {} program {:?}", actual_program, e);
+            return Err(anyhow::anyhow!(
+                "Failed to attach {}: {:?}",
+                actual_program,
+                e
+            ));
+        }
+    };
+
+    info!(
+        "eBPF program {} loaded and attached successfully",
+        program_name
+    );
+    Ok(())
+}

core/src/components/metrics/Cargo.toml

@@ -7,11 +7,21 @@ edition = "2024"
 aya = "0.13.1"
 aya-log = "0.2.1"
 bytes = "1.4"
-tokio = { version = "1.48.0", features = ["rt","macros","time","fs","signal","rt-multi-thread"] }
+tokio = { version = "1.48.0", features = [
+    "rt",
+    "macros",
+    "time",
+    "fs",
+    "signal",
+    "rt-multi-thread",
+] }
 anyhow = "1.0"
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 libc = "0.2.172"
 bytemuck = "1.23.0"
-cortexbrain-common = { path = "../../../common" }
-nix  ={version="0.30.1",features=["net"]}
+cortexbrain-common = { path = "../../../common", features = [
+    "map-handlers",
+    "program-handlers",
+] }
+nix = { version = "0.30.1", features = ["net"] }

core/src/components/metrics/src/helpers.rs

@@ -1,24 +1,23 @@
-use aya::{maps::{
-        perf::PerfEventArrayBuffer, Map, MapData, PerfEventArray
-    }, util::online_cpus};
+use aya::{
+    maps::{Map, MapData, PerfEventArray, perf::PerfEventArrayBuffer},
+    util::online_cpus,
+};
 
 use bytes::BytesMut;
-use tokio::signal;
-use std::{
-    sync::{
-        Arc,
-        atomic::{AtomicBool, Ordering},
-    },
+use std::sync::{
+    Arc,
+    atomic::{AtomicBool, Ordering},
 };
+use tokio::signal;
 
-use tracing::{error, info};
+use tracing::{debug, error, info};
 
 use crate::structs::NetworkMetrics;
 use crate::structs::TimeStampMetrics;
 
 pub async fn display_metrics_map(
     mut perf_buffers: Vec<PerfEventArrayBuffer<MapData>>,
-    running: Arc<AtomicBool>,  // Changed to Arc<AtomicBool>
+    running: Arc<AtomicBool>, // Changed to Arc<AtomicBool>
     mut buffers: Vec<BytesMut>,
 ) {
     info!("Starting metrics event listener...");
@@ -46,10 +45,23 @@ pub async fn display_metrics_map(
                             let sk_receive_buffer_size = net_metrics.sk_receive_buffer_size;
                             info!(
                                 "tgid: {}, comm: {}, ts_us: {}, sk_drops: {}, sk_err: {}, sk_err_soft: {}, sk_backlog_len: {}, sk_write_memory_queued: {}, sk_ack_backlog: {}, sk_receive_buffer_size: {}",
-                                tgid, comm, ts_us, sk_drop_count, sk_err, sk_err_soft, sk_backlog_len, sk_write_memory_queued, sk_ack_backlog, sk_receive_buffer_size
+                                tgid,
+                                comm,
+                                ts_us,
+                                sk_drop_count,
+                                sk_err,
+                                sk_err_soft,
+                                sk_backlog_len,
+                                sk_write_memory_queued,
+                                sk_ack_backlog,
+                                sk_receive_buffer_size
                             );
                         } else {
-                            info!("Received data too small: {} bytes, expected: {}", data.len(), std::mem::size_of::<NetworkMetrics>());
+                            info!(
+                                "Received data too small: {} bytes, expected: {}",
+                                data.len(),
+                                std::mem::size_of::<NetworkMetrics>()
+                            );
                         }
                     }
                 }
@@ -65,7 +77,7 @@ pub async fn display_metrics_map(
 
 pub async fn display_time_stamp_events_map(
     mut perf_buffers: Vec<PerfEventArrayBuffer<MapData>>,
-    running: Arc<AtomicBool>,  // Changed to Arc<AtomicBool>
+    running: Arc<AtomicBool>, // Changed to Arc<AtomicBool>
     mut buffers: Vec<BytesMut>,
 ) {
     info!("Starting timestamp event listener...");
@@ -107,48 +119,67 @@ pub async fn display_time_stamp_events_map(
     info!("Timestamp event listener stopped");
 }
 
-pub async fn event_listener(bpf_maps: (Map, Map)) -> Result<(), anyhow::Error> {
+pub async fn event_listener(bpf_maps: Vec<Map>) -> Result<(), anyhow::Error> {
     info!("Getting CPU count...");
-    let cpu_count = online_cpus().map_err(|e| anyhow::anyhow!("Error {:?}", e))?.len();
-    info!("CPU count: {}", cpu_count);
-    
+
+    let mut perf_event_arrays = Vec::new(); // contains a vector of PerfEventArrays
+    let mut event_buffers = Vec::new(); // contains a vector of buffers
+
     info!("Creating perf buffers...");
-    let mut net_perf_buffer: Vec<PerfEventArrayBuffer<MapData>> = Vec::new();
-    let mut net_perf_array: PerfEventArray<MapData> = PerfEventArray::try_from(bpf_maps.0)?;
-    let mut time_stamp_events_perf_buffer: Vec<PerfEventArrayBuffer<MapData>> = Vec::new();
-    let mut time_stamp_events_perf_array: PerfEventArray<MapData> =
-        PerfEventArray::try_from(bpf_maps.1)?;
-
-    info!("Opening perf buffers for {} CPUs...", cpu_count);
-    for cpu_id in online_cpus().map_err(|e| anyhow::anyhow!("Error {:?}", e))? {
-        let buf: PerfEventArrayBuffer<MapData> = net_perf_array.open(cpu_id, None)?;
-        net_perf_buffer.push(buf);
+    for map in bpf_maps {
+        debug!("Debugging map type:{:?}", map);
+        let perf_event_array = PerfEventArray::try_from(map).map_err(|e| {
+            error!("Cannot create perf_event_array for map.Reason: {}", e);
+            anyhow::anyhow!("Cannot create perf_event_array for map.Reason: {}", e)
+        })?;
+        perf_event_arrays.push(perf_event_array); // this is step 1
+        let perf_event_array_buffer = Vec::new();
+        event_buffers.push(perf_event_array_buffer); //this is step 2 
     }
-    for cpu_id in online_cpus().map_err(|e| anyhow::anyhow!("Error {:?}", e))? {
-        let buf: PerfEventArrayBuffer<MapData> = time_stamp_events_perf_array.open(cpu_id, None)?;
-        time_stamp_events_perf_buffer.push(buf);
+
+    let cpu_count = online_cpus().map_err(|e| anyhow::anyhow!("Error {:?}", e))?;
+
+    //info!("CPU count: {}", cpu_count);
+    for (perf_evt_array, perf_evt_array_buffer) in
+        perf_event_arrays.iter_mut().zip(event_buffers.iter_mut())
+    {
+        for cpu_id in &cpu_count {
+            let single_buffer = perf_evt_array.open(*cpu_id, None)?;
+            perf_evt_array_buffer.push(single_buffer);
+        }
     }
+
+    //info!("Opening perf buffers for {} CPUs...", cpu_count);
     info!("Perf buffers created successfully");
+    let mut event_buffers = event_buffers.into_iter();
+
+    let time_stamp_events_perf_buffer = event_buffers.next().expect("");
+    let net_perf_buffer = event_buffers.next().expect("");
 
     // Create shared running flags
     let net_metrics_running = Arc::new(AtomicBool::new(true));
     let time_stamp_events_running = Arc::new(AtomicBool::new(true));
-    
+
     // Create proper sized buffers
-    let net_metrics_buffers = vec![BytesMut::with_capacity(1024); cpu_count];
-    let time_stamp_events_buffers = vec![BytesMut::with_capacity(1024); cpu_count];
-    
+    let net_metrics_buffers = vec![BytesMut::with_capacity(1024); cpu_count.len()];
+    let time_stamp_events_buffers = vec![BytesMut::with_capacity(1024); cpu_count.len()];
+
     // Clone for the signal handler
     let net_metrics_running_signal = net_metrics_running.clone();
     let time_stamp_events_running_signal = time_stamp_events_running.clone();
-    
+
     info!("Starting event listener tasks...");
     let metrics_map_displayer = tokio::spawn(async move {
         display_metrics_map(net_perf_buffer, net_metrics_running, net_metrics_buffers).await;
     });
 
     let time_stamp_events_displayer = tokio::spawn(async move {
-        display_time_stamp_events_map(time_stamp_events_perf_buffer, time_stamp_events_running, time_stamp_events_buffers).await
+        display_time_stamp_events_map(
+            time_stamp_events_perf_buffer,
+            time_stamp_events_running,
+            time_stamp_events_buffers,
+        )
+        .await
     });
 
     info!("Event listeners started, entering main loop...");
@@ -176,4 +207,4 @@ pub async fn event_listener(bpf_maps: (Map, Map)) -> Result<(), anyhow::Error> {
 
     // return success
     Ok(())
-}
+}

core/src/components/metrics/src/main.rs

@@ -1,27 +1,18 @@
-use aya::{
-    Ebpf
-};
-
+use anyhow::{Context, Ok};
+use aya::Ebpf;
+use cortexbrain_common::{constants, logger};
 use std::{
     env, fs,
     path::Path,
-    sync::{
-        Arc, Mutex,
-    },
+    sync::{Arc, Mutex},
 };
-
-use anyhow::{Context, Ok};
 use tracing::{error, info};
-use cortexbrain_common::{constants, logger};
 
 mod helpers;
-use crate::{helpers::event_listener, maps_handlers::map_pinner, program_handlers::load_and_attach_tcp_programs};
-
-mod maps_handlers;
-use crate::maps_handlers::init_ebpf_maps;
+use crate::helpers::event_listener;
 
-mod program_handlers;
-use crate::program_handlers::load_program;
+use cortexbrain_common::map_handlers::{init_bpf_maps, map_pinner};
+use cortexbrain_common::program_handlers::load_program;
 
 mod structs;
 
@@ -33,41 +24,50 @@ async fn main() -> Result<(), anyhow::Error> {
     info!("Starting metrics service...");
     info!("fetching data");
 
-    let bpf_path = env::var(constants::BPF_PATH).context("BPF_PATH environment variable required")?;
+    let bpf_path =
+        env::var(constants::BPF_PATH).context("BPF_PATH environment variable required")?;
     let data = fs::read(Path::new(&bpf_path)).context("Failed to load file from path")?;
     let bpf = Arc::new(Mutex::new(Ebpf::load(&data)?));
     let tcp_bpf = bpf.clone();
     let tcp_rev_bpf = bpf.clone();
+    let tcp_v6_bpf = bpf.clone();
 
     info!("Running Ebpf logger");
     info!("loading programs");
-    let bpf_map_save_path =
-        std::env::var(constants::PIN_MAP_PATH).context("PIN_MAP_PATH environment variable required")?;
+    let bpf_map_save_path = std::env::var(constants::PIN_MAP_PATH)
+        .context("PIN_MAP_PATH environment variable required")?;
 
-    match init_ebpf_maps(bpf.clone()) {
-        std::result::Result::Ok(maps) => {
+    let map_data = vec!["time_stamp_events".to_string(), "net_metrics".to_string()];
+
+    match init_bpf_maps(bpf.clone(), map_data) {
+        std::result::Result::Ok(bpf_maps) => {
             info!("BPF maps loaded successfully");
             let pin_path = std::path::PathBuf::from(&bpf_map_save_path);
             info!("About to call map_pinner with path: {:?}", pin_path);
-            match map_pinner(&maps, &pin_path).await {
-                std::result::Result::Ok(_) => {
+            match map_pinner(bpf_maps, &pin_path) {
+                std::result::Result::Ok(maps) => {
                     info!("BPF maps pinned successfully to {}", bpf_map_save_path);
 
                     {
                         load_program(bpf.clone(), "metrics_tracer", "tcp_identify_packet_loss")
-                            .context("An error occured during the execution of load_program function")?;
-                    }
-
-                    {
-                        load_and_attach_tcp_programs(tcp_bpf.clone())
-                            .context("An error occured during the execution of load_and_attach_tcp_programs function")?;
+                            .context(
+                                "An error occured during the execution of load_program function",
+                            )?;
+
+                        load_program(tcp_bpf,"tcp_connect","tcp_v4_connect")
+                        .context("An error occured during the execution of load_and_attach_tcp_programs function")?;
+                        load_program(tcp_v6_bpf,"tcp_connect","tcp_v6_connect")
+                        .context("An error occured during the execution of load_and_attach_tcp_programs function")?;
+
+                        load_program(
+                            tcp_rev_bpf,
+                            "tcp_rcv_state_process",
+                            "tcp_rcv_state_process",
+                        )
+                        .context(
+                            "An error occured during the execution of load_program function",
+                        )?;
                     }
-
-                    {
-                        load_program(tcp_rev_bpf.clone(), "tcp_rcv_state_process", "tcp_rcv_state_process")
-                            .context("An error occured during the execution of load_program function")?;
-                    }
-
                     event_listener(maps).await?;
                 }
                 Err(e) => {
@@ -83,4 +83,4 @@ async fn main() -> Result<(), anyhow::Error> {
     }
 
     Ok(())
-}
+}