MySQL tracing without USDT (#1239)

Support tracing MySQL queries even when MySQL is built
without USDT support, by using uprobes on internal functions
responsible for command (query) dispatching.

diff --git a/tools/dbslower.py b/tools/dbslower.py
index e7924e4..b32f6c9 100755 (executable)
--- a/tools/dbslower.py
+++ b/tools/dbslower.py
@@ -2,14 +2,19 @@
 #
 # dbslower      Trace MySQL and PostgreSQL queries slower than a threshold.
 #
-# USAGE: dbslower [-v] [-p PID [PID ...]] [-m THRESHOLD] {mysql,postgres}
+# USAGE: dbslower [-v] [-p PID [PID ...]] [-b PATH_TO_BINARY] [-m THRESHOLD] {mysql,postgres}
 #
 # By default, a threshold of 1ms is used. Set the threshold to 0 to trace all
-# queries (verbose). If no PID is provided, the script attempts to discover
-# all MySQL or PostgreSQL database processes.
-#
-# This tool uses USDT probes, which means it needs MySQL and PostgreSQL built
-# with USDT (DTrace) support.
+# queries (verbose). 
+# 
+# Script works in two different modes: 
+# 1) USDT probes, which means it needs MySQL and PostgreSQL built with 
+# USDT (DTrace) support.
+# 2) uprobe and uretprobe on exported function of binary specified by 
+# PATH_TO_BINARY parameter. (At the moment only MySQL support)
+# 
+# If no PID or PATH_TO_BINARY is provided, the script attempts to discover
+# all MySQL or PostgreSQL database processes and uses USDT probes.
 #
 # Strongly inspired by Brendan Gregg's work on the mysqld_qslower script.
 #
@@ -20,6 +25,7 @@
 
 from bcc import BPF, USDT
 import argparse
+import re
 import ctypes as ct
 import subprocess
 
@@ -28,6 +34,7 @@ examples = """examples:
     dbslower postgres -p 188 322 # trace specific PostgreSQL processes
     dbslower mysql -p 480 -m 30  # trace MySQL queries slower than 30ms
     dbslower mysql -p 480 -v     # trace MySQL queries & print the BPF program
+    dbslower mysql -x $(which mysqld)  # trace MySQL queries with uprobes
 """
 parser = argparse.ArgumentParser(
     description="",
@@ -39,26 +46,53 @@ parser.add_argument("db", choices=["mysql", "postgres"],
     help="the database engine to use")
 parser.add_argument("-p", "--pid", type=int, nargs='*',
     dest="pids", metavar="PID", help="the pid(s) to trace")
+parser.add_argument("-x", "--exe", type=str,
+    dest="path", metavar="PATH", help="path to binary")
 parser.add_argument("-m", "--threshold", type=int, default=1,
     help="trace queries slower than this threshold (ms)")
 args = parser.parse_args()
 
-if not args.pids or len(args.pids) == 0:
-    if args.db == "mysql":
-        args.pids = map(int, subprocess.check_output(
-                                        "pidof mysqld".split()).split())
-    elif args.db == "postgres":
-        args.pids = map(int, subprocess.check_output(
-                                        "pidof postgres".split()).split())
-
 threshold_ns = args.threshold * 1000000
 
+mode = "USDT"
+if args.path and not args.pids:
+    if args.db == "mysql":
+        symbols = BPF.get_user_functions_and_addresses(args.path, "\\w+dispatch_command\\w+")
+
+        if len(symbols) == 0:
+            print("Can't find function 'dispatch_command' in %s" % (args.path))
+            exit(1)
+        
+        (mysql_func_name, addr) = symbols[0]
+
+        if mysql_func_name.find("COM_DATA") >= 0:
+            mode = "MYSQL57"
+        else:
+            mode = "MYSQL56"
+    else:
+        # Placeholder for PostrgeSQL
+        # Look on functions initStringInfo, pgstat_report_activity, EndCommand, NullCommand
+        print("Sorry at the moment PostgreSQL supports only USDT")
+        exit(1)
+
 program = """
 #include <uapi/linux/ptrace.h>
 
+DEFINE_THRESHOLD
+DEFINE_USDT
+DEFINE_MYSQL56
+DEFINE_MYSQL57
+
 struct temp_t {
     u64 timestamp;
+#ifdef USDT
     char *query;
+#else
+    /*
+    MySQL clears query packet before uretprobe call - so copy query in advance
+    */
+    char query[256];
+#endif //USDT
 };
 
 struct data_t {
@@ -71,16 +105,40 @@ struct data_t {
 BPF_HASH(temp, u64, struct temp_t);
 BPF_PERF_OUTPUT(events);
 
-int probe_start(struct pt_regs *ctx) {
+int query_start(struct pt_regs *ctx) {
+
+#if defined(MYSQL56) || defined(MYSQL57)
+/*
+Trace only packets with enum_server_command == COM_QUERY 
+*/
+    #ifdef MYSQL56
+    u64 command  = (u64) PT_REGS_PARM1(ctx);
+    #else //MYSQL57
+    u64 command  = (u64) PT_REGS_PARM3(ctx);
+    #endif
+    if (command != 3) return 0;
+#endif
+
     struct temp_t tmp = {};
     tmp.timestamp = bpf_ktime_get_ns();
+
+#if defined(MYSQL56)
+    bpf_probe_read(&tmp.query, sizeof(tmp.query), (void*) PT_REGS_PARM3(ctx));
+#elif defined(MYSQL57)
+    void* st = (void*) PT_REGS_PARM2(ctx);
+    char* query;
+    bpf_probe_read(&query, sizeof(query), st);
+    bpf_probe_read(&tmp.query, sizeof(tmp.query), query);
+#else //USDT
     bpf_usdt_readarg(1, ctx, &tmp.query);
+#endif
+
     u64 pid = bpf_get_current_pid_tgid();
     temp.update(&pid, &tmp);
     return 0;
 }
 
-int probe_end(struct pt_regs *ctx) {
+int query_end(struct pt_regs *ctx) {
     struct temp_t *tempp;
     u64 pid = bpf_get_current_pid_tgid();
     tempp = temp.lookup(&pid);
@@ -88,30 +146,53 @@ int probe_end(struct pt_regs *ctx) {
         return 0;
 
     u64 delta = bpf_ktime_get_ns() - tempp->timestamp;
-    if (delta >= """ + str(threshold_ns) + """) {
+#ifdef THRESHOLD
+    if (delta >= THRESHOLD) {
+#endif //THRESHOLD     
         struct data_t data = {};
         data.pid = pid >> 32;   // only process id
         data.timestamp = tempp->timestamp;
         data.duration = delta;
         bpf_probe_read(&data.query, sizeof(data.query), tempp->query);
         events.perf_submit(ctx, &data, sizeof(data));
+#ifdef THRESHOLD
     }
+#endif //THRESHOLD
     temp.delete(&pid);
     return 0;
-}
-"""
-
-usdts = map(lambda pid: USDT(pid=pid), args.pids)
-for usdt in usdts:
-    usdt.enable_probe("query__start", "probe_start")
-    usdt.enable_probe("query__done", "probe_end")
+};
+""".replace("DEFINE_USDT", "#define USDT" if mode == "USDT" else "") \
+   .replace("DEFINE_MYSQL56", "#define MYSQL56" if mode == "MYSQL56" else "") \
+   .replace("DEFINE_MYSQL57", "#define MYSQL57" if mode == "MYSQL57" else "") \
+   .replace("DEFINE_THRESHOLD", ("#define THRESHOLD " + str(threshold_ns)) if threshold_ns > 0 else "")
+
+if mode.startswith("MYSQL"):
+    # Uprobes mode
+    bpf = BPF(text=program)
+    bpf.attach_uprobe(name=args.path, sym=mysql_func_name, fn_name="query_start")
+    bpf.attach_uretprobe(name=args.path, sym=mysql_func_name, fn_name="query_end")
+else:
+    # USDT mode
+    if not args.pids or len(args.pids) == 0:
+        if args.db == "mysql":
+            args.pids = map(int, subprocess.check_output(
+                                            "pidof mysqld".split()).split())
+        elif args.db == "postgres":
+            args.pids = map(int, subprocess.check_output(
+                                            "pidof postgres".split()).split())
+
+    usdts = map(lambda pid: USDT(pid=pid), args.pids)
+    for usdt in usdts:
+        usdt.enable_probe("query__start", "query_start")
+        usdt.enable_probe("query__done", "query_end")
+    if args.verbose:
+        print('\n'.join(map(lambda u: u.get_text(), usdts)))
+
+    bpf = BPF(text=program, usdt_contexts=usdts)
 
 if args.verbose:
-    print('\n'.join(map(lambda u: u.get_text(), usdts)))
     print(program)
 
-bpf = BPF(text=program, usdt_contexts=usdts)
-
 class Data(ct.Structure):
     _fields_ = [
         ("pid", ct.c_ulonglong),
@@ -128,8 +209,13 @@ def print_event(cpu, data, size):
         float(event.timestamp - start) / 1000000000,
         event.pid, float(event.delta) / 1000000, event.query))
 
-print("Tracing database queries for pids %s slower than %d ms..." %
-      (', '.join(map(str, args.pids)), args.threshold))
+if mode.startswith("MYSQL"):
+    print("Tracing database queries for application %s slower than %d ms..." %
+        (args.path, args.threshold))
+else:
+    print("Tracing database queries for pids %s slower than %d ms..." %
+        (', '.join(map(str, args.pids)), args.threshold))
+
 print("%-14s %-6s %8s %s" % ("TIME(s)", "PID", "MS", "QUERY"))
 
 bpf["events"].open_perf_buffer(print_event, page_cnt=64)