1#!/usr/bin/python
2#
3# dbslower      Trace MySQL and PostgreSQL queries slower than a threshold.
4#
5# USAGE: dbslower [-v] [-p PID [PID ...]] [-b PATH_TO_BINARY] [-m THRESHOLD]
6#                 {mysql,postgres}
7#
8# By default, a threshold of 1ms is used. Set the threshold to 0 to trace all
9# queries (verbose).
10#
11# Script works in two different modes:
12# 1) USDT probes, which means it needs MySQL and PostgreSQL built with
13# USDT (DTrace) support.
14# 2) uprobe and uretprobe on exported function of binary specified by
15# PATH_TO_BINARY parameter. (At the moment only MySQL support)
16#
17# If no PID or PATH_TO_BINARY is provided, the script attempts to discover
18# all MySQL or PostgreSQL database processes and uses USDT probes.
19#
20# Strongly inspired by Brendan Gregg's work on the mysqld_qslower script.
21#
22# Copyright 2017, Sasha Goldshtein
23# Licensed under the Apache License, Version 2.0
24#
25# 15-Feb-2017   Sasha Goldshtein   Created this.
26
27from bcc import BPF, USDT
28import argparse
29import re
30import ctypes as ct
31import subprocess
32
33examples = """examples:
34    dbslower postgres            # trace PostgreSQL queries slower than 1ms
35    dbslower postgres -p 188 322 # trace specific PostgreSQL processes
36    dbslower mysql -p 480 -m 30  # trace MySQL queries slower than 30ms
37    dbslower mysql -p 480 -v     # trace MySQL queries & print the BPF program
38    dbslower mysql -x $(which mysqld)  # trace MySQL queries with uprobes
39"""
40parser = argparse.ArgumentParser(
41    description="",
42    formatter_class=argparse.RawDescriptionHelpFormatter,
43    epilog=examples)
44parser.add_argument("-v", "--verbose", action="store_true",
45    help="print the BPF program")
46parser.add_argument("db", choices=["mysql", "postgres"],
47    help="the database engine to use")
48parser.add_argument("-p", "--pid", type=int, nargs='*',
49    dest="pids", metavar="PID", help="the pid(s) to trace")
50parser.add_argument("-x", "--exe", type=str,
51    dest="path", metavar="PATH", help="path to binary")
52parser.add_argument("-m", "--threshold", type=int, default=1,
53    help="trace queries slower than this threshold (ms)")
54parser.add_argument("--ebpf", action="store_true",
55    help=argparse.SUPPRESS)
56args = parser.parse_args()
57
58threshold_ns = args.threshold * 1000000
59
60mode = "USDT"
61if args.path and not args.pids:
62    if args.db == "mysql":
63        regex = "\\w+dispatch_command\\w+"
64        symbols = BPF.get_user_functions_and_addresses(args.path, regex)
65
66        if len(symbols) == 0:
67            print("Can't find function 'dispatch_command' in %s" % (args.path))
68            exit(1)
69
70        (mysql_func_name, addr) = symbols[0]
71
72        if mysql_func_name.find("COM_DATA") >= 0:
73            mode = "MYSQL57"
74        else:
75            mode = "MYSQL56"
76    else:
77        # Placeholder for PostrgeSQL
78        # Look on functions initStringInfo, pgstat_report_activity, EndCommand,
79        # NullCommand
80        print("Sorry at the moment PostgreSQL supports only USDT")
81        exit(1)
82
83program = """
84#include <uapi/linux/ptrace.h>
85
86DEFINE_THRESHOLD
87DEFINE_USDT
88DEFINE_MYSQL56
89DEFINE_MYSQL57
90
91struct temp_t {
92    u64 timestamp;
93#ifdef USDT
94    char *query;
95#else
96    /*
97    MySQL clears query packet before uretprobe call - so copy query in advance
98    */
99    char query[256];
100#endif //USDT
101};
102
103struct data_t {
104    u64 pid;
105    u64 timestamp;
106    u64 duration;
107    char query[256];
108};
109
110BPF_HASH(temp, u64, struct temp_t);
111BPF_PERF_OUTPUT(events);
112
113int query_start(struct pt_regs *ctx) {
114
115#if defined(MYSQL56) || defined(MYSQL57)
116/*
117Trace only packets with enum_server_command == COM_QUERY
118*/
119    #ifdef MYSQL56
120    u64 command  = (u64) PT_REGS_PARM1(ctx);
121    #else //MYSQL57
122    u64 command  = (u64) PT_REGS_PARM3(ctx);
123    #endif
124    if (command != 3) return 0;
125#endif
126
127    struct temp_t tmp = {};
128    tmp.timestamp = bpf_ktime_get_ns();
129
130#if defined(MYSQL56)
131    bpf_probe_read(&tmp.query, sizeof(tmp.query), (void*) PT_REGS_PARM3(ctx));
132#elif defined(MYSQL57)
133    void* st = (void*) PT_REGS_PARM2(ctx);
134    char* query;
135    bpf_probe_read(&query, sizeof(query), st);
136    bpf_probe_read(&tmp.query, sizeof(tmp.query), query);
137#else //USDT
138    bpf_usdt_readarg(1, ctx, &tmp.query);
139#endif
140
141    u64 pid = bpf_get_current_pid_tgid();
142    temp.update(&pid, &tmp);
143    return 0;
144}
145
146int query_end(struct pt_regs *ctx) {
147    struct temp_t *tempp;
148    u64 pid = bpf_get_current_pid_tgid();
149    tempp = temp.lookup(&pid);
150    if (!tempp)
151        return 0;
152
153    u64 delta = bpf_ktime_get_ns() - tempp->timestamp;
154#ifdef THRESHOLD
155    if (delta >= THRESHOLD) {
156#endif //THRESHOLD
157        struct data_t data = {};
158        data.pid = pid >> 32;   // only process id
159        data.timestamp = tempp->timestamp;
160        data.duration = delta;
161        bpf_probe_read(&data.query, sizeof(data.query), tempp->query);
162        events.perf_submit(ctx, &data, sizeof(data));
163#ifdef THRESHOLD
164    }
165#endif //THRESHOLD
166    temp.delete(&pid);
167    return 0;
168};
169""".replace("DEFINE_USDT", "#define USDT" if mode == "USDT" else "") \
170   .replace("DEFINE_MYSQL56", "#define MYSQL56" if mode == "MYSQL56" else "") \
171   .replace("DEFINE_MYSQL57", "#define MYSQL57" if mode == "MYSQL57" else "") \
172   .replace("DEFINE_THRESHOLD",
173            "#define THRESHOLD %d" % threshold_ns if threshold_ns > 0 else "")
174
175if mode.startswith("MYSQL"):
176    # Uprobes mode
177    bpf = BPF(text=program)
178    bpf.attach_uprobe(name=args.path, sym=mysql_func_name,
179                      fn_name="query_start")
180    bpf.attach_uretprobe(name=args.path, sym=mysql_func_name,
181                         fn_name="query_end")
182else:
183    # USDT mode
184    if not args.pids or len(args.pids) == 0:
185        if args.db == "mysql":
186            args.pids = map(int, subprocess.check_output(
187                                            "pidof mysqld".split()).split())
188        elif args.db == "postgres":
189            args.pids = map(int, subprocess.check_output(
190                                            "pidof postgres".split()).split())
191
192    usdts = map(lambda pid: USDT(pid=pid), args.pids)
193    for usdt in usdts:
194        usdt.enable_probe("query__start", "query_start")
195        usdt.enable_probe("query__done", "query_end")
196    if args.verbose:
197        print('\n'.join(map(lambda u: u.get_text(), usdts)))
198
199    bpf = BPF(text=program, usdt_contexts=usdts)
200
201if args.verbose or args.ebpf:
202    print(program)
203    if args.ebpf:
204        exit()
205
206class Data(ct.Structure):
207    _fields_ = [
208        ("pid", ct.c_ulonglong),
209        ("timestamp", ct.c_ulonglong),
210        ("delta", ct.c_ulonglong),
211        ("query", ct.c_char * 256)
212    ]
213
214start = BPF.monotonic_time()
215
216def print_event(cpu, data, size):
217    event = ct.cast(data, ct.POINTER(Data)).contents
218    print("%-14.6f %-6d %8.3f %s" % (
219        float(event.timestamp - start) / 1000000000,
220        event.pid, float(event.delta) / 1000000, event.query))
221
222if mode.startswith("MYSQL"):
223    print("Tracing database queries for application %s slower than %d ms..." %
224        (args.path, args.threshold))
225else:
226    print("Tracing database queries for pids %s slower than %d ms..." %
227        (', '.join(map(str, args.pids)), args.threshold))
228
229print("%-14s %-6s %8s %s" % ("TIME(s)", "PID", "MS", "QUERY"))
230
231bpf["events"].open_perf_buffer(print_event, page_cnt=64)
232while True:
233    bpf.perf_buffer_poll()
234