Use plain arrays in ring buf tests

Store ring buffer accounting in plain arrays instead of per-CPU arrays
and use atomic increments in the eBPF test program where concurrent
writers still share the counters.

Co-authored-by: Codex <noreply@openai.com>

Author: tamird

ebpf/aya-ebpf/src/programs/probe.rs

@@ -26,7 +26,7 @@ impl ProbeContext {
     /// unsafe fn try_kprobe_try_to_wake_up(ctx: ProbeContext) -> Result<u32, u32> {
     ///     let tp: *const task_struct = ctx.arg(0).ok_or(1u32)?;
     ///     let pid = unsafe {
-    ///         bpf_probe_read(core::ptr::addr_of!((*tp).pid))
+    ///         bpf_probe_read(&raw const (*tp).pid)
     ///     }.map_err(|err| err as u32)?;
     ///
     ///     // Do something with pid or something else with tp

test/integration-common/src/lib.rs

@@ -62,22 +62,6 @@ pub mod ring_buf {
         pub rejected: u64,
     }
 
-    impl core::ops::Add for Registers {
-        type Output = Self;
-        fn add(self, rhs: Self) -> Self::Output {
-            Self {
-                dropped: self.dropped + rhs.dropped,
-                rejected: self.rejected + rhs.rejected,
-            }
-        }
-    }
-
-    impl<'a> core::iter::Sum<&'a Self> for Registers {
-        fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
-            iter.fold(Default::default(), |a, b| a + *b)
-        }
-    }
-
     #[cfg(feature = "user")]
     unsafe impl aya::Pod for Registers {}
 }

test/integration-ebpf/src/kprobe.rs

@@ -1,18 +1,9 @@
-#![cfg_attr(
-    target_arch = "bpf",
-    expect(
-        internal_features,
-        reason = "core_intrinsics is required for atomic_xadd on BPF targets"
-    ),
-    expect(
-        unstable_features,
-        reason = "core_intrinsics is required for atomic_xadd on BPF targets"
-    ),
-    feature(core_intrinsics)
-)]
 #![no_std]
 #![no_main]
 #![expect(unused_crate_dependencies, reason = "used in other bins")]
+#![expect(internal_features, reason = "atomic_xadd is unstable")]
+#![expect(unstable_features, reason = "atomic_xadd is unstable")]
+#![feature(core_intrinsics)]
 
 use aya_ebpf::{
     EbpfContext as _, Global,
@@ -46,14 +37,11 @@ fn test_kprobe_trigger(ctx: ProbeContext) -> u32 {
         return 0;
     };
 
-    #[cfg(target_arch = "bpf")]
     unsafe {
         core::intrinsics::atomic_xadd::<u64, u64, { core::intrinsics::AtomicOrdering::Relaxed }>(
             hits, 1,
         );
     }
 
-    let _: *mut u64 = hits;
-
     0
 }

test/integration-ebpf/src/ring_buf.rs

@@ -1,11 +1,14 @@
 #![no_std]
 #![no_main]
 #![expect(unused_crate_dependencies, reason = "used in other bins")]
+#![expect(internal_features, reason = "atomic_xadd is unstable")]
+#![expect(unstable_features, reason = "atomic_xadd is unstable")]
+#![feature(core_intrinsics)]
 
 use aya_ebpf::{
     btf_maps::RingBuf as BtfRingBuf,
     macros::{btf_map, map, uprobe},
-    maps::{PerCpuArray, RingBuf as LegacyRingBuf},
+    maps::{Array, RingBuf as LegacyRingBuf},
     programs::ProbeContext,
 };
 use integration_common::ring_buf::Registers;
@@ -21,26 +24,27 @@ static RING_BUF_MISMATCH: BtfRingBuf<u32, 0, 0> = BtfRingBuf::new();
 #[map]
 static RING_BUF_LEGACY: LegacyRingBuf = LegacyRingBuf::with_byte_size(0, 0);
 
-// Use a PerCpuArray to store the registers so that we can update the values from multiple CPUs
-// without needing synchronization. Atomics exist [1], but aren't exposed.
-//
-// [1]: https://lwn.net/Articles/838884/
 #[map]
-static REGISTERS: PerCpuArray<Registers> = PerCpuArray::with_max_entries(1, 0);
+static REGISTERS: Array<Registers> = Array::with_max_entries(1, 0);
 
 #[map]
-static REGISTERS_LEGACY: PerCpuArray<Registers> = PerCpuArray::with_max_entries(1, 0);
+static REGISTERS_LEGACY: Array<Registers> = Array::with_max_entries(1, 0);
 
 macro_rules! define_ring_buf_test {
     ($registers:ident, $name:ident, $reserve:expr) => {
         #[uprobe]
         fn $name(ctx: ProbeContext) {
-            let Registers { dropped, rejected } = match $registers.get_ptr_mut(0) {
-                Some(regs) => unsafe { &mut *regs },
-                None => return,
+            let Some(regs) = $registers.get_ptr_mut(0) else {
+                return;
             };
             let Some(mut entry) = $reserve else {
-                *dropped += 1;
+                unsafe {
+                    core::intrinsics::atomic_xadd::<
+                        u64,
+                        u64,
+                        { core::intrinsics::AtomicOrdering::Relaxed },
+                    >(&raw mut (*regs).dropped, 1);
+                }
                 return;
             };
             // Write the first argument to the function back out to RING_BUF if it is even,
@@ -53,7 +57,13 @@ macro_rules! define_ring_buf_test {
                 entry.write(arg);
                 entry.submit(0);
             } else {
-                *rejected += 1;
+                unsafe {
+                    core::intrinsics::atomic_xadd::<
+                        u64,
+                        u64,
+                        { core::intrinsics::AtomicOrdering::Relaxed },
+                    >(&raw mut (*regs).rejected, 1);
+                }
                 entry.discard(0);
             }
         }

test/integration-test/src/tests/ring_buf.rs

@@ -13,7 +13,7 @@ use anyhow::Context as _;
 use assert_matches::assert_matches;
 use aya::{
     Ebpf, EbpfLoader,
-    maps::{MapData, array::PerCpuArray, ring_buf::RingBuf},
+    maps::{Array, MapData, ring_buf::RingBuf},
     programs::UProbe,
 };
 use aya_obj::generated::BPF_RINGBUF_HDR_SZ;
@@ -25,7 +25,7 @@ use tokio::io::{Interest, unix::AsyncFd};
 struct RingBufTest {
     bpf: Ebpf,
     ring_buf: RingBuf<MapData>,
-    regs: PerCpuArray<MapData, Registers>,
+    regs: Array<MapData, Registers>,
 }
 
 const RING_BUF: &str = "RING_BUF";
@@ -87,7 +87,7 @@ impl RingBufTest {
         let ring_buf = bpf.take_map(variant.map).unwrap();
         let ring_buf = RingBuf::try_from(ring_buf).unwrap();
         let regs = bpf.take_map(variant.regs).unwrap();
-        let regs = PerCpuArray::<_, Registers>::try_from(regs).unwrap();
+        let regs = Array::<_, Registers>::try_from(regs).unwrap();
         let prog: &mut UProbe = bpf.program_mut(variant.prog).unwrap().try_into().unwrap();
         prog.load().unwrap();
         prog.attach("ring_buf_trigger_ebpf_program", "/proc/self/exe", None)
@@ -166,7 +166,7 @@ fn ring_buf(n: usize) {
         // Ensure that the data that was read matches what was passed, and the rejected count was set
         // properly.
         assert_eq!(seen, expected);
-        let Registers { dropped, rejected } = regs.get(&0, 0).unwrap().iter().sum();
+        let Registers { dropped, rejected } = regs.get(&0, 0).unwrap();
         assert_eq!(dropped, expected_dropped);
         assert_eq!(rejected, expected_rejected);
     }
@@ -322,7 +322,7 @@ async fn ring_buf_async_with_drops() {
             u64::try_from(data.len().saturating_sub(RING_BUF_MAX_ENTRIES - 1)).unwrap();
         let max_seen = u64::try_from(data.iter().filter(|v| *v % 2 == 0).count()).unwrap();
         let max_rejected = u64::try_from(data.len()).unwrap() - max_seen;
-        let Registers { dropped, rejected } = regs.get(&0, 0).unwrap().iter().sum();
+        let Registers { dropped, rejected } = regs.get(&0, 0).unwrap();
         let total = u64::try_from(data.len()).unwrap();
         let min_seen = max_seen.saturating_sub(max_dropped);
         let min_rejected = max_rejected.saturating_sub(dropped);
@@ -408,7 +408,7 @@ async fn ring_buf_async_no_drop() {
 
         // Ensure that the data that was read matches what was passed.
         assert_eq!(&seen, &expected);
-        let Registers { dropped, rejected } = regs.get(&0, 0).unwrap().iter().sum();
+        let Registers { dropped, rejected } = regs.get(&0, 0).unwrap();
         assert_eq!(dropped, 0);
         assert_eq!(rejected, (data.len() - expected.len()).try_into().unwrap());
     }