libbpf: ringbuf: Allow to consume up to a certain amount of items

In some cases, instead of always consuming all items from ring buffers
in a greedy way, we may want to consume up to a certain amount of items,
for example when we need to copy items from the BPF ring buffer to a
limited user buffer.

This change allows to set an upper limit to the amount of items consumed
from one or more ring buffers.

Signed-off-by: Andrea Righi <andrea.righi@canonical.com>
Signed-off-by: Andrii Nakryiko <andrii@kernel.org>
Link: https://lore.kernel.org/bpf/20240406092005.92399-3-andrea.righi@canonical.com

diff --git a/tools/lib/bpf/ringbuf.c b/tools/lib/bpf/ringbuf.c
index aacb64278a018c22aab25038f90bccb182577e2d..db05e6f526de751c1c69dedf513b937c83ac66d8 100644 (file)
--- a/tools/lib/bpf/ringbuf.c
+++ b/tools/lib/bpf/ringbuf.c
@@ -231,7 +231,7 @@ static inline int roundup_len(__u32 len)
        return (len + 7) / 8 * 8;
 }
 
-static int64_t ringbuf_process_ring(struct ring *r)
+static int64_t ringbuf_process_ring(struct ring *r, size_t n)
 {
        int *len_ptr, len, err;
        /* 64-bit to avoid overflow in case of extreme application behavior */
@@ -268,6 +268,9 @@ static int64_t ringbuf_process_ring(struct ring *r)
                        }
 
                        smp_store_release(r->consumer_pos, cons_pos);
+
+                       if (cnt >= n)
+                               goto done;
                }
        } while (got_new_data);
 done:
@@ -287,13 +290,15 @@ int ring_buffer__consume(struct ring_buffer *rb)
        for (i = 0; i < rb->ring_cnt; i++) {
                struct ring *ring = rb->rings[i];
 
-               err = ringbuf_process_ring(ring);
+               err = ringbuf_process_ring(ring, INT_MAX);
                if (err < 0)
                        return libbpf_err(err);
                res += err;
+               if (res > INT_MAX) {
+                       res = INT_MAX;
+                       break;
+               }
        }
-       if (res > INT_MAX)
-               return INT_MAX;
        return res;
 }
 
@@ -314,13 +319,13 @@ int ring_buffer__poll(struct ring_buffer *rb, int timeout_ms)
                __u32 ring_id = rb->events[i].data.fd;
                struct ring *ring = rb->rings[ring_id];
 
-               err = ringbuf_process_ring(ring);
+               err = ringbuf_process_ring(ring, INT_MAX);
                if (err < 0)
                        return libbpf_err(err);
                res += err;
        }
        if (res > INT_MAX)
-               return INT_MAX;
+               res = INT_MAX;
        return res;
 }
 
@@ -375,7 +380,7 @@ int ring__consume(struct ring *r)
 {
        int64_t res;
 
-       res = ringbuf_process_ring(r);
+       res = ringbuf_process_ring(r, INT_MAX);
        if (res < 0)
                return libbpf_err(res);