Update Heimdall to also use eBPF iterators

This eliminates another set of syscalls per cycle. The old per_cpu
map system is now removed, in favor of direct iteration of maps.

src/rust/lqos_heimdall/src/flows.rs

@@ -1,43 +1,11 @@
 use crate::{timeline::expire_timeline, FLOW_EXPIRE_SECS};
 use dashmap::DashMap;
 use lqos_bus::{tos_parser, BusResponse, FlowTransport};
-use lqos_sys::bpf_per_cpu_map::BpfPerCpuMap;
+use lqos_sys::heimdall_data::{HeimdallKey, HeimdallData};
 use lqos_utils::{unix_time::time_since_boot, XdpIpAddress};
 use once_cell::sync::Lazy;
 use std::{collections::HashSet, time::Duration};
 
-/// Representation of the eBPF `heimdall_key` type.
-#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
-#[repr(C)]
-pub struct HeimdallKey {
-  /// Mapped `XdpIpAddress` source for the flow.
-  pub src_ip: XdpIpAddress,
-  /// Mapped `XdpIpAddress` destination for the flow
-  pub dst_ip: XdpIpAddress,
-  /// IP protocol (see the Linux kernel!)
-  pub ip_protocol: u8,
-  /// Source port number, or ICMP type.
-  pub src_port: u16,
-  /// Destination port number.
-  pub dst_port: u16,
-}
-
-/// Mapped representation of the eBPF `heimdall_data` type.
-#[derive(Debug, Clone, Default)]
-#[repr(C)]
-pub struct HeimdallData {
-  /// Last seen, in nanoseconds (since boot time).
-  pub last_seen: u64,
-  /// Number of bytes since the flow started being tracked
-  pub bytes: u64,
-  /// Number of packets since the flow started being tracked
-  pub packets: u64,
-  /// IP header TOS value
-  pub tos: u8,
-  /// Reserved to pad the structure
-  pub reserved: [u8; 3],
-}
-
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 struct FlowKey {
   src: XdpIpAddress,
@@ -95,11 +63,12 @@ static FLOW_DATA: Lazy<DashMap<FlowKey, FlowData>> = Lazy::new(DashMap::new);
 fn heimdall_for_each(
   callback: &mut dyn FnMut(&HeimdallKey, &[HeimdallData]),
 ) {
-  if let Ok(heimdall) = BpfPerCpuMap::<HeimdallKey, HeimdallData>::from_path(
+  /*if let Ok(heimdall) = BpfPerCpuMap::<HeimdallKey, HeimdallData>::from_path(
     "/sys/fs/bpf/heimdall",
   ) {
     heimdall.for_each(callback);
-  }
+  }*/
+  lqos_sys::iterate_heimdall(callback);
 }
 
 

src/rust/lqos_sys/src/bpf/common/heimdall.h

@@ -67,6 +67,7 @@ struct heimdall_key
     __u8 ip_protocol;
     __u16 src_port;
     __u16 dst_port;
+    __u8 pad;
 };
 
 struct heimdall_data {
@@ -155,6 +156,7 @@ static __always_inline void update_heimdall(struct dissector_t *dissector, __u32
             {
                 bpf_debug("Failed to insert tracking");
             }
+            //bpf_debug("Inserted tracking");
         }
     } else if (mode == 2) {
         struct heimdall_event event = {0};

src/rust/lqos_sys/src/bpf/lqos_kern.c

@@ -386,4 +386,32 @@ int rtt_reader(struct bpf_iter__bpf_map_elem *ctx)
     return 0;
 }
 
+SEC("iter/bpf_map_elem")
+int heimdall_reader(struct bpf_iter__bpf_map_elem *ctx) {
+    // The sequence file
+    struct seq_file *seq = ctx->meta->seq;
+    void *counter = ctx->value;
+    struct heimdall_key *ip = ctx->key;
+    __u32 num_cpus = NUM_CPUS;
+
+    if (ctx->meta->seq_num == 0) {
+        bpf_seq_write(seq, &num_cpus, sizeof(__u32));
+        bpf_seq_write(seq, &num_cpus, sizeof(__u32)); // Repeat for padding
+    }
+
+    // Bail on end
+    if (counter == NULL || ip == NULL) {
+        return 0;
+    }
+
+    bpf_seq_write(seq, ip, sizeof(struct heimdall_key));
+    for (__u32 i=0; i<NUM_CPUS; i++) {
+        struct heimdall_data * content = counter+(i*sizeof(struct heimdall_data));
+        bpf_seq_write(seq, content, sizeof(struct heimdall_data));
+    }
+
+    //BPF_SEQ_PRINTF(seq, "%d %d\n", counter->download_bytes, counter->upload_bytes);
+    return 0;
+}
+
 char _license[] SEC("license") = "GPL";

src/rust/lqos_sys/src/bpf_iterator.rs

@@ -1,6 +1,6 @@
 use crate::{
   kernel_wrapper::BPF_SKELETON, lqos_kernel::bpf, HostCounter,
-  RttTrackingEntry,
+  RttTrackingEntry, heimdall_data::{HeimdallKey, HeimdallData},
 };
 use lqos_utils::XdpIpAddress;
 use once_cell::sync::Lazy;
@@ -90,7 +90,11 @@ where
         log::error!("{e:?}");
         Err(BpfIteratorError::UnableToCreateIterator)
       }
-      Ok(_) => {
+      Ok(bytes) => {
+        if bytes == 0 {
+          // Not having any data is not an error
+          return Ok(());
+        }
         let first_four_bytes: [u8; 4] = [buf[0], buf[1], buf[2], buf[3]];
         let num_cpus = u32::from_ne_bytes(first_four_bytes) as usize;
         let mut index = 8;
@@ -98,11 +102,13 @@ where
           let key_start = index;
           let key_end = key_start + Self::KEY_SIZE;
           let key_slice = &buf[key_start..key_end];
+          //println!("{:?}", unsafe { &key_slice.align_to::<KEY>() });
           let (_head, key, _tail) = unsafe { &key_slice.align_to::<KEY>() };
 
           let value_start = key_end;
           let value_end = value_start + (num_cpus * Self::VALUE_SIZE);
           let value_slice = &buf[value_start..value_end];
+          //println!("{:?}", unsafe { &value_slice.align_to::<VALUE>() });
           let (_head, values, _tail) =
             unsafe { &value_slice.align_to::<VALUE>() };
           debug_assert_eq!(values.len(), num_cpus);
@@ -181,6 +187,10 @@ static mut RTT_TRACKER: Lazy<
   Option<BpfMapIterator<XdpIpAddress, RttTrackingEntry>>,
 > = Lazy::new(|| None);
 
+static mut HEIMDALL_TRACKER: Lazy<
+  Option<BpfMapIterator<HeimdallKey, HeimdallData>>,
+> = Lazy::new(|| None);
+
 pub unsafe fn iterate_throughput(
   callback: &mut dyn FnMut(&XdpIpAddress, &[HostCounter]),
 ) {
@@ -226,3 +236,28 @@ pub unsafe fn iterate_rtt(
     let _ = iter.for_each(callback);
   }
 }
+
+pub fn iterate_heimdall(
+  callback: &mut dyn FnMut(&HeimdallKey, &[HeimdallData]),
+) {
+  unsafe {
+    if HEIMDALL_TRACKER.is_none() {
+      let lock = BPF_SKELETON.lock().unwrap();
+      if let Some(skeleton) = lock.as_ref() {
+        let skeleton = skeleton.get_ptr();
+        if let Ok(iter) = {
+          BpfMapIterator::new(
+            (*skeleton).progs.heimdall_reader,
+            (*skeleton).maps.heimdall,
+          )
+        } {
+          *HEIMDALL_TRACKER = Some(iter);
+        }
+      }
+    }
+
+    if let Some(iter) = HEIMDALL_TRACKER.as_mut() {
+      let _ = iter.for_each_per_cpu(callback);
+    }
+  }
+}

src/rust/lqos_sys/src/bpf_per_cpu_map.rs

@@ -1,77 +0,0 @@
-use anyhow::{Error, Result};
-use libbpf_sys::{
-  bpf_map_get_next_key, bpf_map_lookup_elem, bpf_obj_get,
-};
-use std::fmt::Debug;
-use std::{
-  ffi::{c_void, CString},
-  marker::PhantomData,
-  ptr::null_mut,
-};
-
-use crate::num_possible_cpus;
-
-/// Represents an underlying BPF map, accessed via the filesystem.
-/// `BpfMap` *only* talks to PER-CPU variants of maps.
-///
-/// `K` is the *key* type, indexing the map.
-/// `V` is the *value* type, and must exactly match the underlying C data type.
-pub struct BpfPerCpuMap<K, V> {
-  fd: i32,
-  _key_phantom: PhantomData<K>,
-  _val_phantom: PhantomData<V>,
-}
-
-impl<K, V> BpfPerCpuMap<K, V>
-where
-  K: Default + Clone,
-  V: Default + Clone + Debug,
-{
-  /// Connect to a PER-CPU BPF map via a filename. Connects the internal
-  /// file descriptor, which is held until the structure is
-  /// dropped. The index of the CPU is *not* specified.
-  pub fn from_path(filename: &str) -> Result<Self> {
-    let filename_c = CString::new(filename)?;
-    let fd = unsafe { bpf_obj_get(filename_c.as_ptr()) };
-    if fd < 0 {
-      Err(Error::msg("Unable to open BPF map"))
-    } else {
-      Ok(Self { fd, _key_phantom: PhantomData, _val_phantom: PhantomData })
-    }
-  }
-
-  /// Instead of clonining into a vector
-  /// and allocating, calls `callback` for each key/value slice
-  /// with references to the data returned from C.
-  pub fn for_each(&self, callback: &mut dyn FnMut(&K, &[V])) {
-    let num_cpus = num_possible_cpus().unwrap();
-    let mut prev_key: *mut K = null_mut();
-    let mut key: K = K::default();
-    let key_ptr: *mut K = &mut key;
-    let mut value = vec![V::default(); num_cpus as usize];
-    let value_ptr = value.as_mut_ptr();
-
-    unsafe {
-      while bpf_map_get_next_key(
-        self.fd,
-        prev_key as *mut c_void,
-        key_ptr as *mut c_void,
-      ) == 0
-      {
-        bpf_map_lookup_elem(
-          self.fd,
-          key_ptr as *mut c_void,
-          value_ptr as *mut c_void,
-        );
-        callback(&key, &value);
-        prev_key = key_ptr;
-      }
-    }
-  }
-}
-
-impl<K, V> Drop for BpfPerCpuMap<K, V> {
-  fn drop(&mut self) {
-    let _ = nix::unistd::close(self.fd);
-  }
-}

src/rust/lqos_sys/src/heimdall_data.rs

@@ -0,0 +1,33 @@
+use lqos_utils::XdpIpAddress;
+use zerocopy::FromBytes;
+
+/// Representation of the eBPF `heimdall_key` type.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, FromBytes)]
+#[repr(C)]
+pub struct HeimdallKey {
+  /// Mapped `XdpIpAddress` source for the flow.
+  pub src_ip: XdpIpAddress,
+  /// Mapped `XdpIpAddress` destination for the flow
+  pub dst_ip: XdpIpAddress,
+  /// IP protocol (see the Linux kernel!)
+  pub ip_protocol: u8,
+  /// Source port number, or ICMP type.
+  pub src_port: u16,
+  /// Destination port number.
+  pub dst_port: u16,
+  _padding: u8,
+}
+
+/// Mapped representation of the eBPF `heimdall_data` type.
+#[derive(Debug, Clone, Default, FromBytes)]
+#[repr(C)]
+pub struct HeimdallData {
+  /// Last seen, in nanoseconds (since boot time).
+  pub last_seen: u64,
+  /// Number of bytes since the flow started being tracked
+  pub bytes: u64,
+  /// Number of packets since the flow started being tracked
+  pub packets: u64,
+  /// IP header TOS value
+  pub tos: u8,
+}

src/rust/lqos_sys/src/lib.rs

@@ -11,10 +11,6 @@ mod bifrost_maps;
 /// built-in, compiled eBPF programs. This is very-low level and should
 /// be handled with caution.
 pub mod bpf_map;
-/// Provides direct access to LibBPF functionality, as exposed by the
-/// built-in, compiled eBPF programs. This is very-low level and should
-/// be handled with caution.
-pub mod bpf_per_cpu_map;
 mod cpu_map;
 mod ip_mapping;
 mod kernel_wrapper;
@@ -23,6 +19,9 @@ mod tcp_rtt;
 mod throughput;
 mod linux;
 mod bpf_iterator;
+/// Data shared between eBPF and Heimdall that needs local access
+/// for map control.
+pub mod heimdall_data;
 
 pub use ip_mapping::{
   add_ip_to_tc, clear_ips_from_tc, del_ip_from_tc, list_mapped_ips,
@@ -32,3 +31,4 @@ pub use linux::num_possible_cpus;
 pub use lqos_kernel::max_tracked_ips;
 pub use tcp_rtt::{rtt_for_each, RttTrackingEntry};
 pub use throughput::{throughput_for_each, HostCounter};
+pub use bpf_iterator::iterate_heimdall;