From f1b769b99e638b392b39fa8bdd15000f82f56901 Mon Sep 17 00:00:00 2001
From: xingfeng2510 <xingfeng25100@163.com>
Date: Tue, 21 Jun 2022 10:41:45 +0800
Subject: [PATCH] tools/biosnoop: Add disk filter support

---
 man/man8/biosnoop.8        |  7 +++++--
 tools/biosnoop.py          | 37 +++++++++++++++++++++++++++++++++++++
 tools/biosnoop_example.txt |  8 +++++---
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/man/man8/biosnoop.8 b/man/man8/biosnoop.8
index 9a0ae2649680..24f19edff200 100644
--- a/man/man8/biosnoop.8
+++ b/man/man8/biosnoop.8
@@ -2,7 +2,7 @@
 .SH NAME
 biosnoop \- Trace block device I/O and print details incl. issuing PID.
 .SH SYNOPSIS
-.B biosnoop [\-hQ]
+.B biosnoop [\-h] [\-Q] [\-d DISK]
 .SH DESCRIPTION
 This tools traces block device I/O (disk I/O), and prints a one-line summary
 for each I/O showing various details. These include the latency from the time of
@@ -29,6 +29,9 @@ Print usage message.
 .TP
 \-Q
 Include a column showing the time spent queued in the OS.
+.TP
+\-d DISK
+Trace this disk only.
 .SH EXAMPLES
 .TP
 Trace all block device I/O and print a summary line per I/O:
@@ -82,6 +85,6 @@ Linux
 .SH STABILITY
 Unstable - in development.
 .SH AUTHOR
-Brendan Gregg
+Brendan Gregg, Rocky Xing
 .SH SEE ALSO
 disksnoop(8), iostat(1)
diff --git a/tools/biosnoop.py b/tools/biosnoop.py
index 3f3ebd2e4f50..e0bdad1cc31d 100755
--- a/tools/biosnoop.py
+++ b/tools/biosnoop.py
@@ -12,16 +12,19 @@
 #
 # 16-Sep-2015   Brendan Gregg   Created this.
 # 11-Feb-2016   Allan McAleavy  updated for BPF_PERF_OUTPUT
+# 21-Jun-2022   Rocky Xing      Added disk filter support.
 
 from __future__ import print_function
 from bcc import BPF
 import re
 import argparse
+import os
 
 # arguments
 examples = """examples:
     ./biosnoop           # trace all block I/O
     ./biosnoop -Q        # include OS queued time
+    ./biolatency -d sdc  # trace sdc only
 """
 parser = argparse.ArgumentParser(
     description="Trace block I/O",
@@ -29,6 +32,8 @@
     epilog=examples)
 parser.add_argument("-Q", "--queue", action="store_true",
     help="include OS queued time")
+parser.add_argument("-d", "--disk", type=str,
+    help="Trace this disk only")
 parser.add_argument("--ebpf", action="store_true",
     help=argparse.SUPPRESS)
 args = parser.parse_args()
@@ -70,6 +75,8 @@
 // cache PID and comm by-req
 int trace_pid_start(struct pt_regs *ctx, struct request *req)
 {
+    DISK_FILTER
+
     struct val_t val = {};
     u64 ts;
 
@@ -86,6 +93,8 @@
 // time block I/O
 int trace_req_start(struct pt_regs *ctx, struct request *req)
 {
+    DISK_FILTER
+
     struct start_req_t start_req = {
         .ts = bpf_ktime_get_ns(),
         .data_len = req->__data_len
@@ -160,6 +169,34 @@
     bpf_text = bpf_text.replace('__RQ_DISK__', 'rq_disk')
 else:
     bpf_text = bpf_text.replace('__RQ_DISK__', 'q->disk')
+
+if args.disk is not None:
+    disk_path = os.path.join('/dev', args.disk)
+    if not os.path.exists(disk_path):
+        print("no such disk '%s'" % args.disk)
+        exit(1)
+
+    stat_info = os.stat(disk_path)
+    major = os.major(stat_info.st_rdev)
+    minor = os.minor(stat_info.st_rdev)
+
+    disk_field_str = ""
+    if BPF.kernel_struct_has_field(b'request', b'rq_disk') == 1:
+        disk_field_str = 'req->rq_disk'
+    else:
+        disk_field_str = 'req->q->disk'
+
+    disk_filter_str = """
+    struct gendisk *disk = %s;
+    if (!(disk->major == %d && disk->first_minor == %d)) {
+        return 0;
+    }
+    """ % (disk_field_str, major, minor)
+
+    bpf_text = bpf_text.replace('DISK_FILTER', disk_filter_str)
+else:
+    bpf_text = bpf_text.replace('DISK_FILTER', '')
+
 if debug or args.ebpf:
     print(bpf_text)
     if args.ebpf:
diff --git a/tools/biosnoop_example.txt b/tools/biosnoop_example.txt
index d8be0624c94e..38b0ca3431f9 100644
--- a/tools/biosnoop_example.txt
+++ b/tools/biosnoop_example.txt
@@ -64,14 +64,16 @@ TIME(s)     COMM           PID    DISK    T SECTOR     BYTES  QUE(ms) LAT(ms)
 
 USAGE message:
 
-usage: biosnoop.py [-h] [-Q]
+usage: biosnoop.py [-h] [-Q] [-d DISK]
 
 Trace block I/O
 
 optional arguments:
-  -h, --help   show this help message and exit
-  -Q, --queue  include OS queued time
+  -h, --help            show this help message and exit
+  -Q, --queue           include OS queued time
+  -d DISK, --disk DISK  Trace this disk only
 
 examples:
     ./biosnoop           # trace all block I/O
     ./biosnoop -Q        # include OS queued time
+    ./biolatency -d sdc  # trace sdc only