PR feedback

Signed-off-by: Alan Jowett <[email protected]>
microsoft · Nov 20, 2024 · b2da4a4 · b2da4a4
1 parent f3bef43
commit b2da4a4
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 19 deletions.
diff --git a/libs/api/ebpf_api.cpp b/libs/api/ebpf_api.cpp
@@ -4266,10 +4266,9 @@ _ebpf_ring_buffer_map_async_query_completion(_Inout_ void* completion_context) N
                 break;
             }
 
-            long size;
             for (;;) {
-                size = ReadAcquire(&record->size);
-                if (size & EBPF_RING_BUFFER_RECORD_FLAG_LOCKED) {
+
+                if (ebpf_ring_buffer_record_is_locked(record)) {
                     // The record is being written to by the producer.
                     // Wait for the producer to finish writing.
                     YieldProcessor();
@@ -4278,16 +4277,16 @@ _ebpf_ring_buffer_map_async_query_completion(_Inout_ void* completion_context) N
                 }
             }
 
-            if (!(size & EBPF_RING_BUFFER_RECORD_FLAG_DISCARDED)) {
+            if (!ebpf_ring_buffer_record_is_discarded(record)) {
                 int callback_result = subscription->sample_callback(
                     subscription->sample_callback_context,
                     const_cast<void*>(reinterpret_cast<const void*>(record->data)),
-                    size - EBPF_OFFSET_OF(ebpf_ring_buffer_record_t, data));
+                    ebpf_ring_buffer_record_size(record) - EBPF_OFFSET_OF(ebpf_ring_buffer_record_t, data));
                 if (callback_result != 0) {
                     break;
                 }
             }
-            consumer += size;
+            consumer += ebpf_ring_buffer_record_size(record);
         }
     }
 

diff --git a/libs/runtime/ebpf_ring_buffer.c b/libs/runtime/ebpf_ring_buffer.c
@@ -157,24 +157,18 @@ ebpf_ring_buffer_reserve(
             (ebpf_ring_buffer_record_t*)(ring->shared_buffer + (producer_offset & ring->mask));
 
         int64_t remaining_space = ring->length - (producer_offset - ReadAcquire64(&ring->consumer_offset));
-        long effective_length = (long)length + 4;
+        long effective_length = (long)length + EBPF_OFFSET_OF(ebpf_ring_buffer_record_t, data);
         if (remaining_space < effective_length) {
             return EBPF_NO_MEMORY;
         }
 
         // Check if locked.
-        if (record->size != 0) {
-            // If locked, pause then try again.
-            // Use _mm_pause() on x86 and __yield() on ARM.
-#if defined(_M_X64) || defined(_M_IX86)
-            _mm_pause();
-#else
-            __yield();
-#endif
+        if (ReadNoFence(&record->size) != 0) {
+            YieldProcessor();
             continue;
         }
 
-        if (InterlockedBitTestAndSetAcquire(&record->size, 31) != 0) {
+        if (InterlockedBitTestAndSetAcquire(&record->size, EBPF_RING_BUFFER_RECORD_FLAG_LOCKED_OFFSET) != 0) {
             continue;
         }
 

diff --git a/libs/runtime/unit/platform_unit_test.cpp b/libs/runtime/unit/platform_unit_test.cpp
@@ -1233,7 +1233,7 @@ TEST_CASE("ring_buffer_stress", "[platform]")
 
     REQUIRE(!bad_record);
     REQUIRE(a_records > 0);
-    // REQUIRE(b_records > 0);
+    REQUIRE(b_records > 0);
 
     ebpf_ring_buffer_destroy(ring_buffer);
 }

diff --git a/libs/shared/ebpf_ring_buffer_record.h b/libs/shared/ebpf_ring_buffer_record.h
@@ -5,8 +5,10 @@
 
 CXPLAT_EXTERN_C_BEGIN
 
-#define EBPF_RING_BUFFER_RECORD_FLAG_LOCKED (long)(0x1ul << 31)
-#define EBPF_RING_BUFFER_RECORD_FLAG_DISCARDED (long)(0x1ul << 30)
+#define EBPF_RING_BUFFER_RECORD_FLAG_LOCKED_OFFSET 31
+#define EBPF_RING_BUFFER_RECORD_FLAG_DISCARDED_OFFSET 30
+#define EBPF_RING_BUFFER_RECORD_FLAG_LOCKED (long)(0x1ul << EBPF_RING_BUFFER_RECORD_FLAG_LOCKED_OFFSET)
+#define EBPF_RING_BUFFER_RECORD_FLAG_DISCARDED (long)(0x1ul << EBPF_RING_BUFFER_RECORD_FLAG_DISCARDED_OFFSET)
 
 typedef struct _ebpf_ring_buffer_record
 {
@@ -34,4 +36,23 @@ ebpf_ring_buffer_next_record(_In_ const uint8_t* buffer, size_t buffer_length, s
     return (ebpf_ring_buffer_record_t*)(buffer + consumer % buffer_length);
 }
 
+inline const bool
+ebpf_ring_buffer_record_is_discarded(_In_ const ebpf_ring_buffer_record_t* record)
+{
+    return (ReadNoFence(&record->size) & EBPF_RING_BUFFER_RECORD_FLAG_DISCARDED) != 0;
+}
+
+inline const bool
+ebpf_ring_buffer_record_is_locked(_In_ const ebpf_ring_buffer_record_t* record)
+{
+    return (ReadNoFence(&record->size) & EBPF_RING_BUFFER_RECORD_FLAG_LOCKED) != 0;
+}
+
+inline const size_t
+ebpf_ring_buffer_record_size(_In_ const ebpf_ring_buffer_record_t* record)
+{
+    return (size_t)(ReadNoFence(&record->size) &
+                    ~(EBPF_RING_BUFFER_RECORD_FLAG_LOCKED | EBPF_RING_BUFFER_RECORD_FLAG_DISCARDED));
+}
+
 CXPLAT_EXTERN_C_END