From 17a57e194e9aa7cf351b63629c979805273f7637 Mon Sep 17 00:00:00 2001
From: Tony Perkins <tony.perkins@delphix.com>
Date: Tue, 10 Jun 2025 10:47:10 -0700
Subject: [PATCH] DLPX-99371: test_stbtrace_io failed because it stbtrace did
 not produce any output

PR URL: https://www.github.com/delphix/performance-diagnostics/pull/109
---
 bpf/estat/backend-io.c | 12 ++++++++----
 bpf/stbtrace/io.st     | 14 ++++++++------
 cmd/estat.py           |  2 ++
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/bpf/estat/backend-io.c b/bpf/estat/backend-io.c
index 85607ea..f93261d 100644
--- a/bpf/estat/backend-io.c
+++ b/bpf/estat/backend-io.c
@@ -30,10 +30,12 @@ typedef struct {
 
 BPF_HASH(io_base_data, u64, io_data_t);
 
-// @@ kprobe|blk_mq_start_request|disk_io_start
+// @@ raw_tracepoint|block_io_start|disk_io_start
 int
-disk_io_start(struct pt_regs *ctx, struct request *reqp)
+disk_io_start(struct bpf_raw_tracepoint_args *ctx)
 {
+	struct request *reqp = (struct request *)ctx->args[0];
+
 	io_data_t data = {};
 	struct gendisk *diskp = reqp->q->disk;
 	data.ts = bpf_ktime_get_ns();
@@ -44,10 +46,12 @@ disk_io_start(struct pt_regs *ctx, struct request *reqp)
 	return (0);
 }
 
-// @@ kprobe|blk_account_io_done|disk_io_done
+// @@ raw_tracepoint|block_io_done|disk_io_done
 int
-disk_io_done(struct pt_regs *ctx, struct request *reqp)
+disk_io_done(struct bpf_raw_tracepoint_args *ctx)
 {
+	struct request *reqp = (struct request *)ctx->args[0];
+
 	u64 ts = bpf_ktime_get_ns();
 	io_data_t *data = io_base_data.lookup((u64 *) &reqp);
 	struct bio *bp = reqp->bio;
diff --git a/bpf/stbtrace/io.st b/bpf/stbtrace/io.st
index ab935be..0647b51 100755
--- a/bpf/stbtrace/io.st
+++ b/bpf/stbtrace/io.st
@@ -63,8 +63,10 @@ $hists:{hist|
 BPF_HASH($hist.name$, io_hist_key_t, u64);
 }$
 
-int disk_io_start(struct pt_regs *ctx, struct request *reqp)
+int disk_io_start(struct bpf_raw_tracepoint_args *ctx)
 {
+    struct request *reqp = (struct request *)ctx->args[0];
+
     io_data_t data = {};
     struct gendisk *diskp = reqp->q->disk;
     data.ts = bpf_ktime_get_ns();
@@ -101,8 +103,10 @@ static int aggregate_data(io_data_t *data, u64 ts, char *opstr)
     return 0;
 }
 
-int disk_io_done(struct pt_regs *ctx, struct request *reqp)
+int disk_io_done(struct bpf_raw_tracepoint_args *ctx)
 {
+    struct request *reqp = (struct request *)ctx->args[0];
+
     u64 ts = bpf_ktime_get_ns();
     io_data_t *data = io_base_data.lookup((u64 *) &reqp);
     struct bio *bp = reqp->bio;
@@ -126,10 +130,8 @@ int disk_io_done(struct pt_regs *ctx, struct request *reqp)
 """  # noqa: W293
 b = BPF(text=bpf_text)
 
-if BPF.get_kprobe_functions(b'blk_start_request'):
-    b.attach_kprobe(event="blk_start_request", fn_name="disk_io_start")
-b.attach_kprobe(event="blk_mq_start_request", fn_name="disk_io_start")
-b.attach_kprobe(event="blk_mq_end_request", fn_name="disk_io_done")
+b.attach_raw_tracepoint("block_io_start", fn_name="disk_io_start")
+b.attach_raw_tracepoint("block_io_done", fn_name="disk_io_done")
 
 
 helper = BCCHelper(b, BCCHelper.ANALYTICS_PRINT_MODE)
diff --git a/cmd/estat.py b/cmd/estat.py
index 47d8f00..18b8c9e 100755
--- a/cmd/estat.py
+++ b/cmd/estat.py
@@ -440,6 +440,8 @@ class Args:
             else:
                 print("WARNING: {}: {} - not found"
                       .format(probe_type, probe_spec[1]))
+        elif probe_type == "raw_tracepoint":
+            b.attach_raw_tracepoint(probe_spec[1], fn_name=probe_spec[2])
         elif probe_type == "kretprobe":
             b.attach_kretprobe(event=probe_spec[1], fn_name=probe_spec[2],
                                maxactive=MAXACTIVE)