From 700041e5c7dac333f1fc3c5617e92418fd4e0fc7 Mon Sep 17 00:00:00 2001
From: Tw <tw19881113@gmail.com>
Date: Wed, 22 Jan 2025 12:15:18 +0800
Subject: [PATCH] tools/trace: add max count support

---
 src/tools/trace/trace.zig | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/tools/trace/trace.zig b/src/tools/trace/trace.zig
index 5e017f1..a13cc7e 100644
--- a/src/tools/trace/trace.zig
+++ b/src/tools/trace/trace.zig
@@ -40,6 +40,7 @@ fn usage() void {
         \\ --testing
         \\ --lbr
         \\ --vmlinux [/path/to/vmlinux]
+        \\ --count [n]
     ++ "\n", .{});
 }
 
@@ -71,6 +72,7 @@ pub fn main() !void {
 
     var seconds: ?usize = null;
     var vmlinux_path: ?[]const u8 = null;
+    var max_count: ?usize = null;
     while (nextArg(args, &arg_idx)) |arg| {
         if (std.mem.eql(u8, arg, "--timeout")) {
             const s = nextArg(args, &arg_idx) orelse {
@@ -90,6 +92,12 @@ pub fn main() !void {
                 usage();
                 return error.ARGS;
             };
+        } else if (std.mem.eql(u8, arg, "--count")) {
+            const s = nextArg(args, &arg_idx) orelse {
+                usage();
+                return error.ARGS;
+            };
+            max_count = try std.fmt.parseUnsigned(usize, s, 0);
         } else {
             print("unknown parameter\n", .{});
             usage();
@@ -157,12 +165,22 @@ pub fn main() !void {
         _ = testing_call(1, 2);
     }
     const begin_ts = std.time.timestamp();
+    var consumed: usize = 0;
     while (!exiting) {
-        ret = libbpf.ring_buffer__poll(ring_buf, 100);
+        if (max_count) |max| {
+            ret = libbpf.ring_buffer__consume_n(ring_buf, max - consumed);
+        } else {
+            ret = libbpf.ring_buffer__poll(ring_buf, 100);
+        }
 
         if (ret < 0) {
             return error.POLL;
         }
+        consumed += @intCast(ret);
+
+        if (max_count) |max| {
+            if (consumed >= max) break;
+        }
 
         if (seconds) |timeout| {
             const cur_ts = std.time.timestamp();