support stack trace

Signed-off-by: Tw <[email protected]>

README.md

@@ -109,5 +109,6 @@ kprobe | [source](samples/kprobe.zig) | [source](src/tests/kprobe.zig)
 kmulprobe | [source](samples/kmulprobe.zig) | [source](src/tests/kmulprobe.zig)
 xdp ping | [source](samples/xdp_ping.zig) | [source](src/tests/xdp_ping.zig)
 kfunc | [source](samples/kfunc.zig) | [source](src/tests/kfunc.zig)
+stack_trace | [source](samples/stacktrace.zig) | [source](src/tests/stacktrace.zig)
 
 **Have fun!**

samples/stacktrace.zig

@@ -0,0 +1,26 @@
+const std = @import("std");
+const bpf = @import("bpf");
+const vmlinux = @import("vmlinux");
+const get_stack = std.os.linux.BPF.kern.helpers.get_stack;
+
+var indexmap = bpf.Map.ArrayMap("indexmap", i32, 1, 0).init();
+var stackmap = bpf.Map.StackTraceMap("stackmap", 16).init();
+var astackmap = bpf.Map.ArrayMap("astackmap", bpf.Map.STACK_TRACE, 1, 0).init();
+
+const tp = bpf.Tracepoint{
+    .category = "sched",
+    .name = "sched_switch",
+};
+
+export fn test_stacktrace(ctx: *tp.Ctx()) linksection(tp.section()) callconv(.C) c_int {
+    if (indexmap.lookup(0)) |i| {
+        if (i.* < 0) {
+            const index = stackmap.get_current_stack(ctx);
+            i.* = @intCast(index);
+            if (astackmap.lookup(0)) |p| {
+                _ = get_stack(ctx, p, @sizeOf(bpf.Map.STACK_TRACE), 0);
+            } else bpf.exit(@src(), @as(c_long, 1));
+        }
+    } else bpf.exit(@src(), @as(c_long, 1));
+    return 0;
+}

src/bpf/map.zig

@@ -275,3 +275,38 @@ pub fn RingBuffer(
         }
     };
 }
+
+pub const STACK_TRACE = [127]u64;
+
+/// Represent `BPF_MAP_TYPE_STACK_TRACE`.
+pub fn StackTraceMap(
+    comptime name: []const u8,
+    comptime max_entries: u32,
+) type {
+    return struct {
+        map: Map(name, .stack_trace, u32, STACK_TRACE, max_entries, 0),
+
+        const Self = @This();
+
+        /// Initialization.
+        pub fn init() Self {
+            return .{ .map = .{} };
+        }
+
+        /// Return the pointer to the entry at index.
+        /// If index out of range, return `null`.
+        /// If any error happens, current program will exit immediately.
+        pub fn lookup(self: *const Self, index: u32) ?*STACK_TRACE {
+            return self.map.lookup(index);
+        }
+
+        /// Get current stack, return the index
+        pub fn get_current_stack(self: *const Self, ctx: *anyopaque) u32 {
+            const rc = helpers.get_stackid(ctx, @ptrCast(&@TypeOf(self.map).def), 0);
+            if (rc < 0) {
+                exit(@src(), @as(c_long, rc));
+            }
+            return @intCast(rc);
+        }
+    };
+}

src/tests/root.zig

@@ -64,4 +64,5 @@ test {
     _ = @import("ksyscall.zig");
     _ = @import("xdp_ping.zig");
     _ = @import("kfunc.zig");
+    _ = @import("stacktrace.zig");
 }

src/tests/stacktrace.zig

@@ -0,0 +1,71 @@
+const std = @import("std");
+const root = @import("root.zig");
+const print = std.debug.print;
+const testing = std.testing;
+const allocator = root.allocator;
+const libbpf = root.libbpf;
+
+const STACK_TRACE = [127]u64;
+
+test "stacktrace" {
+    const bytes = @embedFile("@stacktrace");
+
+    _ = libbpf.libbpf_set_print(root.dbg_printf);
+
+    const obj = libbpf.bpf_object__open_mem(bytes.ptr, bytes.len, null);
+    if (obj == null) {
+        print("failed to open bpf object: {}\n", .{std.posix.errno(-1)});
+        return error.OPEN;
+    }
+    defer libbpf.bpf_object__close(obj);
+
+    var ret = libbpf.bpf_object__load(obj);
+    if (ret != 0) {
+        print("failed to load bpf object: {}\n", .{std.posix.errno(-1)});
+        return error.LOAD;
+    }
+
+    if (libbpf.bpf_object__next_program(obj, null)) |prog| {
+        const stackmap = libbpf.bpf_object__find_map_by_name(obj, "stackmap").?;
+        const astackmap = libbpf.bpf_object__find_map_by_name(obj, "astackmap").?;
+        const indexmap = libbpf.bpf_object__find_map_by_name(obj, "indexmap").?;
+
+        const zero: u32 = 0;
+        var index: i32 = -1;
+        ret = libbpf.bpf_map__update_elem(indexmap, &zero, @sizeOf(u32), &index, @sizeOf(@TypeOf(index)), 0);
+        if (ret != 0) {
+            print("failed update index element: {}\n", .{std.posix.errno(-1)});
+            return error.MAP_UPDATE;
+        }
+
+        const link = libbpf.bpf_program__attach(prog) orelse {
+            print("failed to attach prog {s}: {}\n", .{ libbpf.bpf_program__name(prog), std.posix.errno(-1) });
+            return error.ATTACH;
+        };
+        defer _ = libbpf.bpf_link__destroy(link);
+
+        std.time.sleep(10);
+
+        var v: STACK_TRACE = undefined;
+        var av: STACK_TRACE = undefined;
+
+        ret = libbpf.bpf_map__lookup_elem(indexmap, &zero, @sizeOf(u32), &index, @sizeOf(@TypeOf(index)), 0);
+        if (ret != 0) {
+            print("failed lookup stackmap element: {}\n", .{std.posix.errno(-1)});
+            return error.MAP_LOOKUP;
+        }
+
+        ret = libbpf.bpf_map__lookup_elem(stackmap, &index, @sizeOf(u32), &v, @sizeOf(@TypeOf(v)), 0);
+        if (ret != 0) {
+            print("failed lookup stackmap element: {}\n", .{std.posix.errno(-1)});
+            return error.MAP_LOOKUP;
+        }
+        ret = libbpf.bpf_map__lookup_elem(astackmap, &zero, @sizeOf(u32), &av, @sizeOf(@TypeOf(av)), 0);
+        if (ret != 0) {
+            print("failed lookup astackmap element: {}\n", .{std.posix.errno(-1)});
+            return error.MAP_LOOKUP;
+        }
+
+        try testing.expectEqual(v, av);
+    }
+}